agentex-sdk 0.1.0a6__py3-none-any.whl

agentex/lib/adk/_modules/acp.py

@@ -0,0 +1,247 @@
+from datetime import timedelta
+from typing import Any
+
+from temporalio.common import RetryPolicy
+
+from agentex import AsyncAgentex
+from agentex.lib.core.services.adk.acp.acp import ACPService
+from agentex.lib.core.temporal.activities.activity_helpers import ActivityHelpers
+from agentex.lib.core.temporal.activities.adk.acp.acp_activities import (
+    ACPActivityName,
+    EventSendParams,
+    MessageSendParams,
+    TaskCancelParams,
+    TaskCreateParams,
+)
+from agentex.lib.core.tracing.tracer import AsyncTracer
+from agentex.types.task_message import TaskMessage
+from agentex.types.task import Task
+from agentex.lib.utils.logging import make_logger
+from agentex.lib.utils.temporal import in_temporal_workflow
+from agentex.types.task_message_content import TaskMessageContent
+
+logger = make_logger(__name__)
+
+DEFAULT_RETRY_POLICY = RetryPolicy(maximum_attempts=1)
+
+
+class ACPModule:
+    """
+    Module for managing Agent to Client Protocol (ACP) agent operations in Agentex.
+
+    This interface provides high-level methods for interacting with the agent through the ACP.
+    """
+
+    def __init__(self, acp_service: ACPService | None = None):
+        """
+        Initialize the ACP module.
+
+        Args:
+            acp_activities (Optional[ACPActivities]): Optional pre-configured ACP activities. If None, will be auto-initialized.
+        """
+        if acp_service is None:
+            agentex_client = AsyncAgentex()
+            tracer = AsyncTracer(agentex_client)
+            self._acp_service = ACPService(agentex_client=agentex_client, tracer=tracer)
+        else:
+            self._acp_service = acp_service
+
+    async def create_task(
+        self,
+        name: str | None = None,
+        agent_id: str | None = None,
+        agent_name: str | None = None,
+        params: dict[str, Any] | 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,
+    ):
+        """
+        Create a new task.
+
+        Args:
+            name: The name of the task.
+            agent_id: The ID of the agent to create the task for.
+            agent_name: The name of the agent to create the task for.
+            params: The parameters for the task.
+            start_to_close_timeout: The start to close timeout for the task.
+            heartbeat_timeout: The heartbeat timeout for the task.
+            retry_policy: The retry policy for the task.
+
+        Returns:
+            The task entry.
+        """
+        if in_temporal_workflow():
+            return await ActivityHelpers.execute_activity(
+                activity_name=ACPActivityName.TASK_CREATE,
+                request=TaskCreateParams(
+                    name=name,
+                    agent_id=agent_id,
+                    agent_name=agent_name,
+                    params=params,
+                    trace_id=trace_id,
+                    parent_span_id=parent_span_id,
+                ),
+                response_type=Task,
+                start_to_close_timeout=start_to_close_timeout,
+                retry_policy=retry_policy,  
+                heartbeat_timeout=heartbeat_timeout,
+            )
+        else:
+            return await self._acp_service.task_create(
+                name=name,
+                agent_id=agent_id,
+                agent_name=agent_name,
+                params=params,
+                trace_id=trace_id,
+                parent_span_id=parent_span_id,
+            )
+
+    async def send_event(
+        self,
+        task_id: str,
+        content: TaskMessageContent,
+        agent_id: str | None = None,
+        agent_name: str | 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,
+    ):
+        """
+        Send an event to a task.
+
+        Args:
+            task_id: The ID of the task to send the event to.
+            data: The data to send to the event.
+            agent_id: The ID of the agent to send the event to.
+            agent_name: The name of the agent to send the event to.
+
+        Returns:
+            The event entry.
+        """
+        if in_temporal_workflow():
+            return await ActivityHelpers.execute_activity(
+                activity_name=ACPActivityName.EVENT_SEND,
+                request=EventSendParams(
+                    agent_id=agent_id,
+                    agent_name=agent_name,
+                    task_id=task_id,
+                    content=content,
+                    trace_id=trace_id,
+                    parent_span_id=parent_span_id,
+                ),
+                response_type=None,
+                start_to_close_timeout=start_to_close_timeout,
+                retry_policy=retry_policy,
+                heartbeat_timeout=heartbeat_timeout,
+            )
+        else:
+            return await self._acp_service.event_send(
+                agent_id=agent_id,
+                agent_name=agent_name,
+                task_id=task_id,
+                content=content,
+                trace_id=trace_id,
+                parent_span_id=parent_span_id,
+            )
+
+    async def send_message(
+        self,
+        task_id: str,
+        content: TaskMessageContent,
+        agent_id: str | None = None,
+        agent_name: str | 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,
+    ):
+        """
+        Send a message to a task.
+
+        Args:
+            task_id: The ID of the task to send the message to.
+            content: The task message content to send to the task.
+            agent_id: The ID of the agent to send the message to.
+            agent_name: The name of the agent to send the message to.
+
+        Returns:
+            The message entry.
+        """
+        if in_temporal_workflow():
+            return await ActivityHelpers.execute_activity(
+                activity_name=ACPActivityName.MESSAGE_SEND,
+                request=MessageSendParams(
+                    agent_id=agent_id,
+                    agent_name=agent_name,
+                    task_id=task_id,
+                    content=content,
+                    trace_id=trace_id,
+                    parent_span_id=parent_span_id,
+                ),
+                response_type=TaskMessage,
+                start_to_close_timeout=start_to_close_timeout,
+                retry_policy=retry_policy,
+                heartbeat_timeout=heartbeat_timeout,
+            )
+        else:
+            return await self._acp_service.message_send(
+                agent_id=agent_id,
+                agent_name=agent_name,
+                task_id=task_id,
+                content=content,
+                trace_id=trace_id,
+                parent_span_id=parent_span_id,
+            )
+
+    async def cancel_task(
+        self,
+        task_id: str | None = None,
+        task_name: str | 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,
+    ):
+        """
+        Cancel a task.
+
+        Args:
+            task_id: The ID of the task to cancel.
+            task_name: The name of the task to cancel.
+            trace_id: The trace ID for the task.
+            parent_span_id: The parent span ID for the task.
+            start_to_close_timeout: The start to close timeout for the task.
+            heartbeat_timeout: The heartbeat timeout for the task.
+            retry_policy: The retry policy for the task.
+
+        Returns:
+            The task entry.
+        """
+        if in_temporal_workflow():
+            return await ActivityHelpers.execute_activity(
+                activity_name=ACPActivityName.TASK_CANCEL,
+                request=TaskCancelParams(
+                    task_id=task_id,
+                    task_name=task_name,
+                    trace_id=trace_id,
+                    parent_span_id=parent_span_id,
+                ),
+                response_type=None,
+                start_to_close_timeout=start_to_close_timeout,
+                retry_policy=retry_policy,
+                heartbeat_timeout=heartbeat_timeout,
+            )
+        else:
+            return await self._acp_service.task_cancel(
+                task_id=task_id,
+                task_name=task_name,
+                trace_id=trace_id,
+                parent_span_id=parent_span_id,
+            )

agentex/lib/adk/_modules/agent_task_tracker.py

@@ -0,0 +1,176 @@
+from datetime import timedelta
+
+from temporalio.common import RetryPolicy
+
+from agentex import AsyncAgentex
+from agentex.lib.core.services.adk.agent_task_tracker import AgentTaskTrackerService
+from agentex.lib.core.temporal.activities.activity_helpers import ActivityHelpers
+from agentex.lib.core.temporal.activities.adk.agent_task_tracker_activities import (
+    AgentTaskTrackerActivityName,
+    GetAgentTaskTrackerByTaskAndAgentParams,
+    GetAgentTaskTrackerParams,
+    UpdateAgentTaskTrackerParams,
+)
+from agentex.lib.core.tracing.tracer import AsyncTracer
+from agentex.types.agent_task_tracker import AgentTaskTracker
+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 agent task tracker operations
+DEFAULT_RETRY_POLICY = RetryPolicy(maximum_attempts=1)
+
+
+class AgentTaskTrackerModule:
+    """
+    Module for managing agent task trackers in Agentex.
+    Provides high-level async methods for retrieving, filtering, and updating agent task trackers.
+    """
+
+    def __init__(
+        self,
+        agent_task_tracker_service: AgentTaskTrackerService | None = None,
+    ):
+        if agent_task_tracker_service is None:
+            agentex_client = AsyncAgentex()
+            tracer = AsyncTracer(agentex_client)
+            self._agent_task_tracker_service = AgentTaskTrackerService(
+                agentex_client=agentex_client, tracer=tracer
+            )
+        else:
+            self._agent_task_tracker_service = agent_task_tracker_service
+
+    async def get(
+        self,
+        tracker_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,
+    ) -> AgentTaskTracker:
+        """
+        Get an agent task tracker by ID.
+
+        Args:
+            tracker_id (str): The ID of the tracker.
+            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:
+            AgentTaskTracker: The agent task tracker.
+        """
+        params = GetAgentTaskTrackerParams(
+            tracker_id=tracker_id,
+            trace_id=trace_id,
+            parent_span_id=parent_span_id,
+        )
+        if in_temporal_workflow():
+            return await ActivityHelpers.execute_activity(
+                activity_name=AgentTaskTrackerActivityName.GET_AGENT_TASK_TRACKER,
+                request=params,
+                response_type=AgentTaskTracker,
+                start_to_close_timeout=start_to_close_timeout,
+                retry_policy=retry_policy,
+                heartbeat_timeout=heartbeat_timeout,
+            )
+        else:
+            return await self._agent_task_tracker_service.get_agent_task_tracker(
+                tracker_id=tracker_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,
+    ) -> AgentTaskTracker | None:
+        """
+        Get an agent task tracker by task ID and agent ID.
+        """
+        params = GetAgentTaskTrackerByTaskAndAgentParams(
+            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=AgentTaskTrackerActivityName.GET_AGENT_TASK_TRACKER_BY_TASK_AND_AGENT,
+                request=params,
+                response_type=AgentTaskTracker,
+                start_to_close_timeout=start_to_close_timeout,
+                retry_policy=retry_policy,
+                heartbeat_timeout=heartbeat_timeout,
+            )
+        else:
+            return await self._agent_task_tracker_service.get_by_task_and_agent(
+                task_id=task_id,
+                agent_id=agent_id,
+                trace_id=trace_id,
+                parent_span_id=parent_span_id,
+            )
+
+    async def update(
+        self,
+        tracker_id: str,
+        last_processed_event_id: str | None = None,
+        status: str | None = None,
+        status_reason: str | 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,
+    ) -> AgentTaskTracker:
+        """
+        Update an agent task tracker.
+
+        Args:
+            tracker_id (str): The ID of the tracker to update.
+            request (UpdateAgentTaskTrackerRequest): The update request containing the new values.
+            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:
+            AgentTaskTracker: The updated agent task tracker.
+        """
+        params = UpdateAgentTaskTrackerParams(
+            tracker_id=tracker_id,
+            last_processed_event_id=last_processed_event_id,
+            status=status,
+            status_reason=status_reason,
+            trace_id=trace_id,
+            parent_span_id=parent_span_id,
+        )
+        if in_temporal_workflow():
+            return await ActivityHelpers.execute_activity(
+                activity_name=AgentTaskTrackerActivityName.UPDATE_AGENT_TASK_TRACKER,
+                request=params,
+                response_type=AgentTaskTracker,
+                start_to_close_timeout=start_to_close_timeout,
+                retry_policy=retry_policy,
+                heartbeat_timeout=heartbeat_timeout,
+            )
+        else:
+            return await self._agent_task_tracker_service.update_agent_task_tracker(
+                tracker_id=tracker_id,
+                last_processed_event_id=last_processed_event_id,
+                status=status,
+                status_reason=status_reason,
+                trace_id=trace_id,
+                parent_span_id=parent_span_id,
+            )

agentex/lib/adk/_modules/agents.py

@@ -0,0 +1,77 @@
+from datetime import timedelta
+from typing import Optional
+
+from agentex.lib.core.temporal.activities.adk.agents_activities import AgentsActivityName, GetAgentParams
+from temporalio.common import RetryPolicy
+
+from agentex import AsyncAgentex
+from agentex.lib.core.services.adk.agents import AgentsService
+from agentex.lib.core.temporal.activities.activity_helpers import ActivityHelpers
+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 in_temporal_workflow
+
+logger = make_logger(__name__)
+
+DEFAULT_RETRY_POLICY = RetryPolicy(maximum_attempts=1)
+
+
+class AgentsModule:
+    """
+    Module for managing agents in Agentex.
+    Provides high-level async methods for retrieving, listing, and deleting agents.
+    """
+
+    def __init__(
+        self,
+        agents_service: Optional[AgentsService] = None,
+    ):
+        if agents_service is None:
+            agentex_client = AsyncAgentex()
+            tracer = AsyncTracer(agentex_client)
+            self._agents_service = AgentsService(agentex_client=agentex_client, tracer=tracer)
+        else:
+            self._agents_service = agents_service
+
+    async def get(
+        self,
+        *,
+        agent_id: Optional[str] = None,
+        agent_name: Optional[str] = None,
+        trace_id: Optional[str] = None,
+        parent_span_id: Optional[str] = None,
+        start_to_close_timeout: timedelta = timedelta(seconds=5),
+        heartbeat_timeout: timedelta = timedelta(seconds=5),
+        retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
+    ) -> Agent:
+        """
+        Get an agent by ID or name.
+        Args:
+            agent_id: The ID of the agent to retrieve.
+            agent_name: The name of the agent to retrieve.
+        Returns:
+            The agent entry.
+        """
+        params = GetAgentParams(
+            agent_id=agent_id,
+            agent_name=agent_name,
+            trace_id=trace_id,
+            parent_span_id=parent_span_id,
+        )
+        if in_temporal_workflow():
+            return await ActivityHelpers.execute_activity(
+                activity_name=AgentsActivityName.GET_AGENT,
+                request=params,
+                response_type=Agent,
+                start_to_close_timeout=start_to_close_timeout,
+                retry_policy=retry_policy,
+                heartbeat_timeout=heartbeat_timeout,
+            )
+        else:
+            return await self._agents_service.get_agent(
+                agent_id=agent_id,
+                agent_name=agent_name,
+                trace_id=trace_id,
+                parent_span_id=parent_span_id,
+            )

agentex/lib/adk/_modules/events.py

@@ -0,0 +1,141 @@
+from datetime import timedelta
+
+from temporalio.common import RetryPolicy
+
+from agentex import AsyncAgentex
+from agentex.lib.core.services.adk.events import EventsService
+from agentex.lib.core.temporal.activities.activity_helpers import ActivityHelpers
+from agentex.lib.core.temporal.activities.adk.events_activities import (
+    EventsActivityName,
+    GetEventParams,
+    ListEventsParams,
+)
+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.utils.temporal import in_temporal_workflow
+
+logger = make_logger(__name__)
+
+# Default retry policy for all events operations
+DEFAULT_RETRY_POLICY = RetryPolicy(maximum_attempts=1)
+
+
+class EventsModule:
+    """
+    Module for managing events in Agentex.
+    Provides high-level async methods for retrieving and listing events.
+    """
+
+    def __init__(
+        self,
+        events_service: EventsService | None = None,
+    ):
+        if events_service is None:
+            agentex_client = AsyncAgentex()
+            tracer = AsyncTracer(agentex_client)
+            self._events_service = EventsService(
+                agentex_client=agentex_client, tracer=tracer
+            )
+        else:
+            self._events_service = events_service
+
+    async def get(
+        self,
+        event_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,
+    ) -> Event | None:
+        """
+        Get an event by ID.
+
+        Args:
+            event_id (str): The ID of the event.
+            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[Event]: The event if found, None otherwise.
+        """
+        params = GetEventParams(
+            event_id=event_id,
+            trace_id=trace_id,
+            parent_span_id=parent_span_id,
+        )
+        if in_temporal_workflow():
+            return await ActivityHelpers.execute_activity(
+                activity_name=EventsActivityName.GET_EVENT,
+                request=params,
+                response_type=Event,
+                start_to_close_timeout=start_to_close_timeout,
+                retry_policy=retry_policy,
+                heartbeat_timeout=heartbeat_timeout,
+            )
+        else:
+            return await self._events_service.get_event(
+                event_id=event_id,
+                trace_id=trace_id,
+                parent_span_id=parent_span_id,
+            )
+
+    async def list_events(
+        self,
+        task_id: str,
+        agent_id: str,
+        last_processed_event_id: str | None = None,
+        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[Event]:
+        """
+        List events for a specific task and agent.
+
+        Args:
+            task_id (str): The ID of the task.
+            agent_id (str): The ID of the agent.
+            last_processed_event_id (Optional[str]): Optional event ID to get events after this ID.
+            limit (Optional[int]): Optional limit on number of results.
+            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:
+            List[Event]: List of events ordered by sequence_id.
+        """
+        params = ListEventsParams(
+            task_id=task_id,
+            agent_id=agent_id,
+            last_processed_event_id=last_processed_event_id,
+            limit=limit,
+            trace_id=trace_id,
+            parent_span_id=parent_span_id,
+        )
+        if in_temporal_workflow():
+            return await ActivityHelpers.execute_activity(
+                activity_name=EventsActivityName.LIST_EVENTS,
+                request=params,
+                response_type=list[Event],
+                start_to_close_timeout=start_to_close_timeout,
+                retry_policy=retry_policy,
+                heartbeat_timeout=heartbeat_timeout,
+            )
+        else:
+            return await self._events_service.list_events(
+                task_id=task_id,
+                agent_id=agent_id,
+                last_processed_event_id=last_processed_event_id,
+                limit=limit,
+                trace_id=trace_id,
+                parent_span_id=parent_span_id,
+            )