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

Author: tconley1428

pyproject.toml

@@ -26,7 +26,7 @@ opentelemetry = [
 ]
 pydantic = ["pydantic>=2.0.0,<3"]
 openai-agents = [
-    "openai-agents >= 0.1,<0.2",
+    "openai-agents >= 0.2.3,<0.3",
     "eval-type-backport>=0.2.2; python_version < '3.10'"
 ]
 
@@ -165,6 +165,7 @@ reportAny = "none"
 reportCallInDefaultInitializer = "none"
 reportExplicitAny = "none"
 reportIgnoreCommentWithoutRule = "none"
+reportImplicitAbstractClass = "none"
 reportImplicitOverride = "none"
 reportImplicitStringConcatenation = "none"
 reportImportCycles = "none"
@@ -184,11 +185,6 @@ exclude = [
   "temporalio/bridge/proto",
   "tests/worker/workflow_sandbox/testmodules/proto",
   "temporalio/bridge/worker.py",
-  "temporalio/contrib/opentelemetry.py",
-  "temporalio/contrib/pydantic.py",
-  "temporalio/converter.py",
-  "temporalio/testing/_workflow.py",
-  "temporalio/worker/_activity.py",
   "temporalio/worker/_replayer.py",
   "temporalio/worker/_worker.py",
   "temporalio/worker/workflow_sandbox/_importer.py",
@@ -203,9 +199,7 @@ exclude = [
   "tests/contrib/pydantic/workflows.py",
   "tests/test_converter.py",
   "tests/test_service.py",
-  "tests/test_workflow.py",
   "tests/worker/test_activity.py",
-  "tests/worker/test_workflow.py",
   "tests/worker/workflow_sandbox/test_importer.py",
   "tests/worker/workflow_sandbox/test_restrictions.py",
   # TODO: these pass locally but fail in CI with

temporalio/client.py

@@ -2889,7 +2889,7 @@ def _from_raw_info(
         cls,
         info: temporalio.api.workflow.v1.WorkflowExecutionInfo,
         converter: temporalio.converter.DataConverter,
-        **additional_fields,
+        **additional_fields: Any,
     ) -> WorkflowExecution:
         return cls(
             close_time=info.close_time.ToDatetime().replace(tzinfo=timezone.utc)

temporalio/contrib/openai_agents/__init__.py

@@ -8,12 +8,11 @@
     Use with caution in production environments.
 """
 
-from temporalio.contrib.openai_agents._invoke_model_activity import ModelActivity
 from temporalio.contrib.openai_agents._model_parameters import ModelActivityParameters
 from temporalio.contrib.openai_agents._temporal_openai_agents import (
+    OpenAIAgentsPlugin,
     TestModel,
     TestModelProvider,
-    set_open_ai_agent_temporal_overrides,
 )
 from temporalio.contrib.openai_agents._trace_interceptor import (
     OpenAIAgentsTracingInterceptor,
@@ -22,10 +21,9 @@
 from . import workflow
 
 __all__ = [
-    "ModelActivity",
+    "OpenAIAgentsPlugin",
     "ModelActivityParameters",
-    "set_open_ai_agent_temporal_overrides",
-    "OpenAIAgentsTracingInterceptor",
+    "workflow",
     "TestModel",
     "TestModelProvider",
 ]

temporalio/contrib/openai_agents/_heartbeat_decorator.py

@@ -10,7 +10,7 @@
 def _auto_heartbeater(fn: F) -> F:
     # Propagate type hints from the original callable.
     @wraps(fn)
-    async def wrapper(*args, **kwargs):
+    async def wrapper(*args: Any, **kwargs: Any) -> Any:
         heartbeat_timeout = activity.info().heartbeat_timeout
         heartbeat_task = None
         if heartbeat_timeout:

temporalio/contrib/openai_agents/_invoke_model_activity.py

@@ -173,7 +173,7 @@ def make_tool(tool: ToolInput) -> Tool:
                 raise UserError(f"Unknown tool type: {tool.name}")
 
         tools = [make_tool(x) for x in input.get("tools", [])]
-        handoffs = [
+        handoffs: list[Handoff[Any, Any]] = [
             Handoff(
                 tool_name=x.tool_name,
                 tool_description=x.tool_description,

temporalio/contrib/openai_agents/_openai_runner.py

@@ -1,11 +1,9 @@
 from dataclasses import replace
-from datetime import timedelta
-from typing import Optional, Union
+from typing import Any, Union
 
 from agents import (
     Agent,
     RunConfig,
-    RunHooks,
     RunResult,
     RunResultStreaming,
     TContext,
@@ -14,10 +12,8 @@
 from agents.run import DEFAULT_AGENT_RUNNER, DEFAULT_MAX_TURNS, AgentRunner
 
 from temporalio import workflow
-from temporalio.common import Priority, RetryPolicy
 from temporalio.contrib.openai_agents._model_parameters import ModelActivityParameters
 from temporalio.contrib.openai_agents._temporal_model_stub import _TemporalModelStub
-from temporalio.workflow import ActivityCancellationType, VersioningIntent
 
 
 class TemporalOpenAIRunner(AgentRunner):
@@ -36,7 +32,7 @@ async def run(
         self,
         starting_agent: Agent[TContext],
         input: Union[str, list[TResponseInputItem]],
-        **kwargs,
+        **kwargs: Any,
     ) -> RunResult:
         """Run the agent in a Temporal workflow."""
         if not workflow.in_workflow():
@@ -82,7 +78,7 @@ def run_sync(
         self,
         starting_agent: Agent[TContext],
         input: Union[str, list[TResponseInputItem]],
-        **kwargs,
+        **kwargs: Any,
     ) -> RunResult:
         """Run the agent synchronously (not supported in Temporal workflows)."""
         if not workflow.in_workflow():
@@ -97,7 +93,7 @@ def run_streamed(
         self,
         starting_agent: Agent[TContext],
         input: Union[str, list[TResponseInputItem]],
-        **kwargs,
+        **kwargs: Any,
     ) -> RunResultStreaming:
         """Run the agent with streaming responses (not supported in Temporal workflows)."""
         if not workflow.in_workflow():

temporalio/contrib/openai_agents/_temporal_openai_agents.py

@@ -21,11 +21,20 @@
 from agents.tracing.provider import DefaultTraceProvider
 from openai.types.responses import ResponsePromptParam
 
+import temporalio.client
+import temporalio.worker
+from temporalio.client import ClientConfig
+from temporalio.contrib.openai_agents._invoke_model_activity import ModelActivity
 from temporalio.contrib.openai_agents._model_parameters import ModelActivityParameters
 from temporalio.contrib.openai_agents._openai_runner import TemporalOpenAIRunner
 from temporalio.contrib.openai_agents._temporal_trace_provider import (
     TemporalTraceProvider,
 )
+from temporalio.contrib.openai_agents._trace_interceptor import (
+    OpenAIAgentsTracingInterceptor,
+)
+from temporalio.contrib.pydantic import pydantic_data_converter
+from temporalio.worker import Worker, WorkerConfig
 
 
 @contextmanager
@@ -133,3 +142,121 @@ def stream_response(
     ) -> AsyncIterator[TResponseStreamEvent]:
         """Get a streamed response from the model. Unimplemented."""
         raise NotImplementedError()
+
+
+class OpenAIAgentsPlugin(temporalio.client.Plugin, temporalio.worker.Plugin):
+    """Temporal plugin for integrating OpenAI agents with Temporal workflows.
+
+    .. warning::
+        This class is experimental and may change in future versions.
+        Use with caution in production environments.
+
+    This plugin provides seamless integration between the OpenAI Agents SDK and
+    Temporal workflows. It automatically configures the necessary interceptors,
+    activities, and data converters to enable OpenAI agents to run within
+    Temporal workflows with proper tracing and model execution.
+
+    The plugin:
+    1. Configures the Pydantic data converter for type-safe serialization
+    2. Sets up tracing interceptors for OpenAI agent interactions
+    3. Registers model execution activities
+    4. Manages the OpenAI agent runtime overrides during worker execution
+
+    Args:
+        model_params: Configuration parameters for Temporal activity execution
+            of model calls. If None, default parameters will be used.
+        model_provider: Optional model provider for custom model implementations.
+            Useful for testing or custom model integrations.
+
+    Example:
+        >>> from temporalio.client import Client
+        >>> from temporalio.worker import Worker
+        >>> from temporalio.contrib.openai_agents import OpenAIAgentsPlugin, ModelActivityParameters
+        >>> from datetime import timedelta
+        >>>
+        >>> # Configure model parameters
+        >>> model_params = ModelActivityParameters(
+        ...     start_to_close_timeout=timedelta(seconds=30),
+        ...     retry_policy=RetryPolicy(maximum_attempts=3)
+        ... )
+        >>>
+        >>> # Create plugin
+        >>> plugin = OpenAIAgentsPlugin(model_params=model_params)
+        >>>
+        >>> # Use with client and worker
+        >>> client = await Client.connect(
+        ...     "localhost:7233",
+        ...     plugins=[plugin]
+        ... )
+        >>> worker = Worker(
+        ...     client,
+        ...     task_queue="my-task-queue",
+        ...     workflows=[MyWorkflow],
+        ... )
+    """
+
+    def __init__(
+        self,
+        model_params: Optional[ModelActivityParameters] = None,
+        model_provider: Optional[ModelProvider] = None,
+    ) -> None:
+        """Initialize the OpenAI agents plugin.
+
+        Args:
+            model_params: Configuration parameters for Temporal activity execution
+                of model calls. If None, default parameters will be used.
+            model_provider: Optional model provider for custom model implementations.
+                Useful for testing or custom model integrations.
+        """
+        self._model_params = model_params
+        self._model_provider = model_provider
+
+    def configure_client(self, config: ClientConfig) -> ClientConfig:
+        """Configure the Temporal client for OpenAI agents integration.
+
+        This method sets up the Pydantic data converter to enable proper
+        serialization of OpenAI agent objects and responses.
+
+        Args:
+            config: The client configuration to modify.
+
+        Returns:
+            The modified client configuration.
+        """
+        config["data_converter"] = pydantic_data_converter
+        return super().configure_client(config)
+
+    def configure_worker(self, config: WorkerConfig) -> WorkerConfig:
+        """Configure the Temporal worker for OpenAI agents integration.
+
+        This method adds the necessary interceptors and activities for OpenAI
+        agent execution:
+        - Adds tracing interceptors for OpenAI agent interactions
+        - Registers model execution activities
+
+        Args:
+            config: The worker configuration to modify.
+
+        Returns:
+            The modified worker configuration.
+        """
+        config["interceptors"] = list(config.get("interceptors") or []) + [
+            OpenAIAgentsTracingInterceptor()
+        ]
+        config["activities"] = list(config.get("activities") or []) + [
+            ModelActivity(self._model_provider).invoke_model_activity
+        ]
+        return super().configure_worker(config)
+
+    async def run_worker(self, worker: Worker) -> None:
+        """Run the worker with OpenAI agents temporal overrides.
+
+        This method sets up the necessary runtime overrides for OpenAI agents
+        to work within the Temporal worker context, including custom runners
+        and trace providers.
+
+        Args:
+            worker: The worker instance to run.
+        """
+        with set_open_ai_agent_temporal_overrides(self._model_params):
+            await super().run_worker(worker)

temporalio/contrib/openai_agents/_temporal_trace_provider.py

@@ -1,7 +1,8 @@
 """Provides support for integration with OpenAI Agents SDK tracing across workflows"""
 
 import uuid
-from typing import Any, Optional, Union, cast
+from types import TracebackType
+from typing import Any, Optional, cast
 
 from agents import SpanData, Trace, TracingProcessor
 from agents.tracing import (
@@ -184,6 +185,11 @@ def __enter__(self):
         """Enter the context of the Temporal trace provider."""
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
+    def __exit__(
+        self,
+        exc_type: type[BaseException],
+        exc_val: BaseException,
+        exc_tb: TracebackType,
+    ):
         """Exit the context of the Temporal trace provider."""
         self._multi_processor.shutdown()

temporalio/contrib/openai_agents/workflow.py

@@ -240,4 +240,26 @@ async def run_operation(ctx: RunContextWrapper[Any], input: str) -> Any:
 
 
 class ToolSerializationError(TemporalError):
-    """Error that occurs when a tool output could not be serialized."""
+    """Error that occurs when a tool output could not be serialized.
+
+    .. warning::
+        This exception is experimental and may change in future versions.
+        Use with caution in production environments.
+
+    This exception is raised when a tool (created from an activity or Nexus operation)
+    returns a value that cannot be properly serialized for use by the OpenAI agent.
+    All tool outputs must be convertible to strings for the agent to process them.
+
+    The error typically occurs when:
+    - A tool returns a complex object that doesn't have a meaningful string representation
+    - The returned object cannot be converted using str()
+    - Custom serialization is needed but not implemented
+
+    Example:
+        >>> @activity.defn
+        >>> def problematic_tool() -> ComplexObject:
+        ...     return ComplexObject()  # This might cause ToolSerializationError
+
+    To fix this error, ensure your tool returns string-convertible values or
+    modify the tool to return a string representation of the result.
+    """

temporalio/contrib/opentelemetry.py

@@ -71,7 +71,7 @@ class should return the workflow interceptor subclass from
     custom attributes desired.
     """
 
-    def __init__(
+    def __init__(  # type: ignore[reportMissingSuperCall]
         self,
         tracer: Optional[opentelemetry.trace.Tracer] = None,
         *,
@@ -125,11 +125,10 @@ def workflow_interceptor_class(
         :py:meth:`temporalio.worker.Interceptor.workflow_interceptor_class`.
         """
         # Set the externs needed
-        # TODO(cretz): MyPy works w/ spread kwargs instead of direct passing
         input.unsafe_extern_functions.update(
-            **_WorkflowExternFunctions(
-                __temporal_opentelemetry_completed_span=self._completed_workflow_span,
-            )
+            {
+                "__temporal_opentelemetry_completed_span": self._completed_workflow_span,
+            }
         )
         return TracingWorkflowInboundInterceptor