Move model override from the runconfig to the agents

Author: tconley1428

temporalio/contrib/openai_agents/_model_parameters.py

@@ -2,7 +2,9 @@
 
 from dataclasses import dataclass
 from datetime import timedelta
-from typing import Optional
+from typing import Optional, Union, Callable, Any
+
+from agents import Agent, TResponseInputItem
 
 from temporalio.common import Priority, RetryPolicy
 from temporalio.workflow import ActivityCancellationType, VersioningIntent
@@ -41,7 +43,7 @@ class ModelActivityParameters:
     versioning_intent: Optional[VersioningIntent] = None
     """Versioning intent for the activity."""
 
-    summary_override: Optional[str] = None
+    summary_override: Optional[Union[str, Callable[[Agent[Any], Optional[str], Union[str, list[TResponseInputItem]]], str]]] = None
     """Summary for the activity execution."""
 
     priority: Priority = Priority.default

temporalio/contrib/openai_agents/_openai_runner.py

@@ -11,7 +11,7 @@
     SQLiteSession,
     TContext,
     Tool,
-    TResponseInputItem,
+    TResponseInputItem, Handoff, RunContextWrapper,
 )
 from agents.run import DEFAULT_AGENT_RUNNER, DEFAULT_MAX_TURNS, AgentRunner
 from pydantic_core import to_json
@@ -77,26 +77,50 @@ async def run(
         if run_config is None:
             run_config = RunConfig()
 
-        model_name = run_config.model or starting_agent.model
-        if model_name is not None and not isinstance(model_name, str):
-            raise ValueError(
-                "Temporal workflows require a model name to be a string in the run config and/or agent."
-            )
-        updated_run_config = replace(
-            run_config,
-            model=_TemporalModelStub(
-                model_name=model_name,
+        def model_name(agent: Agent[Any]) -> str:
+            name = run_config.model or agent.model
+            if name is not None and not isinstance(name, str):
+                print("Name: ", name, " Agent: ", agent)
+                raise ValueError(
+                    "Temporal workflows require a model name to be a string in the run config and/or agent."
+                )
+            return name
+
+        def convert_agent(agent: Agent[Any]) -> None:
+            print("Model: ", agent.model)
+
+            # Short circuit if this model was already replaced to prevent looping from circular handoffs
+            if isinstance(agent.model, _TemporalModelStub):
+                return
+
+            name = model_name(agent)
+            agent.model = _TemporalModelStub(
+                model_name=name,
                 model_params=self.model_params,
-            ),
-        )
+                agent=agent,
+            )
+            print("Model after replace: ", agent.model)
+
+            for handoff in agent.handoffs:
+                if isinstance(handoff, Agent):
+                    convert_agent(handoff)
+                elif isinstance(handoff, Handoff):
+                    original_invoke = handoff.on_invoke_handoff
+                    async def on_invoke(context: RunContextWrapper[Any], args: str) -> Agent[Any]:
+                        handoff_agent = await original_invoke(context, args)
+                        convert_agent(handoff_agent)
+                        return handoff_agent
+                    handoff.on_invoke_handoff = on_invoke
+
+        convert_agent(starting_agent)
 
         return await self._runner.run(
             starting_agent=starting_agent,
             input=input,
             context=context,
             max_turns=max_turns,
             hooks=hooks,
-            run_config=updated_run_config,
+            run_config=run_config,
             previous_response_id=previous_response_id,
             session=session,
         )

temporalio/contrib/openai_agents/_temporal_model_stub.py

@@ -25,7 +25,7 @@
     ModelTracing,
     Tool,
     TResponseInputItem,
-    WebSearchTool,
+    WebSearchTool, Agent,
 )
 from agents.items import TResponseStreamEvent
 from openai.types.responses.response_prompt_param import ResponsePromptParam
@@ -50,9 +50,11 @@ def __init__(
         model_name: Optional[str],
         *,
         model_params: ModelActivityParameters,
+        agent: Agent[Any],
     ) -> None:
         self.model_name = model_name
         self.model_params = model_params
+        self.agent = agent
 
     async def get_response(
         self,
@@ -124,7 +126,7 @@ def make_tool_info(tool: Tool) -> ToolInput:
         activity_input = ActivityModelInput(
             model_name=self.model_name,
             system_instructions=system_instructions,
-            input=cast(Union[str, list[TResponseInputItem]], input),
+            input=input,
             model_settings=model_settings,
             tools=tool_infos,
             output_schema=output_schema_input,
@@ -134,10 +136,16 @@ def make_tool_info(tool: Tool) -> ToolInput:
             prompt=prompt,
         )
 
+        if self.model_params.summary_override:
+            summary = self.model_params.summary_override if isinstance(self.model_params.summary_override, str) else (
+                self.model_params.summary_override(self.agent, system_instructions, input))
+        else:
+            summary = self.agent.name
+
         return await workflow.execute_activity_method(
             ModelActivity.invoke_model_activity,
             activity_input,
-            summary=self.model_params.summary_override or _extract_summary(input),
+            summary=summary,
             task_queue=self.model_params.task_queue,
             schedule_to_close_timeout=self.model_params.schedule_to_close_timeout,
             schedule_to_start_timeout=self.model_params.schedule_to_start_timeout,

tests/contrib/openai_agents/test_openai.py

@@ -861,6 +861,7 @@ def init_agents() -> Agent[AirlineAgentContext]:
 
     seat_booking_agent = Agent[AirlineAgentContext](
         name="Seat Booking Agent",
+        model="gpt-4o-mini",
         handoff_description="A helpful agent that can update a seat on a flight.",
         instructions=f"""{RECOMMENDED_PROMPT_PREFIX}
         You are a seat booking agent. If you are speaking to a customer, you probably were transferred to from the triage agent.
@@ -915,6 +916,17 @@ class CustomerServiceModel(StaticTestModel):
     ]
 
 
+class AssertDifferentModelProvider(ModelProvider):
+    model_names = set()
+
+    def __init__(self, model: Model):
+        self._model = model
+
+    def get_model(self, model_name: Union[str, None]) -> Model:
+        self.model_names.add(model_name)
+        return self._model
+
+
 @workflow.defn
 class CustomerServiceWorkflow:
     def __init__(self, input_items: list[TResponseInputItem] = []):
@@ -987,12 +999,13 @@ async def test_customer_service_workflow(client: Client, use_local_model: bool):
     if not use_local_model and not os.environ.get("OPENAI_API_KEY"):
         pytest.skip("No openai API key")
     new_config = client.config()
+    provider = AssertDifferentModelProvider(CustomerServiceModel())
     new_config["plugins"] = [
         openai_agents.OpenAIAgentsPlugin(
             model_params=ModelActivityParameters(
                 start_to_close_timeout=timedelta(seconds=30)
             ),
-            model_provider=TestModelProvider(CustomerServiceModel())
+            model_provider= provider
             if use_local_model
             else None,
         )
@@ -1076,6 +1089,7 @@ async def test_customer_service_workflow(client: Client, use_local_model: bool):
                 .data.decode()
             )
 
+            assert len(provider.model_names) == 2
 
 class InputGuardrailModel(OpenAIResponsesModel):
     __test__ = False