weave icon indicating copy to clipboard operation
weave copied to clipboard

[Feature]: Patch bedrock client method `invoke_model_with_streaming_response`

Open jack-oconnor-lrn opened this issue 1 month ago • 0 comments

Description

I did not find any existing issues when searching "invoke_model_with_streaming_response".

Feature request

The current Bedrock patch client message only patches invoke_model. It would be great if token usage and costs could also be automatically tracked for invoke-model-with-response-stream as well

Suggested Solution

I think a possible first place to look would be patching the generator returned by invoke_model_with_response_stream t check each chunk into and when the final chunk is returned grab / finish collecting the token usage stats.

Weave Project Link

https://wandb.ai/feedback-aide/dev/weave

Screenshots or Mockups

The following two screenshots are from the same overall trace. No automatic spans are created for calls to model_invoke_with_response_stream.

Two calls to model_invoke Image

One call to model_invoke and one call to model_invoke_with_response_stream Image

jack-oconnor-lrn avatar Nov 19 '25 13:11 jack-oconnor-lrn