Repository navigation
fix(bedrock): log token metrics and summarize InvokeModel embeddings - #1081
Open
alexcastgonz (pcbeingused333) wants to merge 1 commit into
Conversation
Read token counts from Bedrock response headers when body usage is absent. Summarize Titan and Cohere embedding vectors as length/count, including Cohere typed dictionaries and Titan binary-only embeddingsByType payloads. Keep existing Converse and Anthropic usage metrics unchanged. Fixes braintrustdata#574
Codex Review Summary: ✅ CompletedThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
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.
Fixes #574
Root cause.
_invoke_model_wrapperlogs the parsedInvokeModelresponse body as span output and only reads token usage from a nestedusageobject (Converse shape) or Anthropic'susage. Amazon Titan Embeddings and Cohere Embed responses have neither, so their spans had noprompt_tokens/tokens, and the output contained the full embedding vectors. Titan v2 also repeats its vector underembeddingsByType.float.Fix. For
invoke_modelthe wrapper now also (1) readsX-Amzn-Bedrock-Input-Token-Count/X-Amzn-Bedrock-Output-Token-Countfrom the response headers, and (2) logsembedding/embeddingsasembedding_length(+embedding_countfor batches), including Cohere typedembeddingsdictionaries and Titan binary-onlyembeddingsByTyperesponses, dropping the raw vector containers, matching what the other provider integrations do. Existing Converse/Anthropic metrics still take precedence. Streaming is unchanged.Test. Two new tests (Titan Embeddings v2, Cohere Embed v3) use cassettes recorded against Bedrock in us-east-1, under
cassettes/latestandcassettes/1.34.116(the second with boto3/botocore 1.34.116). They fail on main (KeyErroron the token metrics) and pass now. Six parameterized unit cases cover typed Cohere embeddings, binary-only Titan embeddings, the list shapes and unchanged non-vector payloads. WithCI=1and no credentials,test_bedrock_runtime.pygives 18 passed on latest, and 17 passed, 1 skipped on 1.34.116.ruff checkandruff format --checkare clean.This change was prepared with the help of an AI coding assistant (Claude Code and Codex), and the author is responsible for it.