@@ -3,23 +3,23 @@ import { SpanAttributes } from "@traceloop/ai-semantic-conventions";
33
44const AI_GENERATE_TEXT_DO_GENERATE = "ai.generateText.doGenerate" ;
55const AI_STREAM_TEXT_DO_STREAM = "ai.streamText.doStream" ;
6+ const HANDLED_SPAN_NAMES : Record < string , string > = {
7+ [ AI_GENERATE_TEXT_DO_GENERATE ] : "ai.generateText.generate" ,
8+ [ AI_STREAM_TEXT_DO_STREAM ] : "ai.streamText.stream" ,
9+ } ;
10+
611const AI_RESPONSE_TEXT = "ai.response.text" ;
712const AI_PROMPT_MESSAGES = "ai.prompt.messages" ;
813const AI_USAGE_PROMPT_TOKENS = "ai.usage.promptTokens" ;
914const AI_USAGE_COMPLETION_TOKENS = "ai.usage.completionTokens" ;
1015const AI_MODEL_PROVIDER = "ai.model.provider" ;
1116
12- export const transformAiSdkSpanName = ( span : ReadableSpan ) : void => {
13- const nameMap : Record < string , string > = {
14- [ AI_GENERATE_TEXT_DO_GENERATE ] : "ai.generateText.generate" ,
15- [ AI_STREAM_TEXT_DO_STREAM ] : "ai.streamText.stream" ,
16- } ;
1717
18- if ( span . name in nameMap ) {
18+
19+ export const transformAiSdkSpanName = ( span : ReadableSpan ) : void => {
1920 // Unfortunately, the span name is not writable as this is not the intended behavior
2021 // but it is a workaround to set the correct span name
21- ( span as any ) . name = nameMap [ span . name ] ;
22- }
22+ ( span as any ) . name = HANDLED_SPAN_NAMES [ span . name ] ;
2323} ;
2424
2525export const transformResponseText = (
@@ -107,7 +107,14 @@ export const transformAiSdkAttributes = (
107107 transformVendor ( attributes ) ;
108108} ;
109109
110+ export const shouldHandleSpan = ( span : ReadableSpan ) : boolean => {
111+ return span . name in HANDLED_SPAN_NAMES ;
112+ } ;
113+
110114export const transformAiSdkSpan = ( span : ReadableSpan ) : void => {
115+ if ( ! shouldHandleSpan ( span ) ) {
116+ return ;
117+ }
111118 transformAiSdkSpanName ( span ) ;
112119 transformAiSdkAttributes ( span . attributes ) ;
113120} ;
0 commit comments