Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
91 changes: 72 additions & 19 deletions packages/traceloop-sdk/src/lib/tracing/ai-sdk-transformations.ts
Original file line number Diff line number Diff line change
@@ -1,15 +1,19 @@
import { ReadableSpan } from "@opentelemetry/sdk-trace-node";
import { ReadableSpan, Span } from "@opentelemetry/sdk-trace-node";
import { SpanAttributes } from "@traceloop/ai-semantic-conventions";

const AI_GENERATE_TEXT = "ai.generateText";
const AI_GENERATE_TEXT_DO_GENERATE = "ai.generateText.doGenerate";
const AI_GENERATE_OBJECT_DO_GENERATE = "ai.generateObject.doGenerate";
const AI_STREAM_TEXT_DO_STREAM = "ai.streamText.doStream";
const HANDLED_SPAN_NAMES: Record<string, string> = {
[AI_GENERATE_TEXT_DO_GENERATE]: "ai.generateText.generate",
[AI_GENERATE_OBJECT_DO_GENERATE]: "ai.generateObject.generate",
[AI_STREAM_TEXT_DO_STREAM]: "ai.streamText.stream",
[AI_GENERATE_TEXT]: "run.ai",
[AI_GENERATE_TEXT_DO_GENERATE]: "text.generate",
[AI_GENERATE_OBJECT_DO_GENERATE]: "object.generate",
[AI_STREAM_TEXT_DO_STREAM]: "text.stream",
Comment on lines +9 to +12
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

these naming changes may break mor's monitors if merged.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@galkleinman the name changes didn't work at all apparently

};

const TOOL_SPAN_NAME = "ai.toolCall";

const AI_RESPONSE_TEXT = "ai.response.text";
const AI_RESPONSE_OBJECT = "ai.response.object";
const AI_RESPONSE_TOOL_CALLS = "ai.response.toolCalls";
Expand All @@ -19,6 +23,7 @@ const AI_USAGE_PROMPT_TOKENS = "ai.usage.promptTokens";
const AI_USAGE_COMPLETION_TOKENS = "ai.usage.completionTokens";
const AI_MODEL_PROVIDER = "ai.model.provider";
const AI_PROMPT_TOOLS = "ai.prompt.tools";
const AI_TELEMETRY_METADATA_PREFIX = "ai.telemetry.metadata.";
const TYPE_TEXT = "text";
const TYPE_TOOL_CALL = "tool_call";
const ROLE_ASSISTANT = "assistant";
Expand Down Expand Up @@ -47,14 +52,6 @@ const VENDOR_MAPPING: Record<string, string> = {
openrouter: "OpenRouter",
};

export const transformAiSdkSpanName = (span: ReadableSpan): void => {
// Unfortunately, the span name is not writable as this is not the intended behavior
// but it is a workaround to set the correct span name
if (span.name in HANDLED_SPAN_NAMES) {
(span as any).name = HANDLED_SPAN_NAMES[span.name];
}
};

const transformResponseText = (attributes: Record<string, any>): void => {
if (AI_RESPONSE_TEXT in attributes) {
attributes[`${SpanAttributes.LLM_COMPLETIONS}.0.content`] =
Expand Down Expand Up @@ -367,9 +364,41 @@ const transformVendor = (attributes: Record<string, any>): void => {
}
};

export const transformAiSdkAttributes = (
attributes: Record<string, any>,
): void => {
const transformTelemetryMetadata = (attributes: Record<string, any>): void => {
const metadataAttributes: Record<string, string> = {};
const keysToDelete: string[] = [];

// Find all ai.telemetry.metadata.* attributes
for (const [key, value] of Object.entries(attributes)) {
if (key.startsWith(AI_TELEMETRY_METADATA_PREFIX)) {
const metadataKey = key.substring(AI_TELEMETRY_METADATA_PREFIX.length);

// Always mark for deletion since it's a telemetry metadata attribute
keysToDelete.push(key);

if (metadataKey && value != null) {
// Convert value to string for association properties
const stringValue = typeof value === "string" ? value : String(value);
metadataAttributes[metadataKey] = stringValue;

// Also set as traceloop association property attribute
attributes[
`${SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.${metadataKey}`
] = stringValue;
}
}
}

// Remove original ai.telemetry.metadata.* attributes
keysToDelete.forEach((key) => {
delete attributes[key];
});

// Note: Context setting for child span inheritance should be done before span creation,
// not during transformation. Use `withTelemetryMetadataContext` function for context propagation.
};

export const transformLLMSpans = (attributes: Record<string, any>): void => {
transformResponseText(attributes);
transformResponseObject(attributes);
transformResponseToolCalls(attributes);
Expand All @@ -379,16 +408,40 @@ export const transformAiSdkAttributes = (
transformCompletionTokens(attributes);
calculateTotalTokens(attributes);
transformVendor(attributes);
transformTelemetryMetadata(attributes);
};

const transformToolCalls = (span: ReadableSpan): void => {
if (
span.attributes["ai.toolCall.args"] &&
span.attributes["ai.toolCall.result"]
) {
span.attributes[SpanAttributes.TRACELOOP_ENTITY_INPUT] =
span.attributes["ai.toolCall.args"];
delete span.attributes["ai.toolCall.args"];
span.attributes[SpanAttributes.TRACELOOP_ENTITY_OUTPUT] =
span.attributes["ai.toolCall.result"];
delete span.attributes["ai.toolCall.result"];
}
};

const shouldHandleSpan = (span: ReadableSpan): boolean => {
return span.name in HANDLED_SPAN_NAMES;
return span.instrumentationScope?.name === "ai";
};

export const transformAiSdkSpanNames = (span: Span): void => {
if (span.name === TOOL_SPAN_NAME) {
span.updateName(`${span.attributes["ai.toolCall.name"] as string}.tool`);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Guard the use of 'ai.toolCall.name' in transformAiSdkSpanNames to avoid formatting to 'undefined.tool' if the attribute is missing.

Suggested change
span.updateName(`${span.attributes["ai.toolCall.name"] as string}.tool`);
span.updateName(`${span.attributes["ai.toolCall.name"] ? span.attributes["ai.toolCall.name"] : "tool"}.tool`);

}
if (span.name in HANDLED_SPAN_NAMES) {
span.updateName(HANDLED_SPAN_NAMES[span.name]);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nice! this updateName is a new api i think so. if that's the case this change will break the otel-v1 tho if we merge it.

just read about it, it exists since v1, makes me wonder why didn't we use it.

}
};
Comment on lines +432 to 439
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Guard against missing ai.toolCall.name to prevent 'undefined.tool'.

Only rename tool-call spans when the name attribute exists and is non-empty.

Apply this diff:

 export const transformAiSdkSpanNames = (span: Span): void => {
-  if (span.name === TOOL_SPAN_NAME) {
-    span.updateName(`${span.attributes["ai.toolCall.name"] as string}.tool`);
-  }
+  if (span.name === TOOL_SPAN_NAME) {
+    const toolName = (span.attributes["ai.toolCall.name"] as unknown) as string | undefined;
+    if (typeof toolName === "string" && toolName.length > 0) {
+      span.updateName(`${toolName}.tool`);
+    }
+  }
   if (span.name in HANDLED_SPAN_NAMES) {
     span.updateName(HANDLED_SPAN_NAMES[span.name]);
   }
 };
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
export const transformAiSdkSpanNames = (span: Span): void => {
if (span.name === TOOL_SPAN_NAME) {
span.updateName(`${span.attributes["ai.toolCall.name"] as string}.tool`);
}
if (span.name in HANDLED_SPAN_NAMES) {
span.updateName(HANDLED_SPAN_NAMES[span.name]);
}
};
export const transformAiSdkSpanNames = (span: Span): void => {
if (span.name === TOOL_SPAN_NAME) {
const toolName = (span.attributes["ai.toolCall.name"] as unknown) as string | undefined;
if (typeof toolName === "string" && toolName.length > 0) {
span.updateName(`${toolName}.tool`);
}
}
if (span.name in HANDLED_SPAN_NAMES) {
span.updateName(HANDLED_SPAN_NAMES[span.name]);
}
};
🤖 Prompt for AI Agents
In packages/traceloop-sdk/src/lib/tracing/ai-sdk-transformations.ts around lines
396 to 403, the tool-call renaming can produce "undefined.tool" when
ai.toolCall.name is missing; change the logic so you only call
span.updateName(`${...}.tool`) if span.attributes["ai.toolCall.name"] is a
non-empty string (e.g. typeof === "string" &&
span.attributes["ai.toolCall.name"].trim() !== ""), otherwise skip that rename;
keep the subsequent HANDLED_SPAN_NAMES check as-is so other mappings still
apply.


export const transformAiSdkSpan = (span: ReadableSpan): void => {
export const transformAiSdkSpanAttributes = (span: ReadableSpan): void => {
if (!shouldHandleSpan(span)) {
return;
}
transformAiSdkSpanName(span);
transformAiSdkAttributes(span.attributes);
transformLLMSpans(span.attributes);
transformToolCalls(span);
Comment on lines +441 to +446
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[STYLE] naming & args aren't consistent, if you changed transformAiSdkSpan -> transformAiSdkSpanAttributes than the methods used inside should also be with attrs suffix imo, and have the same signature

};
12 changes: 9 additions & 3 deletions packages/traceloop-sdk/src/lib/tracing/span-processor.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,10 @@ import {
SimpleSpanProcessor,
BatchSpanProcessor,
SpanProcessor,
Span,
ReadableSpan,
} from "@opentelemetry/sdk-trace-node";
import { Span, context } from "@opentelemetry/api";
import { context } from "@opentelemetry/api";
import { OTLPTraceExporter } from "@opentelemetry/exporter-trace-otlp-proto";
import { SpanExporter } from "@opentelemetry/sdk-trace-base";
import {
Expand All @@ -13,7 +14,10 @@ import {
WORKFLOW_NAME_KEY,
} from "./tracing";
import { SpanAttributes } from "@traceloop/ai-semantic-conventions";
import { transformAiSdkSpan } from "./ai-sdk-transformations";
import {
transformAiSdkSpanAttributes,
transformAiSdkSpanNames,
} from "./ai-sdk-transformations";
import { parseKeyPairsIntoRecord } from "./baggage-utils";

export const ALL_INSTRUMENTATION_LIBRARIES = "all" as const;
Expand Down Expand Up @@ -155,6 +159,8 @@ const onSpanStart = (span: Span): void => {
);
}
}

transformAiSdkSpanNames(span);
};

/**
Expand Down Expand Up @@ -220,7 +226,7 @@ const onSpanEnd = (
}

// Apply AI SDK transformations (if needed)
transformAiSdkSpan(span);
transformAiSdkSpanAttributes(span);

// Ensure OTLP transformer compatibility
const compatibleSpan = ensureSpanCompatibility(span);
Expand Down
19 changes: 7 additions & 12 deletions packages/traceloop-sdk/test/ai-sdk-integration.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -105,17 +105,15 @@ describe("Test AI SDK Integration with Recording", function () {
const spans = memoryExporter.getFinishedSpans();

const generateTextSpan = spans.find(
(span) =>
span.name === "ai.generateText.generate" ||
span.name === "ai.generateText.doGenerate",
(span) => span.name === "text.generate",
);

assert.ok(result);
assert.ok(result.text);
assert.ok(generateTextSpan);

// Verify span name
assert.strictEqual(generateTextSpan.name, "ai.generateText.generate");
// Verify span name (should be transformed from ai.generateText.doGenerate to text.generate)
assert.strictEqual(generateTextSpan.name, "text.generate");

// Verify vendor
assert.strictEqual(generateTextSpan.attributes["gen_ai.system"], "OpenAI");
Expand Down Expand Up @@ -174,17 +172,16 @@ describe("Test AI SDK Integration with Recording", function () {
// Find the Google span specifically (should have workflow name test_google_workflow)
const generateTextSpan = spans.find(
(span) =>
(span.name === "ai.generateText.generate" ||
span.name === "ai.generateText.doGenerate") &&
span.name === "text.generate" &&
span.attributes["traceloop.workflow.name"] === "test_google_workflow",
);

assert.ok(result);
assert.ok(result.text);
assert.ok(generateTextSpan, "Could not find Google generateText span");

// Verify span name
assert.strictEqual(generateTextSpan.name, "ai.generateText.generate");
// Verify span name (should be transformed from ai.generateText.doGenerate to text.generate)
assert.strictEqual(generateTextSpan.name, "text.generate");

// Verify vendor
assert.strictEqual(generateTextSpan.attributes["gen_ai.system"], "Google");
Expand Down Expand Up @@ -236,9 +233,7 @@ describe("Test AI SDK Integration with Recording", function () {
assert.ok(result.text);

const spans = memoryExporter.getFinishedSpans();
const aiSdkSpan = spans.find((span) =>
span.name.startsWith("ai.generateText"),
);
const aiSdkSpan = spans.find((span) => span.name === "text.generate");

assert.ok(aiSdkSpan);

Expand Down
Loading