agentex-sdk 0.1.0a6__py3-none-any.whl

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

@@ -0,0 +1,128 @@
+import asyncio
+import json
+import os
+from collections.abc import AsyncIterator
+from typing import Annotated, Any
+
+import redis.asyncio as redis
+from fastapi import Depends
+
+from agentex.lib.core.adapters.streams.port import EventStreamRepository
+from agentex.lib.utils.logging import make_logger
+
+logger = make_logger(__name__)
+
+
+class RedisEventStreamRepository(EventStreamRepository):
+    """
+    A simplified Redis implementation of the EventStreamRepository interface.
+    Optimized for text/JSON streaming with SSE.
+    """
+
+    def __init__(self, redis_url: str | None = None):
+        # Get Redis URL from environment if not provided
+        self.redis_url = redis_url or os.environ.get(
+            "REDIS_URL", "redis://localhost:6379"
+        )
+        self.redis = redis.from_url(self.redis_url)
+
+    async def send_event(self, topic: str, event: dict[str, Any]) -> str:
+        """
+        Send an event to a Redis stream.
+
+        Args:
+            topic: The stream topic/name
+            event: The event data (will be JSON serialized)
+
+        Returns:
+            The message ID from Redis
+        """
+        try:
+            # Simple JSON serialization
+            event_json = json.dumps(event)
+
+            # # Uncomment to debug
+            # logger.info(f"Sending event to Redis stream {topic}: {event_json}")
+
+            # Add to Redis stream with a reasonable max length
+            message_id = await self.redis.xadd(
+                name=topic,
+                fields={"data": event_json},
+            )
+
+            return message_id
+        except Exception as e:
+            logger.error(f"Error publishing to Redis stream {topic}: {e}")
+            raise
+
+    async def subscribe(
+        self, topic: str, last_id: str = "$"
+    ) -> AsyncIterator[dict[str, Any]]:
+        """
+        Subscribe to a Redis stream and yield events as they come in.
+
+        Args:
+            topic: The stream topic to subscribe to
+            last_id: Where to start reading from:
+                    "$" = only new messages (default)
+                    "0" = all messages from the beginning
+                    "<id>" = messages after the specified ID
+
+        Yields:
+            Parsed event data
+        """
+
+        current_id = last_id
+
+        while True:
+            try:
+                # Read new messages with a reasonable block time
+                streams = {topic: current_id}
+                response = await self.redis.xread(
+                    streams=streams,
+                    count=10,  # Get up to 10 messages at a time (reduces overprocessing)
+                    block=2000,  # Wait up to 2 seconds for new messages
+                )
+
+                if response:
+                    for _, messages in response:
+                        for message_id, fields in messages:
+                            # Update the last_id for next iteration
+                            current_id = message_id
+
+                            # Extract and parse the JSON data
+                            if b"data" in fields:
+                                try:
+                                    data_str = fields[b"data"].decode("utf-8")
+                                    event = json.loads(data_str)
+                                    yield event
+                                except Exception as e:
+                                    logger.warning(
+                                        f"Failed to parse event from Redis stream: {e}"
+                                    )
+
+                # Small sleep to prevent tight loops
+                await asyncio.sleep(0.01)
+
+            except Exception as e:
+                logger.error(f"Error reading from Redis stream: {e}")
+                await asyncio.sleep(1)  # Back off on errors
+
+    async def cleanup_stream(self, topic: str) -> None:
+        """
+        Clean up a Redis stream.
+
+        Args:
+            topic: The stream topic to clean up
+        """
+        try:
+            await self.redis.delete(topic)
+            logger.info(f"Cleaned up Redis stream: {topic}")
+        except Exception as e:
+            logger.error(f"Error cleaning up Redis stream {topic}: {e}")
+            raise
+
+
+DRedisEventStreamRepository = Annotated[
+    RedisEventStreamRepository | None, Depends(RedisEventStreamRepository)
+]

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

@@ -0,0 +1,50 @@
+from abc import ABC, abstractmethod
+from collections.abc import AsyncIterator
+from typing import Any
+
+
+class EventStreamRepository(ABC):
+    """
+    Interface for event streaming repositories.
+    Used to publish and subscribe to event streams.
+    """
+
+    @abstractmethod
+    async def send_event(self, topic: str, event: dict[str, Any]) -> str:
+        """
+        Send an event to a stream.
+
+        Args:
+            topic: The stream topic/name
+            event: The event data
+
+        Returns:
+            The message ID or other identifier
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def subscribe(
+        self, topic: str, last_id: str = "$"
+    ) -> AsyncIterator[dict[str, Any]]:
+        """
+        Subscribe to a stream and yield events as they come in.
+
+        Args:
+            topic: The stream topic to subscribe to
+            last_id: Where to start reading from
+
+        Yields:
+            Event data
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def cleanup_stream(self, topic: str) -> None:
+        """
+        Clean up a stream.
+
+        Args:
+            topic: The stream topic to clean up
+        """
+        raise NotImplementedError

agentex/lib/core/clients/__init__.py

@@ -0,0 +1 @@
+

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

@@ -0,0 +1,181 @@
+from collections.abc import Callable
+from datetime import timedelta
+from typing import Any
+
+from temporalio.client import Client, WorkflowExecutionStatus
+from temporalio.common import RetryPolicy as TemporalRetryPolicy
+from temporalio.common import WorkflowIDReusePolicy
+from temporalio.service import RPCError, RPCStatusCode
+
+from agentex.lib.core.clients.temporal.types import (
+    DuplicateWorkflowPolicy,
+    RetryPolicy,
+    TaskStatus,
+    WorkflowState,
+)
+from agentex.lib.core.clients.temporal.utils import get_temporal_client
+from agentex.lib.utils.logging import make_logger
+from agentex.lib.utils.model_utils import BaseModel
+
+logger = make_logger(__name__)
+
+DEFAULT_RETRY_POLICY = RetryPolicy(maximum_attempts=1)
+
+
+TEMPORAL_STATUS_TO_UPLOAD_STATUS_AND_REASON = {
+    # TODO: Support canceled status
+    WorkflowExecutionStatus.CANCELED: WorkflowState(
+        status=TaskStatus.CANCELED,
+        reason="Task canceled by the user.",
+        is_terminal=True,
+    ),
+    WorkflowExecutionStatus.COMPLETED: WorkflowState(
+        status=TaskStatus.COMPLETED,
+        reason="Task completed successfully.",
+        is_terminal=True,
+    ),
+    WorkflowExecutionStatus.FAILED: WorkflowState(
+        status=TaskStatus.FAILED,
+        reason="Task encountered terminal failure. "
+        "Please contact support if retrying does not resolve the issue.",
+        is_terminal=True,
+    ),
+    WorkflowExecutionStatus.RUNNING: WorkflowState(
+        status=TaskStatus.RUNNING,
+        reason="Task is running.",
+        is_terminal=False,
+    ),
+    WorkflowExecutionStatus.TERMINATED: WorkflowState(
+        status=TaskStatus.CANCELED,
+        reason="Task canceled by the user.",
+        is_terminal=True,
+    ),
+    WorkflowExecutionStatus.TIMED_OUT: WorkflowState(
+        status=TaskStatus.FAILED,
+        reason="Task timed out. Please contact support if retrying does not resolve the issue",
+        is_terminal=True,
+    ),
+    WorkflowExecutionStatus.CONTINUED_AS_NEW: WorkflowState(
+        status=TaskStatus.RUNNING,
+        reason="Task is running.",
+        is_terminal=False,
+    ),
+}
+
+DUPLICATE_POLICY_TO_ID_REUSE_POLICY = {
+    DuplicateWorkflowPolicy.ALLOW_DUPLICATE: WorkflowIDReusePolicy.ALLOW_DUPLICATE,
+    DuplicateWorkflowPolicy.ALLOW_DUPLICATE_FAILED_ONLY: WorkflowIDReusePolicy.ALLOW_DUPLICATE_FAILED_ONLY,
+    DuplicateWorkflowPolicy.REJECT_DUPLICATE: WorkflowIDReusePolicy.REJECT_DUPLICATE,
+    DuplicateWorkflowPolicy.TERMINATE_IF_RUNNING: WorkflowIDReusePolicy.TERMINATE_IF_RUNNING,
+}
+
+
+class TemporalClient:
+    def __init__(self, temporal_client: Client | None = None):
+        self._client: Client = temporal_client
+
+    @classmethod
+    async def create(cls, temporal_address: str):
+        if temporal_address in [
+            "false",
+            "False",
+            "null",
+            "None",
+            "",
+            "undefined",
+            False,
+            None,
+        ]:
+            _client = None
+        else:
+            _client = await get_temporal_client(temporal_address)
+        return cls(_client)
+
+    async def setup(self, temporal_address: str):
+        self._client = await self._get_temporal_client(
+            temporal_address=temporal_address
+        )
+
+    async def _get_temporal_client(self, temporal_address: str) -> Client:
+        if temporal_address in [
+            "false",
+            "False",
+            "null",
+            "None",
+            "",
+            "undefined",
+            False,
+            None,
+        ]:
+            return None
+        else:
+            return await get_temporal_client(temporal_address)
+
+    async def start_workflow(
+        self,
+        *args: Any,
+        duplicate_policy: DuplicateWorkflowPolicy = DuplicateWorkflowPolicy.ALLOW_DUPLICATE,
+        retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
+        task_timeout: timedelta = timedelta(seconds=10),
+        execution_timeout: timedelta = timedelta(seconds=86400),
+        **kwargs: Any,
+    ) -> str:
+        temporal_retry_policy = TemporalRetryPolicy(
+            **retry_policy.model_dump(exclude_unset=True)
+        )
+        workflow_handle = await self._client.start_workflow(
+            *args,
+            retry_policy=temporal_retry_policy,
+            task_timeout=task_timeout,
+            execution_timeout=execution_timeout,
+            id_reuse_policy=DUPLICATE_POLICY_TO_ID_REUSE_POLICY[duplicate_policy],
+            **kwargs,
+        )
+        return workflow_handle.id
+
+    async def send_signal(
+        self,
+        workflow_id: str,
+        signal: str | Callable[[dict[str, Any] | list[Any] | str | int | float | bool | BaseModel], Any],
+        payload: dict[str, Any] | list[Any] | str | int | float | bool | BaseModel,
+    ) -> None:
+        handle = self._client.get_workflow_handle(workflow_id=workflow_id)
+        await handle.signal(signal, payload)
+
+    async def query_workflow(
+        self,
+        workflow_id: str,
+        query: str | Callable[[dict[str, Any] | list[Any] | str | int | float | bool | BaseModel], Any],
+    ) -> Any:
+        """
+        Submit a query to a workflow by name and return the results.
+
+        Args:
+            workflow_id: The ID of the workflow to query
+            query: The name of the query or a callable query function
+
+        Returns:
+            The result of the query
+        """
+        handle = self._client.get_workflow_handle(workflow_id=workflow_id)
+        return await handle.query(query)
+
+    async def get_workflow_status(self, workflow_id: str) -> WorkflowState:
+        try:
+            handle = self._client.get_workflow_handle(workflow_id=workflow_id)
+            description = await handle.describe()
+            return TEMPORAL_STATUS_TO_UPLOAD_STATUS_AND_REASON[description.status]
+        except RPCError as e:
+            if e.status == RPCStatusCode.NOT_FOUND:
+                return WorkflowState(
+                    status="NOT_FOUND",
+                    reason="Workflow not found",
+                    is_terminal=True,
+                )
+            raise
+
+    async def terminate_workflow(self, workflow_id: str) -> None:
+        return await self._client.get_workflow_handle(workflow_id).terminate()
+
+    async def cancel_workflow(self, workflow_id: str) -> None:
+        return await self._client.get_workflow_handle(workflow_id).cancel()

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

@@ -0,0 +1,47 @@
+from datetime import timedelta
+from enum import Enum
+
+from pydantic import Field
+
+from agentex.lib.utils.model_utils import BaseModel
+
+
+class WorkflowState(BaseModel):
+    status: str
+    is_terminal: bool
+    reason: str | None = None
+
+
+class RetryPolicy(BaseModel):
+    initial_interval: timedelta = Field(
+        timedelta(seconds=1),
+        description="Backoff interval for the first retry. Default 1s.",
+    )
+    backoff_coefficient: float = Field(
+        2.0,
+        description="Coefficient to multiply previous backoff interval by to get new interval. Default 2.0.",
+    )
+    maximum_interval: timedelta | None = Field(
+        None,
+        description="Maximum backoff interval between retries. Default 100x :py:attr:`initial_interval`.",
+    )
+    maximum_attempts: int = Field(
+        0,
+        description="Maximum number of attempts. If 0, the default, there is no maximum.",
+    )
+
+
+class DuplicateWorkflowPolicy(str, Enum):
+    ALLOW_DUPLICATE = "ALLOW_DUPLICATE"
+    ALLOW_DUPLICATE_FAILED_ONLY = "ALLOW_DUPLICATE_FAILED_ONLY"
+    REJECT_DUPLICATE = "REJECT_DUPLICATE"
+    TERMINATE_IF_RUNNING = "TERMINATE_IF_RUNNING"
+
+
+class TaskStatus(str, Enum):
+    CANCELED = "CANCELED"
+    COMPLETED = "COMPLETED"
+    FAILED = "FAILED"
+    RUNNING = "RUNNING"
+    TERMINATED = "TERMINATED"
+    TIMED_OUT = "TIMED_OUT"

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

@@ -0,0 +1,56 @@
+from temporalio.client import Client
+from temporalio.contrib.pydantic import pydantic_data_converter
+from temporalio.runtime import OpenTelemetryConfig, Runtime, TelemetryConfig
+
+# class DateTimeJSONEncoder(AdvancedJSONEncoder):
+#     def default(self, o: Any) -> Any:
+#         if isinstance(o, datetime.datetime):
+#             return o.isoformat()
+#         return super().default(o)
+
+
+# class DateTimeJSONTypeConverter(JSONTypeConverter):
+#     def to_typed_value(
+#         self, hint: Type, value: Any
+#     ) -> Union[Optional[Any], _JSONTypeConverterUnhandled]:
+#         if hint == datetime.datetime:
+#             return datetime.datetime.fromisoformat(value)
+#         return JSONTypeConverter.Unhandled
+
+
+# class DateTimePayloadConverter(CompositePayloadConverter):
+#     def __init__(self) -> None:
+#         json_converter = JSONPlainPayloadConverter(
+#             encoder=DateTimeJSONEncoder,
+#             custom_type_converters=[DateTimeJSONTypeConverter()],
+#         )
+#         super().__init__(
+#             *[
+#                 c if not isinstance(c, JSONPlainPayloadConverter) else json_converter
+#                 for c in DefaultPayloadConverter.default_encoding_payload_converters
+#             ]
+#         )
+
+
+# custom_data_converter = dataclasses.replace(
+#     DataConverter.default,
+#     payload_converter_class=DateTimePayloadConverter,
+# )
+
+
+async def get_temporal_client(temporal_address: str, metrics_url: str = None) -> Client:
+    if not metrics_url:
+        client = await Client.connect(
+            target_host=temporal_address,
+            # data_converter=custom_data_converter,
+            data_converter=pydantic_data_converter,
+        )
+    else:
+        runtime = Runtime(telemetry=TelemetryConfig(metrics=OpenTelemetryConfig(url=metrics_url)))
+        client = await Client.connect(
+            target_host=temporal_address,
+            # data_converter=custom_data_converter,
+            data_converter=pydantic_data_converter,
+            runtime=runtime,
+        )
+    return client

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

@@ -0,0 +1,210 @@
+from typing import Any, cast
+
+from agentex import AsyncAgentex
+from agentex.lib.core.tracing.tracer import AsyncTracer
+from agentex.lib.utils.logging import make_logger
+from agentex.lib.utils.temporal import heartbeat_if_in_workflow
+from agentex.types.event import Event
+from agentex.types.task import Task
+from agentex.types.task_message import TaskMessage
+from agentex.types.task_message_content import TaskMessageContent
+from agentex.types.task_message_content_param import TaskMessageContentParam
+
+logger = make_logger(__name__)
+
+
+class ACPService:
+    def __init__(
+        self,
+        agentex_client: AsyncAgentex,
+        tracer: AsyncTracer,
+    ):
+        self._agentex_client = agentex_client
+        self._tracer = tracer
+
+    async def task_create(
+        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,
+    ) -> Task:
+        trace = self._tracer.trace(trace_id=trace_id)
+        async with trace.span(
+            parent_id=parent_span_id,
+            name="task_create",
+            input={
+                "name": name,
+                "agent_id": agent_id,
+                "agent_name": agent_name,
+                "params": params,
+            },
+        ) as span:
+            heartbeat_if_in_workflow("task create")
+            if agent_name:
+                json_rpc_response = await self._agentex_client.agents.rpc_by_name(
+                    agent_name=agent_name,
+                    method="task/create",
+                    params={
+                        "name": name,
+                        "params": params,
+                    },
+                )
+            elif agent_id:
+                json_rpc_response = await self._agentex_client.agents.rpc(
+                    agent_id=agent_id,
+                    method="task/create",
+                    params={
+                        "name": name,
+                        "params": params,
+                    },
+                )
+            else:
+                raise ValueError("Either agent_name or agent_id must be provided")
+
+            task_entry = Task.model_validate(json_rpc_response["result"])
+            if span:
+                span.output = task_entry.model_dump()
+            return task_entry
+
+    async def message_send(
+        self,
+        content: TaskMessageContent,
+        agent_id: str | None = None,
+        agent_name: str | None = None,
+        task_id: str | None = None,
+        task_name: str | None = None,
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+    ) -> TaskMessage:
+        trace = self._tracer.trace(trace_id=trace_id)
+        async with trace.span(
+            parent_id=parent_span_id,
+            name="message_send",
+            input={
+                "agent_id": agent_id,
+                "agent_name": agent_name,
+                "task_id": task_id,
+                "task_name": task_name,
+                "message": content,
+            },
+        ) as span:
+            heartbeat_if_in_workflow("message send")
+            if agent_name:
+                json_rpc_response = await self._agentex_client.agents.rpc_by_name(
+                    agent_name=agent_name,
+                    method="message/send",
+                    params={
+                        "task_id": task_id,
+                        "content": cast(TaskMessageContentParam, content.model_dump()),
+                        "stream": False,
+                    },
+                )
+            elif agent_id:
+                json_rpc_response = await self._agentex_client.agents.rpc(
+                    agent_id=agent_id,
+                    method="message/send",
+                    params={
+                        "task_id": task_id,
+                        "content": cast(TaskMessageContentParam, content.model_dump()),
+                        "stream": False,
+                    },
+                )
+            else:
+                raise ValueError("Either agent_name or agent_id must be provided")
+
+            task_message = TaskMessage.model_validate(json_rpc_response["result"])
+            if span:
+                span.output = task_message.model_dump()
+            return task_message
+
+    async def event_send(
+        self,
+        content: TaskMessageContent,
+        agent_id: str | None = None,
+        agent_name: str | None = None,
+        task_id: str | None = None,
+        task_name: str | None = None,
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+    ) -> Event:
+        trace = self._tracer.trace(trace_id=trace_id)
+        async with trace.span(
+            parent_id=parent_span_id,
+            name="event_send",
+            input={
+                "agent_id": agent_id,
+                "agent_name": agent_name,
+                "task_id": task_id,
+                "content": content,
+            },
+        ) as span:
+            heartbeat_if_in_workflow("event send")
+            if agent_name:
+                json_rpc_response = await self._agentex_client.agents.rpc_by_name(
+                    agent_name=agent_name,
+                    method="event/send",
+                    params={
+                        "task_name": task_name,
+                        "content": cast(TaskMessageContentParam, content.model_dump()),
+                    },
+                )
+            elif agent_id:
+                json_rpc_response = await self._agentex_client.agents.rpc(
+                    agent_id=agent_id,
+                    method="event/send",
+                    params={
+                        "task_id": task_id,
+                        "content": cast(TaskMessageContentParam, content.model_dump()),
+                    },
+                )
+            else:
+                raise ValueError("Either agent_name or agent_id must be provided")
+
+            event_entry = Event.model_validate(json_rpc_response["result"])
+            if span:
+                span.output = event_entry.model_dump()
+            return event_entry
+
+    async def task_cancel(
+        self,
+        task_id: str | None = None,
+        task_name: str | None = None,
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+    ) -> Task:
+        trace = self._tracer.trace(trace_id=trace_id)
+        async with trace.span(
+            parent_id=parent_span_id,
+            name="task_cancel",
+            input={
+                "task_id": task_id,
+                "task_name": task_name,
+            },
+        ) as span:
+            heartbeat_if_in_workflow("task cancel")
+            if task_name:
+                json_rpc_response = await self._agentex_client.agents.rpc_by_name(
+                    agent_name=task_name,
+                    method="task/cancel",
+                    params={
+                        "task_name": task_name,
+                    },
+                )
+            elif task_id:
+                json_rpc_response = await self._agentex_client.agents.rpc(
+                    agent_id=task_id,
+                    method="task/cancel",
+                    params={
+                        "task_id": task_id,
+                    },
+                )
+            else:
+                raise ValueError("Either task_name or task_id must be provided")
+
+            task_entry = Task.model_validate(json_rpc_response["result"])
+            if span:
+                span.output = task_entry.model_dump()
+            return task_entry