diff --git a/examples/samples/agent_hooks_serverless.py b/examples/samples/agent_hooks_serverless.py index 8a958a9b..09983782 100644 --- a/examples/samples/agent_hooks_serverless.py +++ b/examples/samples/agent_hooks_serverless.py @@ -95,11 +95,6 @@ async def with_confirmation( async with my_agent.run(model, messages) as stream: async for event in stream: - # HACK?: When we get a complete assistant message, add it to - # messages so it can get replayed easily. - if isinstance(event, ai.events.StreamEnd): - messages.append(event.message) - if isinstance(event, ai.events.TextDelta): print(event.chunk, end="", flush=True) elif ( @@ -112,6 +107,9 @@ async def with_confirmation( f" Hook pending: {hook_part.hook_id} " f"(metadata={hook_part.metadata})" ) + # Pick up the assistant turn that the loop appended so the + # next run replays from the same point. + messages = stream.messages print("\n Run interrupted; approval will be pre-registered for re-entry.\n") diff --git a/examples/samples/middleware_simple.py b/examples/samples/middleware_simple.py index 146fe7d1..52a98024 100644 --- a/examples/samples/middleware_simple.py +++ b/examples/samples/middleware_simple.py @@ -26,8 +26,8 @@ class PrintMiddleware(ai.Middleware): async def wrap_agent_run( self, - call: ai.middleware.AgentRunContext, - next: Callable[[ai.middleware.AgentRunContext], AsyncGenerator[Any]], + call: ai.Context, + next: Callable[[ai.Context], AsyncGenerator[Any]], ) -> AsyncGenerator[Any]: print(f">>> [run] agent starting tools={len(call.tools)}") t0 = time.perf_counter() diff --git a/src/ai/__init__.py b/src/ai/__init__.py index 3f9ad3fb..ea1389fd 100644 --- a/src/ai/__init__.py +++ b/src/ai/__init__.py @@ -27,7 +27,7 @@ tool_result, yield_from, ) -from .agents.middleware import AgentRunContext, Middleware +from .agents.middleware import Middleware from .models import ( Client, Executor, @@ -119,7 +119,6 @@ "cancel_hook", "TOOL_APPROVAL_HOOK_TYPE", # Middleware - "AgentRunContext", "Middleware", "middleware", # Submodules diff --git a/src/ai/agents/agent.py b/src/ai/agents/agent.py index b8ce72d2..74b8172c 100644 --- a/src/ai/agents/agent.py +++ b/src/ai/agents/agent.py @@ -595,6 +595,55 @@ class LoopFn(Protocol): def __call__(self, context: Context) -> AsyncGenerator[events_.AgentEvent]: ... +class AgentStream: + """Async-iterable wrapper around an agent run's event stream. + + Exposes the run's :class:`Context` via :attr:`context` so callers can + inspect (or use) the live messages/tools without threading them + through their own bookkeeping:: + + async with agent.run(model, messages) as stream: + async for event in stream: + ... + print(stream.context.messages) + + Structurally satisfies the ``AsyncGenerator`` protocol by delegating + to the underlying generator, so it can be passed directly to + :func:`yield_from` and other APIs that expect an async generator. + """ + + def __init__( + self, + gen: AsyncGenerator[events_.AgentEvent], + context: Context, + ) -> None: + self._gen = gen + self._context = context + + def __aiter__(self) -> Self: + return self + + async def __anext__(self) -> events_.AgentEvent: + return await self._gen.__anext__() + + async def asend(self, value: Any) -> events_.AgentEvent: + return await self._gen.asend(value) + + async def athrow(self, *args: Any, **kwargs: Any) -> events_.AgentEvent: + return await self._gen.athrow(*args, **kwargs) + + async def aclose(self) -> None: + await self._gen.aclose() + + @property + def context(self) -> Context: + return self._context + + @property + def messages(self) -> list[types.messages.Message]: + return self._context.messages + + def tool_result( *items: types.messages.Message | types.messages.ToolResultPart @@ -746,14 +795,16 @@ async def run( messages: list[types.messages.Message], *, middleware: list[middleware_.Middleware] | None = None, - ) -> AsyncIterator[AsyncGenerator[events_.AgentEvent]]: + ) -> AsyncIterator[AgentStream]: """Run the agent loop, yielding events to the consumer. - Used as an async context manager whose value is the event stream:: + Used as an async context manager whose value the event stream, + extended with the `context` and `messages` of the stream:: async with agent.run(model, messages) as stream: async for event in stream: ... + print(stream.messages) Args: model: The model to use for LLM calls. @@ -766,37 +817,29 @@ async def run( ``yield_from(..., label=...)`` — the label flows via ``PartialToolCallResult`` rather than on individual messages. """ - call = middleware_.AgentRunContext( + context = Context( model=model, - messages=messages, - tools=self._tools, + messages=list(messages), + tools=[t.tool for t in self._tools], ) + context._agent_tools_by_name = {t.name: t for t in self._tools} + # If the final message is an assistant call with tool + # calls, then probably the situation is that we bailed out + # earlier due to unresolved hooks, and we need to arrange + # to replay the message now. + if ( + context.messages + and context.messages[-1].role == "assistant" + and context.messages[-1].tool_calls + ): + context.messages[-1] = context.messages[-1].model_copy( + update={"replay": True} + ) loop_fn = self._loop_fn or self.default_loop - async def _real( - call: middleware_.AgentRunContext, - ) -> AsyncGenerator[events_.AgentEvent]: - context = Context( - model=call.model, - messages=list(call.messages), - tools=[tool.tool for tool in call.tools], - ) - context._agent_tools_by_name = {tool.name: tool for tool in call.tools} - # If the final message is an assistant call with tool - # calls, then probably the situation is that we bailed out - # earlier due to unresolved hooks, and we need to arrange - # to replay the message now. - if ( - context.messages - and context.messages[-1].role == "assistant" - and context.messages[-1].tool_calls - ): - context.messages[-1] = context.messages[-1].model_copy( - update={"replay": True} - ) - - source = loop_fn(context) + async def _real(call: Context) -> AsyncGenerator[events_.AgentEvent]: + source = loop_fn(call) async for event in runtime.run(source): # Drop replay-flagged events: they're a control-flow # signal for the loop's tool dispatcher (which already @@ -820,13 +863,13 @@ async def _stream() -> AsyncGenerator[events_.AgentEvent]: mw_token = middleware_.activate(parent + middleware) try: chain = middleware_._build_agent_run_chain(_real) - async for event in chain(call): + async for event in chain(context): yield event finally: if mw_token is not None: middleware_.deactivate(mw_token) - yield _stream() + yield AgentStream(_stream(), context) def agent( diff --git a/src/ai/agents/middleware.py b/src/ai/agents/middleware.py index 54aeff4b..0b8bcfcd 100644 --- a/src/ai/agents/middleware.py +++ b/src/ai/agents/middleware.py @@ -45,7 +45,7 @@ if TYPE_CHECKING: from ..models.core.model import Model from ..types import events as events_ - from .agent import AgentTool + from .agent import Context @dataclasses.dataclass(frozen=True) @@ -102,19 +102,6 @@ def __post_init__(self) -> None: object.__setattr__(self, "metadata", dict(self.metadata)) -@dataclasses.dataclass(frozen=True) -class AgentRunContext: - """Context for an agent run.""" - - model: Model[Any] - messages: list[messages_.Message] - tools: list[AgentTool] - - def __post_init__(self) -> None: - object.__setattr__(self, "messages", list(self.messages)) - object.__setattr__(self, "tools", list(self.tools)) - - # --------------------------------------------------------------------------- # Middleware base class — override the methods you care about. # --------------------------------------------------------------------------- @@ -127,7 +114,7 @@ def __post_init__(self) -> None: _Message = messages_.Message # Agent run next-function type: call -> async generator of events. -_AgentRunNext = Callable[[AgentRunContext], AsyncGenerator[_Event]] +_AgentRunNext = Callable[["Context"], AsyncGenerator[_Event]] class Middleware: @@ -138,7 +125,7 @@ class Middleware: async def wrap_agent_run( self, - call: AgentRunContext, + call: Context, next: _AgentRunNext, ) -> AsyncGenerator[_Event]: """Wrap an agent run. @@ -354,7 +341,7 @@ def _build_agent_run_chain( for m in reversed(mw): def _make(m: Middleware, nxt: _AgentRunNext) -> _AgentRunNext: - async def _wrapped(call: AgentRunContext) -> AsyncGenerator[_Event]: + async def _wrapped(call: Context) -> AsyncGenerator[_Event]: async for event in m.wrap_agent_run(call, nxt): yield event @@ -365,7 +352,6 @@ async def _wrapped(call: AgentRunContext) -> AsyncGenerator[_Event]: __all__ = [ - "AgentRunContext", "GenerateContext", "HookContext", "Middleware", diff --git a/tests/test_middleware.py b/tests/test_middleware.py index 14118b70..87c05c38 100644 --- a/tests/test_middleware.py +++ b/tests/test_middleware.py @@ -110,7 +110,7 @@ async def test_wrap_agent_run_ordering() -> None: class Outer(ai.Middleware): async def wrap_agent_run( - self, call: middleware.AgentRunContext, next: Any + self, call: ai.Context, next: Any ) -> AsyncGenerator[ai.events.Event]: order.append("outer-before") async for event in next(call): @@ -119,7 +119,7 @@ async def wrap_agent_run( class Inner(ai.Middleware): async def wrap_agent_run( - self, call: middleware.AgentRunContext, next: Any + self, call: ai.Context, next: Any ) -> AsyncGenerator[ai.events.Event]: order.append("inner-before") async for event in next(call):