diff --git a/examples/temporal-direct/main.py b/examples/temporal-direct/main.py index e9c2ab7b..01d31d41 100644 --- a/examples/temporal-direct/main.py +++ b/examples/temporal-direct/main.py @@ -153,11 +153,10 @@ async def loop( # 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. + # the default loop. ``Stream.replay_message`` replays events + # from llm_msg. async with ( - ai.Stream(ai.events.replay_message_events(llm_msg)) as stream, + ai.Stream.replay_message(llm_msg) as stream, ai.ToolRunner() as tr, ): async for event in ai.util.merge(stream, tr.events()): diff --git a/src/ai/models/core/api.py b/src/ai/models/core/api.py index 42f5c202..3e199cfb 100644 --- a/src/ai/models/core/api.py +++ b/src/ai/models/core/api.py @@ -135,6 +135,53 @@ def __init__( # rather than look like a normal end of turn. self._ended = False + @classmethod + def replay_message( + cls, + message: types.messages.Message, + *, + output_type: type[StreamOutputT] | None = None, + ) -> Stream[StreamOutputT]: + """Synthesize stream events 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.replay_message(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. + + Each part's ``provider_metadata`` (and the message's) rides on its end + event, mirroring the real adapters, so a rebuilt turn keeps it -- + reasoning signatures included, which must survive to replay the turn + back to the provider. + + Parts with no model-layer event analog — ``ToolResultPart``, + ``HookPart`` — are skipped silently; they are agent-layer concerns + and never appear on the model stream. + + ``stream.message`` keeps ``message``'s id; the parts are rebuilt + from the stream. + """ + seed = types.messages.Message( + id=message.id, role=message.role, parts=[] + ) + return cls( + types.events._replay_message_events(message), + seed_message=seed, + output_type=output_type, + ) + async def aclose(self) -> None: await self._gen.aclose() diff --git a/src/ai/types/events.py b/src/ai/types/events.py index 15ee16b6..ea6d1be0 100644 --- a/src/ai/types/events.py +++ b/src/ai/types/events.py @@ -183,47 +183,24 @@ class HookResolution(BaseEvent): ] -async def replay_message_events( +async def _replay_message_events( msg: messages.Message, ) -> AsyncGenerator[Event]: - """Synthesize stream events 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`` — are skipped silently; they are agent-layer concerns - and never appear on the model stream. - """ + """Synthesize stream events for ``msg``.""" + # See Stream.replay_message 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) + yield TextEnd( + block_id=part.id, provider_metadata=part.provider_metadata + ) elif isinstance(part, messages.ReasoningPart): yield ReasoningStart(block_id=part.id) if part.text: yield ReasoningDelta(block_id=part.id, chunk=part.text) - # Carry the signature (and any other reasoning metadata) on the - # end event, mirroring how the real adapters emit it -- otherwise - # a replayed-then-rebuilt turn loses its signature and can't be - # replayed to the provider. yield ReasoningEnd( block_id=part.id, provider_metadata=part.provider_metadata, @@ -238,7 +215,11 @@ async def replay_message_events( tool_call_id=part.tool_call_id, chunk=part.tool_args, ) - yield ToolEnd(tool_call_id=part.tool_call_id, tool_call=part) + yield ToolEnd( + tool_call_id=part.tool_call_id, + tool_call=part, + provider_metadata=part.provider_metadata, + ) elif isinstance(part, messages.BuiltinToolCallPart): yield BuiltinToolStart( tool_call_id=part.tool_call_id, @@ -249,7 +230,11 @@ async def replay_message_events( tool_call_id=part.tool_call_id, chunk=part.tool_args, ) - yield BuiltinToolEnd(tool_call_id=part.tool_call_id, tool_call=part) + yield BuiltinToolEnd( + tool_call_id=part.tool_call_id, + tool_call=part, + provider_metadata=part.provider_metadata, + ) elif isinstance(part, messages.BuiltinToolReturnPart): yield BuiltinToolResult(tool_call_id=part.tool_call_id, result=part) elif isinstance(part, messages.FilePart): @@ -258,8 +243,9 @@ async def replay_message_events( data=part.data, media_type=part.media_type, filename=part.filename, + provider_metadata=part.provider_metadata, ) - yield StreamEnd() + yield StreamEnd(provider_metadata=msg.provider_metadata) # --------------------------------------------------------------------------- diff --git a/tests/models/core/test_api.py b/tests/models/core/test_api.py index 0335d6db..fad22c62 100644 --- a/tests/models/core/test_api.py +++ b/tests/models/core/test_api.py @@ -464,6 +464,43 @@ async def _spy_stream( assert not any(isinstance(e, events_.TextDelta) for e in events) +async def test_replay_message_preserves_id() -> None: + """``Stream.replay_message`` keeps the message id and rebuilds parts.""" + msg = messages_.Message( + role="assistant", + parts=[ + messages_.TextPart(id="t1", text="calling tools"), + messages_.ToolCallPart( + tool_call_id="tc-1", + tool_name="weather", + tool_args='{"city":"SF"}', + ), + ], + ) + + async with ai.Stream.replay_message(msg) as stream: + events: list[events_.Event] = [event async for event in stream] + + # The rebuilt message keeps the original id (a fresh Message, not the + # original object) and the parts are reconstructed from the events -- + # no duplication, same part ids. + assert stream.message is not msg + assert stream.message.id == msg.id + assert len(stream.message.parts) == 2 + assert stream.text == "calling tools" + assert [tc.tool_call_id for tc in stream.tool_calls] == ["tc-1"] + + # The full event set is re-emitted (and is visible to consumers -- + # nothing is flagged ``replay``), enough to drive a tool runner. + assert events + assert not any(e.replay for e in events) + tool_ends = [e for e in events if isinstance(e, events_.ToolEnd)] + assert len(tool_ends) == 1 + assert tool_ends[0].tool_call.tool_call_id == "tc-1" + assert any(isinstance(e, events_.TextDelta) for e in events) + assert any(isinstance(e, events_.ToolStart) for e in events) + + def test_tool_end_replay_flag_excluded_from_json() -> None: """The replay flag is internal — it must not appear in serialized output.""" ev = events_.ToolEnd( diff --git a/tests/types/test_events.py b/tests/types/test_events.py index a54b8c9a..9133872d 100644 --- a/tests/types/test_events.py +++ b/tests/types/test_events.py @@ -23,7 +23,7 @@ async def test_reasoning_signature_survives_replay(self) -> None: ) async with models.Stream( - events.replay_message_events(original) + events._replay_message_events(original) ) as stream: async for _ in stream: pass @@ -54,10 +54,44 @@ async def test_reasoning_signature_on_end_event(self) -> None: reasoning_ends = [ e - async for e in events.replay_message_events(msg) + async for e in events._replay_message_events(msg) if isinstance(e, events.ReasoningEnd) ] assert len(reasoning_ends) == 1 assert reasoning_ends[0].provider_metadata == { "anthropic": {"signature": "sig"} } + + async def test_provider_metadata_survives_replay(self) -> None: + """provider_metadata on every part, and the message itself, round- + trips through the aggregator -- not just reasoning signatures.""" + original = messages.Message( + role="assistant", + parts=[ + messages.TextPart(text="hi", provider_metadata={"p": {"t": 1}}), + messages.ToolCallPart( + tool_call_id="tc-1", + tool_name="weather", + tool_args="{}", + provider_metadata={"p": {"tc": 2}}, + ), + ], + provider_metadata={"p": {"msg": 3}}, + ) + + async with models.Stream( + events._replay_message_events(original) + ) as stream: + async for _ in stream: + pass + + rebuilt = stream.message + assert rebuilt.provider_metadata == {"p": {"msg": 3}} + text = next( + p for p in rebuilt.parts if isinstance(p, messages.TextPart) + ) + assert text.provider_metadata == {"p": {"t": 1}} + tool = next( + p for p in rebuilt.parts if isinstance(p, messages.ToolCallPart) + ) + assert tool.provider_metadata == {"p": {"tc": 2}}