Skip to content

Conversation

@tirth14
Copy link

@tirth14 tirth14 commented Jan 5, 2026

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.

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')
*,
tool_choice: ToolChoice | None = None,
system_prompt_content: Optional[list[SystemContentBlock]] = None,
**kwargs: Any,
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants