diff --git a/src/handlers/message.ts b/src/handlers/message.ts index 248228a..ee3df11 100644 --- a/src/handlers/message.ts +++ b/src/handlers/message.ts @@ -27,6 +27,12 @@ import { TOOL_NAME, TOOL_PARAMETERS, } from "@arizeai/openinference-semantic-conventions" +import { + ATTR_GEN_AI_REQUEST_MODEL, + ATTR_GEN_AI_RESPONSE_MODEL, + ATTR_GEN_AI_USAGE_INPUT_TOKENS, + ATTR_GEN_AI_USAGE_OUTPUT_TOKENS, +} from "@opentelemetry/semantic-conventions/incubating" import { agentAttrs, errorSummary, @@ -42,6 +48,11 @@ import type { HandlerContext } from "../types.ts" const OPENINFERENCE_SPAN_KIND = SemanticConventions.OPENINFERENCE_SPAN_KIND const LLM_FINISH_REASON = "llm.finish_reason" +// Spelled out rather than taken from semconv: the incubating constants emit +// `gen_ai.usage.cache_read.input_tokens`, while Anthropic's API and every GenAI +// consumer we target read the underscored form below. +const ATTR_GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS = "gen_ai.usage.cache_read_input_tokens" +const ATTR_GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS = "gen_ai.usage.cache_creation_input_tokens" type SubtaskPart = { type: "subtask" @@ -130,6 +141,11 @@ export function handleMessageUpdated(e: EventMessageUpdated, ctx: HandlerContext [LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_READ]: assistant.tokens.cache.read, [LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_WRITE]: assistant.tokens.cache.write, [LLM_TOKEN_COUNT_TOTAL]: totalTokens, + [ATTR_GEN_AI_RESPONSE_MODEL]: modelID, + [ATTR_GEN_AI_USAGE_INPUT_TOKENS]: assistant.tokens.input, + [ATTR_GEN_AI_USAGE_OUTPUT_TOKENS]: assistant.tokens.output, + [ATTR_GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS]: assistant.tokens.cache.read, + [ATTR_GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS]: assistant.tokens.cache.write, [LLM_FINISH_REASON]: assistant.error ? "error" : (assistant.finish ?? "stop"), [LLM_COST_TOTAL]: assistant.cost, ...(outputText @@ -454,6 +470,7 @@ export function startMessageSpan( [LLM_PROVIDER]: providerID, "gen_ai.provider.name": genAiProviderName(providerID), [LLM_MODEL_NAME]: modelID, + [ATTR_GEN_AI_REQUEST_MODEL]: modelID, ...(inputText ? { [INPUT_VALUE]: inputText, diff --git a/tests/handlers/spans.test.ts b/tests/handlers/spans.test.ts index 18851d0..ef2bdb3 100644 --- a/tests/handlers/spans.test.ts +++ b/tests/handlers/spans.test.ts @@ -370,6 +370,16 @@ describe("message (LLM) spans", () => { expect(tracer.spans[0]!.attributes[LLM_MODEL_NAME]).toBe("claude-sonnet-4") }) + // Asserted as literal keys, not via the semconv constants the source imports: + // these strings are the wire format Gen AI consumers match on, so the test has + // to fail if a constant is renamed upstream. + test("startMessageSpan sets OTel GenAI model attribute", () => { + const { ctx, tracer } = makeCtx() + startMessageSpan("ses_1", "msg_1", "user_1", "claude-sonnet-4", "amazon-bedrock", 1000, ctx) + expect(tracer.spans[0]!.attributes["gen_ai.request.model"]).toBe("claude-sonnet-4") + expect(tracer.spans[0]!.attributes["gen_ai.provider.name"]).toBe("aws.bedrock") + }) + test("startMessageSpan is a no-op when span already exists for sessionID:messageID", () => { const { ctx, tracer } = makeCtx() startMessageSpan("ses_1", "msg_1", "user_1", "claude", "anthropic", 1000, ctx) @@ -434,6 +444,44 @@ describe("message (LLM) spans", () => { expect(span.attributes["agent.type"]).toBe("subagent") }) + test("handleMessageUpdated sets OTel GenAI token attributes on span", () => { + const { ctx, tracer } = makeCtx() + startMessageSpan("ses_1", "msg_1", "user_1", "claude-3-5-sonnet", "anthropic", 1000, ctx) + handleMessageUpdated( + makeAssistantMessageUpdated({ + id: "msg_1", + modelID: "claude-3-5-sonnet", + tokens: { input: 200, output: 80, reasoning: 10, cache: { read: 30, write: 5 } }, + }), + ctx, + ) + const span = tracer.spans[0]! + expect(span.attributes["gen_ai.response.model"]).toBe("claude-3-5-sonnet") + expect(span.attributes["gen_ai.usage.input_tokens"]).toBe(200) + expect(span.attributes["gen_ai.usage.output_tokens"]).toBe(80) + // Underscored, matching Anthropic's API and Gen AI consumers. The semconv + // constants spell these with a dot (gen_ai.usage.cache_read.input_tokens). + expect(span.attributes["gen_ai.usage.cache_read_input_tokens"]).toBe(30) + expect(span.attributes["gen_ai.usage.cache_creation_input_tokens"]).toBe(5) + }) + + test("GenAI token attributes do not displace the OpenInference ones", () => { + const { ctx, tracer } = makeCtx() + startMessageSpan("ses_1", "msg_1", "user_1", "claude-3-5-sonnet", "anthropic", 1000, ctx) + handleMessageUpdated( + makeAssistantMessageUpdated({ + id: "msg_1", + tokens: { input: 200, output: 80, reasoning: 10, cache: { read: 30, write: 5 } }, + }), + ctx, + ) + const span = tracer.spans[0]! + expect(span.attributes[LLM_TOKEN_COUNT_PROMPT]).toBe(200) + expect(span.attributes["gen_ai.usage.input_tokens"]).toBe(200) + expect(span.attributes[LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_WRITE]).toBe(5) + expect(span.attributes["gen_ai.usage.cache_creation_input_tokens"]).toBe(5) + }) + test("handleMessageUpdated no-ops span handling when no span exists for messageID", () => { const { ctx, tracer } = makeCtx() const spansBefore = tracer.spans.length