feat(weave): instrument invoke_model_with_response_stream for token tracking - #7678
Open
Aftabbs wants to merge 1 commit into
Open
feat(weave): instrument invoke_model_with_response_stream for token tracking#7678Aftabbs wants to merge 1 commit into
Aftabbs wants to merge 1 commit into
Conversation
… tracking Adds weave tracing for bedrock-runtime's invoke_model_with_response_stream API, which was previously untracked (issue wandb#5767). The implementation follows the same pattern as invoke_model and invoke_agent: - postprocess_output_invoke_stream buffers the EventStream body events, provides callers with a generator of the original events, and returns a logged copy with the concatenated content for the weave trace. - bedrock_on_finish_invoke_stream reads prompt/completion token counts from the x-amzn-bedrock-input-token-count and x-amzn-bedrock-output-token-count HTTP response headers (same header names as invoke_model). - _patch_invoke_model_stream wires the op into the bedrock-runtime client. - patch_client now calls _patch_invoke_model_stream when the method is present on the client.
Contributor
|
CLA Assistant Lite bot All contributors have signed the CLA ✍️ ✅ |
Aftabbs
changed the title
Author
|
I have read the CLA Document and I hereby sign the CLA |
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
1 participant
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 #5767
Summary
invoke_model_with_response_streamwas the only bedrock-runtime methodwith no weave tracing. This PR closes that gap with the same pattern
already used for
invoke_modelandinvoke_agent.What changed
weave/integrations/bedrock/bedrock_sdk.pypostprocess_output_invoke_stream— consumes thebodyEventStreambefore the call returns, buffers all events, puts a generator of the
original events back on
outputs["body"](so callers can stilliterate it), and returns a logged copy with the concatenated content
and event count.
ResponseMetadatais preserved on the originaloutputsdict for the finish handler.bedrock_on_finish_invoke_stream— reads token counts from thex-amzn-bedrock-input-token-count/x-amzn-bedrock-output-token-countHTTP response headers, matching the pattern in
bedrock_on_finish_invoke._patch_invoke_model_stream— wires the op withpostprocess_inputs_invoke(reused from
invoke_model), the newpostprocess_outputand finish handler.patch_client— calls_patch_invoke_model_streamwhen the methodis present on the client (guarded with
hasattrso it degradesgracefully on older boto3 versions).
tests/integrations/bedrock/bedrock_test.pyMOCK_INVOKE_STREAM_EVENTS— two chunk events with JSON text bytes.mock_invoke_stream_make_api_call— returns a dict withResponseMetadata(including the token-count headers) andbodyas an iterator over the mock events.
test_bedrock_invoke_model_stream— verifies that after callinginvoke_model_with_response_streamthe caller can still consume thestream, the trace is captured, and token counts are populated in the
call summary.
Test plan
All 13 existing tests still pass; the 3 ARN-parametrized errors are a
pre-existing Windows path-length limitation unrelated to this change.