Naming and imports cleanup.

Author: mfateev

custom_decorator/activity_utils.py

@@ -8,7 +8,7 @@
 F = TypeVar("F", bound=Callable[..., Awaitable[Any]])
 
 
-def auto_heartbeater(fn: F) -> F:
+def _auto_heartbeater(fn: F) -> F:
     # We want to ensure that the type hints from the original callable are
     # available via our wrapper, so we use the functools wraps decorator
     @wraps(fn)

custom_decorator/worker.py

@@ -6,13 +6,13 @@
 from temporalio.common import RetryPolicy
 from temporalio.worker import Worker
 
-from custom_decorator.activity_utils import auto_heartbeater
+from custom_decorator.activity_utils import _auto_heartbeater
 
 
 # Here we use our automatic heartbeater decorator. If this wasn't present, our
 # activity would timeout since it isn't heartbeating.
 @activity.defn
-@auto_heartbeater
+@_auto_heartbeater
 async def wait_for_cancel_activity() -> str:
     # Wait forever, catch the cancel, and return some value
     try:

…i_agents/adapters/heartbeat_decorator.py …_agents/adapters/_heartbeat_decorator.pyopenai_agents/adapters/heartbeat_decorator.py renamed to openai_agents/adapters/_heartbeat_decorator.py openai_agents/adapters/heartbeat_decorator.py renamed to openai_agents/adapters/_heartbeat_decorator.py

@@ -1,5 +1,4 @@
 import asyncio
-from datetime import datetime
 from functools import wraps
 from typing import Any, Awaitable, Callable, TypeVar, cast
 

openai_agents/adapters/model_activity.py …agents/adapters/invoke_model_activity.pyopenai_agents/adapters/model_activity.py renamed to openai_agents/adapters/invoke_model_activity.py openai_agents/adapters/model_activity.py renamed to openai_agents/adapters/invoke_model_activity.py

@@ -9,7 +9,7 @@
 
 from temporalio import activity
 
-from openai_agents.adapters.heartbeat_decorator import auto_heartbeater
+from custom_decorator.activity_utils import _auto_heartbeater
 
 
 class OpenAIActivityInput(TypedDict, total=False):
@@ -39,7 +39,7 @@ class OpenAIActivityInput(TypedDict, total=False):
 
 
 @activity.defn
-@auto_heartbeater
+@_auto_heartbeater
 async def invoke_open_ai_model(input: OpenAIActivityInput) -> Response:
     client = AsyncOpenAI()
     return await client.responses.create(**input)

openai_agents/adapters/open_ai_converter.py

@@ -4,16 +4,12 @@
 """
 from __future__ import annotations
 
-import typing_extensions
-from openai import BaseModel, NOT_GIVEN
-from openai._types import NotGiven
-from typing import Any, Optional, Type, Generic, TypeVar, Required
-
-from pydantic import TypeAdapter, RootModel
-import pydantic
 import json
+from typing import Any, Optional, Type, TypeVar
 
 import temporalio.api.common.v1
+from openai import BaseModel, NOT_GIVEN
+from pydantic import TypeAdapter, RootModel
 from temporalio.converter import (
     CompositePayloadConverter,
     DataConverter,
@@ -24,26 +20,13 @@
 
 T = TypeVar("T", bound=BaseModel)
 
-class WrapperModel(RootModel[T]):
+class _WrapperModel(RootModel[T]):
     model_config = {
         "arbitrary_types_allowed": True,
     }
 
 
-def strip_not_given(obj: Any) -> Any:
-    if isinstance(obj, dict):
-        return {
-            k: strip_not_given(v)
-            for k, v in obj.items()
-            if v != NOT_GIVEN and v is not None
-        }
-    elif isinstance(obj, list):
-        return [strip_not_given(v) for v in obj if v != NOT_GIVEN and v is not None]
-    else:
-        return obj
-
-
-class OpenAIJSONPlainPayloadConverter(EncodingPayloadConverter):
+class _OpenAIJSONPlainPayloadConverter(EncodingPayloadConverter):
     """Pydantic JSON payload converter.
 
     Supports conversion of all types supported by Pydantic to and from JSON.
@@ -70,8 +53,19 @@ def to_payload(self, value: Any) -> Optional[temporalio.api.common.v1.Payload]:
         https://docs.pydantic.dev/latest/api/pydantic_core/#pydantic_core.to_json.
         """
 
-        # json = Wrapper(value=value).model_dump_json().encode()
-        wrapper = WrapperModel[type(value)](root=value)
+        def strip_not_given(obj: Any) -> Any:
+            if isinstance(obj, dict):
+                return {
+                    k: strip_not_given(v)
+                    for k, v in obj.items()
+                    if v != NOT_GIVEN and v is not None
+                }
+            elif isinstance(obj, list):
+                return [strip_not_given(v) for v in obj if v != NOT_GIVEN and v is not None]
+            else:
+                return obj
+
+        wrapper = _WrapperModel[type(value)](root=value)
         dump = wrapper.model_dump(mode="python", by_alias=True)
         # NotGiven values are not JSON serializable, so we need to strip them out
         dump = strip_not_given(dump)
@@ -87,7 +81,7 @@ def from_payload(
             type_hint: Optional[Type] = None,
     ) -> Any:
         _type_hint = type_hint if type_hint is not None else Any
-        wrapper = WrapperModel[_type_hint]
+        wrapper = _WrapperModel[_type_hint]
         return TypeAdapter(wrapper).validate_json(payload.data.decode()).root
 
 
@@ -100,7 +94,7 @@ class OpenAIPayloadConverter(CompositePayloadConverter):
 
     def __init__(self) -> None:
         """Initialize object"""
-        json_payload_converter = OpenAIJSONPlainPayloadConverter()
+        json_payload_converter = _OpenAIJSONPlainPayloadConverter()
         super().__init__(
             *(
                 c

openai_agents/adapters/activity_model.py …ents/adapters/temporal_model_provider.pyopenai_agents/adapters/activity_model.py renamed to openai_agents/adapters/temporal_model_provider.py openai_agents/adapters/activity_model.py renamed to openai_agents/adapters/temporal_model_provider.py

@@ -2,14 +2,15 @@
 
 from temporalio import workflow
 
+from openai_agents.adapters.invoke_model_activity import OpenAIActivityInput, invoke_open_ai_model
+
 with workflow.unsafe.imports_passed_through():
     from datetime import timedelta
     from idlelib.query import Query
     from typing import Union, Optional, List, Literal, Iterable, Callable, Any
     from wsgiref.headers import Headers
     from agents.function_schema import function_schema
     from agents.models.openai_provider import DEFAULT_MODEL
-    from openai_agents.adapters.model_activity import OpenAIActivityInput, invoke_open_ai_model
     from agents import ModelProvider, Model, OpenAIResponsesModel, Tool, RunContextWrapper, FunctionTool
     import httpx
     from fastapi import Body
@@ -79,7 +80,7 @@ def get_summary(input: Any) -> str:
     return client
 
 
-class ModelStubProvider(ModelProvider):
+class TemporalModelProvider(ModelProvider):
     def get_model(self, model_name: str | None) -> Model:
         if model_name is None:
             model_name = DEFAULT_MODEL

openai_agents/run_agents_as_tools_workflow.py

@@ -1,12 +1,10 @@
 import asyncio
-import logging
 
 from temporalio.client import Client
 from temporalio.common import WorkflowIDReusePolicy
 
-from openai_agents.workflows.agents_as_tools_workflow import AgentsAsToolsWorkflow
 from openai_agents.adapters.open_ai_converter import open_ai_data_converter
-# Import the workflow from the previous code
+from openai_agents.workflows.agents_as_tools_workflow import AgentsAsToolsWorkflow
 
 
 async def main():

openai_agents/run_hello_world_workflow.py

@@ -1,11 +1,9 @@
 import asyncio
-import logging
 
 from temporalio.client import Client
 from temporalio.common import WorkflowIDReusePolicy
 
 from openai_agents.adapters.open_ai_converter import open_ai_data_converter
-# Import the workflow from the previous code
 from openai_agents.workflows.hello_world_workflow import HelloWorldAgent
 
 

openai_agents/run_research_workflow.py

@@ -1,12 +1,12 @@
 import asyncio
-import logging
 
 from temporalio.client import Client
 from temporalio.common import WorkflowIDReusePolicy
 
 from openai_agents.adapters.open_ai_converter import open_ai_data_converter
 from openai_agents.workflows.research_bot_workflow import ResearchWorkflow
 
+
 async def main():
     # Create client connected to server at the given address
     client = await Client.connect(

openai_agents/run_tools_workflow.py

@@ -1,5 +1,4 @@
 import asyncio
-import logging
 
 from temporalio.client import Client
 from temporalio.common import WorkflowIDReusePolicy