diff --git a/.changeset/bright-otters-stream.md b/.changeset/bright-otters-stream.md new file mode 100644 index 000000000..8bd07b5d3 --- /dev/null +++ b/.changeset/bright-otters-stream.md @@ -0,0 +1,5 @@ +--- +"braintrust": patch +--- + +fix(bedrock): Log text and token usage for streamed InvokeModel calls to OpenAI models such as gpt-oss diff --git a/js/src/instrumentation/plugins/bedrock-runtime-plugin.test.ts b/js/src/instrumentation/plugins/bedrock-runtime-plugin.test.ts index 5016574bb..1d09680ce 100644 --- a/js/src/instrumentation/plugins/bedrock-runtime-plugin.test.ts +++ b/js/src/instrumentation/plugins/bedrock-runtime-plugin.test.ts @@ -594,6 +594,209 @@ describe("aggregateInvokeModelResponseStreamChunks", () => { }); }); + it("aggregates OpenAI (gpt-oss) stream text and invocation metrics from InvokeModel response stream chunks", () => { + // gpt-oss sends Chat Completions chunks with no `usage`; the token counts + // only arrive in the last chunk's amazon-bedrock-invocationMetrics. + expect( + aggregateInvokeModelResponseStreamChunks( + [ + { + choices: [ + { + delta: { content: "", role: "assistant" }, + finish_reason: null, + index: 0, + }, + ], + object: "chat.completion.chunk", + }, + { + choices: [ + { delta: { content: "Hello" }, finish_reason: null, index: 0 }, + ], + object: "chat.completion.chunk", + }, + { + choices: [ + { delta: { content: " world" }, finish_reason: null, index: 0 }, + ], + object: "chat.completion.chunk", + }, + { + choices: [{ delta: {}, finish_reason: "stop", index: 0 }], + object: "chat.completion.chunk", + "amazon-bedrock-invocationMetrics": { + inputTokenCount: 126, + outputTokenCount: 125, + invocationLatency: 2006, + firstByteLatency: 685, + }, + }, + ].map((event) => ({ + chunk: { bytes: new TextEncoder().encode(JSON.stringify(event)) }, + })), + ), + ).toEqual({ + output: [ + { + index: 0, + message: { role: "assistant", content: "Hello world" }, + finish_reason: "stop", + logprobs: null, + }, + ], + metrics: { + completion_tokens: 125, + prompt_tokens: 126, + tokens: 251, + }, + }); + }); + + it.each([undefined, "Checking the weather."])( + "preserves OpenAI tool calls alongside content %j", + (content) => { + const events = [ + { + choices: [ + { + index: 0, + delta: { + role: "assistant", + content, + reasoning_content: "Look up the weather.", + }, + }, + ], + }, + { + choices: [ + { + index: 0, + delta: { + tool_calls: [ + { + index: 0, + id: "call_weather", + type: "function", + function: { + name: "get_weather", + arguments: '{"city":', + }, + }, + ], + }, + }, + ], + }, + { + choices: [ + { + index: 0, + delta: { + tool_calls: [ + { + index: 0, + function: { arguments: '"Berlin"}' }, + }, + ], + }, + }, + ], + }, + { + choices: [{ index: 0, delta: {}, finish_reason: "tool_calls" }], + "amazon-bedrock-invocationMetrics": { + inputTokenCount: 12, + outputTokenCount: 20, + }, + }, + ]; + + expect( + aggregateInvokeModelResponseStreamChunks( + events.map((event) => ({ + chunk: { bytes: new TextEncoder().encode(JSON.stringify(event)) }, + })), + ), + ).toEqual({ + output: [ + { + index: 0, + message: { + role: "assistant", + content, + reasoning_content: "Look up the weather.", + tool_calls: [ + { + id: "call_weather", + type: "function", + function: { + name: "get_weather", + arguments: '{"city":"Berlin"}', + }, + }, + ], + }, + finish_reason: "tool_calls", + logprobs: null, + }, + ], + metrics: { + prompt_tokens: 12, + completion_tokens: 20, + tokens: 32, + }, + }); + }, + ); + + it("prefers Nova metadata usage over its zeroed invocation metrics in InvokeModel response stream chunks", () => { + // Nova reports inputTokenCount/outputTokenCount as 0 in + // amazon-bedrock-invocationMetrics; the real counts are in metadata.usage. + expect( + aggregateInvokeModelResponseStreamChunks( + [ + { messageStart: { role: "assistant" } }, + { + contentBlockDelta: { + delta: { text: "Hello" }, + contentBlockIndex: 0, + }, + }, + { contentBlockStop: { contentBlockIndex: 0 } }, + { messageStop: { stopReason: "end_turn" } }, + { + metadata: { + usage: { + inputTokens: 6, + outputTokens: 5, + cacheReadInputTokenCount: 0, + cacheWriteInputTokenCount: 0, + }, + metrics: {}, + trace: {}, + }, + "amazon-bedrock-invocationMetrics": { + inputTokenCount: 0, + outputTokenCount: 0, + invocationLatency: 428, + firstByteLatency: 365, + }, + }, + ].map((event) => ({ + chunk: { bytes: new TextEncoder().encode(JSON.stringify(event)) }, + })), + ).metrics, + ).toEqual({ + completion_tokens: 5, + prompt_cache_creation_tokens: 0, + prompt_cached_tokens: 0, + prompt_tokens: 6, + tokens: 11, + }); + }); + it("aggregates Titan stream outputText from InvokeModel response stream chunks", () => { expect( aggregateInvokeModelResponseStreamChunks([ diff --git a/js/src/instrumentation/plugins/bedrock-runtime-plugin.ts b/js/src/instrumentation/plugins/bedrock-runtime-plugin.ts index f9a04d300..5e4d3987b 100644 --- a/js/src/instrumentation/plugins/bedrock-runtime-plugin.ts +++ b/js/src/instrumentation/plugins/bedrock-runtime-plugin.ts @@ -2,6 +2,7 @@ import type { ReturnOf } from "../core/interceptor-definitions"; import type { StartOf } from "../core/tracing-types"; import { SpanTypeAttribute, isObject } from "../../../util/index"; import type { Span } from "../../logger"; +import type { OpenAIChatCompletionChunk } from "../../vendor-sdk-types/openai"; import { INSTRUMENTATION_NAMES } from "../../span-origin"; import { getCurrentUnixTimestamp } from "../../util"; import { BasePlugin } from "../core"; @@ -27,6 +28,7 @@ import { getBedrockRuntimeCommandName, getBedrockRuntimeOperation, } from "./bedrock-runtime-common"; +import { aggregateChatCompletionChunks } from "./openai-plugin"; export class BedrockRuntimePlugin extends BasePlugin { protected onEnable(): void { @@ -530,6 +532,9 @@ export function aggregateInvokeModelResponseStreamChunks( .map((chunk) => parseJsonBody(chunk.chunk?.bytes)) .filter((chunk) => chunk !== undefined); const jsonLikeChunks = parsedChunks.filter(isObject); + const chatCompletionChunks = jsonLikeChunks.filter((chunk) => + Array.isArray(chunk.choices), + ) as OpenAIChatCompletionChunk[]; const text = parsedChunks.map(extractTextFromJsonLike).join(""); const lastMetadataChunk = jsonLikeChunks .slice() @@ -556,25 +561,42 @@ export function aggregateInvokeModelResponseStreamChunks( return { output: - text.length > 0 - ? { text } - : { - chunk_count: chunks.length, - chunks: sanitizeBedrockValue( - (jsonLikeChunks.length > 0 ? jsonLikeChunks : chunks).slice( - 0, - 20, + chatCompletionChunks.length > 0 + ? aggregateChatCompletionChunks(chatCompletionChunks).output + : text.length > 0 + ? { text } + : { + chunk_count: chunks.length, + chunks: sanitizeBedrockValue( + (jsonLikeChunks.length > 0 ? jsonLikeChunks : chunks).slice( + 0, + 20, + ), ), - ), - }, + }, metrics: parseBedrockRuntimeMetrics( - usage ?? (isObject(metadata) ? metadata.usage : undefined), + usage ?? + (isObject(metadata) ? metadata.usage : undefined) ?? + lastInvocationMetrics(jsonLikeChunks), isObject(metadata) ? metadata.metrics : undefined, ), ...(metadata ? { metadata } : {}), }; } +// Models whose stream carries no usage of its own (gpt-oss, for example) only +// report token counts in the amazon-bedrock-invocationMetrics of the last +// chunk. Nova sends zeros there, so it is only a fallback after metadata.usage. +function lastInvocationMetrics( + chunks: Record[], +): Record | undefined { + const metrics = chunks + .map((chunk) => chunk["amazon-bedrock-invocationMetrics"]) + .reverse() + .find(isObject); + return metrics ? sanitizeRecord(metrics) : undefined; +} + function logBedrockStreamAggregation( span: Span, aggregated: BedrockRuntimeStreamAggregation, @@ -819,6 +841,5 @@ function extractTextFromJsonLike(value: unknown): string { if (Array.isArray(value.outputs)) { return value.outputs.map(extractTextFromJsonLike).join(""); } - return ""; }