fix(instrumentation-cohere): record chat token usage and use .content for the user turn

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()) {

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<string, unknown>)
+      : undefined;
+
+  const readCounts = (source: unknown) => {
+    if (typeof source !== "object" || source === null) {
+      return;
+    }
+    const counts = source as Record<string, unknown>;
+    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<string, unknown>;
+      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;
+};

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,
     );

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), {});
+  });
+});