diff --git a/strands-py/src/strands/_middleware/stages.py b/strands-py/src/strands/_middleware/stages.py index ad8c65838c..69efde00c8 100644 --- a/strands-py/src/strands/_middleware/stages.py +++ b/strands-py/src/strands/_middleware/stages.py @@ -13,6 +13,7 @@ from ..agent.agent import Agent from ..experimental.bidi import BidiAgent from ..interrupt import _InterruptState + from ..models.model import Model from ..types._events import ModelStopReason, ToolResultEvent, TypedEvent from ..types.content import Messages, SystemPrompt from ..types.tools import AgentTool, ToolChoice, ToolSpec, ToolUse @@ -25,6 +26,9 @@ class InvokeModelContext: All collection fields (messages, system_prompt, tool_specs, tool_choice) are defensive copies — middleware cannot accidentally mutate agent state. invocation_state is shared by reference (hooks and tools write to it during streaming). + + ``model`` is the model this call invokes; it starts as ``agent.model`` and middleware + may replace it per call. """ agent: Agent @@ -33,6 +37,7 @@ class InvokeModelContext: tool_specs: list[ToolSpec] tool_choice: ToolChoice | None invocation_state: dict[str, Any] + model: Model projected_input_tokens: int | None = None diff --git a/strands-py/src/strands/event_loop/event_loop.py b/strands-py/src/strands/event_loop/event_loop.py index aa9536c9c3..93ac76a76b 100644 --- a/strands-py/src/strands/event_loop/event_loop.py +++ b/strands-py/src/strands/event_loop/event_loop.py @@ -537,6 +537,7 @@ async def _handle_model_execution( tool_specs=copy.deepcopy(tool_specs), tool_choice=copy.deepcopy(structured_output_context.tool_choice), invocation_state=invocation_state, + model=agent.model, projected_input_tokens=projected_input_tokens, ) @@ -559,8 +560,7 @@ async def _handle_model_execution( if last_event is None: raise RuntimeError( - "Middleware chain did not yield a result event. " - "Ensure middleware forwards events from next()." + "Middleware chain did not yield a result event. Ensure middleware forwards events from next()." ) # Write the post-stream model state back to the agent. Skipped on error @@ -663,7 +663,7 @@ def _make_invoke_model_terminal( async def terminal(ctx: InvokeModelContext) -> AsyncGenerator[Any, None]: system_prompt_str, system_prompt_content = split_system_prompt(ctx.system_prompt) - model_id = agent.model.config.get("model_id") if hasattr(agent.model, "config") else None + model_id = ctx.model.config.get("model_id") if hasattr(ctx.model, "config") else None model_invoke_span = tracer.start_model_invoke_span( messages=ctx.messages, parent_span=cycle_span, @@ -675,7 +675,7 @@ async def terminal(ctx: InvokeModelContext) -> AsyncGenerator[Any, None]: with trace_api.use_span(model_invoke_span, end_on_exit=False): try: async for event in stream_messages( - agent.model, + ctx.model, system_prompt_str, ctx.messages, ctx.tool_specs, diff --git a/strands-py/tests/strands/agent/conversation_manager/test_token_usage_middleware.py b/strands-py/tests/strands/agent/conversation_manager/test_token_usage_middleware.py index 1c55a3c2ad..8057b20d60 100644 --- a/strands-py/tests/strands/agent/conversation_manager/test_token_usage_middleware.py +++ b/strands-py/tests/strands/agent/conversation_manager/test_token_usage_middleware.py @@ -25,6 +25,7 @@ def make_context(**overrides) -> InvokeModelContext: tool_specs=[], tool_choice=None, invocation_state={}, + model=Mock(), ) defaults.update(overrides) return InvokeModelContext(**defaults) diff --git a/strands-py/tests/strands/injection/test_message_injection.py b/strands-py/tests/strands/injection/test_message_injection.py index 4a13144749..bf3386e27f 100644 --- a/strands-py/tests/strands/injection/test_message_injection.py +++ b/strands-py/tests/strands/injection/test_message_injection.py @@ -51,6 +51,7 @@ def invoke_ctx(messages: list[dict], agent: Any = None) -> InvokeModelContext: tool_specs=[], tool_choice=None, invocation_state={}, + model=MagicMock(), ) diff --git a/strands-py/tests/strands/memory/test_memory_manager.py b/strands-py/tests/strands/memory/test_memory_manager.py index 94ab713e00..591e961bc9 100644 --- a/strands-py/tests/strands/memory/test_memory_manager.py +++ b/strands-py/tests/strands/memory/test_memory_manager.py @@ -1192,6 +1192,7 @@ def _invoke_ctx(messages: list[dict], agent: Any) -> Any: tool_specs=[], tool_choice=None, invocation_state={}, + model=getattr(agent, "model", None), ) diff --git a/strands-py/tests/strands/middleware/test_agent_middleware.py b/strands-py/tests/strands/middleware/test_agent_middleware.py index 9bd2a7ce3b..260575ce20 100644 --- a/strands-py/tests/strands/middleware/test_agent_middleware.py +++ b/strands-py/tests/strands/middleware/test_agent_middleware.py @@ -720,3 +720,38 @@ async def retry_middleware(context, next_fn): result = agent("test") assert result.message["content"][0]["text"] == "Success!" assert call_count == 3 + + +# --- per-call model on context (routing plumbing) --- + + +def test_invoke_model_context_exposes_agent_model(agent): + """InvokeModelContext.model defaults to agent.model.""" + captured: list = [] + + async def capture(context, next_fn): + captured.append(context.model) + async for event in next_fn(context): + yield event + + agent._middleware_registry.add_middleware(InvokeModelStage, capture) + agent("test") + + assert captured == [agent.model] + + +def test_terminal_streams_context_model_override(): + """The terminal streams the model set on the context, not agent.model.""" + model_a = MockedModelProvider([{"role": "assistant", "content": [{"text": "A"}]}]) + model_b = MockedModelProvider([{"role": "assistant", "content": [{"text": "B"}]}]) + agent = Agent(model=model_a, callback_handler=None) + model_a.stream = AsyncMock(wraps=model_a.stream) + + def route_to_b(context): + return replace(context, model=model_b) + + agent._middleware_registry.add_middleware(InvokeModelStage.Input, route_to_b) + result = agent("test") + + assert result.message["content"][0]["text"] == "B" + model_a.stream.assert_not_called() diff --git a/strands-py/tests/strands/vended_plugins/context_injector/test_plugin.py b/strands-py/tests/strands/vended_plugins/context_injector/test_plugin.py index 372e62cdbf..6d9ce3b575 100644 --- a/strands-py/tests/strands/vended_plugins/context_injector/test_plugin.py +++ b/strands-py/tests/strands/vended_plugins/context_injector/test_plugin.py @@ -40,6 +40,7 @@ def invoke_ctx(messages: list[dict], agent: Any) -> InvokeModelContext: tool_specs=[], tool_choice=None, invocation_state={}, + model=MagicMock(), )