diff --git a/docs/ai-python/content/docs/reference/ai/agent.mdx b/docs/ai-python/content/docs/reference/ai/agent.mdx index 02692e6e..476f6bd5 100644 --- a/docs/ai-python/content/docs/reference/ai/agent.mdx +++ b/docs/ai-python/content/docs/reference/ai/agent.mdx @@ -107,6 +107,24 @@ stream.output text. When `output_type` is set, the text is validated as JSON against that Pydantic model and the parsed instance is returned. +## after_step + +Override `async def after_step(self, context: ai.Context) -> None` to inspect +or modify the context between steps without replacing the default loop. +The default implementation does nothing. + +The default loop calls it once after each completed model response and all +of its tool results have been appended to history, after the step's stream, +tool runner, and event merge have closed. This includes a final response +without tool calls. Changes to the context affect the next `keep_running()` +check; for example, appending a pending user message continues the run. +An exception or cancellation propagates and prevents the next model request. +A failed step does not call `after_step`. + +The hook does not wait for the consumer to drain buffered events. Set +`LOOP_BUFFER = 0` if it must run after consumer side effects for the step's +events. A custom `loop` must call `after_step` itself if needed. + ## loop Override `Agent.loop(context)` to customize control flow. The default loop uses diff --git a/src/ai/agents/agent.py b/src/ai/agents/agent.py index ec53583c..fe77aeb2 100644 --- a/src/ai/agents/agent.py +++ b/src/ai/agents/agent.py @@ -1460,6 +1460,19 @@ def resolve( """Resolve ToolCallPart(s) into callable ToolCall object(s).""" return self._tool_resolver.resolve(tool_part) + async def after_step(self, context: Context) -> None: + """Run after a completed step, before deciding whether to continue. + + The assistant message and all tool results are in ``context``, and + the step's stream and tool runner are closed. This includes final + assistant-only steps. Override to inspect or modify the context; + the default does nothing. Exceptions and cancellation propagate. + + This does not wait for the consumer to drain buffered events. Use + ``LOOP_BUFFER = 0`` when sequencing against consumer side effects. + Custom ``loop`` implementations must call this method themselves. + """ + async def loop( self, context: Context ) -> AsyncGenerator[events_.AgentEvent]: @@ -1490,6 +1503,8 @@ async def loop( # the loop. context.add(tr.get_tool_message()) + await self.after_step(context) + @overload def run( self, diff --git a/tests/agents/test_after_step.py b/tests/agents/test_after_step.py new file mode 100644 index 00000000..b044ec28 --- /dev/null +++ b/tests/agents/test_after_step.py @@ -0,0 +1,94 @@ +"""Default-loop step boundaries through the public Agent.run API.""" + +from __future__ import annotations + +import asyncio + +import ai +from ai.types import events, messages + +from ..conftest import ( + MOCK_MODEL, + collect_messages, + mock_llm, + text_msg, + tool_call_msg, +) + + +async def test_after_step_can_continue_final_response() -> None: + snapshots: list[list[str]] = [] + + class PendingInputAgent(ai.Agent): + async def after_step(self, context: ai.Context) -> None: + snapshots.append([message.text for message in context.messages]) + if len(snapshots) == 1: + context.add(ai.user_message("One more question")) + + llm = mock_llm([[text_msg("First answer")], [text_msg("Second answer")]]) + async with PendingInputAgent().run( + MOCK_MODEL, [ai.user_message("Question")] + ) as stream: + await collect_messages(stream) + assert stream.output == "Second answer" + + assert llm.call_count == 2 + assert snapshots == [ + ["Question", "First answer"], + ["Question", "First answer", "One more question", "Second answer"], + ] + + +async def test_after_step_sees_complete_concurrent_batch() -> None: + started: list[int] = [] + both_started = asyncio.Event() + second_finished = asyncio.Event() + snapshots: list[list[messages.Message]] = [] + consumed_results: set[str] = set() + + @ai.tool + async def complete(number: int) -> int: + """Start both tools, then complete the second one first.""" + started.append(number) + if len(started) == 2: + both_started.set() + await both_started.wait() + if number == 1: + await second_finished.wait() + else: + second_finished.set() + return number * 10 + + class BatchAgent(ai.Agent): + LOOP_BUFFER = 0 + + async def after_step(self, context: ai.Context) -> None: + assert consumed_results == {"tc-1", "tc-2"} + snapshots.append(list(context.messages)) + + calls = [ + tool_call_msg( + tc_id=f"tc-{number}", + name="complete", + args=f'{{"number": {number}}}', + ) + for number in (1, 2) + ] + llm = mock_llm([calls, [text_msg("Done")]]) + async with asyncio.timeout(5): + async with BatchAgent(tools=[complete]).run( + MOCK_MODEL, [ai.user_message("Run both")] + ) as stream: + async for event in stream: + if isinstance(event, events.ToolCallResult): + consumed_results.update( + result.tool_call_id for result in event.results + ) + + assert llm.call_count == 2 + assert len(snapshots) == 2 + assert { + result.tool_call_id: result.result + for result in snapshots[0][-1].tool_results + } == {"tc-1": 10, "tc-2": 20} + assert snapshots[1][-1].text == "Done"