Introduce summary provider, don't mutate user provided agents

Author: tconley1428

temporalio/contrib/openai_agents/_model_parameters.py

@@ -1,5 +1,6 @@
 """Parameters for configuring Temporal activity execution for model calls."""
 
+from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from datetime import timedelta
 from typing import Any, Callable, Optional, Union
@@ -10,6 +11,22 @@
 from temporalio.workflow import ActivityCancellationType, VersioningIntent
 
 
+class ModelSummaryProvider(ABC):
+    """Abstract base class for providing model summaries. Essentially just a callable,
+    but the arguments are sufficiently complex to benefit from names.
+    """
+
+    @abstractmethod
+    def provide(
+        self,
+        agent: Optional[Agent[Any]],
+        instructions: Optional[str],
+        input: Union[str, list[TResponseInputItem]],
+    ) -> str:
+        """Given the provided information, produce a summary for the model invocation activity."""
+        pass
+
+
 @dataclass
 class ModelActivityParameters:
     """Parameters for configuring Temporal activity execution for model calls.
@@ -46,14 +63,7 @@ class ModelActivityParameters:
     summary_override: Optional[
         Union[
             str,
-            Callable[
-                [
-                    Optional[Agent[Any]],
-                    Optional[str],
-                    Union[str, list[TResponseInputItem]],
-                ],
-                str,
-            ],
+            ModelSummaryProvider,
         ]
     ] = None
     """Summary for the activity execution."""

temporalio/contrib/openai_agents/_openai_runner.py

@@ -92,37 +92,50 @@ async def run(
             )
 
         # Recursively replace models in all agents
-        def convert_agent(agent: Agent[Any]) -> None:
-            # Short circuit if this model was already replaced to prevent looping from circular handoffs
-            if isinstance(agent.model, _TemporalModelStub):
-                return
+        def convert_agent(
+            agent: Agent[Any], seen: Optional[set[int]] = None
+        ) -> Agent[Any]:
+            if seen is None:
+                seen = set()
+
+            # Short circuit if this model was already seen to prevent looping from circular handoffs
+            if id(agent) in seen:
+                return agent
+            seen.add(id(agent))
 
             name = _model_name(agent)
-            agent.model = _TemporalModelStub(
-                model_name=name,
-                model_params=self.model_params,
-                agent=agent,
-            )
 
+            new_handoffs: list[Union[Agent, Handoff]] = []
             for handoff in agent.handoffs:
                 if isinstance(handoff, Agent):
-                    convert_agent(handoff)
+                    new_handoffs.append(convert_agent(handoff))
                 elif isinstance(handoff, Handoff):
                     original_invoke = handoff.on_invoke_handoff
 
                     async def on_invoke(
                         context: RunContextWrapper[Any], args: str
-                    ) -> Agent[Any]:
+                    ) -> Agent:
                         handoff_agent = await original_invoke(context, args)
-                        convert_agent(handoff_agent)
-                        return handoff_agent
+                        return convert_agent(handoff_agent, seen)
 
-                    handoff.on_invoke_handoff = on_invoke
+                    new_handoffs.append(
+                        dataclasses.replace(handoff, on_invoke_handoff=on_invoke)
+                    )
+                else:
+                    raise ValueError(f"Unknown handoff type: {type(handoff)}")
 
-        convert_agent(starting_agent)
+            return dataclasses.replace(
+                agent,
+                model=_TemporalModelStub(
+                    model_name=name,
+                    model_params=self.model_params,
+                    agent=agent,
+                ),
+                handoffs=new_handoffs,
+            )
 
         return await self._runner.run(
-            starting_agent=starting_agent,
+            starting_agent=convert_agent(starting_agent),
             input=input,
             context=context,
             max_turns=max_turns,

temporalio/contrib/openai_agents/_temporal_model_stub.py

@@ -142,7 +142,7 @@ def make_tool_info(tool: Tool) -> ToolInput:
                 self.model_params.summary_override
                 if isinstance(self.model_params.summary_override, str)
                 else (
-                    self.model_params.summary_override(
+                    self.model_params.summary_override.provide(
                         self.agent, system_instructions, input
                     )
                 )

tests/contrib/openai_agents/test_openai.py

@@ -86,6 +86,7 @@
     TestModel,
     TestModelProvider,
 )
+from temporalio.contrib.openai_agents._model_parameters import ModelSummaryProvider
 from temporalio.contrib.openai_agents._temporal_model_stub import _extract_summary
 from temporalio.contrib.pydantic import pydantic_data_converter
 from temporalio.exceptions import ApplicationError, CancelledError
@@ -2144,3 +2145,42 @@ async def test_run_config_models(client: Client):
 
         # Only the model from the runconfig override is used
         assert provider.model_names == {"gpt-4o"}
+
+
+async def test_summary_provider(client: Client):
+    class SummaryProvider(ModelSummaryProvider):
+        def provide(
+            self,
+            agent: Optional[Agent[Any]],
+            instructions: Optional[str],
+            input: Union[str, list[TResponseInputItem]],
+        ) -> str:
+            return "My summary"
+
+    new_config = client.config()
+    new_config["plugins"] = [
+        openai_agents.OpenAIAgentsPlugin(
+            model_params=ModelActivityParameters(
+                start_to_close_timeout=timedelta(seconds=120),
+                summary_override=SummaryProvider(),
+            ),
+            model_provider=TestModelProvider(TestHelloModel()),
+        )
+    ]
+    client = Client(**new_config)
+
+    async with new_worker(
+        client,
+        HelloWorldAgent,
+    ) as worker:
+        workflow_handle = await client.start_workflow(
+            HelloWorldAgent.run,
+            "Prompt",
+            id=f"summary-provider-model-{uuid.uuid4()}",
+            task_queue=worker.task_queue,
+            execution_timeout=timedelta(seconds=10),
+        )
+        result = await workflow_handle.result()
+        async for e in workflow_handle.fetch_history_events():
+            if e.HasField("activity_task_scheduled_event_attributes"):
+                assert e.user_metadata.summary.data == b'"My summary"'