From 1dc41ce457cf3ebbf4138e551960fc04e213e291 Mon Sep 17 00:00:00 2001 From: Stephen Crowley Date: Wed, 15 Jul 2026 14:46:36 +0100 Subject: [PATCH] fix(instrumentation-bedrock): avoid NaN total_tokens for Titan embedding responses Titan embedding responses carry inputTextTokenCount but no output-token field, so computing total_tokens as input + undefined produced NaN, which was serialized as the string "NaN" on exported spans. Guard both the streaming and non-streaming Amazon paths: omit output_tokens when the response has none and compute total_tokens as input + (output ?? 0). Fixes #1031 Co-Authored-By: Claude Fable 5 --- .../src/instrumentation.ts | 22 ++-- .../tests/titan-embedding-tokens.test.ts | 100 ++++++++++++++++++ 2 files changed, 116 insertions(+), 6 deletions(-) create mode 100644 packages/instrumentation-bedrock/tests/titan-embedding-tokens.test.ts diff --git a/packages/instrumentation-bedrock/src/instrumentation.ts b/packages/instrumentation-bedrock/src/instrumentation.ts index 299de20b..ee201e86 100644 --- a/packages/instrumentation-bedrock/src/instrumentation.ts +++ b/packages/instrumentation-bedrock/src/instrumentation.ts @@ -599,16 +599,22 @@ export class BedrockInstrumentation extends InstrumentationBase { ], } : {}), - // Titan includes token counts on the final chunk + // Titan includes token counts on the final chunk. Embedding + // responses have no output tokens, so guard against + // `input + undefined` producing NaN. ...(response["inputTextTokenCount"] != null ? { [ATTR_GEN_AI_USAGE_INPUT_TOKENS]: response["inputTextTokenCount"], - [ATTR_GEN_AI_USAGE_OUTPUT_TOKENS]: - response["totalOutputTextTokenCount"], + ...(response["totalOutputTextTokenCount"] != null + ? { + [ATTR_GEN_AI_USAGE_OUTPUT_TOKENS]: + response["totalOutputTextTokenCount"], + } + : {}), [SpanAttributes.GEN_AI_USAGE_TOTAL_TOKENS]: response["inputTextTokenCount"] + - response["totalOutputTextTokenCount"], + (response["totalOutputTextTokenCount"] ?? 0), } : {}), }; @@ -642,9 +648,13 @@ export class BedrockInstrumentation extends InstrumentationBase { ...(titanInputTokens != null ? { [ATTR_GEN_AI_USAGE_INPUT_TOKENS]: titanInputTokens, - [ATTR_GEN_AI_USAGE_OUTPUT_TOKENS]: titanOutputTokens, + // Titan embedding responses have no `results` array, so + // output tokens are undefined; guard against NaN totals. + ...(titanOutputTokens != null + ? { [ATTR_GEN_AI_USAGE_OUTPUT_TOKENS]: titanOutputTokens } + : {}), [SpanAttributes.GEN_AI_USAGE_TOTAL_TOKENS]: - titanInputTokens + titanOutputTokens, + titanInputTokens + (titanOutputTokens ?? 0), } : novaUsage != null ? { diff --git a/packages/instrumentation-bedrock/tests/titan-embedding-tokens.test.ts b/packages/instrumentation-bedrock/tests/titan-embedding-tokens.test.ts new file mode 100644 index 00000000..a3035fd9 --- /dev/null +++ b/packages/instrumentation-bedrock/tests/titan-embedding-tokens.test.ts @@ -0,0 +1,100 @@ +/* + * 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 { BedrockInstrumentation } from "../src/instrumentation"; +import { BedrockVendor } from "../src/types"; +import { SpanAttributes } from "@traceloop/ai-semantic-conventions"; +import { + ATTR_GEN_AI_USAGE_INPUT_TOKENS, + ATTR_GEN_AI_USAGE_OUTPUT_TOKENS, +} from "@opentelemetry/semantic-conventions/incubating"; + +// Titan embedding responses carry inputTextTokenCount but no output-token +// field (there is no generated text). total_tokens must not be computed as +// `input + undefined`, which evaluates to NaN and ends up serialized as the +// string "NaN" on the exported span. +// https://github.com/traceloop/openllmetry-js/issues/1031 + +describe("Bedrock Titan embedding token accounting", () => { + const instrumentation = new BedrockInstrumentation(); + + const setResponseAttrs = (response: Record, isStream = false) => + (instrumentation as any)._setResponseAttributes( + BedrockVendor.AMAZON, + response, + isStream, + ); + + it("emits numeric total_tokens for embedding responses (no output tokens)", () => { + const attrs = setResponseAttrs({ + embedding: [0.1, 0.2, 0.3], + inputTextTokenCount: 117, + }); + assert.strictEqual(attrs[ATTR_GEN_AI_USAGE_INPUT_TOKENS], 117); + assert.strictEqual( + attrs[SpanAttributes.GEN_AI_USAGE_TOTAL_TOKENS], + 117, + "total_tokens should equal input tokens when there is no output", + ); + assert.ok( + !Number.isNaN(attrs[SpanAttributes.GEN_AI_USAGE_TOTAL_TOKENS]), + "total_tokens must never be NaN", + ); + assert.strictEqual( + ATTR_GEN_AI_USAGE_OUTPUT_TOKENS in attrs, + false, + "output_tokens should be omitted when the response has none", + ); + }); + + it("still sums input + output for Titan text responses", () => { + const attrs = setResponseAttrs({ + inputTextTokenCount: 10, + results: [ + { tokenCount: 25, outputText: "hello", completionReason: "FINISH" }, + ], + }); + assert.strictEqual(attrs[ATTR_GEN_AI_USAGE_INPUT_TOKENS], 10); + assert.strictEqual(attrs[ATTR_GEN_AI_USAGE_OUTPUT_TOKENS], 25); + assert.strictEqual(attrs[SpanAttributes.GEN_AI_USAGE_TOTAL_TOKENS], 35); + }); + + it("guards the streaming final-chunk path the same way", () => { + const attrs = setResponseAttrs({ inputTextTokenCount: 42 }, true); + assert.strictEqual(attrs[ATTR_GEN_AI_USAGE_INPUT_TOKENS], 42); + assert.strictEqual( + attrs[SpanAttributes.GEN_AI_USAGE_TOTAL_TOKENS], + 42, + "streaming total_tokens should not be NaN when output count is absent", + ); + assert.strictEqual(ATTR_GEN_AI_USAGE_OUTPUT_TOKENS in attrs, false); + }); + + it("still sums both counts on the streaming final chunk when present", () => { + const attrs = setResponseAttrs( + { + inputTextTokenCount: 42, + totalOutputTextTokenCount: 8, + completionReason: "FINISH", + }, + true, + ); + assert.strictEqual(attrs[ATTR_GEN_AI_USAGE_INPUT_TOKENS], 42); + assert.strictEqual(attrs[ATTR_GEN_AI_USAGE_OUTPUT_TOKENS], 8); + assert.strictEqual(attrs[SpanAttributes.GEN_AI_USAGE_TOTAL_TOKENS], 50); + }); +});