From 6ea103458cee8d7b7e44320162942f68385b4062 Mon Sep 17 00:00:00 2001 From: "yogeshpatil0314@gmail.com" Date: Mon, 5 Oct 2026 20:45:18 +0530 Subject: [PATCH] fix(instrumentation-bedrock): safely compute titan total tokens when output tokens are missing --- .../src/instrumentation.ts | 16 ++- .../tests/titan-tokens.test.ts | 122 ++++++++++++++++++ 2 files changed, 133 insertions(+), 5 deletions(-) create mode 100644 packages/instrumentation-bedrock/tests/titan-tokens.test.ts diff --git a/packages/instrumentation-bedrock/src/instrumentation.ts b/packages/instrumentation-bedrock/src/instrumentation.ts index 299de20b..31e4bfca 100644 --- a/packages/instrumentation-bedrock/src/instrumentation.ts +++ b/packages/instrumentation-bedrock/src/instrumentation.ts @@ -604,11 +604,15 @@ export class BedrockInstrumentation extends InstrumentationBase { ? { [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 +646,11 @@ export class BedrockInstrumentation extends InstrumentationBase { ...(titanInputTokens != null ? { [ATTR_GEN_AI_USAGE_INPUT_TOKENS]: titanInputTokens, - [ATTR_GEN_AI_USAGE_OUTPUT_TOKENS]: titanOutputTokens, + ...(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-tokens.test.ts b/packages/instrumentation-bedrock/tests/titan-tokens.test.ts new file mode 100644 index 00000000..2531e50f --- /dev/null +++ b/packages/instrumentation-bedrock/tests/titan-tokens.test.ts @@ -0,0 +1,122 @@ +/* + * 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"; + +describe("Amazon Titan token counting semantics", () => { + const instrumentation = new BedrockInstrumentation(); + + const setResponseAttrs = ( + response: Record, + isStream = false, + ) => + // eslint-disable-next-line @typescript-eslint/no-explicit-any + (instrumentation as any)._setResponseAttributes( + BedrockVendor.AMAZON, + response, + isStream, + ); + + describe("Non-streaming responses", () => { + it("safely handles Titan embedding responses without output tokens", () => { + // Amazon Titan embedding models (e.g. amazon.titan-embed-text-v2:0) + // return inputTextTokenCount and an embedding array, but no results[0].tokenCount + const attrs = setResponseAttrs({ + inputTextTokenCount: 15, + embedding: [0.1, 0.2, 0.3], + }); + + assert.strictEqual( + attrs[ATTR_GEN_AI_USAGE_INPUT_TOKENS], + 15, + "input_tokens should equal inputTextTokenCount", + ); + assert.strictEqual( + attrs[ATTR_GEN_AI_USAGE_OUTPUT_TOKENS], + undefined, + "output_tokens should not be present when undefined", + ); + assert.strictEqual( + attrs[SpanAttributes.GEN_AI_USAGE_TOTAL_TOKENS], + 15, + "total_tokens should equal inputTextTokenCount (not NaN)", + ); + assert.ok( + !Number.isNaN(attrs[SpanAttributes.GEN_AI_USAGE_TOTAL_TOKENS]), + "total_tokens must not be NaN", + ); + }); + + it("correctly computes total tokens when both input and output tokens are present", () => { + const attrs = setResponseAttrs({ + inputTextTokenCount: 20, + results: [ + { + tokenCount: 45, + outputText: "Hello world", + completionReason: "FINISH", + }, + ], + }); + + assert.strictEqual(attrs[ATTR_GEN_AI_USAGE_INPUT_TOKENS], 20); + assert.strictEqual(attrs[ATTR_GEN_AI_USAGE_OUTPUT_TOKENS], 45); + assert.strictEqual(attrs[SpanAttributes.GEN_AI_USAGE_TOTAL_TOKENS], 65); + }); + }); + + describe("Streaming responses", () => { + it("safely handles Titan streaming chunks when totalOutputTextTokenCount is absent", () => { + const attrs = setResponseAttrs( + { + inputTextTokenCount: 30, + completionReason: "FINISH", + }, + true, + ); + + assert.strictEqual(attrs[ATTR_GEN_AI_USAGE_INPUT_TOKENS], 30); + assert.strictEqual(attrs[ATTR_GEN_AI_USAGE_OUTPUT_TOKENS], undefined); + assert.strictEqual(attrs[SpanAttributes.GEN_AI_USAGE_TOTAL_TOKENS], 30); + assert.ok( + !Number.isNaN(attrs[SpanAttributes.GEN_AI_USAGE_TOTAL_TOKENS]), + "total_tokens must not be NaN in streaming", + ); + }); + + it("correctly computes total tokens in streaming when both counts are present", () => { + const attrs = setResponseAttrs( + { + inputTextTokenCount: 30, + totalOutputTextTokenCount: 70, + completionReason: "FINISH", + }, + true, + ); + + assert.strictEqual(attrs[ATTR_GEN_AI_USAGE_INPUT_TOKENS], 30); + assert.strictEqual(attrs[ATTR_GEN_AI_USAGE_OUTPUT_TOKENS], 70); + assert.strictEqual(attrs[SpanAttributes.GEN_AI_USAGE_TOTAL_TOKENS], 100); + }); + }); +});