Repository navigation
fix(bedrock): Log text and usage for streamed InvokeModel calls to OpenAI models - #2632
Open
alexcastgonz (pcbeingused333) wants to merge 1 commit into
Conversation
…enAI models OpenAI models on Bedrock (gpt-oss, for example) stream Chat Completions chunks from InvokeModelWithResponseStream. Those chunks carry no `usage` field; the token counts only arrive in the last chunk's amazon-bedrock-invocationMetrics. aggregateInvokeModelResponseStreamChunks never read it, and extractTextFromJsonLike had no case for choices[].delta.content, so the span logged no token metrics and fell back to dumping the first 20 raw chunks as output. Read text from choices[].delta.content, and use the last chunk's invocationMetrics as a fallback after message/metadata usage. Nova reports zeros there, so its metadata.usage keeps taking precedence; a test pins that. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
OpenAI models on Bedrock (
openai.gpt-oss-20b-1:0,openai.gpt-oss-120b-1:0, …) stream Chat Completions chunks fromInvokeModelWithResponseStream. The span for such a call ends up with no token metrics, and its output is a dump of the first 20 raw chunks instead of the text.Two things in
aggregateInvokeModelResponseStreamChunkscause this:usagefield (unless the caller asks forstream_options.include_usage). The counts only arrive inamazon-bedrock-invocationMetricson the last chunk, and the aggregator readsusage,message.usageandmetadata.usage, but never that.extractTextFromJsonLikehas no case forchoices[].delta.content, sotextis empty and the output falls back to{chunk_count, chunks}.Here's a real
openai.gpt-oss-20b-1:0stream (27 chunks, us-east-1) replayed throughwrapBedrockRuntimefrom braintrust 3.37.2:Changes
extractTextFromJsonLikereadschoices[].delta.content.usage/metadata.usageis present, the last chunk'samazon-bedrock-invocationMetricsis used.parseBedrockRuntimeMetricsalready mapsinputTokenCount/outputTokenCount.amazon-bedrock-invocationMetrics, but withinputTokenCount: 0, outputTokenCount: 0(seen on nova-micro, lite, pro and 2-lite), so itsmetadata.usagehas to keep winning. A test pins that.Tests
The two new tests in
bedrock-runtime-plugin.test.tsuse chunk shapes taken from real Bedrock streams:prompt_tokens: 126, completion_tokens: 125, tokens: 251.Includes a patch changeset.
Not in this PR: tool calls in streamed InvokeModel output (gpt-oss
delta.tool_calls, NovatoolUseblocks) still aren't reconstructed; only the text is logged. Happy to follow up if that's wanted.🤖 Generated with Claude Code