From 75ee384b14997ffac89c9aa147143de60d62f9d6 Mon Sep 17 00:00:00 2001 From: "Michael J. Sullivan" Date: Fri, 8 May 2026 10:17:07 -0700 Subject: [PATCH] Convert temporal-direct loop to use ToolRunner The custom @agent.loop now mirrors Agent.default_loop closely: a streaming model call feeds a ToolRunner, tool calls are scheduled as they arrive, and results are folded back into the context via context.add(stream.message) and context.add(tr.get_tool_message()). To support this we add a `replay_message_events` function to `ai.events`. --- examples/temporal-direct/main.py | 101 ++++++++++++++++++++----------- src/ai/models/core/api.py | 12 ++++ src/ai/types/events.py | 74 +++++++++++++++++++++- 3 files changed, 150 insertions(+), 37 deletions(-) diff --git a/examples/temporal-direct/main.py b/examples/temporal-direct/main.py index 23ef9fc7..cfdaf7fd 100644 --- a/examples/temporal-direct/main.py +++ b/examples/temporal-direct/main.py @@ -6,7 +6,14 @@ without re-executing. This is the "lego bricks" approach: the framework gives you ``Agent``, -``Context``, ``@tool``, and the message types. You compose them yourself. +``Context``, ``@tool``, the message types, and ``ToolRunner``. You compose +them yourself. + +The loop mirrors :py:meth:`ai.Agent.default_loop`: a streaming model call +feeds a ``ToolRunner``, tool calls are scheduled as they're emitted, and +results are folded back into the context. The only difference is that +the model "stream" is synthesized from the result of an LLM activity and +each tool call is dispatched as a Temporal activity. Prerequisites: 1. Temporal dev server: temporal server start-dev @@ -19,9 +26,9 @@ from __future__ import annotations -import asyncio import dataclasses import datetime +import json import sys import uuid from collections.abc import AsyncGenerator @@ -117,9 +124,10 @@ async def llm_call_activity(params: LLMParams) -> LLMResult: # ── Agent with custom loop ─────────────────────────────────────── # -# The loop replaces ai.models.stream() and tool execution with -# Temporal activity calls. The structure mirrors the default loop: -# call LLM → add message to context → execute tool calls → repeat. +# The loop mirrors ai.Agent.default_loop: stream → schedule tools as +# they arrive → fold results back into the context. The only twist +# is that the "stream" comes from a Temporal activity (not the LLM +# directly) and each scheduled tool dispatches via another activity. weather_agent = ai.agent(tools=[get_weather, get_population]) @@ -130,8 +138,8 @@ async def temporal_loop(context: ai.Context) -> AsyncGenerator[ai.events.AgentEv {"name": t.name, "args": t.args.model_dump(mode="json")} for t in context.tools ] - while True: - # 1. LLM call via activity + while context.keep_running(): + # 1. LLM call via activity → complete message result = await temporalio.workflow.execute_activity( llm_call_activity, LLMParams( @@ -141,34 +149,53 @@ async def temporal_loop(context: ai.Context) -> AsyncGenerator[ai.events.AgentEv start_to_close_timeout=datetime.timedelta(minutes=5), retry_policy=temporalio.common.RetryPolicy(maximum_attempts=3), ) - msg = ai.messages.Message.model_validate(result.message) - yield ai.events.StreamEnd(message=msg) - context.add(msg) - - # 2. No tool calls → done - if not msg.tool_calls: - break - - # 3. Execute each tool call as a Temporal activity (parallel) - async def run_tool(tc: ai.messages.ToolCallPart) -> ai.messages.ToolResultPart: - import json - - activity_fn = TOOL_ACTIVITIES[tc.tool_name] - kwargs = json.loads(tc.tool_args) if tc.tool_args else {} - result = await temporalio.workflow.execute_activity( - activity_fn, - args=list(kwargs.values()), - start_to_close_timeout=datetime.timedelta(minutes=2), - ) - return ai.tool_result_part( - tc.tool_call_id, tool_name=tc.tool_name, result=result - ) - - tasks = [asyncio.ensure_future(run_tool(tc)) for tc in msg.tool_calls] - parts = await asyncio.gather(*tasks) - result_event = ai.tool_result(*parts) - yield result_event - context.add(result_event.message) + llm_msg = ai.messages.Message.model_validate(result.message) + + # 2. Wrap the complete message in a synthetic stream so we can + # drive the rest of the loop with ToolRunner — same shape as + # the default loop. ``replay_message_events`` is the framework + # helper that decomposes a complete ``Message`` back into the + # events a streaming adapter would have produced. + async with ( + ai.Stream(ai.events.replay_message_events(llm_msg)) as stream, + ai.ToolRunner() as tr, + ): + async for event in ai.util.merge(stream, tr.events()): + yield event + + if isinstance(event, ai.events.ToolEnd): + tr.schedule(_activity_tool_call(event.tool_call)) + + context.add(stream.message) + context.add(tr.get_tool_message()) + + +def _activity_tool_call( + tc: ai.messages.ToolCallPart, +) -> ai.ToolCallLike: + """Build a ``ToolCallLike`` that runs the tool as a Temporal activity. + + ``ToolRunner.schedule`` accepts any zero-arg callable that returns + a coroutine resolving to a ``ToolCallResult``. This lets us route + tool execution through a Temporal activity (durable!) while keeping + the rest of the loop identical to ``Agent.default_loop``. + """ + + async def _call() -> ai.events.ToolCallResult: + activity_fn = TOOL_ACTIVITIES[tc.tool_name] + kwargs = json.loads(tc.tool_args) if tc.tool_args else {} + result = await temporalio.workflow.execute_activity( + activity_fn, + args=list(kwargs.values()), + start_to_close_timeout=datetime.timedelta(minutes=2), + ) + return ai.tool_result( + tool_call_id=tc.tool_call_id, + tool_name=tc.tool_name, + result=result, + ) + + return _call # ── Workflow ───────────────────────────────────────────────────── @@ -189,7 +216,7 @@ async def run(self, user_query: str) -> str: final_text = "" async with weather_agent.run(model, messages) as stream: async for event in stream: - if isinstance(event, ai.events.TerminalEvent): + if isinstance(event, ai.events.StreamEnd): final_text = event.message.text return final_text @@ -226,6 +253,8 @@ async def main(user_query: str) -> None: if __name__ == "__main__": + import asyncio + query = ( sys.argv[1] if len(sys.argv) > 1 diff --git a/src/ai/models/core/api.py b/src/ai/models/core/api.py index 18c581a9..384fc404 100644 --- a/src/ai/models/core/api.py +++ b/src/ai/models/core/api.py @@ -141,6 +141,18 @@ def __init__( async def aclose(self) -> None: await self._gen.aclose() + async def __aenter__(self) -> Self: + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: object, + ) -> bool: + await self.aclose() + return False + def __aiter__(self) -> Self: return self diff --git a/src/ai/types/events.py b/src/ai/types/events.py index f192cb71..019485a9 100644 --- a/src/ai/types/events.py +++ b/src/ai/types/events.py @@ -1,5 +1,5 @@ import abc -from collections.abc import Callable, Sequence +from collections.abc import AsyncGenerator, Callable, Sequence from typing import Annotated, Any, Literal import pydantic @@ -183,6 +183,78 @@ class HookResolution(BaseEvent): ] +async def replay_message_events( + msg: messages.Message, +) -> AsyncGenerator[Event]: + """Synthesize the events ``ai.models.stream`` would have emitted for ``msg``. + + Use when you have a complete ``Message`` from a non-streaming source — + e.g., the result of a Temporal activity, a cached LLM response, or an + offline test fixture — and want to feed it through code that consumes + an async event stream (``ai.Stream``, ``ai.ToolRunner``, custom loops + that mirror the default loop's shape, etc.):: + + async with ai.Stream(ai.events.replay_message_events(msg)) as stream: + async with ai.ToolRunner() as tr: + async for event in ai.util.merge(stream, tr.events()): + ... + + Each part is emitted as the start/delta/end triple a streaming adapter + would have produced, in part order, bracketed by ``StreamStart`` and + ``StreamEnd``. The full body of text/reasoning/tool-args is sent as a + single delta — the granularity of the model's original chunking is + not recoverable from a complete message. + + Parts with no model-layer event analog — ``ToolResultPart``, + ``HookPart``, ``StructuredOutputPart`` — are skipped silently; they + are agent-layer concerns and never appear on the model stream. + """ + yield StreamStart() + for part in msg.parts: + if isinstance(part, messages.TextPart): + yield TextStart(block_id=part.id) + if part.text: + yield TextDelta(block_id=part.id, chunk=part.text) + yield TextEnd(block_id=part.id) + elif isinstance(part, messages.ReasoningPart): + yield ReasoningStart(block_id=part.id) + if part.text: + yield ReasoningDelta(block_id=part.id, chunk=part.text) + yield ReasoningEnd(block_id=part.id) + elif isinstance(part, messages.ToolCallPart): + yield ToolStart( + tool_call_id=part.tool_call_id, + tool_name=part.tool_name, + ) + if part.tool_args: + yield ToolDelta( + tool_call_id=part.tool_call_id, + chunk=part.tool_args, + ) + yield ToolEnd(tool_call_id=part.tool_call_id, tool_call=part) + elif isinstance(part, messages.BuiltinToolCallPart): + yield BuiltinToolStart( + tool_call_id=part.tool_call_id, + tool_name=part.tool_name, + ) + if part.tool_args: + yield BuiltinToolDelta( + tool_call_id=part.tool_call_id, + chunk=part.tool_args, + ) + yield BuiltinToolEnd(tool_call_id=part.tool_call_id, tool_call=part) + elif isinstance(part, messages.BuiltinToolReturnPart): + yield BuiltinToolResult(tool_call_id=part.tool_call_id, result=part) + elif isinstance(part, messages.FilePart): + yield FileEvent( + block_id=part.id, + data=part.data, + media_type=part.media_type, + filename=part.filename, + ) + yield StreamEnd() + + # --------------------------------------------------------------------------- # Agent-layer event types #