Streaming with Temporal

Author: DouweM

pydantic_ai_slim/pydantic_ai/temporal/__init__.py

@@ -11,10 +11,8 @@
 
 
 class _TemporalRunContext(RunContext[AgentDepsT]):
-    _data: dict[str, Any]
-
     def __init__(self, **kwargs: Any):
-        self._data = kwargs
+        self.__dict__ = kwargs
         setattr(
             self,
             '__dataclass_fields__',
@@ -25,10 +23,12 @@ def __getattribute__(self, name: str) -> Any:
         try:
             return super().__getattribute__(name)
         except AttributeError as e:
-            data = super().__getattribute__('_data')
-            if name in data:
-                return data[name]
-            raise e  # TODO: Explain how to make a new run context attribute available
+            if name in RunContext.__dataclass_fields__:
+                raise AttributeError(
+                    f'Temporalized {RunContext.__name__!r} object has no attribute {name!r}. To make the attribute available, pass a `TemporalSettings` object to `temporalize_agent` that has a custom `serialize_run_context` function that returns a dictionary that includes the attribute.'
+                )
+            else:
+                raise e
 
     @classmethod
     def serialize_run_context(cls, ctx: RunContext[AgentDepsT]) -> dict[str, Any]:
@@ -75,7 +75,7 @@ def for_tool(self, toolset_id: str, tool_id: str) -> TemporalSettings:
     deserialize_run_context: Callable[[dict[str, Any]], RunContext] = _TemporalRunContext.deserialize_run_context
 
     @property
-    def execute_activity_kwargs(self) -> dict[str, Any]:
+    def execute_activity_options(self) -> dict[str, Any]:
         return {
             'task_queue': self.task_queue,
             'schedule_to_close_timeout': self.schedule_to_close_timeout,

pydantic_ai_slim/pydantic_ai/temporal/agent.py

@@ -30,7 +30,7 @@ def temporalize_toolset(toolset: AbstractToolset, settings: TemporalSettings | N
 
 
 def temporalize_agent(
-    agent: Agent,
+    agent: Agent[Any, Any],
     settings: TemporalSettings | None = None,
     temporalize_toolset_func: Callable[
         [AbstractToolset, TemporalSettings | None], list[Callable[..., Any]]
@@ -52,7 +52,7 @@ def temporalize_agent(
 
     activities: list[Callable[..., Any]] = []
     if isinstance(agent.model, Model):
-        activities.extend(temporalize_model(agent.model, settings))
+        activities.extend(temporalize_model(agent.model, settings, agent._event_stream_handler))  # pyright: ignore[reportPrivateUsage]
 
     def temporalize_toolset(toolset: AbstractToolset) -> None:
         activities.extend(temporalize_toolset_func(toolset, settings))

pydantic_ai_slim/pydantic_ai/temporal/function_toolset.py

@@ -57,44 +57,11 @@ async def call_tool(name: str, tool_args: dict[str, Any], ctx: RunContext, tool:
         return await workflow.execute_activity(  # pyright: ignore[reportUnknownMemberType]
             activity=call_tool_activity,
             arg=_CallToolParams(name=name, tool_args=tool_args, serialized_run_context=serialized_run_context),
-            **tool_settings.execute_activity_kwargs,
+            **tool_settings.execute_activity_options,
         )
 
     toolset.call_tool = call_tool
 
     activities = [call_tool_activity]
     setattr(toolset, '__temporal_activities', activities)
     return activities
-
-
-# class TemporalFunctionToolset(FunctionToolset[AgentDepsT]):
-#     def __init__(
-#         self,
-#         tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] = [],
-#         max_retries: int = 1,
-#         temporal_settings: TemporalSettings | None = None,
-#         serialize_run_context: Callable[[RunContext[AgentDepsT]], Any] | None = None,
-#         deserialize_run_context: Callable[[Any], RunContext[AgentDepsT]] | None = None,
-#     ):
-#         super().__init__(tools, max_retries)
-#         self.temporal_settings = temporal_settings or TemporalSettings()
-#         self.serialize_run_context = serialize_run_context or TemporalRunContext[AgentDepsT].serialize_run_context
-#         self.deserialize_run_context = deserialize_run_context or TemporalRunContext[AgentDepsT].deserialize_run_context
-
-#         @activity.defn(name='function_toolset_call_tool')
-#         async def call_tool_activity(params: FunctionCallToolParams) -> Any:
-#             ctx = self.deserialize_run_context(params.serialized_run_context)
-#             tool = (await self.get_tools(ctx))[params.name]
-#             return await FunctionToolset[AgentDepsT].call_tool(self, params.name, params.tool_args, ctx, tool)
-
-#         self.call_tool_activity = call_tool_activity
-
-#     async def call_tool(
-#         self, name: str, tool_args: dict[str, Any], ctx: RunContext[AgentDepsT], tool: ToolsetTool[AgentDepsT]
-#     ) -> Any:
-#         serialized_run_context = self.serialize_run_context(ctx)
-#         return await workflow.execute_activity(
-#             activity=self.call_tool_activity,
-#             arg=FunctionCallToolParams(name=name, tool_args=tool_args, serialized_run_context=serialized_run_context),
-#             **self.temporal_settings.__dict__,
-#         )

pydantic_ai_slim/pydantic_ai/temporal/mcp_server.py

@@ -55,7 +55,7 @@ async def call_tool_activity(params: _CallToolParams) -> ToolResult:
     async def list_tools() -> list[mcp_types.Tool]:
         return await workflow.execute_activity(  # pyright: ignore[reportUnknownVariableType,reportUnknownMemberType]
             activity=list_tools_activity,
-            **settings.execute_activity_kwargs,
+            **settings.execute_activity_options,
         )
 
     async def direct_call_tool(
@@ -66,7 +66,7 @@ async def direct_call_tool(
         return await workflow.execute_activity(  # pyright: ignore[reportUnknownMemberType]
             activity=call_tool_activity,
             arg=_CallToolParams(name=name, tool_args=args, metadata=metadata),
-            **settings.for_tool(id, name).execute_activity_kwargs,
+            **settings.for_tool(id, name).execute_activity_options,
         )
 
     server.list_tools = list_tools
@@ -75,47 +75,3 @@ async def direct_call_tool(
     activities = [list_tools_activity, call_tool_activity]
     setattr(server, '__temporal_activities', activities)
     return activities
-
-
-# class TemporalMCPServer(WrapperToolset[Any]):
-#     temporal_settings: TemporalSettings
-
-#     @property
-#     def wrapped_server(self) -> MCPServer:
-#         assert isinstance(self.wrapped, MCPServer)
-#         return self.wrapped
-
-#     def __init__(self, wrapped: MCPServer, temporal_settings: TemporalSettings | None = None):
-#         assert isinstance(self.wrapped, MCPServer)
-#         super().__init__(wrapped)
-#         self.temporal_settings = temporal_settings or TemporalSettings()
-
-#         @activity.defn(name='mcp_server_list_tools')
-#         async def list_tools_activity() -> list[mcp_types.Tool]:
-#             return await self.wrapped_server.list_tools()
-
-#         self.list_tools_activity = list_tools_activity
-
-#         @activity.defn(name='mcp_server_call_tool')
-#         async def call_tool_activity(params: MCPCallToolParams) -> ToolResult:
-#             return await self.wrapped_server.direct_call_tool(params.name, params.tool_args, params.metadata)
-
-#         self.call_tool_activity = call_tool_activity
-
-#     async def list_tools(self) -> list[mcp_types.Tool]:
-#         return await workflow.execute_activity(
-#             activity=self.list_tools_activity,
-#             **self.temporal_settings.__dict__,
-#         )
-
-#     async def direct_call_tool(
-#         self,
-#         name: str,
-#         args: dict[str, Any],
-#         metadata: dict[str, Any] | None = None,
-#     ) -> ToolResult:
-#         return await workflow.execute_activity(
-#             activity=self.call_tool_activity,
-#             arg=MCPCallToolParams(name=name, tool_args=args, metadata=metadata),
-#             **self.temporal_settings.__dict__,
-#         )