feat: pass invocation_state to model providers

[tirth14]

Enables custom model providers to access invocation_state from agent calls via kwargs. This supports use cases like custom request metadata, tracing context, and provider-specific configuration.

Changes:

• Add invocation_state parameter to stream_messages()
• Pass invocation_state through event_loop to model.stream()
• Update BedrockModel to extract and forward kwargs

Custom providers can access via: kwargs.get('invocation_state')

Description

This PR enables custom model providers to receive invocation_state from agent invocations. The invocation_state parameter flows from the agent call through the event loop to the model's stream() method, where custom providers can access it via kwargs.get('invocation_state').

This supports use cases such as:

• Passing custom request metadata to model providers
• Propagating tracing/observability context
• Provider-specific configuration per invocation

The implementation is fully backward compatible - the abstract Model.stream() already accepts **kwargs, so no changes are needed to the base interface.

Related Issues

N/A

Documentation PR

N/A - No documentation changes needed as this extends existing **kwargs functionality

Type of Change

New feature

Testing

How have you tested the change?

• I ran hatch run prepare

Checklist

• I have read the CONTRIBUTING document
• I have added any necessary tests that prove my fix is effective or my feature works
• I have updated the documentation accordingly
• I have added an appropriate example to the documentation to outline the feature, or no new docs are needed
• My changes generate no new warnings
• Any dependent changes have been merged and published

---

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

[dbschmigelski]

Hi, I am generally inclined.

The one change I would ask for is for us to make it explicit that invocation_state is being passed. So that would mean adding the invocation_state: Optional[dict[str, Any]] = None on stream and then passing it explicitly to the private methods.

I do not think we need kwargs on the private methods. Th reason we have it on stream is for extensibility without breaking backwards compatibilty. For private methods, we are not breaking anyone by adding a new parameter.

Also, this PR seems fairly small would you be able to update the other mode providers too?

Thanks

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[dbschmigelski]

@tirth14, I made some modifications here as I play around with addressing your request. I have reverted back to the original state though

Passing invocation state to the models makes sense for me in terms of extensibility. This allows custom implementations of the interface to leverage it as they see fit.

But, what was proposed in 230e6d8 seems like it might not be the right fix. You are passing the invocation_state to _format_request, what I imagine you are doing is overriding the _format_request and want access to invocation state as you are doing that.

This seems like the wrong approach.

For 1, this is a private method. So if you are overriding this it means we probably want some other mechanism for extensibility.

For 2, passing unrelated and otherwise unused parameters to a private method doesn't make sense. The only reason to do that would be if you believe that method might be overriden, which I see as problematic and related to #1.

---

I think I have a good compromise proposal though. If we begin passing in invocation_state, and I'm correct about what you are doing with _format_request you can override bedrock.py as you need. If we want to explore more extensibility, specifically in the Bedrock provider, I think that could be worth exploring but lets do it in a new GitHub issue.

Does this more limited change work for you?

[tirth14]

Hi @dbschmigelski, thank you for the review. Your changes makes sense to me. That would be a good compromise for now which would unblock us to access invocation_state from our custom model provider.

Sharing more context behind passing it through _format_request, you guessed it correctly. We are attempting override that method and pass additional arguments from invocation_state to the request. We have a proxy service on top of Bedrock calls which handles token management and some additional features, it follows the same request pattern as Bedrock converse + request metadata. Hence, our use case is to override bedrock's boto3 client with ours, and pass request metadata through invocation_state. We still want to reuse all other implementation for BedrockModel's stream as it is.

I understand overriding the private method _format_request might not be the best approach. I would appreciate if you have any thoughts / suggestions which can minimize rewriting most of BedrockModel implementation for our use case. I can create a separate Github issue for this to avoid mixing with current set of changes.

[dbschmigelski]

Hi regarding,

I understand overriding the private method _format_request might not be the best approach. I would appreciate if you have any thoughts / suggestions which can minimize rewriting most of BedrockModel implementation for our use case. I can create a separate Github issue for this to avoid mixing with current set of changes.

Yes, please create a new issue and lets discuss there! The overriding of the private method is flaky because we may unintentionally break you if we, for example, removed or refactored that method. But, the fact that you feel the need to override it in the first place means we likely a use case, beyond just you, where customizing the model implementation would be useful.

Let me get a commit out with the interface-only change and happy to work through what a more customizable BedrockModel looks like!