From 548e21c75c184feb4890da32c046a0caf856b08e Mon Sep 17 00:00:00 2001 From: "Michael J. Sullivan" Date: Mon, 11 May 2026 10:15:03 -0700 Subject: [PATCH] Let `stream` accept an `agents.Context` instead of model/messages/tools Adds a `StreamContext` Protocol (so the model layer doesn't have to import the agents package) and overloads `stream` to accept a kw-only `context=` in addition to the existing positional `model, messages` (plus `tools=`). Mixing positional model/messages with `context=` is a runtime TypeError. Updates the default agent loop, all custom-loop examples, and the context-bearing tests to use the shorthand. --- README.md | 2 +- examples/fastapi-vite/backend/agent.py | 4 +- examples/multiagent-textual/server.py | 2 +- examples/samples/agent_custom_loop.py | 4 +- examples/samples/agent_hooks.py | 2 +- examples/samples/agent_hooks_inline.py | 2 +- examples/samples/agent_hooks_serverless.py | 4 +- src/ai/agents/agent.py | 4 +- src/ai/models/core/api.py | 85 +++++++++++++++++++--- tests/agents/test_hooks.py | 10 +-- tests/models/core/test_api.py | 37 ++++++++++ tests/test_middleware.py | 2 +- 12 files changed, 127 insertions(+), 31 deletions(-) diff --git a/README.md b/README.md index 56496df7..ccd4ccff 100644 --- a/README.md +++ b/README.md @@ -107,7 +107,7 @@ Override the default loop when you need approval gates, routing, or custom orche @agent.loop async def custom(context: ai.Context): while True: - async with ai.stream(context.model, context.messages, tools=context.tools) as s: + async with ai.stream(context=context) as s: async for event in s: yield event context.add(s.message) diff --git a/examples/fastapi-vite/backend/agent.py b/examples/fastapi-vite/backend/agent.py index e89e1c08..2dab3b6a 100644 --- a/examples/fastapi-vite/backend/agent.py +++ b/examples/fastapi-vite/backend/agent.py @@ -56,9 +56,7 @@ async def graph(context: ai.Context) -> AsyncGenerator[ai.events.AgentEvent]: Reject buttons and sends the decision back on the next request. """ while context.keep_running(): - async with ai.models.stream( - context.model, context.messages, tools=context.tools - ) as s: + async with ai.models.stream(context=context) as s: async for event in s: yield event context.add(s.message) diff --git a/examples/multiagent-textual/server.py b/examples/multiagent-textual/server.py index 43919d4d..c12200b5 100644 --- a/examples/multiagent-textual/server.py +++ b/examples/multiagent-textual/server.py @@ -110,7 +110,7 @@ def _gated_agent( async def gated_loop(context: ai.Context) -> AsyncGenerator[ai.events.AgentEvent]: while context.keep_running(): async with ( - ai.stream(context.model, context.messages, tools=context.tools) as s, + ai.stream(context=context) as s, ai.agents.ToolRunner() as tr, ): async for event in ai.util.merge(s, tr.events()): diff --git a/examples/samples/agent_custom_loop.py b/examples/samples/agent_custom_loop.py index 1ffe5fe3..f9bd7fe0 100644 --- a/examples/samples/agent_custom_loop.py +++ b/examples/samples/agent_custom_loop.py @@ -27,9 +27,7 @@ async def default_loop( """Stream, execute tools with logging, repeat.""" while context.keep_running(): async with ( - ai.models.stream( - context.model, context.messages, tools=context.tools - ) as stream, + ai.models.stream(context=context) as stream, ai.ToolRunner() as tr, ): async for event in ai.util.merge(stream, tr.events()): diff --git a/examples/samples/agent_hooks.py b/examples/samples/agent_hooks.py index dc50c9a6..86ebfd76 100644 --- a/examples/samples/agent_hooks.py +++ b/examples/samples/agent_hooks.py @@ -70,7 +70,7 @@ async def with_approval( ) -> AsyncGenerator[ai.events.AgentEvent]: while context.keep_running(): async with ( - ai.stream(context.model, context.messages, tools=context.tools) as s, + ai.stream(context=context) as s, ai.ToolRunner() as tr, ): async for event in ai.util.merge(s, tr.events()): diff --git a/examples/samples/agent_hooks_inline.py b/examples/samples/agent_hooks_inline.py index aa9f0cf0..09548248 100644 --- a/examples/samples/agent_hooks_inline.py +++ b/examples/samples/agent_hooks_inline.py @@ -46,7 +46,7 @@ async def with_approval( ) -> AsyncGenerator[ai.events.AgentEvent]: while context.keep_running(): async with ( - ai.stream(context.model, context.messages, tools=context.tools) as s, + ai.stream(context=context) as s, ai.ToolRunner() as tr, ): async for event in ai.util.merge(s, tr.events()): diff --git a/examples/samples/agent_hooks_serverless.py b/examples/samples/agent_hooks_serverless.py index 09983782..5bd917f9 100644 --- a/examples/samples/agent_hooks_serverless.py +++ b/examples/samples/agent_hooks_serverless.py @@ -47,9 +47,7 @@ async def with_confirmation( context: ai.Context, ) -> AsyncGenerator[ai.events.AgentEvent]: while context.keep_running(): - async with ai.models.stream( - context.model, context.messages, tools=context.tools - ) as s: + async with ai.models.stream(context=context) as s: async for event in s: yield event diff --git a/src/ai/agents/agent.py b/src/ai/agents/agent.py index 011df48d..9b3d3fec 100644 --- a/src/ai/agents/agent.py +++ b/src/ai/agents/agent.py @@ -770,9 +770,7 @@ async def default_loop( """Stream, execute tools, repeat.""" while context.keep_running(): async with ( - models.stream( - context.model, context.messages, tools=context.tools - ) as stream, + models.stream(context=context) as stream, ToolRunner() as tr, ): async for event in util.merge(stream, tr.events()): diff --git a/src/ai/models/core/api.py b/src/ai/models/core/api.py index 697d5940..bc672913 100644 --- a/src/ai/models/core/api.py +++ b/src/ai/models/core/api.py @@ -1,7 +1,8 @@ import contextlib import dataclasses from collections.abc import AsyncGenerator, AsyncIterator, Sequence -from typing import Any, Protocol, Self, runtime_checkable +from contextlib import AbstractAsyncContextManager +from typing import Any, Protocol, Self, overload, runtime_checkable import pydantic @@ -302,8 +303,32 @@ async def _replay_tool_calls( ) -@contextlib.asynccontextmanager -async def stream( +@runtime_checkable +class StreamContext(Protocol): + """Anything that exposes ``model``/``messages``/``tools``. + + Used to let callers pass an ``agents.Context`` to :func:`stream` + without an import-time circular dependency. + """ + + @property + def model(self) -> model_.Model: ... + @property + def messages(self) -> list[types.messages.Message]: ... + @property + def tools(self) -> list[types.tools.Tool]: ... + + +@overload +def stream( + *, + context: StreamContext, + output_type: type[pydantic.BaseModel] | None = None, + params: Any = None, + executor: StreamExecutor = _default_executor, +) -> AbstractAsyncContextManager[Stream]: ... +@overload +def stream[ProviderParamsT: pydantic.BaseModel]( model: model_.Model, messages: list[types.messages.Message], *, @@ -311,19 +336,61 @@ async def stream( output_type: type[pydantic.BaseModel] | None = None, params: Any = None, executor: StreamExecutor = _default_executor, -) -> AsyncIterator[Stream]: +) -> AbstractAsyncContextManager[Stream]: ... +def stream( + model: model_.Model | None = None, + messages: list[types.messages.Message] | None = None, + *, + context: StreamContext | None = None, + tools: Sequence[types.tools.Tool] | None = None, + output_type: type[pydantic.BaseModel] | None = None, + params: Any = None, + executor: StreamExecutor = _default_executor, +) -> AbstractAsyncContextManager[Stream]: """Stream an LLM response. - Used as an async context manager whose value is the :class:`Stream`:: + Used as an async context manager whose value is the :class:`Stream`. + Pass either positional ``model, messages`` (plus optional ``tools=``) + or ``context=`` (an ``agents.Context`` or anything matching + :class:`StreamContext`):: - async with ai.stream(model, messages) as s: - async for event in s: - ... - print(s.message) + async with ai.stream(model, messages) as s: ... + async with ai.stream(context=context) as s: ... If the last message is marked ``replay=True``, replay that turn as synthetic stream events instead of calling the model. """ + if context is not None: + if model is not None or messages is not None or tools is not None: + raise TypeError( + "stream() takes either model/messages/tools or context=, not both" + ) + model = context.model + messages = context.messages + tools = context.tools + elif model is None or messages is None: + raise TypeError("stream() requires either model and messages or context=") + + return _stream( + model=model, + messages=messages, + tools=tools, + output_type=output_type, + params=params, + executor=executor, + ) + + +@contextlib.asynccontextmanager +async def _stream( + *, + model: model_.Model, + messages: list[types.messages.Message], + tools: Sequence[types.tools.Tool] | None, + output_type: type[pydantic.BaseModel] | None, + params: Any, + executor: StreamExecutor, +) -> AsyncIterator[Stream]: if messages and messages[-1].replay: last = messages[-1] s = Stream(_replay_tool_calls(last), seed_message=last.model_copy(deep=True)) diff --git a/tests/agents/test_hooks.py b/tests/agents/test_hooks.py index 4c9623d7..eedc197f 100644 --- a/tests/agents/test_hooks.py +++ b/tests/agents/test_hooks.py @@ -31,7 +31,7 @@ async def test_resolve_live_future() -> None: @my_agent.loop async def custom(context: ai.Context) -> AsyncGenerator[ai.events.Event]: nonlocal resolved_value - async with ai.models.stream(context.model, context.messages) as stream: + async with ai.models.stream(context=context) as stream: async for event in stream: yield event result = await ai.hook("confirm_1", payload=Confirmation) @@ -63,7 +63,7 @@ async def test_cancel_live_hook() -> None: @my_agent.loop async def custom(context: ai.Context) -> AsyncGenerator[ai.events.Event]: nonlocal was_cancelled - async with ai.models.stream(context.model, context.messages) as stream: + async with ai.models.stream(context=context) as stream: async for event in stream: yield event try: @@ -102,7 +102,7 @@ async def test_pre_registered_resolution_consumed() -> None: @my_agent.loop async def custom(context: ai.Context) -> AsyncGenerator[ai.events.Event]: nonlocal resolved_value - async with ai.models.stream(context.model, context.messages) as stream: + async with ai.models.stream(context=context) as stream: async for event in stream: yield event resolved_value = await ai.hook("pre_reg_1", payload=Confirmation) @@ -142,7 +142,7 @@ async def test_resolved_hook_emits_message() -> None: @my_agent.loop async def custom(context: ai.Context) -> AsyncGenerator[ai.events.Event]: - async with ai.models.stream(context.model, context.messages) as stream: + async with ai.models.stream(context=context) as stream: async for event in stream: yield event await ai.hook("emit_test", payload=Confirmation) @@ -171,7 +171,7 @@ async def test_hook_metadata_in_pending() -> None: @my_agent.loop async def custom(context: ai.Context) -> AsyncGenerator[ai.events.Event]: - async with ai.models.stream(context.model, context.messages) as stream: + async with ai.models.stream(context=context) as stream: async for event in stream: yield event await ai.hook( diff --git a/tests/models/core/test_api.py b/tests/models/core/test_api.py index 162f2c99..ed4c43d3 100644 --- a/tests/models/core/test_api.py +++ b/tests/models/core/test_api.py @@ -4,6 +4,7 @@ from typing import Any import pydantic +import pytest import ai from ai import models @@ -227,6 +228,42 @@ async def _spy_stream( assert received_params == [params] +async def test_stream_accepts_context() -> None: + """``stream(context=ctx)`` reads model/messages/tools off the context.""" + mock = mock_llm([[text_msg("ok")]]) + ctx = ai.Context( + model=MOCK_MODEL, + messages=[ai.user_message("Hi")], + tools=[], + ) + async with models.stream(context=ctx) as s: + async for _ in s: + pass + assert mock.call_count == 1 + assert s.text == "ok" + + +async def test_stream_rejects_context_with_positional_args() -> None: + """Passing both positional model/messages and ``context=`` is a TypeError.""" + ctx = ai.Context( + model=MOCK_MODEL, + messages=[ai.user_message("Hi")], + tools=[], + ) + with pytest.raises(TypeError, match="either model/messages/tools or context="): + async with models.stream( # type: ignore[call-overload] + MOCK_MODEL, [ai.user_message("Hi")], context=ctx + ): + pass + + +async def test_stream_requires_model_messages_or_context() -> None: + """Passing nothing is a TypeError.""" + with pytest.raises(TypeError, match="either model and messages or context="): + async with models.stream(): # type: ignore[call-overload] + pass + + async def test_generate_dispatches_to_registered_adapter() -> None: provider = MockProvider(adapter="mock-generate") model = models.Model( diff --git a/tests/test_middleware.py b/tests/test_middleware.py index 87c05c38..360a300c 100644 --- a/tests/test_middleware.py +++ b/tests/test_middleware.py @@ -79,7 +79,7 @@ async def wrap_hook(self, call: middleware.HookContext, next: Any) -> Any: @my_agent.loop async def custom(context: ai.Context) -> AsyncGenerator[ai.events.Event]: - async with ai.models.stream(context.model, context.messages) as stream: + async with ai.models.stream(context=context) as stream: async for event in stream: yield event await ai.hook("test_hook", payload=Confirmation)