agentex-sdk 0.4.19__py3-none-any.whl → 0.4.21__py3-none-any.whl

agentex/lib/core/adapters/llm/adapter_sgp.py

@@ -1,11 +1,12 @@
 import os
-from collections.abc import AsyncGenerator, Generator
+from typing import override
+from collections.abc import Generator, AsyncGenerator
 
-from scale_gp import AsyncSGPClient, SGPClient
+from scale_gp import SGPClient, AsyncSGPClient
 
-from agentex.lib.core.adapters.llm.port import LLMGateway
-from agentex.lib.types.llm_messages import Completion
 from agentex.lib.utils.logging import make_logger
+from agentex.lib.types.llm_messages import Completion
+from agentex.lib.core.adapters.llm.port import LLMGateway
 
 logger = make_logger(__name__)
 
@@ -17,6 +18,7 @@ class SGPLLMGateway(LLMGateway):
             api_key=os.environ.get("SGP_API_KEY", sgp_api_key)
         )
 
+    @override
     def completion(self, *args, **kwargs) -> Completion:
         if kwargs.get("stream", True):
             raise ValueError(
@@ -26,6 +28,7 @@ class SGPLLMGateway(LLMGateway):
         response = self.sync_client.beta.chat.completions.create(*args, **kwargs)
         return Completion.model_validate(response)
 
+    @override
     def completion_stream(self, *args, **kwargs) -> Generator[Completion, None, None]:
         if not kwargs.get("stream"):
             raise ValueError("To use streaming, please set stream=True in the kwargs")
@@ -33,6 +36,7 @@ class SGPLLMGateway(LLMGateway):
         for chunk in self.sync_client.beta.chat.completions.create(*args, **kwargs):
             yield Completion.model_validate(chunk)
 
+    @override
     async def acompletion(self, *args, **kwargs) -> Completion:
         if kwargs.get("stream", True):
             raise ValueError(
@@ -43,13 +47,12 @@ class SGPLLMGateway(LLMGateway):
         response = await self.async_client.beta.chat.completions.create(*args, **kwargs)
         return Completion.model_validate(response)
 
+    @override
     async def acompletion_stream(
         self, *args, **kwargs
     ) -> AsyncGenerator[Completion, None]:
         if not kwargs.get("stream"):
             raise ValueError("To use streaming, please set stream=True in the kwargs")
 
-        async for chunk in await self.async_client.beta.chat.completions.create(
-            *args, **kwargs
-        ):
+        async for chunk in self.async_client.beta.chat.completions.create(*args, **kwargs):  # type: ignore[misc]
             yield Completion.model_validate(chunk)

agentex/lib/core/adapters/llm/port.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from collections.abc import AsyncGenerator, Generator
+from collections.abc import Generator, AsyncGenerator
 
 from agentex.lib.types.llm_messages import Completion
 

agentex/lib/core/adapters/streams/adapter_redis.py

@@ -1,14 +1,14 @@
-import asyncio
-import json
 import os
+import json
+import asyncio
+from typing import Any, Annotated, override
 from collections.abc import AsyncIterator
-from typing import Annotated, Any
 
 import redis.asyncio as redis
 from fastapi import Depends
 
-from agentex.lib.core.adapters.streams.port import StreamRepository
 from agentex.lib.utils.logging import make_logger
+from agentex.lib.core.adapters.streams.port import StreamRepository
 
 logger = make_logger(__name__)
 
@@ -26,6 +26,7 @@ class RedisStreamRepository(StreamRepository):
         )
         self.redis = redis.from_url(self.redis_url)
 
+    @override
     async def send_event(self, topic: str, event: dict[str, Any]) -> str:
         """
         Send an event to a Redis stream.
@@ -55,6 +56,7 @@ class RedisStreamRepository(StreamRepository):
             logger.error(f"Error publishing to Redis stream {topic}: {e}")
             raise
 
+    @override
     async def subscribe(
         self, topic: str, last_id: str = "$"
     ) -> AsyncIterator[dict[str, Any]]:
@@ -108,6 +110,7 @@ class RedisStreamRepository(StreamRepository):
                 logger.error(f"Error reading from Redis stream: {e}")
                 await asyncio.sleep(1)  # Back off on errors
 
+    @override
     async def cleanup_stream(self, topic: str) -> None:
         """
         Clean up a Redis stream.

agentex/lib/core/adapters/streams/port.py

@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
-from collections.abc import AsyncIterator
 from typing import Any
+from collections.abc import AsyncIterator
 
 
 class StreamRepository(ABC):

agentex/lib/core/clients/temporal/temporal_client.py

@@ -1,25 +1,29 @@
-from collections.abc import Callable
-from datetime import timedelta
 from typing import Any
+from datetime import timedelta
+from collections.abc import Callable
 
 from temporalio.client import Client, WorkflowExecutionStatus
-from temporalio.common import RetryPolicy as TemporalRetryPolicy
-from temporalio.common import WorkflowIDReusePolicy
+from temporalio.common import RetryPolicy as TemporalRetryPolicy, WorkflowIDReusePolicy
 from temporalio.service import RPCError, RPCStatusCode
 
+from agentex.lib.utils.logging import make_logger
+from agentex.lib.utils.model_utils import BaseModel
 from agentex.lib.core.clients.temporal.types import (
-    DuplicateWorkflowPolicy,
-    RetryPolicy,
     TaskStatus,
+    RetryPolicy,
     WorkflowState,
+    DuplicateWorkflowPolicy,
 )
 from agentex.lib.core.clients.temporal.utils import get_temporal_client
-from agentex.lib.utils.logging import make_logger
-from agentex.lib.utils.model_utils import BaseModel
 
 logger = make_logger(__name__)
 
-DEFAULT_RETRY_POLICY = RetryPolicy(maximum_attempts=1)
+DEFAULT_RETRY_POLICY = RetryPolicy(
+    maximum_attempts=1,
+    initial_interval=timedelta(seconds=1),
+    backoff_coefficient=2.0,
+    maximum_interval=timedelta(minutes=10),
+)
 
 
 TEMPORAL_STATUS_TO_UPLOAD_STATUS_AND_REASON = {
@@ -36,8 +40,7 @@ TEMPORAL_STATUS_TO_UPLOAD_STATUS_AND_REASON = {
     ),
     WorkflowExecutionStatus.FAILED: WorkflowState(
         status=TaskStatus.FAILED,
-        reason="Task encountered terminal failure. "
-        "Please contact support if retrying does not resolve the issue.",
+        reason="Task encountered terminal failure. Please contact support if retrying does not resolve the issue.",
         is_terminal=True,
     ),
     WorkflowExecutionStatus.RUNNING: WorkflowState(
@@ -72,9 +75,16 @@ DUPLICATE_POLICY_TO_ID_REUSE_POLICY = {
 
 class TemporalClient:
     def __init__(self, temporal_client: Client | None = None, plugins: list[Any] = []):
-        self._client: Client = temporal_client
+        self._client: Client | None = temporal_client
         self._plugins = plugins
 
+    @property
+    def client(self) -> Client:
+        """Get the temporal client, raising an error if not initialized."""
+        if self._client is None:
+            raise RuntimeError("Temporal client not initialized - ensure temporal_address is properly configured")
+        return self._client
+
     @classmethod
     async def create(cls, temporal_address: str, plugins: list[Any] = []):
         if temporal_address in [
@@ -89,18 +99,13 @@ class TemporalClient:
         ]:
             _client = None
         else:
-            _client = await get_temporal_client(
-                temporal_address, 
-                plugins=plugins
-            )
+            _client = await get_temporal_client(temporal_address, plugins=plugins)
         return cls(_client, plugins)
 
     async def setup(self, temporal_address: str):
-        self._client = await self._get_temporal_client(
-            temporal_address=temporal_address
-        )
+        self._client = await self._get_temporal_client(temporal_address=temporal_address)
 
-    async def _get_temporal_client(self, temporal_address: str) -> Client:
+    async def _get_temporal_client(self, temporal_address: str) -> Client | None:
         if temporal_address in [
             "false",
             "False",
@@ -113,10 +118,7 @@ class TemporalClient:
         ]:
             return None
         else:
-            return await get_temporal_client(
-                temporal_address, 
-                plugins=self._plugins
-            )
+            return await get_temporal_client(temporal_address, plugins=self._plugins)
 
     async def start_workflow(
         self,
@@ -127,10 +129,8 @@ class TemporalClient:
         execution_timeout: timedelta = timedelta(seconds=86400),
         **kwargs: Any,
     ) -> str:
-        temporal_retry_policy = TemporalRetryPolicy(
-            **retry_policy.model_dump(exclude_unset=True)
-        )
-        workflow_handle = await self._client.start_workflow(
+        temporal_retry_policy = TemporalRetryPolicy(**retry_policy.model_dump(exclude_unset=True))
+        workflow_handle = await self.client.start_workflow(
             *args,
             retry_policy=temporal_retry_policy,
             task_timeout=task_timeout,
@@ -146,8 +146,8 @@ class TemporalClient:
         signal: str | Callable[[dict[str, Any] | list[Any] | str | int | float | bool | BaseModel], Any],
         payload: dict[str, Any] | list[Any] | str | int | float | bool | BaseModel,
     ) -> None:
-        handle = self._client.get_workflow_handle(workflow_id=workflow_id)
-        await handle.signal(signal, payload)
+        handle = self.client.get_workflow_handle(workflow_id=workflow_id)
+        await handle.signal(signal, payload)  # type: ignore[misc]
 
     async def query_workflow(
         self,
@@ -164,12 +164,12 @@ class TemporalClient:
         Returns:
             The result of the query
         """
-        handle = self._client.get_workflow_handle(workflow_id=workflow_id)
+        handle = self.client.get_workflow_handle(workflow_id=workflow_id)
         return await handle.query(query)
 
     async def get_workflow_status(self, workflow_id: str) -> WorkflowState:
         try:
-            handle = self._client.get_workflow_handle(workflow_id=workflow_id)
+            handle = self.client.get_workflow_handle(workflow_id=workflow_id)
             description = await handle.describe()
             return TEMPORAL_STATUS_TO_UPLOAD_STATUS_AND_REASON[description.status]
         except RPCError as e:
@@ -182,7 +182,7 @@ class TemporalClient:
             raise
 
     async def terminate_workflow(self, workflow_id: str) -> None:
-        return await self._client.get_workflow_handle(workflow_id).terminate()
+        return await self.client.get_workflow_handle(workflow_id).terminate()
 
     async def cancel_workflow(self, workflow_id: str) -> None:
-        return await self._client.get_workflow_handle(workflow_id).cancel()
+        return await self.client.get_workflow_handle(workflow_id).cancel()

agentex/lib/core/clients/temporal/types.py

@@ -1,5 +1,5 @@
-from datetime import timedelta
 from enum import Enum
+from datetime import timedelta
 
 from pydantic import Field
 

agentex/lib/core/clients/temporal/utils.py

@@ -1,7 +1,9 @@
 from typing import Any
+
 from temporalio.client import Client, Plugin as ClientPlugin
+from temporalio.runtime import Runtime, TelemetryConfig, OpenTelemetryConfig
 from temporalio.contrib.pydantic import pydantic_data_converter
-from temporalio.runtime import OpenTelemetryConfig, Runtime, TelemetryConfig
+from temporalio.contrib.openai_agents import OpenAIAgentsPlugin
 
 # class DateTimeJSONEncoder(AdvancedJSONEncoder):
 #     def default(self, o: Any) -> Any:
@@ -42,10 +44,10 @@ from temporalio.runtime import OpenTelemetryConfig, Runtime, TelemetryConfig
 def validate_client_plugins(plugins: list[Any]) -> None:
     """
     Validate that all items in the plugins list are valid Temporal client plugins.
-    
+
     Args:
         plugins: List of plugins to validate
-        
+
     Raises:
         TypeError: If any plugin is not a valid ClientPlugin instance
     """
@@ -57,40 +59,40 @@ def validate_client_plugins(plugins: list[Any]) -> None:
             )
 
 
-async def get_temporal_client(
-    temporal_address: str, 
-    metrics_url: str = None,
-    plugins: list[Any] = []
-) -> Client:
+async def get_temporal_client(temporal_address: str, metrics_url: str | None = None, plugins: list[Any] = []) -> Client:
     """
     Create a Temporal client with plugin integration.
-    
+
     Args:
         temporal_address: Temporal server address
-        metrics_url: Optional metrics endpoint URL  
+        metrics_url: Optional metrics endpoint URL
         plugins: List of Temporal plugins to include
-        
+
     Returns:
         Configured Temporal client
     """
     # Validate plugins if any are provided
     if plugins:
         validate_client_plugins(plugins)
-        
+
+    # Check if OpenAI plugin is present - it needs to configure its own data converter
+    has_openai_plugin = any(
+        isinstance(p, OpenAIAgentsPlugin) for p in (plugins or [])
+    )
+
+    # Only set data_converter if OpenAI plugin is not present
+    connect_kwargs = {
+        "target_host": temporal_address,
+        "plugins": plugins,
+    }
+
+    if not has_openai_plugin:
+        connect_kwargs["data_converter"] = pydantic_data_converter
+
     if not metrics_url:
-        client = await Client.connect(
-            target_host=temporal_address,
-            # data_converter=custom_data_converter,
-            data_converter=pydantic_data_converter,
-            plugins=plugins,
-        )
+        client = await Client.connect(**connect_kwargs)
     else:
         runtime = Runtime(telemetry=TelemetryConfig(metrics=OpenTelemetryConfig(url=metrics_url)))
-        client = await Client.connect(
-            target_host=temporal_address,
-            # data_converter=custom_data_converter,
-            data_converter=pydantic_data_converter,
-            runtime=runtime,
-            plugins=plugins,
-        )
+        connect_kwargs["runtime"] = runtime
+        client = await Client.connect(**connect_kwargs)
     return client

agentex/lib/core/services/adk/acp/acp.py

@@ -1,18 +1,18 @@
 from typing import Any, List, cast
 
 from agentex import AsyncAgentex
-from agentex.lib.core.tracing.tracer import AsyncTracer
+from agentex.types.task import Task
+from agentex.types.event import Event
 from agentex.lib.utils.logging import make_logger
 from agentex.lib.utils.temporal import heartbeat_if_in_workflow
-from agentex.types.event import Event
-from agentex.types.task import Task
 from agentex.types.task_message import TaskMessage
-from agentex.types.task_message_content import TaskMessageContent
-from agentex.types.task_message_content_param import TaskMessageContentParam
 from agentex.types.agent_rpc_params import (
-    ParamsCancelTaskRequest as RpcParamsCancelTaskRequest,
     ParamsSendEventRequest as RpcParamsSendEventRequest,
+    ParamsCancelTaskRequest as RpcParamsCancelTaskRequest,
 )
+from agentex.lib.core.tracing.tracer import AsyncTracer
+from agentex.types.task_message_content import TaskMessageContent
+from agentex.types.task_message_content_param import TaskMessageContentParam
 
 logger = make_logger(__name__)
 

agentex/lib/core/services/adk/agent_task_tracker.py

@@ -1,6 +1,6 @@
 from agentex import AsyncAgentex
-from agentex.lib.core.tracing.tracer import AsyncTracer
 from agentex.lib.utils.logging import make_logger
+from agentex.lib.core.tracing.tracer import AsyncTracer
 from agentex.types.agent_task_tracker import AgentTaskTracker
 
 logger = make_logger(__name__)

agentex/lib/core/services/adk/agents.py

@@ -1,10 +1,10 @@
 from typing import Optional
 
 from agentex import AsyncAgentex
-from agentex.lib.core.tracing.tracer import AsyncTracer
 from agentex.types.agent import Agent
 from agentex.lib.utils.logging import make_logger
 from agentex.lib.utils.temporal import heartbeat_if_in_workflow
+from agentex.lib.core.tracing.tracer import AsyncTracer
 
 logger = make_logger(__name__)
 

agentex/lib/core/services/adk/events.py

@@ -1,7 +1,7 @@
 from agentex import AsyncAgentex
-from agentex.lib.core.tracing.tracer import AsyncTracer
 from agentex.types.event import Event
 from agentex.lib.utils.logging import make_logger
+from agentex.lib.core.tracing.tracer import AsyncTracer
 
 logger = make_logger(__name__)
 

agentex/lib/core/services/adk/messages.py

@@ -1,14 +1,13 @@
 import asyncio
-from typing import Any, Coroutine, cast
+from typing import Any, Coroutine
 
 from agentex import AsyncAgentex
-from agentex.lib.core.services.adk.streaming import StreamingService
-from agentex.lib.core.tracing.tracer import AsyncTracer
-from agentex.types.task_message_update import StreamTaskMessageFull, TaskMessageUpdate
-from agentex.types.task_message import TaskMessage, TaskMessageContent
 from agentex.lib.utils.logging import make_logger
 from agentex.lib.utils.temporal import heartbeat_if_in_workflow
-from agentex.types.task_message_content_param import TaskMessageContentParam
+from agentex.types.task_message import TaskMessage, TaskMessageContent
+from agentex.lib.core.tracing.tracer import AsyncTracer
+from agentex.types.task_message_update import TaskMessageUpdate, StreamTaskMessageFull
+from agentex.lib.core.services.adk.streaming import StreamingService
 
 logger = make_logger(__name__)
 

agentex/lib/core/services/adk/providers/litellm.py

@@ -1,23 +1,23 @@
 from collections.abc import AsyncGenerator
 
 from agentex import AsyncAgentex
-from agentex.lib.core.adapters.llm.adapter_litellm import LiteLLMGateway
-from agentex.lib.core.services.adk.streaming import StreamingService
-from agentex.lib.core.tracing.tracer import AsyncTracer
+from agentex.lib.utils import logging
+from agentex.lib.utils.temporal import heartbeat_if_in_workflow
+from agentex.types.task_message import TaskMessage
+from agentex.lib.utils.completions import concat_completion_chunks
 from agentex.lib.types.llm_messages import (
-    Completion,
     LLMConfig,
+    Completion,
 )
+from agentex.lib.core.tracing.tracer import AsyncTracer
+from agentex.types.task_message_delta import TextDelta
 from agentex.types.task_message_update import (
-    StreamTaskMessageDelta,
     StreamTaskMessageFull,
+    StreamTaskMessageDelta,
 )
-from agentex.types.task_message_delta import TextDelta
-from agentex.types.task_message import TaskMessage
 from agentex.types.task_message_content import TextContent
-from agentex.lib.utils import logging
-from agentex.lib.utils.completions import concat_completion_chunks
-from agentex.lib.utils.temporal import heartbeat_if_in_workflow
+from agentex.lib.core.services.adk.streaming import StreamingService
+from agentex.lib.core.adapters.llm.adapter_litellm import LiteLLMGateway
 
 logger = logging.make_logger(__name__)