From 383bf36d9cd8aacdbbac4733d72739e1bd773f6a Mon Sep 17 00:00:00 2001 From: survivor-agent <326275588+survivor-agent@users.noreply.github.com> Date: Mon, 21 Sep 2026 13:47:14 +0000 Subject: [PATCH] fix(instrumentation-cohere): record chat token usage and use .content for the user turn Chat spans read token usage only from a top-level `token_count` object, which current Cohere chat responses no longer return, so gen_ai.usage.prompt_tokens, gen_ai.usage.completion_tokens and llm.usage.total_tokens were silently absent. Usage is now read from `meta.tokens` first, falling back to the legacy `token_count` and then to `meta.billedUnits` (the field generate spans already use), so older responses keep reporting the same values. The current user message was also recorded under gen_ai.prompt.N.user while chat history used gen_ai.prompt.N.content, so consumers reading .content missed the turn that triggered the request. Fixes #1035 --- .../src/instrumentation.ts | 50 +++++------- packages/instrumentation-cohere/src/usage.ts | 80 +++++++++++++++++++ .../instrumentation-cohere/tests/chat.test.ts | 4 +- .../tests/usage.test.ts | 60 ++++++++++++++ 4 files changed, 160 insertions(+), 34 deletions(-) create mode 100644 packages/instrumentation-cohere/src/usage.ts create mode 100644 packages/instrumentation-cohere/tests/usage.test.ts diff --git a/packages/instrumentation-cohere/src/instrumentation.ts b/packages/instrumentation-cohere/src/instrumentation.ts index 8fa51db8..d41b2b92 100644 --- a/packages/instrumentation-cohere/src/instrumentation.ts +++ b/packages/instrumentation-cohere/src/instrumentation.ts @@ -13,6 +13,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ +import { extractChatTokenUsage } from "./usage"; import { Span, Attributes, @@ -266,7 +267,7 @@ export class CohereInstrumentation extends InstrumentationBase { `${ATTR_GEN_AI_PROMPT}.${params.chatHistory?.length ?? 0}.role` ] = "user"; attributes[ - `${ATTR_GEN_AI_PROMPT}.${params.chatHistory?.length ?? 0}.user` + `${ATTR_GEN_AI_PROMPT}.${params.chatHistory?.length ?? 0}.content` ] = params.message; } else if (type === "rerank" && "query" in params) { attributes[`${ATTR_GEN_AI_PROMPT}.0.role`] = "user"; @@ -409,39 +410,24 @@ export class CohereInstrumentation extends InstrumentationBase { result: cohere.Cohere.NonStreamedChatResponse, ) { try { - if ("token_count" in result && typeof result.token_count === "object") { - if ( - result.token_count && - "prompt_tokens" in result.token_count && - typeof result.token_count.prompt_tokens === "number" - ) { - span.setAttribute( - ATTR_GEN_AI_USAGE_PROMPT_TOKENS, - result.token_count?.prompt_tokens, - ); - } + const usage = extractChatTokenUsage(result); - if ( - result.token_count && - "response_tokens" in result.token_count && - typeof result.token_count.response_tokens === "number" - ) { - span.setAttribute( - ATTR_GEN_AI_USAGE_COMPLETION_TOKENS, - result.token_count?.response_tokens, - ); - } + if (usage.promptTokens !== undefined) { + span.setAttribute(ATTR_GEN_AI_USAGE_PROMPT_TOKENS, usage.promptTokens); + } - if ( - result.token_count && - "total_tokens" in result.token_count && - typeof result.token_count.total_tokens === "number" - ) { - span.setAttribute( - SpanAttributes.LLM_USAGE_TOTAL_TOKENS, - result.token_count?.total_tokens, - ); - } + if (usage.completionTokens !== undefined) { + span.setAttribute( + ATTR_GEN_AI_USAGE_COMPLETION_TOKENS, + usage.completionTokens, + ); + } + + if (usage.totalTokens !== undefined) { + span.setAttribute( + SpanAttributes.LLM_USAGE_TOTAL_TOKENS, + usage.totalTokens, + ); } if (this._shouldSendPrompts()) { diff --git a/packages/instrumentation-cohere/src/usage.ts b/packages/instrumentation-cohere/src/usage.ts new file mode 100644 index 00000000..3b352206 --- /dev/null +++ b/packages/instrumentation-cohere/src/usage.ts @@ -0,0 +1,80 @@ +/* + * Copyright Traceloop + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +export interface TokenUsage { + promptTokens?: number; + completionTokens?: number; + totalTokens?: number; +} + +const asNumber = (value: unknown): number | undefined => + typeof value === "number" ? value : undefined; + +/** + * Chat responses report token usage in more than one place depending on the + * API version that produced them, so they are read in order of accuracy: + * + * 1. `meta.tokens` - actual token counts on current responses. + * 2. a top-level `token_count` object - present on older responses only. It is + * not part of `NonStreamedChatResponse`, but the SDK passes unrecognized + * keys through, so it is still readable at runtime. + * 3. `meta.billedUnits` - billed rather than actual tokens, used as a last + * resort (this is the field `generate` spans already fall back to). + */ +export const extractChatTokenUsage = (result: unknown): TokenUsage => { + if (typeof result !== "object" || result === null) { + return {}; + } + + const usage: TokenUsage = {}; + const meta = + "meta" in result && typeof result.meta === "object" && result.meta !== null + ? (result.meta as Record) + : undefined; + + const readCounts = (source: unknown) => { + if (typeof source !== "object" || source === null) { + return; + } + const counts = source as Record; + usage.promptTokens ??= asNumber(counts.inputTokens); + usage.completionTokens ??= asNumber(counts.outputTokens); + }; + + readCounts(meta?.tokens); + + if ("token_count" in result) { + const legacy = result.token_count; + if (typeof legacy === "object" && legacy !== null) { + const counts = legacy as Record; + usage.promptTokens ??= asNumber(counts.prompt_tokens); + usage.completionTokens ??= asNumber(counts.response_tokens); + usage.totalTokens ??= asNumber(counts.total_tokens); + } + } + + readCounts(meta?.billedUnits); + + if ( + usage.totalTokens === undefined && + usage.promptTokens !== undefined && + usage.completionTokens !== undefined + ) { + usage.totalTokens = usage.promptTokens + usage.completionTokens; + } + + return usage; +}; diff --git a/packages/instrumentation-cohere/tests/chat.test.ts b/packages/instrumentation-cohere/tests/chat.test.ts index 5b6ddf20..e21ac022 100644 --- a/packages/instrumentation-cohere/tests/chat.test.ts +++ b/packages/instrumentation-cohere/tests/chat.test.ts @@ -124,7 +124,7 @@ describe.skip("Test Chat with Cohere Instrumentation", () => { ); assert.strictEqual( attributes[ - `${ATTR_GEN_AI_PROMPT}.${params.chatHistory?.length ?? 0}.user` + `${ATTR_GEN_AI_PROMPT}.${params.chatHistory?.length ?? 0}.content` ], params.message, ); @@ -221,7 +221,7 @@ describe.skip("Test Chat with Cohere Instrumentation", () => { ); assert.strictEqual( attributes[ - `${ATTR_GEN_AI_PROMPT}.${params.chatHistory?.length ?? 0}.user` + `${ATTR_GEN_AI_PROMPT}.${params.chatHistory?.length ?? 0}.content` ], params.message, ); diff --git a/packages/instrumentation-cohere/tests/usage.test.ts b/packages/instrumentation-cohere/tests/usage.test.ts new file mode 100644 index 00000000..2752c9b7 --- /dev/null +++ b/packages/instrumentation-cohere/tests/usage.test.ts @@ -0,0 +1,60 @@ +/* + * Copyright Traceloop + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import * as assert from "assert"; +import { extractChatTokenUsage } from "../src/usage"; + +describe("Test chat token usage extraction", () => { + it("should read actual token counts from meta.tokens", () => { + assert.deepStrictEqual( + extractChatTokenUsage({ + meta: { + tokens: { inputTokens: 11, outputTokens: 22 }, + billedUnits: { inputTokens: 5, outputTokens: 6 }, + }, + }), + { promptTokens: 11, completionTokens: 22, totalTokens: 33 }, + ); + }); + + it("should read the legacy top-level token_count when meta.tokens is absent", () => { + assert.deepStrictEqual( + extractChatTokenUsage({ + token_count: { + prompt_tokens: 1376, + response_tokens: 60, + total_tokens: 1436, + }, + meta: { billedUnits: { inputTokens: 23, outputTokens: 60 } }, + }), + { promptTokens: 1376, completionTokens: 60, totalTokens: 1436 }, + ); + }); + + it("should fall back to meta.billedUnits when nothing else is reported", () => { + assert.deepStrictEqual( + extractChatTokenUsage({ + meta: { billedUnits: { inputTokens: 23, outputTokens: 60 } }, + }), + { promptTokens: 23, completionTokens: 60, totalTokens: 83 }, + ); + }); + + it("should return nothing when the response carries no usage", () => { + assert.deepStrictEqual(extractChatTokenUsage({ text: "hello" }), {}); + assert.deepStrictEqual(extractChatTokenUsage(undefined), {}); + }); +});