agentex-sdk 0.4.11__py3-none-any.whl → 0.4.13__py3-none-any.whl

agentex/lib/core/temporal/activities/adk/providers/openai_activities.py

@@ -1,34 +1,42 @@
 # Standard library imports
 import base64
-from collections.abc import Callable
-from contextlib import AsyncExitStack, asynccontextmanager
 from enum import Enum
 from typing import Any, Literal, Optional
-
-from pydantic import Field, PrivateAttr
+from contextlib import AsyncExitStack, asynccontextmanager
+from collections.abc import Callable
 
 import cloudpickle
-from agents import RunContextWrapper, RunResult, RunResultStreaming
+from mcp import StdioServerParameters
+from agents import RunResult, RunContextWrapper, RunResultStreaming
+from pydantic import Field, PrivateAttr
+from agents.mcp import MCPServerStdio, MCPServerStdioParams
+from temporalio import activity
+from agents.tool import (
+    ComputerTool as OAIComputerTool,
+    FunctionTool as OAIFunctionTool,
+    WebSearchTool as OAIWebSearchTool,
+    FileSearchTool as OAIFileSearchTool,
+    LocalShellTool as OAILocalShellTool,
+    CodeInterpreterTool as OAICodeInterpreterTool,
+    ImageGenerationTool as OAIImageGenerationTool,
+)
 from agents.guardrail import InputGuardrail, OutputGuardrail
 from agents.exceptions import InputGuardrailTripwireTriggered, OutputGuardrailTripwireTriggered
-from agents.mcp import MCPServerStdio, MCPServerStdioParams
 from agents.model_settings import ModelSettings as OAIModelSettings
-from agents.tool import FunctionTool as OAIFunctionTool
-from mcp import StdioServerParameters
-from openai.types.responses.response_includable import ResponseIncludable
 from openai.types.shared.reasoning import Reasoning
-from temporalio import activity
-from agentex.lib.core.services.adk.providers.openai import OpenAIService
+from openai.types.responses.response_includable import ResponseIncludable
+
+from agentex.lib.utils import logging
+
+# Third-party imports
+from agentex.lib.types.tracing import BaseModelWithTraceParams
 
 # Local imports
 from agentex.lib.types.agent_results import (
     SerializableRunResult,
     SerializableRunResultStreaming,
 )
-
-# Third-party imports
-from agentex.lib.types.tracing import BaseModelWithTraceParams
-from agentex.lib.utils import logging
+from agentex.lib.core.services.adk.providers.openai import OpenAIService
 
 logger = logging.make_logger(__name__)
 
@@ -42,6 +50,147 @@ class OpenAIActivityName(str, Enum):
     RUN_AGENT_STREAMED_AUTO_SEND = "run_agent_streamed_auto_send"
 
 
+class WebSearchTool(BaseModelWithTraceParams):
+    """Temporal-compatible wrapper for WebSearchTool."""
+
+    user_location: Optional[dict[str, Any]] = None  # UserLocation object
+    search_context_size: Optional[Literal["low", "medium", "high"]] = "medium"
+
+    def to_oai_function_tool(self) -> OAIWebSearchTool:
+        kwargs = {}
+        if self.user_location is not None:
+            kwargs["user_location"] = self.user_location
+        if self.search_context_size is not None:
+            kwargs["search_context_size"] = self.search_context_size
+        return OAIWebSearchTool(**kwargs)
+
+
+class FileSearchTool(BaseModelWithTraceParams):
+    """Temporal-compatible wrapper for FileSearchTool."""
+
+    vector_store_ids: list[str]
+    max_num_results: Optional[int] = None
+    include_search_results: bool = False
+    ranking_options: Optional[dict[str, Any]] = None
+    filters: Optional[dict[str, Any]] = None
+
+    def to_oai_function_tool(self):
+        return OAIFileSearchTool(
+            vector_store_ids=self.vector_store_ids,
+            max_num_results=self.max_num_results,
+            include_search_results=self.include_search_results,
+            ranking_options=self.ranking_options,
+            filters=self.filters,
+        )
+
+
+class ComputerTool(BaseModelWithTraceParams):
+    """Temporal-compatible wrapper for ComputerTool."""
+
+    # We need to serialize the computer object and safety check function
+    computer_serialized: str = Field(default="", description="Serialized computer object")
+    on_safety_check_serialized: str = Field(default="", description="Serialized safety check function")
+
+    _computer: Any = PrivateAttr()
+    _on_safety_check: Optional[Callable] = PrivateAttr()
+
+    def __init__(
+        self,
+        *,
+        computer: Any = None,
+        on_safety_check: Optional[Callable] = None,
+        **data,
+    ):
+        super().__init__(**data)
+        if computer is not None:
+            self.computer_serialized = self._serialize_callable(computer)
+            self._computer = computer
+        elif self.computer_serialized:
+            self._computer = self._deserialize_callable(self.computer_serialized)
+
+        if on_safety_check is not None:
+            self.on_safety_check_serialized = self._serialize_callable(on_safety_check)
+            self._on_safety_check = on_safety_check
+        elif self.on_safety_check_serialized:
+            self._on_safety_check = self._deserialize_callable(self.on_safety_check_serialized)
+
+    @classmethod
+    def _deserialize_callable(cls, serialized: str) -> Any:
+        encoded = serialized.encode()
+        serialized_bytes = base64.b64decode(encoded)
+        return cloudpickle.loads(serialized_bytes)
+
+    @classmethod
+    def _serialize_callable(cls, func: Any) -> str:
+        serialized_bytes = cloudpickle.dumps(func)
+        encoded = base64.b64encode(serialized_bytes)
+        return encoded.decode()
+
+    def to_oai_function_tool(self):
+        return OAIComputerTool(
+            computer=self._computer,
+            on_safety_check=self._on_safety_check,
+        )
+
+
+class CodeInterpreterTool(BaseModelWithTraceParams):
+    """Temporal-compatible wrapper for CodeInterpreterTool."""
+
+    tool_config: dict[str, Any] = Field(
+        default_factory=lambda: {"type": "code_interpreter"}, description="Tool configuration dict"
+    )
+
+    def to_oai_function_tool(self):
+        return OAICodeInterpreterTool(tool_config=self.tool_config)
+
+
+class ImageGenerationTool(BaseModelWithTraceParams):
+    """Temporal-compatible wrapper for ImageGenerationTool."""
+
+    tool_config: dict[str, Any] = Field(
+        default_factory=lambda: {"type": "image_generation"}, description="Tool configuration dict"
+    )
+
+    def to_oai_function_tool(self):
+        return OAIImageGenerationTool(tool_config=self.tool_config)
+
+
+class LocalShellTool(BaseModelWithTraceParams):
+    """Temporal-compatible wrapper for LocalShellTool."""
+
+    executor_serialized: str = Field(default="", description="Serialized LocalShellExecutor object")
+
+    _executor: Any = PrivateAttr()
+
+    def __init__(
+        self,
+        *,
+        executor: Any = None,
+        **data,
+    ):
+        super().__init__(**data)
+        if executor is not None:
+            self.executor_serialized = self._serialize_callable(executor)
+            self._executor = executor
+        elif self.executor_serialized:
+            self._executor = self._deserialize_callable(self.executor_serialized)
+
+    @classmethod
+    def _deserialize_callable(cls, serialized: str) -> Any:
+        encoded = serialized.encode()
+        serialized_bytes = base64.b64decode(encoded)
+        return cloudpickle.loads(serialized_bytes)
+
+    @classmethod
+    def _serialize_callable(cls, func: Any) -> str:
+        serialized_bytes = cloudpickle.dumps(func)
+        encoded = base64.b64encode(serialized_bytes)
+        return encoded.decode()
+
+    def to_oai_function_tool(self):
+        return OAILocalShellTool(executor=self._executor)
+
+
 class FunctionTool(BaseModelWithTraceParams):
     name: str
     description: str
@@ -79,22 +228,16 @@ class FunctionTool(BaseModelWithTraceParams):
         super().__init__(**data)
         if not on_invoke_tool:
             if not self.on_invoke_tool_serialized:
-                raise ValueError(
-                    "One of `on_invoke_tool` or `on_invoke_tool_serialized` should be set"
-                )
+                raise ValueError("One of `on_invoke_tool` or `on_invoke_tool_serialized` should be set")
             else:
-                on_invoke_tool = self._deserialize_callable(
-                    self.on_invoke_tool_serialized
-                )
+                on_invoke_tool = self._deserialize_callable(self.on_invoke_tool_serialized)
         else:
             self.on_invoke_tool_serialized = self._serialize_callable(on_invoke_tool)
 
         self._on_invoke_tool = on_invoke_tool
 
     @classmethod
-    def _deserialize_callable(
-        cls, serialized: str
-    ) -> Callable[[RunContextWrapper, str], Any]:
+    def _deserialize_callable(cls, serialized: str) -> Callable[[RunContextWrapper, str], Any]:
         encoded = serialized.encode()
         serialized_bytes = base64.b64decode(encoded)
         return cloudpickle.loads(serialized_bytes)
@@ -108,11 +251,9 @@ class FunctionTool(BaseModelWithTraceParams):
     @property
     def on_invoke_tool(self) -> Callable[[RunContextWrapper, str], Any]:
         if self._on_invoke_tool is None and self.on_invoke_tool_serialized:
-            self._on_invoke_tool = self._deserialize_callable(
-                self.on_invoke_tool_serialized
-            )
+            self._on_invoke_tool = self._deserialize_callable(self.on_invoke_tool_serialized)
         return self._on_invoke_tool
-    
+
     @on_invoke_tool.setter
     def on_invoke_tool(self, value: Callable[[RunContextWrapper, str], Any]):
         self.on_invoke_tool_serialized = self._serialize_callable(value)
@@ -135,8 +276,9 @@ class FunctionTool(BaseModelWithTraceParams):
 
 
 class TemporalInputGuardrail(BaseModelWithTraceParams):
-    """Temporal-compatible wrapper for InputGuardrail with function 
+    """Temporal-compatible wrapper for InputGuardrail with function
     serialization."""
+
     name: str
     _guardrail_function: Callable = PrivateAttr()
     guardrail_function_serialized: str = Field(
@@ -157,19 +299,12 @@ class TemporalInputGuardrail(BaseModelWithTraceParams):
         super().__init__(**data)
         if not guardrail_function:
             if not self.guardrail_function_serialized:
-                raise ValueError(
-                    "One of `guardrail_function` or "
-                    "`guardrail_function_serialized` should be set"
-                )
+                raise ValueError("One of `guardrail_function` or `guardrail_function_serialized` should be set")
             else:
-                guardrail_function = self._deserialize_callable(
-                    self.guardrail_function_serialized
-                )
+                guardrail_function = self._deserialize_callable(self.guardrail_function_serialized)
         else:
-            self.guardrail_function_serialized = self._serialize_callable(
-                guardrail_function
-            )
-        
+            self.guardrail_function_serialized = self._serialize_callable(guardrail_function)
+
         self._guardrail_function = guardrail_function
 
     @classmethod
@@ -186,13 +321,10 @@ class TemporalInputGuardrail(BaseModelWithTraceParams):
 
     @property
     def guardrail_function(self) -> Callable:
-        if (self._guardrail_function is None and 
-                self.guardrail_function_serialized):
-            self._guardrail_function = self._deserialize_callable(
-                self.guardrail_function_serialized
-            )
+        if self._guardrail_function is None and self.guardrail_function_serialized:
+            self._guardrail_function = self._deserialize_callable(self.guardrail_function_serialized)
         return self._guardrail_function
-    
+
     @guardrail_function.setter
     def guardrail_function(self, value: Callable):
         self.guardrail_function_serialized = self._serialize_callable(value)
@@ -200,15 +332,13 @@ class TemporalInputGuardrail(BaseModelWithTraceParams):
 
     def to_oai_input_guardrail(self) -> InputGuardrail:
         """Convert to OpenAI InputGuardrail."""
-        return InputGuardrail(
-            guardrail_function=self.guardrail_function,
-            name=self.name
-        )
+        return InputGuardrail(guardrail_function=self.guardrail_function, name=self.name)
 
 
 class TemporalOutputGuardrail(BaseModelWithTraceParams):
-    """Temporal-compatible wrapper for OutputGuardrail with function 
+    """Temporal-compatible wrapper for OutputGuardrail with function
     serialization."""
+
     name: str
     _guardrail_function: Callable = PrivateAttr()
     guardrail_function_serialized: str = Field(
@@ -229,19 +359,12 @@ class TemporalOutputGuardrail(BaseModelWithTraceParams):
         super().__init__(**data)
         if not guardrail_function:
             if not self.guardrail_function_serialized:
-                raise ValueError(
-                    "One of `guardrail_function` or "
-                    "`guardrail_function_serialized` should be set"
-                )
+                raise ValueError("One of `guardrail_function` or `guardrail_function_serialized` should be set")
             else:
-                guardrail_function = self._deserialize_callable(
-                    self.guardrail_function_serialized
-                )
+                guardrail_function = self._deserialize_callable(self.guardrail_function_serialized)
         else:
-            self.guardrail_function_serialized = self._serialize_callable(
-                guardrail_function
-            )
-        
+            self.guardrail_function_serialized = self._serialize_callable(guardrail_function)
+
         self._guardrail_function = guardrail_function
 
     @classmethod
@@ -258,13 +381,10 @@ class TemporalOutputGuardrail(BaseModelWithTraceParams):
 
     @property
     def guardrail_function(self) -> Callable:
-        if (self._guardrail_function is None and 
-                self.guardrail_function_serialized):
-            self._guardrail_function = self._deserialize_callable(
-                self.guardrail_function_serialized
-            )
+        if self._guardrail_function is None and self.guardrail_function_serialized:
+            self._guardrail_function = self._deserialize_callable(self.guardrail_function_serialized)
         return self._guardrail_function
-    
+
     @guardrail_function.setter
     def guardrail_function(self, value: Callable):
         self.guardrail_function_serialized = self._serialize_callable(value)
@@ -272,10 +392,7 @@ class TemporalOutputGuardrail(BaseModelWithTraceParams):
 
     def to_oai_output_guardrail(self) -> OutputGuardrail:
         """Convert to OpenAI OutputGuardrail."""
-        return OutputGuardrail(
-            guardrail_function=self.guardrail_function,
-            name=self.name
-        )
+        return OutputGuardrail(guardrail_function=self.guardrail_function, name=self.name)
 
 
 class ModelSettings(BaseModelWithTraceParams):
@@ -297,9 +414,7 @@ class ModelSettings(BaseModelWithTraceParams):
     extra_args: dict[str, Any] | None = None
 
     def to_oai_model_settings(self) -> OAIModelSettings:
-        return OAIModelSettings(
-            **self.model_dump(exclude=["trace_id", "parent_span_id"])
-        )
+        return OAIModelSettings(**self.model_dump(exclude=["trace_id", "parent_span_id"]))
 
 
 class RunAgentParams(BaseModelWithTraceParams):
@@ -313,12 +428,24 @@ class RunAgentParams(BaseModelWithTraceParams):
     handoffs: list["RunAgentParams"] | None = None
     model: str | None = None
     model_settings: ModelSettings | None = None
-    tools: list[FunctionTool] | None = None
+    tools: (
+        list[
+            FunctionTool
+            | WebSearchTool
+            | FileSearchTool
+            | ComputerTool
+            | CodeInterpreterTool
+            | ImageGenerationTool
+            | LocalShellTool
+        ]
+        | None
+    ) = None
     output_type: Any = None
     tool_use_behavior: Literal["run_llm_again", "stop_on_first_tool"] = "run_llm_again"
     mcp_timeout_seconds: int | None = None
     input_guardrails: list[TemporalInputGuardrail] | None = None
     output_guardrails: list[TemporalOutputGuardrail] | None = None
+    max_turns: int | None = None
 
 
 class RunAgentAutoSendParams(RunAgentParams):
@@ -365,11 +492,11 @@ class OpenAIActivities:
         input_guardrails = None
         if params.input_guardrails:
             input_guardrails = [g.to_oai_input_guardrail() for g in params.input_guardrails]
-        
+
         output_guardrails = None
         if params.output_guardrails:
             output_guardrails = [g.to_oai_output_guardrail() for g in params.output_guardrails]
-        
+
         result = await self._openai_service.run_agent(
             input_list=params.input_list,
             mcp_server_params=params.mcp_server_params,
@@ -386,23 +513,23 @@ class OpenAIActivities:
             tool_use_behavior=params.tool_use_behavior,
             input_guardrails=input_guardrails,
             output_guardrails=output_guardrails,
+            mcp_timeout_seconds=params.mcp_timeout_seconds,
+            max_turns=params.max_turns,
         )
         return self._to_serializable_run_result(result)
 
     @activity.defn(name=OpenAIActivityName.RUN_AGENT_AUTO_SEND)
-    async def run_agent_auto_send(
-        self, params: RunAgentAutoSendParams
-    ) -> SerializableRunResult:
+    async def run_agent_auto_send(self, params: RunAgentAutoSendParams) -> SerializableRunResult:
         """Run an agent with automatic TaskMessage creation."""
         # Convert Temporal guardrails to OpenAI guardrails
         input_guardrails = None
         if params.input_guardrails:
             input_guardrails = [g.to_oai_input_guardrail() for g in params.input_guardrails]
-        
+
         output_guardrails = None
         if params.output_guardrails:
             output_guardrails = [g.to_oai_output_guardrail() for g in params.output_guardrails]
-        
+
         try:
             result = await self._openai_service.run_agent_auto_send(
                 task_id=params.task_id,
@@ -421,65 +548,60 @@ class OpenAIActivities:
                 tool_use_behavior=params.tool_use_behavior,
                 input_guardrails=input_guardrails,
                 output_guardrails=output_guardrails,
+                mcp_timeout_seconds=params.mcp_timeout_seconds,
+                max_turns=params.max_turns,
             )
             return self._to_serializable_run_result(result)
         except InputGuardrailTripwireTriggered as e:
             # Handle guardrail trigger gracefully
-            rejection_message = "I'm sorry, but I cannot process this request due to a guardrail. Please try a different question."
-            
+            rejection_message = (
+                "I'm sorry, but I cannot process this request due to a guardrail. Please try a different question."
+            )
+
             # Try to extract rejection message from the guardrail result
-            if hasattr(e, 'guardrail_result') and hasattr(e.guardrail_result, 'output'):
-                output_info = getattr(e.guardrail_result.output, 'output_info', {})
-                if isinstance(output_info, dict) and 'rejection_message' in output_info:
-                    rejection_message = output_info['rejection_message']
-            
+            if hasattr(e, "guardrail_result") and hasattr(e.guardrail_result, "output"):
+                output_info = getattr(e.guardrail_result.output, "output_info", {})
+                if isinstance(output_info, dict) and "rejection_message" in output_info:
+                    rejection_message = output_info["rejection_message"]
+
             # Build the final input list with the rejection message
             final_input_list = list(params.input_list or [])
-            final_input_list.append({
-                "role": "assistant",
-                "content": rejection_message
-            })
-            
-            return SerializableRunResult(
-                final_output=rejection_message,
-                final_input_list=final_input_list
-            )
+            final_input_list.append({"role": "assistant", "content": rejection_message})
+
+            return SerializableRunResult(final_output=rejection_message, final_input_list=final_input_list)
         except OutputGuardrailTripwireTriggered as e:
             # Handle output guardrail trigger gracefully
-            rejection_message = "I'm sorry, but I cannot provide this response due to a guardrail. Please try a different question."
-            
+            rejection_message = (
+                "I'm sorry, but I cannot provide this response due to a guardrail. Please try a different question."
+            )
+
             # Try to extract rejection message from the guardrail result
-            if hasattr(e, 'guardrail_result') and hasattr(e.guardrail_result, 'output'):
-                output_info = getattr(e.guardrail_result.output, 'output_info', {})
-                if isinstance(output_info, dict) and 'rejection_message' in output_info:
-                    rejection_message = output_info['rejection_message']
-            
+            if hasattr(e, "guardrail_result") and hasattr(e.guardrail_result, "output"):
+                output_info = getattr(e.guardrail_result.output, "output_info", {})
+                if isinstance(output_info, dict) and "rejection_message" in output_info:
+                    rejection_message = output_info["rejection_message"]
+
             # Build the final input list with the rejection message
             final_input_list = list(params.input_list or [])
-            final_input_list.append({
-                "role": "assistant",
-                "content": rejection_message
-            })
-            
-            return SerializableRunResult(
-                final_output=rejection_message,
-                final_input_list=final_input_list
-            )
+            final_input_list.append({"role": "assistant", "content": rejection_message})
+
+            return SerializableRunResult(final_output=rejection_message, final_input_list=final_input_list)
 
     @activity.defn(name=OpenAIActivityName.RUN_AGENT_STREAMED_AUTO_SEND)
     async def run_agent_streamed_auto_send(
         self, params: RunAgentStreamedAutoSendParams
     ) -> SerializableRunResultStreaming:
         """Run an agent with streaming and automatic TaskMessage creation."""
+
         # Convert Temporal guardrails to OpenAI guardrails
         input_guardrails = None
         if params.input_guardrails:
             input_guardrails = [g.to_oai_input_guardrail() for g in params.input_guardrails]
-        
+
         output_guardrails = None
         if params.output_guardrails:
             output_guardrails = [g.to_oai_output_guardrail() for g in params.output_guardrails]
-        
+
         try:
             result = await self._openai_service.run_agent_streamed_auto_send(
                 task_id=params.task_id,
@@ -498,50 +620,44 @@ class OpenAIActivities:
                 tool_use_behavior=params.tool_use_behavior,
                 input_guardrails=input_guardrails,
                 output_guardrails=output_guardrails,
+                mcp_timeout_seconds=params.mcp_timeout_seconds,
+                max_turns=params.max_turns,
             )
             return self._to_serializable_run_result_streaming(result)
         except InputGuardrailTripwireTriggered as e:
             # Handle guardrail trigger gracefully
-            rejection_message = "I'm sorry, but I cannot process this request due to a guardrail. Please try a different question."
-            
+            rejection_message = (
+                "I'm sorry, but I cannot process this request due to a guardrail. Please try a different question."
+            )
+
             # Try to extract rejection message from the guardrail result
-            if hasattr(e, 'guardrail_result') and hasattr(e.guardrail_result, 'output'):
-                output_info = getattr(e.guardrail_result.output, 'output_info', {})
-                if isinstance(output_info, dict) and 'rejection_message' in output_info:
-                    rejection_message = output_info['rejection_message']
-            
+            if hasattr(e, "guardrail_result") and hasattr(e.guardrail_result, "output"):
+                output_info = getattr(e.guardrail_result.output, "output_info", {})
+                if isinstance(output_info, dict) and "rejection_message" in output_info:
+                    rejection_message = output_info["rejection_message"]
+
             # Build the final input list with the rejection message
             final_input_list = list(params.input_list or [])
-            final_input_list.append({
-                "role": "assistant",
-                "content": rejection_message
-            })
-            
-            return SerializableRunResultStreaming(
-                final_output=rejection_message,
-                final_input_list=final_input_list
-            )
+            final_input_list.append({"role": "assistant", "content": rejection_message})
+
+            return SerializableRunResultStreaming(final_output=rejection_message, final_input_list=final_input_list)
         except OutputGuardrailTripwireTriggered as e:
             # Handle output guardrail trigger gracefully
-            rejection_message = "I'm sorry, but I cannot provide this response due to a guardrail. Please try a different question."
-            
+            rejection_message = (
+                "I'm sorry, but I cannot provide this response due to a guardrail. Please try a different question."
+            )
+
             # Try to extract rejection message from the guardrail result
-            if hasattr(e, 'guardrail_result') and hasattr(e.guardrail_result, 'output'):
-                output_info = getattr(e.guardrail_result.output, 'output_info', {})
-                if isinstance(output_info, dict) and 'rejection_message' in output_info:
-                    rejection_message = output_info['rejection_message']
-            
+            if hasattr(e, "guardrail_result") and hasattr(e.guardrail_result, "output"):
+                output_info = getattr(e.guardrail_result.output, "output_info", {})
+                if isinstance(output_info, dict) and "rejection_message" in output_info:
+                    rejection_message = output_info["rejection_message"]
+
             # Build the final input list with the rejection message
             final_input_list = list(params.input_list or [])
-            final_input_list.append({
-                "role": "assistant",
-                "content": rejection_message
-            })
-            
-            return SerializableRunResultStreaming(
-                final_output=rejection_message,
-                final_input_list=final_input_list
-            )
+            final_input_list.append({"role": "assistant", "content": rejection_message})
+
+            return SerializableRunResultStreaming(final_output=rejection_message, final_input_list=final_input_list)
 
     @staticmethod
     def _to_serializable_run_result(result: RunResult) -> SerializableRunResult:

agentex/lib/core/temporal/workers/worker.py

@@ -4,7 +4,7 @@ import os
 import uuid
 from collections.abc import Callable
 from concurrent.futures import ThreadPoolExecutor
-from typing import Any
+from typing import Any, overload
 
 from aiohttp import web
 from temporalio.client import Client
@@ -99,10 +99,28 @@ class AgentexWorker:
         self.healthy = False
         self.health_check_port = health_check_port
 
+    @overload
     async def run(
         self,
         activities: list[Callable],
+        *,
         workflow: type,
+    ) -> None: ...
+    
+    @overload
+    async def run(
+        self,
+        activities: list[Callable],
+        *,
+        workflows: list[type],
+    ) -> None: ...
+
+    async def run(
+        self,
+        activities: list[Callable],
+        *,
+        workflow: type | None = None,
+        workflows: list[type] | None = None,
     ):
         await self.start_health_check_server()
         await self._register_agent()
@@ -115,11 +133,14 @@ class AgentexWorker:
         if debug_enabled:
             logger.info("🐛 [WORKER] Temporal debug mode enabled - deadlock detection disabled")
         
+        if workflow is None and workflows is None:
+            raise ValueError("Either workflow or workflows must be provided")
+        
         worker = Worker(
             client=temporal_client,
             task_queue=self.task_queue,
             activity_executor=ThreadPoolExecutor(max_workers=self.max_workers),
-            workflows=[workflow],
+            workflows=[workflow] if workflows is None else workflows,
             activities=activities,
             workflow_runner=UnsandboxedWorkflowRunner(),
             max_concurrent_activities=self.max_concurrent_activities,