Merge remote-tracking branch 'origin/main' into openai/mcp

Author: tconley1428

temporalio/contrib/openai_agents/_model_parameters.py

@@ -1,13 +1,32 @@
 """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 Optional
+from typing import Any, Callable, Optional, Union
+
+from agents import Agent, TResponseInputItem
 
 from temporalio.common import Priority, RetryPolicy
 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.
@@ -41,7 +60,12 @@ class ModelActivityParameters:
     versioning_intent: Optional[VersioningIntent] = None
     """Versioning intent for the activity."""
 
-    summary_override: Optional[str] = None
+    summary_override: Optional[
+        Union[
+            str,
+            ModelSummaryProvider,
+        ]
+    ] = None
     """Summary for the activity execution."""
 
     priority: Priority = Priority.default

temporalio/contrib/openai_agents/_openai_runner.py

@@ -1,12 +1,13 @@
+import dataclasses
 import json
 import typing
-import warnings
-from dataclasses import replace
-from typing import Any, Union
+from typing import Any, Optional, Union
 
 from agents import (
     Agent,
+    Handoff,
     RunConfig,
+    RunContextWrapper,
     RunResult,
     RunResultStreaming,
     SQLiteSession,
@@ -91,26 +92,70 @@ 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."
+        if run_config.model:
+            if not isinstance(run_config.model, str):
+                raise ValueError(
+                    "Temporal workflows require a model name to be a string in the run config."
+                )
+            run_config = dataclasses.replace(
+                run_config,
+                model=_TemporalModelStub(
+                    run_config.model, model_params=self.model_params, agent=None
+                ),
+            )
+
+        # Recursively replace models in all agents
+        def convert_agent(agent: Agent[Any], seen: Optional[set[int]]) -> 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))
+
+            # This agent has already been processed in some other run
+            if isinstance(agent.model, _TemporalModelStub):
+                return agent
+
+            name = _model_name(agent)
+
+            new_handoffs: list[Union[Agent, Handoff]] = []
+            for handoff in agent.handoffs:
+                if isinstance(handoff, Agent):
+                    new_handoffs.append(convert_agent(handoff, seen))
+                elif isinstance(handoff, Handoff):
+                    original_invoke = handoff.on_invoke_handoff
+
+                    async def on_invoke(
+                        context: RunContextWrapper[Any], args: str
+                    ) -> Agent:
+                        handoff_agent = await original_invoke(context, args)
+                        return convert_agent(handoff_agent, seen)
+
+                    new_handoffs.append(
+                        dataclasses.replace(handoff, on_invoke_handoff=on_invoke)
+                    )
+                else:
+                    raise ValueError(f"Unknown handoff type: {type(handoff)}")
+
+            return dataclasses.replace(
+                agent,
+                model=_TemporalModelStub(
+                    model_name=name,
+                    model_params=self.model_params,
+                    agent=agent,
+                ),
+                handoffs=new_handoffs,
             )
-        updated_run_config = replace(
-            run_config,
-            model=_TemporalModelStub(
-                model_name=model_name,
-                model_params=self.model_params,
-            ),
-        )
 
         return await self._runner.run(
-            starting_agent=starting_agent,
+            starting_agent=convert_agent(starting_agent, None),
             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,
         )
@@ -144,3 +189,12 @@ def run_streamed(
                 **kwargs,
             )
         raise RuntimeError("Temporal workflows do not support streaming.")
+
+
+def _model_name(agent: Agent[Any]) -> Optional[str]:
+    name = agent.model
+    if name is not None and not isinstance(name, str):
+        raise ValueError(
+            "Temporal workflows require a model name to be a string in the agent."
+        )
+    return name

temporalio/contrib/openai_agents/_temporal_model_stub.py

@@ -11,6 +11,7 @@
 from typing import Any, AsyncIterator, Union, cast
 
 from agents import (
+    Agent,
     AgentOutputSchema,
     AgentOutputSchemaBase,
     CodeInterpreterTool,
@@ -50,9 +51,11 @@ def __init__(
         model_name: Optional[str],
         *,
         model_params: ModelActivityParameters,
+        agent: Optional[Agent[Any]],
     ) -> None:
         self.model_name = model_name
         self.model_params = model_params
+        self.agent = agent
 
     async def get_response(
         self,
@@ -124,7 +127,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 +137,25 @@ 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.provide(
+                        self.agent, system_instructions, input
+                    )
+                )
+            )
+        elif self.agent:
+            summary = self.agent.name
+        else:
+            summary = None
+
         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,

temporalio/contrib/openai_agents/_trace_interceptor.py

@@ -5,14 +5,14 @@
 import random
 import uuid
 from contextlib import contextmanager
-from typing import Any, Mapping, Protocol, Type
+from typing import Any, Mapping, Optional, Protocol, Type
 
 from agents import CustomSpanData, custom_span, get_current_span, trace
 from agents.tracing import (
     get_trace_provider,
 )
 from agents.tracing.scope import Scope
-from agents.tracing.spans import NoOpSpan
+from agents.tracing.spans import NoOpSpan, Span
 
 import temporalio.activity
 import temporalio.api.common.v1
@@ -370,55 +370,78 @@ class _ContextPropagationWorkflowOutboundInterceptor(
     async def signal_child_workflow(
         self, input: temporalio.worker.SignalChildWorkflowInput
     ) -> None:
-        with custom_span(
-            name="temporal:signalChildWorkflow",
-            data={"workflowId": input.child_workflow_id},
-        ):
+        trace = get_trace_provider().get_current_trace()
+        if trace:
+            with custom_span(
+                name="temporal:signalChildWorkflow",
+                data={"workflowId": input.child_workflow_id},
+            ):
+                set_header_from_context(input, temporalio.workflow.payload_converter())
+                await self.next.signal_child_workflow(input)
+        else:
             set_header_from_context(input, temporalio.workflow.payload_converter())
             await self.next.signal_child_workflow(input)
 
     async def signal_external_workflow(
         self, input: temporalio.worker.SignalExternalWorkflowInput
     ) -> None:
-        with custom_span(
-            name="temporal:signalExternalWorkflow",
-            data={"workflowId": input.workflow_id},
-        ):
+        trace = get_trace_provider().get_current_trace()
+        if trace:
+            with custom_span(
+                name="temporal:signalExternalWorkflow",
+                data={"workflowId": input.workflow_id},
+            ):
+                set_header_from_context(input, temporalio.workflow.payload_converter())
+                await self.next.signal_external_workflow(input)
+        else:
             set_header_from_context(input, temporalio.workflow.payload_converter())
             await self.next.signal_external_workflow(input)
 
     def start_activity(
         self, input: temporalio.worker.StartActivityInput
     ) -> temporalio.workflow.ActivityHandle:
-        span = custom_span(
-            name="temporal:startActivity", data={"activity": input.activity}
-        )
-        span.start(mark_as_current=True)
+        trace = get_trace_provider().get_current_trace()
+        span: Optional[Span] = None
+        if trace:
+            span = custom_span(
+                name="temporal:startActivity", data={"activity": input.activity}
+            )
+            span.start(mark_as_current=True)
+
         set_header_from_context(input, temporalio.workflow.payload_converter())
         handle = self.next.start_activity(input)
-        handle.add_done_callback(lambda _: span.finish())
+        if span:
+            handle.add_done_callback(lambda _: span.finish())  # type: ignore
         return handle
 
     async def start_child_workflow(
         self, input: temporalio.worker.StartChildWorkflowInput
     ) -> temporalio.workflow.ChildWorkflowHandle:
-        span = custom_span(
-            name="temporal:startChildWorkflow", data={"workflow": input.workflow}
-        )
-        span.start(mark_as_current=True)
+        trace = get_trace_provider().get_current_trace()
+        span: Optional[Span] = None
+        if trace:
+            span = custom_span(
+                name="temporal:startChildWorkflow", data={"workflow": input.workflow}
+            )
+            span.start(mark_as_current=True)
         set_header_from_context(input, temporalio.workflow.payload_converter())
         handle = await self.next.start_child_workflow(input)
-        handle.add_done_callback(lambda _: span.finish())
+        if span:
+            handle.add_done_callback(lambda _: span.finish())  # type: ignore
         return handle
 
     def start_local_activity(
         self, input: temporalio.worker.StartLocalActivityInput
     ) -> temporalio.workflow.ActivityHandle:
-        span = custom_span(
-            name="temporal:startLocalActivity", data={"activity": input.activity}
-        )
-        span.start(mark_as_current=True)
+        trace = get_trace_provider().get_current_trace()
+        span: Optional[Span] = None
+        if trace:
+            span = custom_span(
+                name="temporal:startLocalActivity", data={"activity": input.activity}
+            )
+            span.start(mark_as_current=True)
         set_header_from_context(input, temporalio.workflow.payload_converter())
         handle = self.next.start_local_activity(input)
-        handle.add_done_callback(lambda _: span.finish())
+        if span:
+            handle.add_done_callback(lambda _: span.finish())  # type: ignore
         return handle