Skip to content

[Feature]: Patch bedrock client method invoke_model_with_streaming_response #5767

@jack-oconnor-lrn

Description

@jack-oconnor-lrn

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

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions