Skip to content

Commit d2921fc

Browse files
nina-kollmannina-kollman
committed
pretty
1 parent c331f0b commit d2921fc

File tree

1 file changed

+38
-50
lines changed

1 file changed

+38
-50
lines changed

‎packages/traceloop-sdk/src/lib/tracing/ai-sdk-transformations.ts‎

Lines changed: 38 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import {
1111
ATTR_GEN_AI_OPERATION_NAME,
1212
ATTR_GEN_AI_OUTPUT_MESSAGES,
1313
ATTR_GEN_AI_PROMPT,
14+
ATTR_GEN_AI_PROVIDER_NAME,
1415
ATTR_GEN_AI_REQUEST_MODEL,
1516
ATTR_GEN_AI_RESPONSE_FINISH_REASONS,
1617
ATTR_GEN_AI_RESPONSE_ID,
@@ -95,8 +96,7 @@ const transformResponseText = (attributes: Record<string, any>): void => {
9596
if (AI_RESPONSE_TEXT in attributes) {
9697
attributes[`${ATTR_GEN_AI_COMPLETION}.0.content`] =
9798
attributes[AI_RESPONSE_TEXT];
98-
attributes[`${ATTR_GEN_AI_COMPLETION}.0.role`] =
99-
ROLE_ASSISTANT;
99+
attributes[`${ATTR_GEN_AI_COMPLETION}.0.role`] = ROLE_ASSISTANT;
100100

101101
const outputMessage = {
102102
role: ROLE_ASSISTANT,
@@ -107,9 +107,7 @@ const transformResponseText = (attributes: Record<string, any>): void => {
107107
},
108108
],
109109
};
110-
attributes[ATTR_GEN_AI_OUTPUT_MESSAGES] = JSON.stringify([
111-
outputMessage,
112-
]);
110+
attributes[ATTR_GEN_AI_OUTPUT_MESSAGES] = JSON.stringify([outputMessage]);
113111

114112
delete attributes[AI_RESPONSE_TEXT];
115113
}
@@ -119,8 +117,7 @@ const transformResponseObject = (attributes: Record<string, any>): void => {
119117
if (AI_RESPONSE_OBJECT in attributes) {
120118
attributes[`${ATTR_GEN_AI_COMPLETION}.0.content`] =
121119
attributes[AI_RESPONSE_OBJECT];
122-
attributes[`${ATTR_GEN_AI_COMPLETION}.0.role`] =
123-
ROLE_ASSISTANT;
120+
attributes[`${ATTR_GEN_AI_COMPLETION}.0.role`] = ROLE_ASSISTANT;
124121

125122
const outputMessage = {
126123
role: ROLE_ASSISTANT,
@@ -131,9 +128,7 @@ const transformResponseObject = (attributes: Record<string, any>): void => {
131128
},
132129
],
133130
};
134-
attributes[ATTR_GEN_AI_OUTPUT_MESSAGES] = JSON.stringify([
135-
outputMessage,
136-
]);
131+
attributes[ATTR_GEN_AI_OUTPUT_MESSAGES] = JSON.stringify([outputMessage]);
137132

138133
delete attributes[AI_RESPONSE_OBJECT];
139134
}
@@ -146,15 +141,13 @@ const transformResponseToolCalls = (attributes: Record<string, any>): void => {
146141
attributes[AI_RESPONSE_TOOL_CALLS] as string,
147142
);
148143

149-
attributes[`${ATTR_GEN_AI_COMPLETION}.0.role`] =
150-
ROLE_ASSISTANT;
144+
attributes[`${ATTR_GEN_AI_COMPLETION}.0.role`] = ROLE_ASSISTANT;
151145

152146
const toolCallParts: any[] = [];
153147
toolCalls.forEach((toolCall: any, index: number) => {
154148
if (toolCall.toolCallType === "function") {
155-
attributes[
156-
`${ATTR_GEN_AI_COMPLETION}.0.tool_calls.${index}.name`
157-
] = toolCall.toolName;
149+
attributes[`${ATTR_GEN_AI_COMPLETION}.0.tool_calls.${index}.name`] =
150+
toolCall.toolName;
158151
attributes[
159152
`${ATTR_GEN_AI_COMPLETION}.0.tool_calls.${index}.arguments`
160153
] = toolCall.args;
@@ -174,9 +167,9 @@ const transformResponseToolCalls = (attributes: Record<string, any>): void => {
174167
role: ROLE_ASSISTANT,
175168
parts: toolCallParts,
176169
};
177-
attributes[ATTR_GEN_AI_OUTPUT_MESSAGES] = JSON.stringify(
178-
[outputMessage],
179-
);
170+
attributes[ATTR_GEN_AI_OUTPUT_MESSAGES] = JSON.stringify([
171+
outputMessage,
172+
]);
180173
}
181174

182175
delete attributes[AI_RESPONSE_TOOL_CALLS];
@@ -304,8 +297,7 @@ const transformPrompts = (attributes: Record<string, any>): void => {
304297
const processedContent = processMessageContent(msg.content);
305298
const contentKey = `${ATTR_GEN_AI_PROMPT}.${index}.content`;
306299
attributes[contentKey] = processedContent;
307-
attributes[`${ATTR_GEN_AI_PROMPT}.${index}.role`] =
308-
msg.role;
300+
attributes[`${ATTR_GEN_AI_PROMPT}.${index}.role`] = msg.role;
309301

310302
// Add to OpenTelemetry standard gen_ai.input.messages format
311303
inputMessages.push({
@@ -321,8 +313,7 @@ const transformPrompts = (attributes: Record<string, any>): void => {
321313

322314
// Set the OpenTelemetry standard input messages attribute
323315
if (inputMessages.length > 0) {
324-
attributes[ATTR_GEN_AI_INPUT_MESSAGES] =
325-
JSON.stringify(inputMessages);
316+
attributes[ATTR_GEN_AI_INPUT_MESSAGES] = JSON.stringify(inputMessages);
326317
}
327318

328319
delete attributes[AI_PROMPT_MESSAGES];
@@ -343,8 +334,7 @@ const transformPrompts = (attributes: Record<string, any>): void => {
343334
const processedContent = processMessageContent(msg.content);
344335
const contentKey = `${ATTR_GEN_AI_PROMPT}.${index}.content`;
345336
attributes[contentKey] = processedContent;
346-
attributes[`${ATTR_GEN_AI_PROMPT}.${index}.role`] =
347-
msg.role;
337+
attributes[`${ATTR_GEN_AI_PROMPT}.${index}.role`] = msg.role;
348338

349339
inputMessages.push({
350340
role: msg.role,
@@ -365,8 +355,7 @@ const transformPrompts = (attributes: Record<string, any>): void => {
365355

366356
delete attributes[AI_PROMPT];
367357
} else if (promptData.prompt && typeof promptData.prompt === "string") {
368-
attributes[`${ATTR_GEN_AI_PROMPT}.0.content`] =
369-
promptData.prompt;
358+
attributes[`${ATTR_GEN_AI_PROMPT}.0.content`] = promptData.prompt;
370359
attributes[`${ATTR_GEN_AI_PROMPT}.0.role`] = ROLE_USER;
371360

372361
const inputMessage = {
@@ -378,9 +367,7 @@ const transformPrompts = (attributes: Record<string, any>): void => {
378367
},
379368
],
380369
};
381-
attributes[ATTR_GEN_AI_INPUT_MESSAGES] = JSON.stringify([
382-
inputMessage,
383-
]);
370+
attributes[ATTR_GEN_AI_INPUT_MESSAGES] = JSON.stringify([inputMessage]);
384371

385372
delete attributes[AI_PROMPT];
386373
}
@@ -465,8 +452,7 @@ const transformProviderMetadata = (attributes: Record<string, any>): void => {
465452

466453
const calculateTotalTokens = (attributes: Record<string, any>): void => {
467454
const inputTokens = attributes[ATTR_GEN_AI_USAGE_INPUT_TOKENS];
468-
const outputTokens =
469-
attributes[ATTR_GEN_AI_USAGE_OUTPUT_TOKENS];
455+
const outputTokens = attributes[ATTR_GEN_AI_USAGE_OUTPUT_TOKENS];
470456

471457
if (inputTokens && outputTokens) {
472458
attributes[`${SpanAttributes.LLM_USAGE_TOTAL_TOKENS}`] =
@@ -480,16 +466,23 @@ const transformVendor = (attributes: Record<string, any>): void => {
480466

481467
// Find matching vendor prefix in mapping
482468
let mappedVendor = null;
469+
let providerName = vendor;
483470
if (typeof vendor === "string" && vendor.length > 0) {
471+
// Extract provider name (part before first dot, or entire string if no dot)
472+
const dotIndex = vendor.indexOf(".");
473+
providerName = dotIndex > 0 ? vendor.substring(0, dotIndex) : vendor;
474+
484475
for (const prefix of Object.keys(VENDOR_MAPPING)) {
485476
if (vendor.startsWith(prefix)) {
486477
mappedVendor = VENDOR_MAPPING[prefix];
487478
break;
488479
}
489480
}
481+
482+
attributes[ATTR_GEN_AI_SYSTEM] = mappedVendor || vendor;
483+
attributes[ATTR_GEN_AI_PROVIDER_NAME] = providerName;
490484
}
491485

492-
attributes[ATTR_GEN_AI_SYSTEM] = mappedVendor || vendor;
493486
delete attributes[AI_MODEL_PROVIDER];
494487
}
495488
};
@@ -520,8 +513,7 @@ const transformOperationName = (
520513
const transformModelId = (attributes: Record<string, any>): void => {
521514
const AI_MODEL_ID = "ai.model.id";
522515
if (AI_MODEL_ID in attributes) {
523-
attributes[ATTR_GEN_AI_REQUEST_MODEL] =
524-
attributes[AI_MODEL_ID];
516+
attributes[ATTR_GEN_AI_REQUEST_MODEL] = attributes[AI_MODEL_ID];
525517
delete attributes[AI_MODEL_ID];
526518
}
527519
};
@@ -530,22 +522,23 @@ const transformFinishReason = (attributes: Record<string, any>): void => {
530522
const AI_RESPONSE_FINISH_REASON = "ai.response.finishReason";
531523
if (AI_RESPONSE_FINISH_REASON in attributes) {
532524
const finishReason = attributes[AI_RESPONSE_FINISH_REASON];
533-
attributes[ATTR_GEN_AI_RESPONSE_FINISH_REASONS] =
534-
Array.isArray(finishReason) ? finishReason : [finishReason];
525+
attributes[ATTR_GEN_AI_RESPONSE_FINISH_REASONS] = Array.isArray(
526+
finishReason,
527+
)
528+
? finishReason
529+
: [finishReason];
535530
delete attributes[AI_RESPONSE_FINISH_REASON];
536531
}
537532
};
538533

539534
const transformToolCallAttributes = (attributes: Record<string, any>): void => {
540535
if ("ai.toolCall.name" in attributes) {
541-
attributes[ATTR_GEN_AI_TOOL_NAME] =
542-
attributes["ai.toolCall.name"];
536+
attributes[ATTR_GEN_AI_TOOL_NAME] = attributes["ai.toolCall.name"];
543537
// Keep ai.toolCall.name for now, will be deleted in transformToolCalls
544538
}
545539

546540
if ("ai.toolCall.id" in attributes) {
547-
attributes[ATTR_GEN_AI_TOOL_CALL_ID] =
548-
attributes["ai.toolCall.id"];
541+
attributes[ATTR_GEN_AI_TOOL_CALL_ID] = attributes["ai.toolCall.id"];
549542
delete attributes["ai.toolCall.id"];
550543
}
551544

@@ -556,8 +549,7 @@ const transformToolCallAttributes = (attributes: Record<string, any>): void => {
556549
}
557550

558551
if ("ai.toolCall.result" in attributes) {
559-
attributes[ATTR_GEN_AI_TOOL_CALL_RESULT] =
560-
attributes["ai.toolCall.result"];
552+
attributes[ATTR_GEN_AI_TOOL_CALL_RESULT] = attributes["ai.toolCall.result"];
561553
// Don't delete yet - transformToolCalls will handle entity input/output
562554
}
563555
};
@@ -578,14 +570,12 @@ const transformResponseMetadata = (attributes: Record<string, any>): void => {
578570
const AI_RESPONSE_ID = "ai.response.id";
579571

580572
if (AI_RESPONSE_MODEL in attributes) {
581-
attributes[ATTR_GEN_AI_RESPONSE_MODEL] =
582-
attributes[AI_RESPONSE_MODEL];
573+
attributes[ATTR_GEN_AI_RESPONSE_MODEL] = attributes[AI_RESPONSE_MODEL];
583574
delete attributes[AI_RESPONSE_MODEL];
584575
}
585576

586577
if (AI_RESPONSE_ID in attributes) {
587-
attributes[ATTR_GEN_AI_RESPONSE_ID] =
588-
attributes[AI_RESPONSE_ID];
578+
attributes[ATTR_GEN_AI_RESPONSE_ID] = attributes[AI_RESPONSE_ID];
589579
delete attributes[AI_RESPONSE_ID];
590580
}
591581
};
@@ -639,10 +629,8 @@ const transformTelemetryMetadata = (
639629
TraceloopSpanKindValues.AGENT;
640630
attributes[SpanAttributes.TRACELOOP_ENTITY_NAME] = agentName;
641631

642-
const inputMessages =
643-
attributes[ATTR_GEN_AI_INPUT_MESSAGES];
644-
const outputMessages =
645-
attributes[ATTR_GEN_AI_OUTPUT_MESSAGES];
632+
const inputMessages = attributes[ATTR_GEN_AI_INPUT_MESSAGES];
633+
const outputMessages = attributes[ATTR_GEN_AI_OUTPUT_MESSAGES];
646634
const toolArgs = attributes["ai.toolCall.args"];
647635
const toolResult = attributes["ai.toolCall.result"];
648636

0 commit comments

Comments
 (0)