Support event_stream_handler

igalshilman · igalshilman · commit 01dd2dd03a22 · 2025-11-03T15:18:56.000+01:00
diff --git a/pydantic_ai_slim/pydantic_ai/durable_exec/restate/_agent.py b/pydantic_ai_slim/pydantic_ai/durable_exec/restate/_agent.py
@@ -1,8 +1,8 @@
 from __future__ import annotations
 
-from collections.abc import Iterator, Sequence
-from contextlib import contextmanager
-from typing import Any, overload
+from collections.abc import AsyncIterable, AsyncIterator, Iterator, Sequence
+from contextlib import AbstractAsyncContextManager, asynccontextmanager, contextmanager
+from typing import Any, Never, overload
 
 from restate import Context, TerminalError
 
@@ -11,12 +11,14 @@
 from pydantic_ai.agent.abstract import AbstractAgent, EventStreamHandler, RunOutputDataT
 from pydantic_ai.agent.wrapper import WrapperAgent
 from pydantic_ai.builtin_tools import AbstractBuiltinTool
-from pydantic_ai.messages import ModelMessage, UserContent
+from pydantic_ai.exceptions import UserError
+from pydantic_ai.messages import AgentStreamEvent, ModelMessage, UserContent
 from pydantic_ai.models import Model
 from pydantic_ai.output import OutputDataT, OutputSpec
+from pydantic_ai.result import StreamedRunResult
 from pydantic_ai.run import AgentRunResult
 from pydantic_ai.settings import ModelSettings
-from pydantic_ai.tools import DeferredToolResults
+from pydantic_ai.tools import DeferredToolResults, RunContext
 from pydantic_ai.toolsets.abstract import AbstractToolset
 from pydantic_ai.toolsets.function import FunctionToolset
 from pydantic_ai.usage import RunUsage, UsageLimits
@@ -87,14 +89,21 @@ def __init__(
         wrapped: AbstractAgent[AgentDepsT, OutputDataT],
         restate_context: Context,
         *,
+        event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
         disable_auto_wrapping_tools: bool = False,
     ):
         super().__init__(wrapped)
         if not isinstance(wrapped.model, Model):
             raise TerminalError(
                 'An agent needs to have a `model` in order to be used with Restate, it cannot be set at agent run time.'
             )
-        self._model = RestateModelWrapper(wrapped.model, restate_context, max_attempts=3)
+
+        self.restate_context = restate_context
+        self._event_stream_handler = event_stream_handler
+        self._disable_auto_wrapping_tools = disable_auto_wrapping_tools
+        self._model = RestateModelWrapper(
+            wrapped.model, restate_context, event_stream_handler=event_stream_handler, max_attempts=3
+        )
 
         def set_context(toolset: AbstractToolset[AgentDepsT]) -> AbstractToolset[AgentDepsT]:
             """Set the Restate context for the toolset, wrapping tools if needed."""
@@ -122,6 +131,37 @@ def _restate_overrides(self) -> Iterator[None]:
         ):
             yield
 
+    @property
+    def model(self) -> models.Model | models.KnownModelName | str | None:
+        return self._model
+
+    @property
+    def event_stream_handler(self) -> EventStreamHandler[AgentDepsT] | None:
+        handler = self._event_stream_handler or super().event_stream_handler
+        if handler is None:
+            return None
+        if self._disable_auto_wrapping_tools:
+            return handler
+        return self.wrapped_event_stream_handler
+
+    async def wrapped_event_stream_handler(
+        self, ctx: RunContext[AgentDepsT], stream: AsyncIterable[AgentStreamEvent]
+    ) -> None:
+        fn = self._event_stream_handler
+        if fn is None:
+            return
+        async for event in stream:
+
+            async def single_event():
+                yield event
+
+            await self.restate_context.run_typed('run event', lambda: fn(ctx, single_event()))
+
+    @property
+    def toolsets(self) -> Sequence[AbstractToolset[AgentDepsT]]:
+        with self._restate_overrides():
+            return super().toolsets
+
     @overload
     async def run(
         self,
@@ -232,3 +272,100 @@ async def main():
                 toolsets=toolsets,
                 event_stream_handler=event_stream_handler,
             )
+
+    @overload
+    def run_stream(
+        self,
+        user_prompt: str | Sequence[UserContent] | None = None,
+        *,
+        output_type: None = None,
+        message_history: Sequence[ModelMessage] | None = None,
+        deferred_tool_results: DeferredToolResults | None = None,
+        model: models.Model | models.KnownModelName | str | None = None,
+        deps: AgentDepsT = None,
+        model_settings: ModelSettings | None = None,
+        usage_limits: UsageLimits | None = None,
+        usage: RunUsage | None = None,
+        infer_name: bool = True,
+        toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
+        builtin_tools: Sequence[AbstractBuiltinTool] | None = None,
+        event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
+    ) -> AbstractAsyncContextManager[StreamedRunResult[AgentDepsT, OutputDataT]]: ...
+
+    @overload
+    def run_stream(
+        self,
+        user_prompt: str | Sequence[UserContent] | None = None,
+        *,
+        output_type: OutputSpec[RunOutputDataT],
+        message_history: Sequence[ModelMessage] | None = None,
+        deferred_tool_results: DeferredToolResults | None = None,
+        model: models.Model | models.KnownModelName | str | None = None,
+        deps: AgentDepsT = None,
+        model_settings: ModelSettings | None = None,
+        usage_limits: UsageLimits | None = None,
+        usage: RunUsage | None = None,
+        infer_name: bool = True,
+        toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
+        builtin_tools: Sequence[AbstractBuiltinTool] | None = None,
+        event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
+    ) -> AbstractAsyncContextManager[StreamedRunResult[AgentDepsT, RunOutputDataT]]: ...
+
+    @asynccontextmanager
+    async def run_stream(
+        self,
+        user_prompt: str | Sequence[UserContent] | None = None,
+        *,
+        output_type: OutputSpec[RunOutputDataT] | None = None,
+        message_history: Sequence[ModelMessage] | None = None,
+        deferred_tool_results: DeferredToolResults | None = None,
+        model: models.Model | models.KnownModelName | str | None = None,
+        deps: AgentDepsT = None,
+        model_settings: ModelSettings | None = None,
+        usage_limits: UsageLimits | None = None,
+        usage: RunUsage | None = None,
+        infer_name: bool = True,
+        toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
+        builtin_tools: Sequence[AbstractBuiltinTool] | None = None,
+        event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
+        **_deprecated_kwargs: Never,
+    ) -> AsyncIterator[StreamedRunResult[AgentDepsT, Any]]:
+        """Run the agent with a user prompt in async mode, returning a streamed response.
+
+        Example:
+        ```python
+        from pydantic_ai import Agent
+
+        agent = Agent('openai:gpt-4o')
+
+        async def main():
+            async with agent.run_stream('What is the capital of the UK?') as response:
+                print(await response.get_output())
+                #> The capital of the UK is London.
+        ```
+
+        Args:
+            user_prompt: User input to start/continue the conversation.
+            output_type: Custom output type to use for this run, `output_type` may only be used if the agent has no
+                output validators since output validators would expect an argument that matches the agent's output type.
+            message_history: History of the conversation so far.
+            deferred_tool_results: Optional results for deferred tool calls in the message history.
+            model: Optional model to use for this run, required if `model` was not set when creating the agent.
+            deps: Optional dependencies to use for this run.
+            model_settings: Optional settings to use for this model's request.
+            usage_limits: Optional limits on model request count or token usage.
+            usage: Optional usage to start with, useful for resuming a conversation or agents used in tools.
+            infer_name: Whether to try to infer the agent name from the call frame if it's not set.
+            toolsets: Optional additional toolsets for this run.
+            builtin_tools: Optional additional builtin tools for this run.
+            event_stream_handler: Optional event stream handler to use for this run. It will receive all the events up until the final result is found, which you can then read or stream from inside the context manager.
+
+        Returns:
+            The result of the run.
+        """
+        raise UserError(
+            '`agent.run_stream()` cannot be used inside a restate handler. '
+            'Set an `event_stream_handler` on the agent and use `agent.run()` instead.'
+        )
+
+        yield
diff --git a/pydantic_ai_slim/pydantic_ai/durable_exec/restate/_model.py b/pydantic_ai_slim/pydantic_ai/durable_exec/restate/_model.py
@@ -1,20 +1,99 @@
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from datetime import datetime
 from typing import Any
 
 from restate import Context, RunOptions
 
+from pydantic_ai.agent.abstract import EventStreamHandler
 from pydantic_ai.durable_exec.restate._serde import PydanticTypeAdapter
-from pydantic_ai.messages import ModelResponse
-from pydantic_ai.models import Model
+from pydantic_ai.exceptions import UserError
+from pydantic_ai.messages import ModelMessage, ModelResponse, ModelResponseStreamEvent
+from pydantic_ai.models import Model, ModelRequestParameters, StreamedResponse
 from pydantic_ai.models.wrapper import WrapperModel
+from pydantic_ai.settings import ModelSettings
+from pydantic_ai.tools import AgentDepsT, RunContext
+from pydantic_ai.usage import RequestUsage
 
 MODEL_RESPONSE_SERDE = PydanticTypeAdapter(ModelResponse)
 
 
+class RestateStreamedResponse(StreamedResponse):
+    def __init__(self, model_request_parameters: ModelRequestParameters, response: ModelResponse):
+        super().__init__(model_request_parameters)
+        self.response = response
+
+    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
+        return
+        # noinspection PyUnreachableCode
+        yield
+
+    def get(self) -> ModelResponse:
+        return self.response
+
+    def usage(self) -> RequestUsage:
+        return self.response.usage  # pragma: no cover
+
+    @property
+    def model_name(self) -> str:
+        return self.response.model_name or ''  # pragma: no cover
+
+    @property
+    def provider_name(self) -> str:
+        return self.response.provider_name or ''  # pragma: no cover
+
+    @property
+    def timestamp(self) -> datetime:
+        return self.response.timestamp  # pragma: no cover
+
+
 class RestateModelWrapper(WrapperModel):
-    def __init__(self, wrapped: Model, context: Context, max_attempts: int | None = None):
+    def __init__(
+        self,
+        wrapped: Model,
+        context: Context,
+        event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
+        max_attempts: int | None = None,
+    ):
         super().__init__(wrapped)
-        self.options = RunOptions(serde=MODEL_RESPONSE_SERDE, max_attempts=max_attempts)
-        self.context = context
+        self._options = RunOptions(serde=MODEL_RESPONSE_SERDE, max_attempts=max_attempts)
+        self._context = context
+        self._event_stream_handler = event_stream_handler
 
     async def request(self, *args: Any, **kwargs: Any) -> ModelResponse:
-        return await self.context.run_typed('Model call', self.wrapped.request, self.options, *args, **kwargs)
+        return await self._context.run_typed('Model call', self.wrapped.request, self._options, *args, **kwargs)
+
+    @asynccontextmanager
+    async def request_stream(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+        run_context: RunContext[AgentDepsT] | None = None,
+    ) -> AsyncIterator[StreamedResponse]:
+        if run_context is None:
+            raise UserError(
+                'A model cannot be used with `pydantic_ai.direct.model_request_stream()` as it requires a `run_context`. Set an `event_stream_handler` on the agent and use `agent.run()` instead.'
+            )
+
+        # We can never get here without an `event_stream_handler`, as `TemporalAgent.run_stream` and `TemporalAgent.iter` raise an error saying to use `TemporalAgent.run` instead,
+        # and that only calls `request_stream` if `event_stream_handler` is set.
+        fn = self._event_stream_handler
+        assert fn is not None
+
+        async def request_stream_run():
+            async with self.wrapped.request_stream(
+                messages,
+                model_settings,
+                model_request_parameters,
+                run_context,
+            ) as streamed_response:
+                await fn(run_context, streamed_response)
+
+                async for _ in streamed_response:
+                    pass
+            return streamed_response.get()
+
+        response = await self._context.run_typed('Model stream call', request_stream_run, self._options)
+
+        yield RestateStreamedResponse(model_request_parameters, response)
