Skip to content

fix(bedrock): Log text and usage for streamed InvokeModel calls to OpenAI models - #2632

Open
alexcastgonz (pcbeingused333) wants to merge 1 commit into
braintrustdata:mainfrom
pcbeingused333:fix/bedrock-invoke-stream-openai-usage
Open

alexcastgonz (pcbeingused333) wants to merge 1 commit into
braintrustdata:mainfrom
pcbeingused333:fix/bedrock-invoke-stream-openai-usage

Conversation

@pcbeingused333

Copy link
Copy Markdown

OpenAI models on Bedrock (openai.gpt-oss-20b-1:0, openai.gpt-oss-120b-1:0, …) stream Chat Completions chunks from InvokeModelWithResponseStream. 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 aggregateInvokeModelResponseStreamChunks cause this:

  • Usage. gpt-oss chunks have no usage field (unless the caller asks for stream_options.include_usage). The counts only arrive in amazon-bedrock-invocationMetrics on the last chunk, and the aggregator reads usage, message.usage and metadata.usage, but never that.
  • Text. extractTextFromJsonLike has no case for choices[].delta.content, so text is empty and the output falls back to {chunk_count, chunks}.

Here's a real openai.gpt-oss-20b-1:0 stream (27 chunks, us-east-1) replayed through wrapBedrockRuntime from braintrust 3.37.2:

logged output:  {"chunk_count":27,"chunks":[{"choices":[{"delta":{"content":"","role":"assistant"}, ...
logged metrics: {"time_to_first_token":0.014}
last chunk:     "amazon-bedrock-invocationMetrics": {"inputTokenCount":126,"outputTokenCount":125, ...}

Changes

  • extractTextFromJsonLike reads choices[].delta.content.
  • When no usage / metadata.usage is present, the last chunk's amazon-bedrock-invocationMetrics is used. parseBedrockRuntimeMetrics already maps inputTokenCount / outputTokenCount.
  • It's only a fallback: Nova also sends amazon-bedrock-invocationMetrics, but with inputTokenCount: 0, outputTokenCount: 0 (seen on nova-micro, lite, pro and 2-lite), so its metadata.usage has to keep winning. A test pins that.

Tests

The two new tests in bedrock-runtime-plugin.test.ts use chunk shapes taken from real Bedrock streams:

  • gpt-oss: text + invocationMetrics → prompt_tokens: 126, completion_tokens: 125, tokens: 251.
  • Nova: metadata.usage 6/5 alongside zeroed invocationMetrics → 6/5.
$ cd js && npx vitest run src/instrumentation/plugins/bedrock-runtime-plugin.test.ts
Tests  19 passed (19)
# main 5a8b3c5 with only the tests added: 1 failed (the gpt-oss test), 18 passed
$ npx vitest run src/instrumentation
Tests  1006 passed (1006)
$ npx tsc --noEmit -p .        # no errors
$ npx eslint <changed files>   # 0 errors (8 existing warnings, none on changed lines)
$ npx prettier --check ...     # clean

Includes a patch changeset.

Not in this PR: tool calls in streamed InvokeModel output (gpt-oss delta.tool_calls, Nova toolUse blocks) still aren't reconstructed; only the text is logged. Happy to follow up if that's wanted.

🤖 Generated with Claude Code

…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

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant