agentex-sdk 0.1.0a6__py3-none-any.whl

agentex/lib/sdk/state_machine/state_machine.py

@@ -0,0 +1,189 @@
+from abc import ABC, abstractmethod
+from typing import Any, Generic, TypeVar
+
+from agentex.lib import adk
+from agentex.lib.sdk.state_machine.state import State
+from agentex.lib.sdk.state_machine.state_workflow import StateWorkflow
+from agentex.lib.utils.model_utils import BaseModel
+
+T = TypeVar("T", bound=BaseModel)
+
+
+class StateMachine(ABC, Generic[T]):
+    def __init__(
+        self,
+        initial_state: str,
+        states: list[State],
+        task_id: str | None = None,
+        state_machine_data: T | None = None,
+        trace_transitions: bool = False,
+    ):
+        self._task_id = task_id
+        self._state_map: dict[str, State] = {state.name: state for state in states}
+        self.state_machine_data = state_machine_data
+        self._initial_state = initial_state
+        self._trace_transitions = trace_transitions
+
+        # Validate that initial state exists
+        if initial_state not in self._state_map:
+            raise ValueError(f"Initial state '{initial_state}' not found in states")
+        self._current_state = self._state_map[initial_state]
+
+    def set_task_id(self, task_id: str):
+        self._task_id = task_id
+
+    def get_current_state(self) -> str:
+        return self._current_state.name
+
+    def get_current_workflow(self) -> StateWorkflow:
+        """
+        Get the workflow of the current state.
+
+        Returns:
+            The workflow of the current state
+
+        Raises:
+            ValueError: If the current state is not found in the state map
+        """
+        current_state = self._state_map.get(self.get_current_state())
+        if not current_state:
+            raise ValueError(f"State {self.get_current_state()} not found")
+        return current_state.workflow
+
+    async def transition(self, target_state_name: str):
+        if not self._state_map.get(target_state_name):
+            raise ValueError(f"State {target_state_name} not found")
+        self._current_state = self._state_map[target_state_name]
+
+    def get_state_machine_data(self) -> T:
+        return self.state_machine_data
+
+    @abstractmethod
+    async def terminal_condition(self) -> bool:
+        pass
+
+    # Overwrite this if you want to add more logic to the state machine
+    async def run(self):
+        while not await self.terminal_condition():
+            await self.step()
+
+    async def step(self) -> str:
+        current_state_name = self.get_current_state()
+        current_state = self._state_map.get(current_state_name)
+
+        if self._trace_transitions:
+            if self._task_id is None:
+                raise ValueError(
+                    "Task ID is must be set before tracing can be enabled"
+                )
+            span = await adk.tracing.start_span(
+                trace_id=self._task_id,
+                name="state_transition",
+                input=self.state_machine_data.model_dump(),
+                data={"input_state": current_state_name},
+            )
+
+        next_state_name = await current_state.workflow.execute(
+            state_machine=self, state_machine_data=self.state_machine_data
+        )
+
+        if self._trace_transitions:
+            if self._task_id is None:
+                raise ValueError(
+                    "Task ID is must be set before tracing can be enabled"
+                )
+            span.output = self.state_machine_data.model_dump()
+            span.data["output_state"] = next_state_name
+            await adk.tracing.end_span(trace_id=self._task_id, span=span)
+
+        await self.transition(next_state_name)
+
+        return next_state_name
+
+    async def reset_to_initial_state(self):
+        """
+        Reset the state machine to its initial state.
+        """
+        if self._trace_transitions:
+            if self._task_id is None:
+                raise ValueError(
+                    "Task ID is must be set before tracing can be enabled"
+                )
+            span = await adk.tracing.start_span(
+                trace_id=self._task_id,
+                name="state_transition_reset",
+                input={"input_state": self.get_current_state()},
+            )
+
+        await self.transition(self._initial_state)
+
+        if self._trace_transitions:
+            span.output = {"output_state": self._initial_state}
+            await adk.tracing.end_span(trace_id=self._task_id, span=span)
+
+    def dump(self) -> dict[str, Any]:
+        """
+        Save the current state of the state machine to a serializable dictionary.
+        This includes the current state, task_id, state machine data, and initial state.
+
+        Returns:
+            Dict[str, Any]: A dictionary containing the serialized state machine state
+        """
+        return {
+            "task_id": self._task_id,
+            "current_state": self.get_current_state(),
+            "initial_state": self._initial_state,
+            "state_machine_data": self.state_machine_data.model_dump(mode="json")
+            if self.state_machine_data
+            else None,
+            "trace_transitions": self._trace_transitions,
+        }
+
+    @classmethod
+    async def load(cls, data: dict[str, Any], states: list[State]) -> "StateMachine[T]":
+        """
+        Load a state machine from a previously saved dictionary.
+
+        Args:
+            data: The dictionary containing the saved state machine state
+            states: List of all possible states
+
+        Returns:
+            StateMachine: A new state machine instance restored to the saved state
+
+        Raises:
+            ValueError: If the data is invalid or missing required fields
+        """
+        try:
+            task_id = data.get("task_id")
+            current_state_name = data.get("current_state")
+            initial_state = data.get("initial_state")
+            state_machine_data_dict = data.get("state_machine_data")
+            trace_transitions = data.get("trace_transitions")
+
+            if initial_state is None:
+                raise ValueError("Initial state not found in saved data")
+
+            # Reconstruct the state machine data into its Pydantic model
+            state_machine_data = None
+            if state_machine_data_dict is not None:
+                # Get the actual model type from the class's type parameters
+                model_type = cls.__orig_bases__[0].__args__[0]
+                state_machine_data = model_type.model_validate(state_machine_data_dict)
+
+            # Create a new instance
+            instance = cls(
+                initial_state=initial_state,
+                states=states,
+                task_id=task_id,
+                state_machine_data=state_machine_data,
+                trace_transitions=trace_transitions,
+            )
+
+            # If there's a saved state, transition to it
+            if current_state_name:
+                await instance.transition(target_state_name=current_state_name)
+
+            return instance
+        except Exception as e:
+            raise ValueError(f"Failed to restore state machine: {str(e)}") from e

agentex/lib/sdk/state_machine/state_workflow.py

@@ -0,0 +1,16 @@
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+from pydantic import BaseModel
+
+# Import StateMachine only for type checking to avoid circular imports
+if TYPE_CHECKING:
+    from agentex.lib.sdk.state_machine import StateMachine
+
+
+class StateWorkflow(ABC):
+    @abstractmethod
+    async def execute(
+        self, state_machine: "StateMachine", state_machine_data: BaseModel | None = None
+    ) -> str:
+        pass

agentex/lib/sdk/utils/messages.py

@@ -0,0 +1,223 @@
+import json
+from abc import ABC, abstractmethod
+from typing import Any, Literal, override
+
+from agentex.lib.types.llm_messages import (
+    AssistantMessage,
+    Message,
+    ToolCall,
+    ToolCallRequest,
+    ToolMessage,
+    UserMessage,
+)
+from agentex.types.data_content import DataContent
+from agentex.types.task_message import TaskMessage
+from agentex.types.text_content import TextContent
+from agentex.types.tool_request_content import ToolRequestContent
+from agentex.types.tool_response_content import ToolResponseContent
+
+
+class TaskMessageConverter(ABC):
+    """
+    Abstract base class for converting a specific type of TaskMessage to an LLM Message.
+
+    Each converter should be responsible for one content type.
+    """
+
+    @abstractmethod
+    def convert(self, task_message: TaskMessage) -> Message:
+        """
+        Convert a TaskMessage to an LLM Message.
+
+        Args:
+            task_message: The TaskMessage to convert
+
+        Returns:
+            A Message (Pydantic model)
+        """
+        pass
+
+
+class DefaultTextContentConverter(TaskMessageConverter):
+    """Converter for TEXT content type."""
+
+    @override
+    def convert(self, task_message: TaskMessage) -> Message:
+        """Convert TEXT content to UserMessage or AssistantMessage based on author."""
+        if not isinstance(task_message.content, TextContent):
+            raise ValueError(f"Expected TextContent, got {type(task_message.content)}")
+        content = task_message.content
+        if content.author == "user":
+            return UserMessage(content=content.content)
+        else:  # AGENT or custom author
+            return AssistantMessage(content=content.content)
+
+
+class DefaultToolRequestConverter(TaskMessageConverter):
+    """Converter for TOOL_REQUEST content type."""
+
+    @override
+    def convert(self, task_message: TaskMessage) -> Message:
+        """Convert TOOL_REQUEST content to AssistantMessage with tool_calls."""
+        if not isinstance(task_message.content, ToolRequestContent):
+            raise ValueError(f"Expected ToolRequestContent, got {type(task_message.content)}")
+        
+        content = task_message.content
+
+        # Ensure arguments are properly JSON serialized
+        arguments_str = json.dumps(content.arguments)
+
+        tool_call = ToolCallRequest(
+            id=content.tool_call_id,
+            function=ToolCall(name=content.name, arguments=arguments_str),
+        )
+        return AssistantMessage(content=None, tool_calls=[tool_call])
+
+
+class DefaultToolResponseConverter(TaskMessageConverter):
+    """Converter for TOOL_RESPONSE content type."""
+
+    @override
+    def convert(self, task_message: TaskMessage) -> Message:
+        """Convert TOOL_RESPONSE content to ToolMessage."""
+        if not isinstance(task_message.content, ToolResponseContent):
+            raise ValueError(f"Expected ToolResponseContent, got {type(task_message.content)}")
+            
+        content = task_message.content
+        return ToolMessage(
+            content=str(content.content),
+            tool_call_id=content.tool_call_id,
+            name=content.name,
+        )
+
+
+class DefaultDataContentConverter(TaskMessageConverter):
+    """Converter for DATA content type."""
+
+    @override
+    def convert(self, task_message: TaskMessage) -> Message:
+        """Convert DATA content to UserMessage or AssistantMessage based on author."""
+        if not isinstance(task_message.content, DataContent):
+            raise ValueError(f"Expected DataContent, got {type(task_message.content)}")
+        
+        content = task_message.content
+        content_str = str(content.data)
+        if content.author == "user":
+            return UserMessage(content=content_str)
+        else:  # AGENT or custom author
+            return AssistantMessage(content=content_str)
+
+
+class DefaultUnknownContentConverter(TaskMessageConverter):
+    """Converter for unknown content types."""
+
+    @override
+    def convert(self, task_message: TaskMessage) -> Message:
+        """Convert unknown content types to AssistantMessage with fallback text."""
+        
+        content = task_message.content
+        fallback_content = f"Unknown message type: {content.type}"
+        return AssistantMessage(content=fallback_content)
+
+
+def convert_task_message_to_llm_messages(
+    task_message: TaskMessage,
+    output_mode: Literal["pydantic", "dict"] = "pydantic",
+    text_converter: TaskMessageConverter | None = None,
+    tool_request_converter: TaskMessageConverter | None = None,
+    tool_response_converter: TaskMessageConverter | None = None,
+    data_converter: TaskMessageConverter | None = None,
+    unknown_converter: TaskMessageConverter | None = None,
+) -> Message | dict[str, Any]:
+    """
+    Convert a TaskMessage to an LLM Message format.
+
+    Args:
+        task_message: The TaskMessage to convert
+        output_mode: Whether to return a Pydantic model or dict
+        text_converter: Optional converter for TEXT content. Uses DefaultTextContentConverter if None.
+        tool_request_converter: Optional converter for TOOL_REQUEST content. Uses DefaultToolRequestConverter if None.
+        tool_response_converter: Optional converter for TOOL_RESPONSE content. Uses DefaultToolResponseConverter if None.
+        data_converter: Optional converter for DATA content. Uses DefaultDataContentConverter if None.
+        unknown_converter: Optional converter for unknown content. Uses DefaultUnknownContentConverter if None.
+
+    Returns:
+        Either a Message (Pydantic model) or dict representation
+    """
+    content = task_message.content
+
+    # Get the appropriate converter for this content type
+    if content.type == "text":
+        converter = (
+            text_converter
+            if text_converter is not None
+            else DefaultTextContentConverter()
+        )
+    elif content.type == "tool_request":
+        converter = (
+            tool_request_converter
+            if tool_request_converter is not None
+            else DefaultToolRequestConverter()
+        )
+    elif content.type == "tool_response":
+        converter = (
+            tool_response_converter
+            if tool_response_converter is not None
+            else DefaultToolResponseConverter()
+        )
+    elif content.type == "data":
+        converter = (
+            data_converter
+            if data_converter is not None
+            else DefaultDataContentConverter()
+        )
+    else:
+        converter = (
+            unknown_converter
+            if unknown_converter is not None
+            else DefaultUnknownContentConverter()
+        )
+
+    message = converter.convert(task_message)
+
+    if output_mode == "dict":
+        return message.model_dump()
+    return message
+
+
+def convert_task_messages_to_llm_messages(
+    task_messages: list[TaskMessage],
+    output_mode: Literal["pydantic", "dict"] = "pydantic",
+    text_converter: TaskMessageConverter | None = None,
+    tool_request_converter: TaskMessageConverter | None = None,
+    tool_response_converter: TaskMessageConverter | None = None,
+    data_converter: TaskMessageConverter | None = None,
+    unknown_converter: TaskMessageConverter | None = None,
+) -> list[Message | dict[str, Any]]:
+    """
+    Convert a list of TaskMessages to LLM Message format.
+
+    Args:
+        task_messages: List of TaskMessages to convert
+        output_mode: Whether to return Pydantic models or dicts
+        text_converter: Optional converter for TEXT content. Uses DefaultTextContentConverter if None.
+        tool_request_converter: Optional converter for TOOL_REQUEST content. Uses DefaultToolRequestConverter if None.
+        tool_response_converter: Optional converter for TOOL_RESPONSE content. Uses DefaultToolResponseConverter if None.
+        data_converter: Optional converter for DATA content. Uses DefaultDataContentConverter if None.
+        unknown_converter: Optional converter for unknown content. Uses DefaultUnknownContentConverter if None.
+
+    Returns:
+        List of either Messages (Pydantic models) or dicts
+    """
+    return [
+        convert_task_message_to_llm_messages(
+            task_message,
+            output_mode,
+            text_converter,
+            tool_request_converter,
+            tool_response_converter,
+            data_converter,
+            unknown_converter,
+        )
+        for task_message in task_messages
+    ]

agentex/lib/types/acp.py

@@ -0,0 +1,94 @@
+from enum import Enum
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+from agentex.types.agent import Agent
+from agentex.types.event import Event
+from agentex.types.task_message_content import TaskMessageContent
+from agentex.types.task import Task
+
+
+class RPCMethod(str, Enum):
+    """Available JSON-RPC methods for agent communication."""
+
+    EVENT_SEND = "event/send"
+    MESSAGE_SEND = "message/send"
+    TASK_CANCEL = "task/cancel"
+    TASK_CREATE = "task/create"
+
+
+class CreateTaskParams(BaseModel):
+    """Parameters for task/create method.
+
+    Attributes:
+        agent: The agent that the task was sent to.
+        task: The task to be created.
+        params: The parameters for the task as inputted by the user.
+    """
+
+    agent: Agent = Field(..., description="The agent that the task was sent to")
+    task: Task = Field(..., description="The task to be created")
+    params: dict[str, Any] | None = Field(
+        None,
+        description="The parameters for the task as inputted by the user",
+    )
+
+
+class SendMessageParams(BaseModel):
+    """Parameters for message/send method.
+
+    Attributes:
+        agent: The agent that the message was sent to.
+        task: The task that the message was sent to.
+        content: The message that was sent to the agent.
+        stream: Whether to stream the message back to the agentex server from the agent.
+    """
+
+    agent: Agent = Field(..., description="The agent that the message was sent to")
+    task: Task = Field(..., description="The task that the message was sent to")
+    content: TaskMessageContent = Field(
+        ..., description="The message that was sent to the agent"
+    )
+    stream: bool = Field(
+        False,
+        description="Whether to stream the message back to the agentex server from the agent",
+    )
+
+
+class SendEventParams(BaseModel):
+    """Parameters for event/send method.
+
+    Attributes:
+        agent: The agent that the event was sent to.
+        task: The task that the message was sent to.
+        event: The event that was sent to the agent.
+    """
+
+    agent: Agent = Field(..., description="The agent that the event was sent to")
+    task: Task = Field(..., description="The task that the message was sent to")
+    event: Event = Field(..., description="The event that was sent to the agent")
+
+
+class CancelTaskParams(BaseModel):
+    """Parameters for task/cancel method.
+
+    Attributes:
+        agent: The agent that the task was sent to.
+        task: The task that was cancelled.
+    """
+
+    agent: Agent = Field(..., description="The agent that the task was sent to")
+    task: Task = Field(..., description="The task that was cancelled")
+
+
+RPC_SYNC_METHODS = [
+    RPCMethod.MESSAGE_SEND,
+]
+
+PARAMS_MODEL_BY_METHOD: dict[RPCMethod, type[BaseModel]] = {
+    RPCMethod.EVENT_SEND: SendEventParams,
+    RPCMethod.TASK_CANCEL: CancelTaskParams,
+    RPCMethod.MESSAGE_SEND: SendMessageParams,
+    RPCMethod.TASK_CREATE: CreateTaskParams,
+}

agentex/lib/types/agent_configs.py

@@ -0,0 +1,79 @@
+from pydantic import BaseModel, Field, model_validator, validator
+
+
+class TemporalWorkflowConfig(BaseModel):
+    """
+    Configuration for the temporal workflow that defines the agent.
+
+    Attributes:
+        name: The name of the temporal workflow that defines the agent.
+        queue_name: The name of the temporal queue to send tasks to.
+    """
+
+    name: str = Field(
+        ..., description="The name of the temporal workflow that defines the agent."
+    )
+    queue_name: str = Field(
+        ..., description="The name of the temporal queue to send tasks to."
+    )
+
+
+# TODO: Remove this class when we remove the agentex agents create
+class TemporalWorkerConfig(BaseModel):
+    """
+    Configuration for temporal worker deployment
+
+    Attributes:
+        image: The image to use for the temporal worker
+        workflow: The temporal workflow configuration
+    """
+
+    image: str | None = Field(
+        default=None, description="Image to use for the temporal worker"
+    )
+    workflow: TemporalWorkflowConfig | None = Field(
+        default=None,
+        description="Configuration for the temporal workflow that defines the agent. Only required for agents that leverage Temporal.",
+    )
+
+
+class TemporalConfig(BaseModel):
+    """
+    Simplified temporal configuration for agents
+
+    Attributes:
+        enabled: Whether this agent uses Temporal workflows
+        workflow: The temporal workflow configuration
+        workflows: The list of temporal workflow configurations
+    """
+
+    enabled: bool = Field(
+        default=False, description="Whether this agent uses Temporal workflows"
+    )
+    workflow: TemporalWorkflowConfig | None = Field(
+        default=None,
+        description="Temporal workflow configuration. Required when enabled=True. (deprecated: use workflows instead)",
+    )
+    workflows: list[TemporalWorkflowConfig] | None = Field(
+        default=None,
+        description="List of temporal workflow configurations. Used when enabled=true.",
+    )
+
+    @validator("workflows")
+    def validate_workflows_not_empty(cls, v):
+        """Ensure workflows list is not empty when provided"""
+        if v is not None and len(v) == 0:
+            raise ValueError("workflows list cannot be empty when provided")
+        return v
+
+    @model_validator(mode="after")
+    def validate_temporal_config_when_enabled(self):
+        """Validate that workflow configuration exists when enabled=true"""
+        if self.enabled:
+            # Must have either workflow (legacy) or workflows (new)
+            if not self.workflow and (not self.workflows or len(self.workflows) == 0):
+                raise ValueError(
+                    "When temporal.enabled=true, either 'workflow' or 'workflows' must be provided and non-empty"
+                )
+
+        return self

agentex/lib/types/agent_results.py

@@ -0,0 +1,29 @@
+from typing import Any
+
+from pydantic import BaseModel
+
+
+class SerializableRunResult(BaseModel):
+    """
+    Serializable version of RunResult.
+
+    Attributes:
+        final_output: The final output of the run.
+        final_input_list: The final input list of the run.
+    """
+
+    final_output: Any
+    final_input_list: list[dict[str, Any]]
+
+
+class SerializableRunResultStreaming(BaseModel):
+    """
+    Serializable version of RunResultStreaming.
+
+    Attributes:
+        final_output: The final output of the run.
+        final_input_list: The final input list of the run.
+    """
+
+    final_output: Any
+    final_input_list: list[dict[str, Any]]

agentex/lib/types/credentials.py

@@ -0,0 +1,34 @@
+from pydantic import BaseModel, Field
+
+
+class CredentialMapping(BaseModel):
+    """Maps a Kubernetes secret to an environment variable in the agent container.
+
+    This allows agents to securely access credentials stored in Kubernetes secrets
+    by mapping them to environment variables. For example, you can map a secret
+    containing an API key to an environment variable that your agent code expects.
+
+    Example:
+        A mapping of {"env_var_name": "OPENAI_API_KEY",
+                     "secret_name": "ai-credentials",
+                     "secret_key": "openai-key"}
+        will make the value from the "openai-key" field in the "ai-credentials"
+        Kubernetes secret available to the agent as OPENAI_API_KEY environment variable.
+
+    Attributes:
+        env_var_name: The name of the environment variable that will be available to the agent
+        secret_name: The name of the Kubernetes secret containing the credential
+        secret_key: The key within the Kubernetes secret that contains the credential value
+    """
+
+    env_var_name: str = Field(
+        ...,
+        description="Name of the environment variable that will be available to the agent",
+    )
+    secret_name: str = Field(
+        ..., description="Name of the Kubernetes secret containing the credential"
+    )
+    secret_key: str = Field(
+        ...,
+        description="Key within the Kubernetes secret that contains the credential value",
+    )