fix(server-utils): Read Bedrock Mistral and Cohere stream token counts from invocation metrics

packages/server-utils/src/integrations/aws-sdk/services/bedrock-runtime.ts

@@ -542,28 +542,36 @@ function recordLlamaAttributes(parsedChunk: ParsedChunk, span: Span): void {
   }
 }
 
-function recordMistralAttributes(parsedChunk: ParsedChunk, span: Span): void {
-  if (parsedChunk.outputs?.[0]?.text !== undefined) {
-    span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, Math.ceil(parsedChunk.outputs[0].text.length / 6));
+/**
+ * Bedrock adds the token counts to the last chunk of a stream. For models that send no counts of their own,
+ * this is the only real count.
+ */
+function recordInvocationMetrics(parsedChunk: ParsedChunk, span: Span): void {
+  const metrics = parsedChunk['amazon-bedrock-invocationMetrics'];
+  if (typeof metrics?.inputTokenCount === 'number') {
+    span.setAttribute(GEN_AI_USAGE_INPUT_TOKENS, metrics.inputTokenCount);
   }
+  if (typeof metrics?.outputTokenCount === 'number') {
+    span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, metrics.outputTokenCount);
+  }
+}
+
+function recordMistralAttributes(parsedChunk: ParsedChunk, span: Span): void {
+  recordInvocationMetrics(parsedChunk, span);
   if (parsedChunk.outputs?.[0]?.stop_reason !== undefined) {
     span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, [parsedChunk.outputs[0].stop_reason]);
   }
 }
 
 function recordCohereAttributes(parsedChunk: ParsedChunk, span: Span): void {
-  if (parsedChunk.generations?.[0]?.text !== undefined) {
-    span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, Math.ceil(parsedChunk.generations[0].text.length / 6));
-  }
+  recordInvocationMetrics(parsedChunk, span);
   if (parsedChunk.generations?.[0]?.finish_reason !== undefined) {
     span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, [parsedChunk.generations[0].finish_reason]);
   }
 }
 
 function recordCohereRAttributes(parsedChunk: ParsedChunk, span: Span): void {
-  if (parsedChunk.text !== undefined) {
-    span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, Math.ceil(parsedChunk.text.length / 6));
-  }
+  recordInvocationMetrics(parsedChunk, span);
   if (parsedChunk.finish_reason !== undefined) {
     span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, [parsedChunk.finish_reason]);
   }

packages/server-utils/test/integrations/aws-sdk-bedrock.test.ts

@@ -172,6 +172,54 @@ describe('BedrockRuntimeServiceExtension stream spans', () => {
     expect(spanToJSON(span).attributes?.[GEN_AI_USAGE_OUTPUT_TOKENS]).toBe(42);
   });
 
+  const INVOCATION_METRICS = {
+    inputTokenCount: 11,
+    outputTokenCount: 3,
+    invocationLatency: 412,
+    firstByteLatency: 198,
+  };
+
+  it.each(
+    [
+      { modelId: 'mistral.mistral-large-2402-v1:0', toChunk: (text: string) => ({ outputs: [{ text }] }) },
+      { modelId: 'cohere.command-text-v14', toChunk: (text: string) => ({ generations: [{ text }] }) },
+      { modelId: 'cohere.command-r-plus-v1:0', toChunk: (text: string) => ({ text }) },
+    ].flatMap(family => [
+      {
+        ...family,
+        when: 'has invocation metrics',
+        metrics: INVOCATION_METRICS,
+        expected: { [GEN_AI_USAGE_INPUT_TOKENS]: 11, [GEN_AI_USAGE_OUTPUT_TOKENS]: 3 },
+      },
+      {
+        ...family,
+        when: 'has no invocation metrics',
+        metrics: undefined,
+        expected: { [GEN_AI_USAGE_INPUT_TOKENS]: undefined, [GEN_AI_USAGE_OUTPUT_TOKENS]: undefined },
+      },
+    ]),
+  )('records only real token counts when a $modelId stream $when', async ({ modelId, toChunk, metrics, expected }) => {
+    const span = startInactiveSpan({ name: 'bedrock' });
+    const encoder = new TextEncoder();
+    const chunks = [toChunk('The capital'), { ...toChunk(' is Paris.'), 'amazon-bedrock-invocationMetrics': metrics }];
+    const data = {
+      body: (async function* () {
+        for (const chunk of chunks) {
+          yield { chunk: { bytes: encoder.encode(JSON.stringify(chunk)) } };
+        }
+      })(),
+    };
+
+    respond(span, 'InvokeModelWithResponseStream', { modelId }, data);
+    await drain(data.body);
+
+    const attributes = spanToJSON(span).attributes;
+    expect({
+      [GEN_AI_USAGE_INPUT_TOKENS]: attributes?.[GEN_AI_USAGE_INPUT_TOKENS],
+      [GEN_AI_USAGE_OUTPUT_TOKENS]: attributes?.[GEN_AI_USAGE_OUTPUT_TOKENS],
+    }).toEqual(expected);
+  });
+
   it('ends an InvokeModelWithResponseStream span without a modelId when the stream is drained', async () => {
     const span = startInactiveSpan({ name: 'bedrock' });
     const data = { body: (async function* () {})() };