agentex-sdk 0.1.0__py3-none-any.whl

agentex/lib/adk/_modules/messages.py

@@ -0,0 +1,285 @@
+from datetime import timedelta
+
+from temporalio.common import RetryPolicy
+
+from agentex import AsyncAgentex
+from agentex.lib.core.adapters.streams.adapter_redis import RedisEventStreamRepository
+from agentex.lib.core.services.adk.messages import MessagesService
+from agentex.lib.core.services.adk.streaming import StreamingService
+from agentex.lib.core.temporal.activities.activity_helpers import ActivityHelpers
+from agentex.lib.core.temporal.activities.adk.messages_activities import (
+    CreateMessageParams,
+    CreateMessagesBatchParams,
+    ListMessagesParams,
+    MessagesActivityName,
+    UpdateMessageParams,
+    UpdateMessagesBatchParams,
+)
+from agentex.lib.core.tracing.tracer import AsyncTracer
+from agentex.types.task_message import TaskMessage, TaskMessageContent
+from agentex.lib.utils.logging import make_logger
+from agentex.lib.utils.temporal import in_temporal_workflow
+
+logger = make_logger(__name__)
+
+# Default retry policy for all message operations
+DEFAULT_RETRY_POLICY = RetryPolicy(maximum_attempts=1)
+
+
+class MessagesModule:
+    """
+    Module for managing task messages in Agentex.
+    Provides high-level async methods for creating, retrieving, updating, and deleting messages.
+    """
+
+    def __init__(
+        self,
+        messages_service: MessagesService | None = None,
+    ):
+        if messages_service is None:
+            agentex_client = AsyncAgentex()
+            stream_repository = RedisEventStreamRepository()
+            streaming_service = StreamingService(
+                agentex_client=agentex_client,
+                stream_repository=stream_repository,
+            )
+            tracer = AsyncTracer(agentex_client)
+            self._messages_service = MessagesService(
+                agentex_client=agentex_client,
+                streaming_service=streaming_service,
+                tracer=tracer,
+            )
+        else:
+            self._messages_service = messages_service
+
+    async def create(
+        self,
+        task_id: str,
+        content: TaskMessageContent,
+        emit_updates: bool = True,
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+        start_to_close_timeout: timedelta = timedelta(seconds=5),
+        heartbeat_timeout: timedelta = timedelta(seconds=5),
+        retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
+    ) -> TaskMessage:
+        """
+        Create a new message for a task.
+
+        Args:
+            task_id (str): The ID of the task.
+            message (TaskMessage): The message to create.
+            trace_id (Optional[str]): The trace ID for tracing.
+            parent_span_id (Optional[str]): The parent span ID for tracing.
+            start_to_close_timeout (timedelta): The start to close timeout.
+            heartbeat_timeout (timedelta): The heartbeat timeout.
+            retry_policy (RetryPolicy): The retry policy.
+
+        Returns:
+            TaskMessageEntity: The created message.
+        """
+        params = CreateMessageParams(
+            trace_id=trace_id,
+            parent_span_id=parent_span_id,
+            task_id=task_id,
+            content=content,
+            emit_updates=emit_updates,
+        )
+        if in_temporal_workflow():
+            return await ActivityHelpers.execute_activity(
+                activity_name=MessagesActivityName.CREATE_MESSAGE,
+                request=params,
+                response_type=TaskMessage,
+                start_to_close_timeout=start_to_close_timeout,
+                retry_policy=retry_policy,
+                heartbeat_timeout=heartbeat_timeout,
+            )
+        else:
+            return await self._messages_service.create_message(
+                task_id=task_id,
+                content=content,
+                emit_updates=emit_updates,
+            )
+
+    async def update(
+        self,
+        task_id: str,
+        message_id: str,
+        content: TaskMessageContent,
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+        start_to_close_timeout: timedelta = timedelta(seconds=5),
+        heartbeat_timeout: timedelta = timedelta(seconds=5),
+        retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
+    ) -> TaskMessage:
+        """
+        Update a message for a task.
+
+        Args:
+            task_id (str): The ID of the task.
+            message_id (str): The ID of the message.
+            message (TaskMessage): The message to update.
+            start_to_close_timeout (timedelta): The start to close timeout.
+            heartbeat_timeout (timedelta): The heartbeat timeout.
+            retry_policy (RetryPolicy): The retry policy.
+
+        Returns:
+            TaskMessageEntity: The updated message.
+        """
+        params = UpdateMessageParams(
+            task_id=task_id,
+            message_id=message_id,
+            content=content,
+            trace_id=trace_id,
+            parent_span_id=parent_span_id,
+        )
+        if in_temporal_workflow():
+            return await ActivityHelpers.execute_activity(
+                activity_name=MessagesActivityName.UPDATE_MESSAGE,
+                request=params,
+                response_type=TaskMessage,
+                start_to_close_timeout=start_to_close_timeout,
+                retry_policy=retry_policy,
+                heartbeat_timeout=heartbeat_timeout,
+            )
+        else:
+            return await self._messages_service.update_message(
+                task_id=task_id,
+                message_id=message_id,
+                content=content,
+            )
+
+    async def create_batch(
+        self,
+        task_id: str,
+        contents: list[TaskMessageContent],
+        emit_updates: bool = True,
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+        start_to_close_timeout: timedelta = timedelta(seconds=5),
+        heartbeat_timeout: timedelta = timedelta(seconds=5),
+        retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
+    ) -> list[TaskMessage]:
+        """
+        Create a batch of messages for a task.
+
+        Args:
+            task_id (str): The ID of the task.
+            messages (List[TaskMessage]): The messages to create.
+            start_to_close_timeout (timedelta): The start to close timeout.
+            heartbeat_timeout (timedelta): The heartbeat timeout.
+            retry_policy (RetryPolicy): The retry policy.
+
+        Returns:
+            List[TaskMessageEntity]: The created messages.
+        """
+        params = CreateMessagesBatchParams(
+            task_id=task_id,
+            contents=contents,
+            emit_updates=emit_updates,
+            trace_id=trace_id,
+            parent_span_id=parent_span_id,
+        )
+        if in_temporal_workflow():
+            return await ActivityHelpers.execute_activity(
+                activity_name=MessagesActivityName.CREATE_MESSAGES_BATCH,
+                request=params,
+                response_type=list[TaskMessage],
+                start_to_close_timeout=start_to_close_timeout,
+                retry_policy=retry_policy,
+                heartbeat_timeout=heartbeat_timeout,
+            )
+        else:
+            return await self._messages_service.create_messages_batch(
+                task_id=task_id,
+                contents=contents,
+                emit_updates=emit_updates,
+            )
+
+    async def update_batch(
+        self,
+        task_id: str,
+        updates: dict[str, TaskMessageContent],
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+        start_to_close_timeout: timedelta = timedelta(seconds=5),
+        heartbeat_timeout: timedelta = timedelta(seconds=5),
+        retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
+    ) -> list[TaskMessage]:
+        """
+        Update a batch of messages for a task.
+
+        Args:
+            task_id (str): The ID of the task.
+            updates (Dict[str, TaskMessage]): The updates to apply to the messages.
+            start_to_close_timeout (timedelta): The start to close timeout.
+            heartbeat_timeout (timedelta): The heartbeat timeout.
+            retry_policy (RetryPolicy): The retry policy.
+
+        Returns:
+            List[TaskMessageEntity]: The updated messages.
+        """
+        params = UpdateMessagesBatchParams(
+            task_id=task_id,
+            updates=updates,
+            trace_id=trace_id,
+            parent_span_id=parent_span_id,
+        )
+        if in_temporal_workflow():
+            return await ActivityHelpers.execute_activity(
+                activity_name=MessagesActivityName.UPDATE_MESSAGES_BATCH,
+                request=params,
+                response_type=list[TaskMessage],
+                start_to_close_timeout=start_to_close_timeout,
+                retry_policy=retry_policy,
+                heartbeat_timeout=heartbeat_timeout,
+            )
+        else:
+            return await self._messages_service.update_messages_batch(
+                task_id=task_id,
+                updates=updates,
+            )
+
+    async def list(
+        self,
+        task_id: str,
+        limit: int | None = None,
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+        start_to_close_timeout: timedelta = timedelta(seconds=5),
+        heartbeat_timeout: timedelta = timedelta(seconds=5),
+        retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
+    ) -> list[TaskMessage]:
+        """
+        List messages for a task.
+
+        Args:
+            task_id (str): The ID of the task.
+            limit (Optional[int]): The maximum number of messages to return.
+            start_to_close_timeout (timedelta): The start to close timeout.
+            heartbeat_timeout (timedelta): The heartbeat timeout.
+            retry_policy (RetryPolicy): The retry policy.
+
+        Returns:
+            List[TaskMessageEntity]: The list of messages.
+        """
+        params = ListMessagesParams(
+            task_id=task_id,
+            limit=limit,
+            trace_id=trace_id,
+            parent_span_id=parent_span_id,
+        )
+        if in_temporal_workflow():
+            return await ActivityHelpers.execute_activity(
+                activity_name=MessagesActivityName.LIST_MESSAGES,
+                request=params,
+                response_type=list[TaskMessage],
+                start_to_close_timeout=start_to_close_timeout,
+                retry_policy=retry_policy,
+                heartbeat_timeout=heartbeat_timeout,
+            )
+        else:
+            return await self._messages_service.list_messages(
+                task_id=task_id,
+                limit=limit,
+            )

agentex/lib/adk/_modules/state.py

@@ -0,0 +1,291 @@
+from datetime import timedelta
+from typing import Any
+
+from pydantic import BaseModel
+from temporalio.common import RetryPolicy
+
+from agentex import AsyncAgentex
+from agentex.lib.core.services.adk.state import StateService
+from agentex.lib.core.temporal.activities.activity_helpers import ActivityHelpers
+from agentex.lib.core.temporal.activities.adk.state_activities import (
+    CreateStateParams,
+    DeleteStateParams,
+    GetStateParams,
+    StateActivityName,
+    UpdateStateParams,
+)
+from agentex.lib.core.tracing.tracer import AsyncTracer
+from agentex.types.state import State
+from agentex.lib.utils.logging import make_logger
+from agentex.lib.utils.temporal import in_temporal_workflow
+
+logger = make_logger(__name__)
+
+# Default retry policy for all state operations
+DEFAULT_RETRY_POLICY = RetryPolicy(maximum_attempts=1)
+
+
+class StateModule:
+    """
+    Module for managing task state in Agentex.
+    Provides high-level async methods for creating, retrieving, updating, and deleting state.
+    """
+
+    def __init__(
+        self,
+        state_service: StateService | None = None,
+    ):
+        if state_service is None:
+            agentex_client = AsyncAgentex()
+            tracer = AsyncTracer(agentex_client)
+            self._state_service = StateService(
+                agentex_client=agentex_client, tracer=tracer
+            )
+        else:
+            self._state_service = state_service
+
+    async def create(
+        self,
+        task_id: str,
+        agent_id: str,
+        state: dict[str, Any] | BaseModel,
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+        start_to_close_timeout: timedelta = timedelta(seconds=5),
+        heartbeat_timeout: timedelta = timedelta(seconds=5),
+        retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
+    ) -> State:
+        """
+        Create a new state for a task and agent.
+
+        Args:
+            task_id (str): The ID of the task.
+            agent_id (str): The ID of the agent.
+            state (Dict[str, Any]): The state to create.
+            trace_id (Optional[str]): The trace ID for tracing.
+            parent_span_id (Optional[str]): The parent span ID for tracing.
+            start_to_close_timeout (timedelta): The start to close timeout.
+            heartbeat_timeout (timedelta): The heartbeat timeout.
+            retry_policy (RetryPolicy): The retry policy.
+
+        Returns:
+            State: The created state.
+        """
+        state_dict = state.model_dump() if isinstance(state, BaseModel) else state
+        params = CreateStateParams(
+            task_id=task_id,
+            agent_id=agent_id,
+            state=state_dict,
+            trace_id=trace_id,
+            parent_span_id=parent_span_id,
+        )
+        if in_temporal_workflow():
+            return await ActivityHelpers.execute_activity(
+                activity_name=StateActivityName.CREATE_STATE,
+                request=params,
+                response_type=State,
+                start_to_close_timeout=start_to_close_timeout,
+                retry_policy=retry_policy,
+                heartbeat_timeout=heartbeat_timeout,
+            )
+        else:
+            return await self._state_service.create_state(
+                task_id=task_id,
+                agent_id=agent_id,
+                state=state_dict,
+                trace_id=trace_id,
+                parent_span_id=parent_span_id,
+            )
+
+    async def get(
+        self,
+        state_id: str,
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+        start_to_close_timeout: timedelta = timedelta(seconds=5),
+        heartbeat_timeout: timedelta = timedelta(seconds=5),
+        retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
+    ) -> State | None:
+        """
+        Get a state by ID.
+
+        Args:
+            state_id (str): The ID of the state.
+            trace_id (Optional[str]): The trace ID for tracing.
+            parent_span_id (Optional[str]): The parent span ID for tracing.
+            start_to_close_timeout (timedelta): The start to close timeout.
+            heartbeat_timeout (timedelta): The heartbeat timeout.
+            retry_policy (RetryPolicy): The retry policy.
+
+        Returns:
+            Optional[State]: The state if found, None otherwise.
+        """
+        params = GetStateParams(
+            state_id=state_id,
+            trace_id=trace_id,
+            parent_span_id=parent_span_id,
+        )
+        if in_temporal_workflow():
+            return await ActivityHelpers.execute_activity(
+                activity_name=StateActivityName.GET_STATE,
+                request=params,
+                response_type=State,
+                start_to_close_timeout=start_to_close_timeout,
+                retry_policy=retry_policy,
+                heartbeat_timeout=heartbeat_timeout,
+            )
+        else:
+            return await self._state_service.get_state(
+                state_id=state_id,
+                trace_id=trace_id,
+                parent_span_id=parent_span_id,
+            )
+
+    async def get_by_task_and_agent(
+        self,
+        task_id: str,
+        agent_id: str,
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+        start_to_close_timeout: timedelta = timedelta(seconds=5),
+        heartbeat_timeout: timedelta = timedelta(seconds=5),
+        retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
+    ) -> State | None:
+        """
+        Get a state by task and agent ID. A state is uniquely identified by task and the agent that created it.
+
+        Args:
+            task_id (str): The ID of the task.
+            agent_id (str): The ID of the agent.
+            trace_id (Optional[str]): The trace ID for tracing.
+            parent_span_id (Optional[str]): The parent span ID for tracing.
+            start_to_close_timeout (timedelta): The start to close timeout.
+            heartbeat_timeout (timedelta): The heartbeat timeout.
+            retry_policy (RetryPolicy): The retry policy.
+
+        Returns:
+            Optional[State]: The state if found, None otherwise.
+        """
+        params = GetStateParams(
+            task_id=task_id,
+            agent_id=agent_id,
+            trace_id=trace_id,
+            parent_span_id=parent_span_id,
+        )
+        if in_temporal_workflow():
+            return await ActivityHelpers.execute_activity(
+                activity_name=StateActivityName.GET_STATE,
+                request=params,
+                response_type=State,
+                start_to_close_timeout=start_to_close_timeout,
+                retry_policy=retry_policy,
+                heartbeat_timeout=heartbeat_timeout,
+            )
+        else:
+            return await self._state_service.get_state(
+                task_id=task_id,
+                agent_id=agent_id,
+                trace_id=trace_id,
+                parent_span_id=parent_span_id,
+            )
+
+    async def update(
+        self,
+        state_id: str,
+        task_id: str,
+        agent_id: str,
+        state: dict[str, Any] | BaseModel,
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+        start_to_close_timeout: timedelta = timedelta(seconds=5),
+        heartbeat_timeout: timedelta = timedelta(seconds=5),
+        retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
+    ) -> State:
+        """
+        Update a state by ID.
+
+        Args:
+            state_id (str): The ID of the state.
+            task_id (str): The ID of the task.
+            agent_id (str): The ID of the agent.
+            state (Dict[str, Any]): The state to update.
+            trace_id (Optional[str]): The trace ID for tracing.
+            parent_span_id (Optional[str]): The parent span ID for tracing.
+            start_to_close_timeout (timedelta): The start to close timeout.
+            heartbeat_timeout (timedelta): The heartbeat timeout.
+            retry_policy (RetryPolicy): The retry policy.
+
+        Returns:
+            State: The updated state.
+        """
+        state_dict = state.model_dump() if isinstance(state, BaseModel) else state
+        params = UpdateStateParams(
+            state_id=state_id,
+            task_id=task_id,
+            agent_id=agent_id,
+            state=state_dict,
+            trace_id=trace_id,
+            parent_span_id=parent_span_id,
+        )
+        if in_temporal_workflow():
+            return await ActivityHelpers.execute_activity(
+                activity_name=StateActivityName.UPDATE_STATE,
+                request=params,
+                response_type=State,
+                start_to_close_timeout=start_to_close_timeout,
+                retry_policy=retry_policy,
+                heartbeat_timeout=heartbeat_timeout,
+            )
+        else:
+            return await self._state_service.update_state(
+                state_id=state_id,
+                task_id=task_id,
+                agent_id=agent_id,
+                state=state_dict,
+                trace_id=trace_id,
+                parent_span_id=parent_span_id,
+            )
+
+    async def delete(
+        self,
+        state_id: str,
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+        start_to_close_timeout: timedelta = timedelta(seconds=5),
+        heartbeat_timeout: timedelta = timedelta(seconds=5),
+        retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
+    ) -> State:
+        """
+        Delete a state by ID.
+
+        Args:
+            state_id (str): The ID of the state.
+            trace_id (Optional[str]): The trace ID for tracing.
+            parent_span_id (Optional[str]): The parent span ID for tracing.
+            start_to_close_timeout (timedelta): The start to close timeout.
+            heartbeat_timeout (timedelta): The heartbeat timeout.
+            retry_policy (RetryPolicy): The retry policy.
+
+        Returns:
+            State: The deleted state.
+        """
+        params = DeleteStateParams(
+            state_id=state_id,
+            trace_id=trace_id,
+            parent_span_id=parent_span_id,
+        )
+        if in_temporal_workflow():
+            return await ActivityHelpers.execute_activity(
+                activity_name=StateActivityName.DELETE_STATE,
+                request=params,
+                response_type=State,
+                start_to_close_timeout=start_to_close_timeout,
+                retry_policy=retry_policy,
+                heartbeat_timeout=heartbeat_timeout,
+            )
+        else:
+            return await self._state_service.delete_state(
+                state_id=state_id,
+                trace_id=trace_id,
+                parent_span_id=parent_span_id,
+            )

agentex/lib/adk/_modules/streaming.py

@@ -0,0 +1,75 @@
+from temporalio.common import RetryPolicy
+
+from agentex import AsyncAgentex
+from agentex.lib.core.adapters.streams.adapter_redis import RedisEventStreamRepository
+from agentex.lib.core.services.adk.streaming import (
+    StreamingService,
+    StreamingTaskMessageContext,
+)
+from agentex.types.task_message_content import TaskMessageContent
+from agentex.lib.utils.logging import make_logger
+from agentex.lib.utils.temporal import in_temporal_workflow
+
+logger = make_logger(__name__)
+
+DEFAULT_RETRY_POLICY = RetryPolicy(maximum_attempts=1)
+
+
+class StreamingModule:
+    """
+    Module for streaming content to clients in Agentex.
+
+    This interface wraps around the StreamingService and provides a high-level API
+    for streaming events to clients, supporting both synchronous and asynchronous
+    (Temporal workflow) contexts.
+    """
+
+    def __init__(self, streaming_service: StreamingService | None = None):
+        """
+        Initialize the streaming interface.
+
+        Args:
+            streaming_service (Optional[StreamingService]): Optional StreamingService instance. If not provided,
+                a new service will be created with default parameters.
+        """
+        if streaming_service is None:
+            stream_repository = RedisEventStreamRepository()
+            agentex_client = AsyncAgentex()
+            self._streaming_service = StreamingService(
+                agentex_client=agentex_client,
+                stream_repository=stream_repository,
+            )
+        else:
+            self._streaming_service = streaming_service
+
+    def streaming_task_message_context(
+        self,
+        task_id: str,
+        initial_content: TaskMessageContent,
+    ) -> StreamingTaskMessageContext:
+        """
+        Create a streaming context for managing TaskMessage lifecycle.
+
+        This is a context manager that automatically creates a TaskMessage, sends START event,
+        and sends DONE event when the context exits. Perfect for simple streaming scenarios.
+
+        Args:
+            task_id: The ID of the task
+            initial_content: The initial content for the TaskMessage
+            agentex_client: The agentex client for creating/updating messages
+
+        Returns:
+            StreamingTaskMessageContext: Context manager for streaming operations
+        """
+        # Note: We don't support Temporal activities for streaming context methods yet
+        # since they involve complex state management across multiple activity calls
+        if in_temporal_workflow():
+            logger.warning(
+                "Streaming context methods are not yet supported in Temporal workflows. "
+                "You should wrap the entire streaming context in an activity. All nondeterministic network calls should be wrapped in an activity and generators cannot operate across activities and workflows."
+            )
+
+        return self._streaming_service.streaming_task_message_context(
+            task_id=task_id,
+            initial_content=initial_content,
+        )