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

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,25 +8,22 @@
     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._trace_interceptor import (
-    OpenAIAgentsTracingInterceptor,
-)
-from temporalio.contrib.openai_agents.temporal_openai_agents import (
+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,
 )
 
 from . import workflow
 
 __all__ = [
-    "ModelActivity",
+    "OpenAIAgentsPlugin",
     "ModelActivityParameters",
     "workflow",
-    "set_open_ai_agent_temporal_overrides",
-    "OpenAIAgentsTracingInterceptor",
     "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

@@ -6,6 +6,7 @@
 import enum
 import json
 from dataclasses import dataclass
+from datetime import timedelta
 from typing import Any, Optional, Union, cast
 
 from agents import (
@@ -17,17 +18,25 @@
     ModelResponse,
     ModelSettings,
     ModelTracing,
+    OpenAIProvider,
     RunContextWrapper,
     Tool,
     TResponseInputItem,
     UserError,
     WebSearchTool,
 )
 from agents.models.multi_provider import MultiProvider
+from openai import (
+    APIStatusError,
+    AsyncOpenAI,
+    AuthenticationError,
+    PermissionDeniedError,
+)
 from typing_extensions import Required, TypedDict
 
 from temporalio import activity
 from temporalio.contrib.openai_agents._heartbeat_decorator import _auto_heartbeater
+from temporalio.exceptions import ApplicationError
 
 
 @dataclass
@@ -117,11 +126,15 @@ class ActivityModelInput(TypedDict, total=False):
 
 
 class ModelActivity:
-    """Class wrapper for model invocation activities to allow model customization."""
+    """Class wrapper for model invocation activities to allow model customization. By default, we use an OpenAIProvider with retries disabled.
+    Disabling retries in your model of choice is recommended to allow activity retries to define the retry model.
+    """
 
     def __init__(self, model_provider: Optional[ModelProvider] = None):
         """Initialize the activity with a model provider."""
-        self._model_provider = model_provider or MultiProvider()
+        self._model_provider = model_provider or OpenAIProvider(
+            openai_client=AsyncOpenAI(max_retries=0)
+        )
 
     @activity.defn
     @_auto_heartbeater
@@ -160,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,
@@ -171,14 +184,51 @@ def make_tool(tool: ToolInput) -> Tool:
             )
             for x in input.get("handoffs", [])
         ]
-        return await model.get_response(
-            system_instructions=input.get("system_instructions"),
-            input=input_input,
-            model_settings=input["model_settings"],
-            tools=tools,
-            output_schema=input.get("output_schema"),
-            handoffs=handoffs,
-            tracing=ModelTracing(input["tracing"]),
-            previous_response_id=input.get("previous_response_id"),
-            prompt=input.get("prompt"),
-        )
+
+        try:
+            return await model.get_response(
+                system_instructions=input.get("system_instructions"),
+                input=input_input,
+                model_settings=input["model_settings"],
+                tools=tools,
+                output_schema=input.get("output_schema"),
+                handoffs=handoffs,
+                tracing=ModelTracing(input["tracing"]),
+                previous_response_id=input.get("previous_response_id"),
+                prompt=input.get("prompt"),
+            )
+        except APIStatusError as e:
+            # Listen to server hints
+            retry_after = None
+            retry_after_ms_header = e.response.headers.get("retry-after-ms")
+            if retry_after_ms_header is not None:
+                retry_after = timedelta(milliseconds=float(retry_after_ms_header))
+
+            if retry_after is None:
+                retry_after_header = e.response.headers.get("retry-after")
+                if retry_after_header is not None:
+                    retry_after = timedelta(seconds=float(retry_after_header))
+
+            should_retry_header = e.response.headers.get("x-should-retry")
+            if should_retry_header == "true":
+                raise e
+            if should_retry_header == "false":
+                raise ApplicationError(
+                    "Non retryable OpenAI error",
+                    non_retryable=True,
+                    next_retry_delay=retry_after,
+                ) from e
+
+            # Specifically retryable status codes
+            if e.response.status_code in [408, 409, 429, 500]:
+                raise ApplicationError(
+                    "Retryable OpenAI status code",
+                    non_retryable=False,
+                    next_retry_delay=retry_after,
+                ) from e
+
+            raise ApplicationError(
+                "Non retryable OpenAI status code",
+                non_retryable=True,
+                next_retry_delay=retry_after,
+            ) from e

temporalio/contrib/openai_agents/_model_parameters.py

@@ -20,13 +20,13 @@ class ModelActivityParameters:
     task_queue: Optional[str] = None
     """Specific task queue to use for model activities."""
 
-    schedule_to_close_timeout: Optional[timedelta] = timedelta(seconds=60)
+    schedule_to_close_timeout: Optional[timedelta] = None
     """Maximum time from scheduling to completion."""
 
     schedule_to_start_timeout: Optional[timedelta] = None
     """Maximum time from scheduling to starting."""
 
-    start_to_close_timeout: Optional[timedelta] = None
+    start_to_close_timeout: Optional[timedelta] = timedelta(seconds=60)
     """Maximum time for the activity to complete."""
 
     heartbeat_timeout: Optional[timedelta] = None

temporalio/contrib/openai_agents/_openai_runner.py

@@ -1,23 +1,21 @@
+import typing
 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,
+    Tool,
     TResponseInputItem,
 )
 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 +34,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():
@@ -46,6 +44,13 @@ async def run(
                 **kwargs,
             )
 
+        tool_types = typing.get_args(Tool)
+        for t in starting_agent.tools:
+            if not isinstance(t, tool_types):
+                raise ValueError(
+                    "Provided tool is not a tool type. If using an activity, make sure to wrap it with openai_agents.workflow.activity_as_tool."
+                )
+
         context = kwargs.get("context")
         max_turns = kwargs.get("max_turns", DEFAULT_MAX_TURNS)
         hooks = kwargs.get("hooks")
@@ -67,22 +72,21 @@ async def run(
             ),
         )
 
-        with workflow.unsafe.imports_passed_through():
-            return await self._runner.run(
-                starting_agent=starting_agent,
-                input=input,
-                context=context,
-                max_turns=max_turns,
-                hooks=hooks,
-                run_config=updated_run_config,
-                previous_response_id=previous_response_id,
-            )
+        return await self._runner.run(
+            starting_agent=starting_agent,
+            input=input,
+            context=context,
+            max_turns=max_turns,
+            hooks=hooks,
+            run_config=updated_run_config,
+            previous_response_id=previous_response_id,
+        )
 
     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 +101,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():