pygeai 0.1.6__py3-none-any.whl → 0.6.0b15__py3-none-any.whl

pygeai/vendor/a2a/server/request_handlers/jsonrpc_handler.py

@@ -0,0 +1,327 @@
+import logging
+
+from collections.abc import AsyncIterable
+
+from a2a.server.context import ServerCallContext
+from a2a.server.request_handlers.request_handler import RequestHandler
+from a2a.server.request_handlers.response_helpers import prepare_response_object
+from a2a.types import (
+    AgentCard,
+    CancelTaskRequest,
+    CancelTaskResponse,
+    CancelTaskSuccessResponse,
+    GetTaskPushNotificationConfigRequest,
+    GetTaskPushNotificationConfigResponse,
+    GetTaskPushNotificationConfigSuccessResponse,
+    GetTaskRequest,
+    GetTaskResponse,
+    GetTaskSuccessResponse,
+    InternalError,
+    JSONRPCErrorResponse,
+    Message,
+    SendMessageRequest,
+    SendMessageResponse,
+    SendMessageSuccessResponse,
+    SendStreamingMessageRequest,
+    SendStreamingMessageResponse,
+    SendStreamingMessageSuccessResponse,
+    SetTaskPushNotificationConfigRequest,
+    SetTaskPushNotificationConfigResponse,
+    SetTaskPushNotificationConfigSuccessResponse,
+    Task,
+    TaskArtifactUpdateEvent,
+    TaskNotFoundError,
+    TaskPushNotificationConfig,
+    TaskResubscriptionRequest,
+    TaskStatusUpdateEvent,
+)
+from a2a.utils.errors import ServerError
+from a2a.utils.helpers import validate
+from a2a.utils.telemetry import SpanKind, trace_class
+
+
+logger = logging.getLogger(__name__)
+
+
+@trace_class(kind=SpanKind.SERVER)
+class JSONRPCHandler:
+    """Maps incoming JSON-RPC requests to the appropriate request handler method and formats responses."""
+
+    def __init__(
+        self,
+        agent_card: AgentCard,
+        request_handler: RequestHandler,
+    ):
+        """Initializes the JSONRPCHandler.
+
+        Args:
+            agent_card: The AgentCard describing the agent's capabilities.
+            request_handler: The underlying `RequestHandler` instance to delegate requests to.
+        """
+        self.agent_card = agent_card
+        self.request_handler = request_handler
+
+    async def on_message_send(
+        self,
+        request: SendMessageRequest,
+        context: ServerCallContext | None = None,
+    ) -> SendMessageResponse:
+        """Handles the 'message/send' JSON-RPC method.
+
+        Args:
+            request: The incoming `SendMessageRequest` object.
+            context: Context provided by the server.
+
+        Returns:
+            A `SendMessageResponse` object containing the result (Task or Message)
+            or a JSON-RPC error response if a `ServerError` is raised by the handler.
+        """
+        # TODO: Wrap in error handler to return error states
+        try:
+            task_or_message = await self.request_handler.on_message_send(
+                request.params, context
+            )
+            return prepare_response_object(
+                request.id,
+                task_or_message,
+                (Task, Message),
+                SendMessageSuccessResponse,
+                SendMessageResponse,
+            )
+        except ServerError as e:
+            return SendMessageResponse(
+                root=JSONRPCErrorResponse(
+                    id=request.id, error=e.error if e.error else InternalError()
+                )
+            )
+
+    @validate(
+        lambda self: self.agent_card.capabilities.streaming,
+        'Streaming is not supported by the agent',
+    )
+    async def on_message_send_stream(
+        self,
+        request: SendStreamingMessageRequest,
+        context: ServerCallContext | None = None,
+    ) -> AsyncIterable[SendStreamingMessageResponse]:
+        """Handles the 'message/stream' JSON-RPC method.
+
+        Yields response objects as they are produced by the underlying handler's stream.
+
+        Args:
+            request: The incoming `SendStreamingMessageRequest` object.
+            context: Context provided by the server.
+
+        Yields:
+            `SendStreamingMessageResponse` objects containing streaming events
+            (Task, Message, TaskStatusUpdateEvent, TaskArtifactUpdateEvent)
+            or JSON-RPC error responses if a `ServerError` is raised.
+        """
+        try:
+            async for event in self.request_handler.on_message_send_stream(
+                request.params, context
+            ):
+                yield prepare_response_object(
+                    request.id,
+                    event,
+                    (
+                        Task,
+                        Message,
+                        TaskArtifactUpdateEvent,
+                        TaskStatusUpdateEvent,
+                    ),
+                    SendStreamingMessageSuccessResponse,
+                    SendStreamingMessageResponse,
+                )
+        except ServerError as e:
+            yield SendStreamingMessageResponse(
+                root=JSONRPCErrorResponse(
+                    id=request.id, error=e.error if e.error else InternalError()
+                )
+            )
+
+    async def on_cancel_task(
+        self,
+        request: CancelTaskRequest,
+        context: ServerCallContext | None = None,
+    ) -> CancelTaskResponse:
+        """Handles the 'tasks/cancel' JSON-RPC method.
+
+        Args:
+            request: The incoming `CancelTaskRequest` object.
+            context: Context provided by the server.
+
+        Returns:
+            A `CancelTaskResponse` object containing the updated Task or a JSON-RPC error.
+        """
+        try:
+            task = await self.request_handler.on_cancel_task(
+                request.params, context
+            )
+            if task:
+                return prepare_response_object(
+                    request.id,
+                    task,
+                    (Task,),
+                    CancelTaskSuccessResponse,
+                    CancelTaskResponse,
+                )
+            raise ServerError(error=TaskNotFoundError())
+        except ServerError as e:
+            return CancelTaskResponse(
+                root=JSONRPCErrorResponse(
+                    id=request.id, error=e.error if e.error else InternalError()
+                )
+            )
+
+    async def on_resubscribe_to_task(
+        self,
+        request: TaskResubscriptionRequest,
+        context: ServerCallContext | None = None,
+    ) -> AsyncIterable[SendStreamingMessageResponse]:
+        """Handles the 'tasks/resubscribe' JSON-RPC method.
+
+        Yields response objects as they are produced by the underlying handler's stream.
+
+        Args:
+            request: The incoming `TaskResubscriptionRequest` object.
+            context: Context provided by the server.
+
+        Yields:
+            `SendStreamingMessageResponse` objects containing streaming events
+            or JSON-RPC error responses if a `ServerError` is raised.
+        """
+        try:
+            async for event in self.request_handler.on_resubscribe_to_task(
+                request.params, context
+            ):
+                yield prepare_response_object(
+                    request.id,
+                    event,
+                    (
+                        Task,
+                        Message,
+                        TaskArtifactUpdateEvent,
+                        TaskStatusUpdateEvent,
+                    ),
+                    SendStreamingMessageSuccessResponse,
+                    SendStreamingMessageResponse,
+                )
+        except ServerError as e:
+            yield SendStreamingMessageResponse(
+                root=JSONRPCErrorResponse(
+                    id=request.id, error=e.error if e.error else InternalError()
+                )
+            )
+
+    async def get_push_notification(
+        self,
+        request: GetTaskPushNotificationConfigRequest,
+        context: ServerCallContext | None = None,
+    ) -> GetTaskPushNotificationConfigResponse:
+        """Handles the 'tasks/pushNotificationConfig/get' JSON-RPC method.
+
+        Args:
+            request: The incoming `GetTaskPushNotificationConfigRequest` object.
+            context: Context provided by the server.
+
+        Returns:
+            A `GetTaskPushNotificationConfigResponse` object containing the config or a JSON-RPC error.
+        """
+        try:
+            config = (
+                await self.request_handler.on_get_task_push_notification_config(
+                    request.params, context
+                )
+            )
+            return prepare_response_object(
+                request.id,
+                config,
+                (TaskPushNotificationConfig,),
+                GetTaskPushNotificationConfigSuccessResponse,
+                GetTaskPushNotificationConfigResponse,
+            )
+        except ServerError as e:
+            return GetTaskPushNotificationConfigResponse(
+                root=JSONRPCErrorResponse(
+                    id=request.id, error=e.error if e.error else InternalError()
+                )
+            )
+
+    @validate(
+        lambda self: self.agent_card.capabilities.pushNotifications,
+        'Push notifications are not supported by the agent',
+    )
+    async def set_push_notification(
+        self,
+        request: SetTaskPushNotificationConfigRequest,
+        context: ServerCallContext | None = None,
+    ) -> SetTaskPushNotificationConfigResponse:
+        """Handles the 'tasks/pushNotificationConfig/set' JSON-RPC method.
+
+        Requires the agent to support push notifications.
+
+        Args:
+            request: The incoming `SetTaskPushNotificationConfigRequest` object.
+            context: Context provided by the server.
+
+        Returns:
+            A `SetTaskPushNotificationConfigResponse` object containing the config or a JSON-RPC error.
+
+        Raises:
+            ServerError: If push notifications are not supported by the agent
+                (due to the `@validate` decorator).
+        """
+        try:
+            config = (
+                await self.request_handler.on_set_task_push_notification_config(
+                    request.params, context
+                )
+            )
+            return prepare_response_object(
+                request.id,
+                config,
+                (TaskPushNotificationConfig,),
+                SetTaskPushNotificationConfigSuccessResponse,
+                SetTaskPushNotificationConfigResponse,
+            )
+        except ServerError as e:
+            return SetTaskPushNotificationConfigResponse(
+                root=JSONRPCErrorResponse(
+                    id=request.id, error=e.error if e.error else InternalError()
+                )
+            )
+
+    async def on_get_task(
+        self,
+        request: GetTaskRequest,
+        context: ServerCallContext | None = None,
+    ) -> GetTaskResponse:
+        """Handles the 'tasks/get' JSON-RPC method.
+
+        Args:
+            request: The incoming `GetTaskRequest` object.
+            context: Context provided by the server.
+
+        Returns:
+            A `GetTaskResponse` object containing the Task or a JSON-RPC error.
+        """
+        try:
+            task = await self.request_handler.on_get_task(
+                request.params, context
+            )
+            if task:
+                return prepare_response_object(
+                    request.id,
+                    task,
+                    (Task,),
+                    GetTaskSuccessResponse,
+                    GetTaskResponse,
+                )
+            raise ServerError(error=TaskNotFoundError())
+        except ServerError as e:
+            return GetTaskResponse(
+                root=JSONRPCErrorResponse(
+                    id=request.id, error=e.error if e.error else InternalError()
+                )
+            )

pygeai/vendor/a2a/server/request_handlers/request_handler.py

@@ -0,0 +1,161 @@
+from abc import ABC, abstractmethod
+from collections.abc import AsyncGenerator
+
+from a2a.server.context import ServerCallContext
+from a2a.server.events.event_queue import Event
+from a2a.types import (
+    Message,
+    MessageSendParams,
+    Task,
+    TaskIdParams,
+    TaskPushNotificationConfig,
+    TaskQueryParams,
+    UnsupportedOperationError,
+)
+from a2a.utils.errors import ServerError
+
+
+class RequestHandler(ABC):
+    """A2A request handler interface.
+
+    This interface defines the methods that an A2A server implementation must
+    provide to handle incoming JSON-RPC requests.
+    """
+
+    @abstractmethod
+    async def on_get_task(
+        self,
+        params: TaskQueryParams,
+        context: ServerCallContext | None = None,
+    ) -> Task | None:
+        """Handles the 'tasks/get' method.
+
+        Retrieves the state and history of a specific task.
+
+        Args:
+            params: Parameters specifying the task ID and optionally history length.
+            context: Context provided by the server.
+
+        Returns:
+            The `Task` object if found, otherwise `None`.
+        """
+
+    @abstractmethod
+    async def on_cancel_task(
+        self,
+        params: TaskIdParams,
+        context: ServerCallContext | None = None,
+    ) -> Task | None:
+        """Handles the 'tasks/cancel' method.
+
+        Requests the agent to cancel an ongoing task.
+
+        Args:
+            params: Parameters specifying the task ID.
+            context: Context provided by the server.
+
+        Returns:
+            The `Task` object with its status updated to canceled, or `None` if the task was not found.
+        """
+
+    @abstractmethod
+    async def on_message_send(
+        self,
+        params: MessageSendParams,
+        context: ServerCallContext | None = None,
+    ) -> Task | Message:
+        """Handles the 'message/send' method (non-streaming).
+
+        Sends a message to the agent to create, continue, or restart a task,
+        and waits for the final result (Task or Message).
+
+        Args:
+            params: Parameters including the message and configuration.
+            context: Context provided by the server.
+
+        Returns:
+            The final `Task` object or a final `Message` object.
+        """
+
+    @abstractmethod
+    async def on_message_send_stream(
+        self,
+        params: MessageSendParams,
+        context: ServerCallContext | None = None,
+    ) -> AsyncGenerator[Event]:
+        """Handles the 'message/stream' method (streaming).
+
+        Sends a message to the agent and yields stream events as they are
+        produced (Task updates, Message chunks, Artifact updates).
+
+        Args:
+            params: Parameters including the message and configuration.
+            context: Context provided by the server.
+
+        Yields:
+            `Event` objects from the agent's execution.
+
+        Raises:
+             ServerError(UnsupportedOperationError): By default, if not implemented.
+        """
+        raise ServerError(error=UnsupportedOperationError())
+        yield
+
+    @abstractmethod
+    async def on_set_task_push_notification_config(
+        self,
+        params: TaskPushNotificationConfig,
+        context: ServerCallContext | None = None,
+    ) -> TaskPushNotificationConfig:
+        """Handles the 'tasks/pushNotificationConfig/set' method.
+
+        Sets or updates the push notification configuration for a task.
+
+        Args:
+            params: Parameters including the task ID and push notification configuration.
+            context: Context provided by the server.
+
+        Returns:
+            The provided `TaskPushNotificationConfig` upon success.
+        """
+
+    @abstractmethod
+    async def on_get_task_push_notification_config(
+        self,
+        params: TaskIdParams,
+        context: ServerCallContext | None = None,
+    ) -> TaskPushNotificationConfig:
+        """Handles the 'tasks/pushNotificationConfig/get' method.
+
+        Retrieves the current push notification configuration for a task.
+
+        Args:
+            params: Parameters including the task ID.
+            context: Context provided by the server.
+
+        Returns:
+            The `TaskPushNotificationConfig` for the task.
+        """
+
+    @abstractmethod
+    async def on_resubscribe_to_task(
+        self,
+        params: TaskIdParams,
+        context: ServerCallContext | None = None,
+    ) -> AsyncGenerator[Event]:
+        """Handles the 'tasks/resubscribe' method.
+
+        Allows a client to re-subscribe to a running streaming task's event stream.
+
+        Args:
+            params: Parameters including the task ID.
+            context: Context provided by the server.
+
+        Yields:
+             `Event` objects from the agent's ongoing execution for the specified task.
+
+        Raises:
+             ServerError(UnsupportedOperationError): By default, if not implemented.
+        """
+        raise ServerError(error=UnsupportedOperationError())
+        yield

pygeai/vendor/a2a/server/request_handlers/response_helpers.py

@@ -0,0 +1,133 @@
+"""Helper functions for building A2A JSON-RPC responses."""
+
+# response types
+from typing import TypeVar
+
+from a2a.types import (
+    A2AError,
+    CancelTaskResponse,
+    CancelTaskSuccessResponse,
+    GetTaskPushNotificationConfigResponse,
+    GetTaskPushNotificationConfigSuccessResponse,
+    GetTaskResponse,
+    GetTaskSuccessResponse,
+    InvalidAgentResponseError,
+    JSONRPCError,
+    JSONRPCErrorResponse,
+    Message,
+    SendMessageResponse,
+    SendMessageSuccessResponse,
+    SendStreamingMessageResponse,
+    SendStreamingMessageSuccessResponse,
+    SetTaskPushNotificationConfigResponse,
+    SetTaskPushNotificationConfigSuccessResponse,
+    Task,
+    TaskArtifactUpdateEvent,
+    TaskPushNotificationConfig,
+    TaskStatusUpdateEvent,
+)
+
+
+RT = TypeVar(
+    'RT',
+    GetTaskResponse,
+    CancelTaskResponse,
+    SendMessageResponse,
+    SetTaskPushNotificationConfigResponse,
+    GetTaskPushNotificationConfigResponse,
+    SendStreamingMessageResponse,
+)
+"""Type variable for RootModel response types."""
+
+# success types
+SPT = TypeVar(
+    'SPT',
+    GetTaskSuccessResponse,
+    CancelTaskSuccessResponse,
+    SendMessageSuccessResponse,
+    SetTaskPushNotificationConfigSuccessResponse,
+    GetTaskPushNotificationConfigSuccessResponse,
+    SendStreamingMessageSuccessResponse,
+)
+"""Type variable for SuccessResponse types."""
+
+# result types
+EventTypes = (
+    Task
+    | Message
+    | TaskArtifactUpdateEvent
+    | TaskStatusUpdateEvent
+    | TaskPushNotificationConfig
+    | A2AError
+    | JSONRPCError
+)
+"""Type alias for possible event types produced by handlers."""
+
+
+def build_error_response(
+    request_id: str | int | None,
+    error: A2AError | JSONRPCError,
+    response_wrapper_type: type[RT],
+) -> RT:
+    """Helper method to build a JSONRPCErrorResponse wrapped in the appropriate response type.
+
+    Args:
+        request_id: The ID of the request that caused the error.
+        error: The A2AError or JSONRPCError object.
+        response_wrapper_type: The Pydantic RootModel type that wraps the response
+                                for the specific RPC method (e.g., `SendMessageResponse`).
+
+    Returns:
+        A Pydantic model representing the JSON-RPC error response,
+        wrapped in the specified response type.
+    """
+    return response_wrapper_type(
+        JSONRPCErrorResponse(
+            id=request_id,
+            error=error.root if isinstance(error, A2AError) else error,
+        )
+    )
+
+
+def prepare_response_object(
+    request_id: str | int | None,
+    response: EventTypes,
+    success_response_types: tuple[type, ...],
+    success_payload_type: type[SPT],
+    response_type: type[RT],
+) -> RT:
+    """Helper method to build appropriate JSONRPCResponse object for RPC methods.
+
+    Based on the type of the `response` object received from the handler,
+    it constructs either a success response wrapped in the appropriate payload type
+    or an error response.
+
+    Args:
+        request_id: The ID of the request.
+        response: The object received from the request handler.
+        success_response_types: A tuple of expected Pydantic model types for a successful result.
+        success_payload_type: The Pydantic model type for the success payload
+                                (e.g., `SendMessageSuccessResponse`).
+        response_type: The Pydantic RootModel type that wraps the final response
+                       (e.g., `SendMessageResponse`).
+
+    Returns:
+        A Pydantic model representing the final JSON-RPC response (success or error).
+    """
+    if isinstance(response, success_response_types):
+        return response_type(
+            root=success_payload_type(id=request_id, result=response)  # type:ignore
+        )
+
+    if isinstance(response, A2AError | JSONRPCError):
+        return build_error_response(request_id, response, response_type)
+
+    # If consumer_data is not an expected success type and not an error,
+    # it's an invalid type of response from the agent for this specific method.
+    response = A2AError(
+        root=InvalidAgentResponseError(
+            message='Agent returned invalid type response for this method'
+        )
+    )
+
+    return build_error_response(request_id, response, response_type)

pygeai/vendor/a2a/server/tasks/__init__.py

@@ -0,0 +1,20 @@
+"""Components for managing tasks within the A2A server."""
+
+from a2a.server.tasks.inmemory_push_notifier import InMemoryPushNotifier
+from a2a.server.tasks.inmemory_task_store import InMemoryTaskStore
+from a2a.server.tasks.push_notifier import PushNotifier
+from a2a.server.tasks.result_aggregator import ResultAggregator
+from a2a.server.tasks.task_manager import TaskManager
+from a2a.server.tasks.task_store import TaskStore
+from a2a.server.tasks.task_updater import TaskUpdater
+
+
+__all__ = [
+    'InMemoryPushNotifier',
+    'InMemoryTaskStore',
+    'PushNotifier',
+    'ResultAggregator',
+    'TaskManager',
+    'TaskStore',
+    'TaskUpdater',
+]

pygeai/vendor/a2a/server/tasks/inmemory_push_notifier.py

@@ -0,0 +1,62 @@
+import asyncio
+import logging
+
+import httpx
+
+from a2a.server.tasks.push_notifier import PushNotifier
+from a2a.types import PushNotificationConfig, Task
+
+
+logger = logging.getLogger(__name__)
+
+
+class InMemoryPushNotifier(PushNotifier):
+    """In-memory implementation of PushNotifier interface.
+
+    Stores push notification configurations in memory and uses an httpx client
+    to send notifications.
+    """
+
+    def __init__(self, httpx_client: httpx.AsyncClient) -> None:
+        """Initializes the InMemoryPushNotifier.
+
+        Args:
+            httpx_client: An async HTTP client instance to send notifications.
+        """
+        self._client = httpx_client
+        self.lock = asyncio.Lock()
+        self._push_notification_infos: dict[str, PushNotificationConfig] = {}
+
+    async def set_info(
+        self, task_id: str, notification_config: PushNotificationConfig
+    ) -> None:
+        """Sets or updates the push notification configuration for a task in memory."""
+        async with self.lock:
+            self._push_notification_infos[task_id] = notification_config
+
+    async def get_info(self, task_id: str) -> PushNotificationConfig | None:
+        """Retrieves the push notification configuration for a task from memory."""
+        async with self.lock:
+            return self._push_notification_infos.get(task_id)
+
+    async def delete_info(self, task_id: str) -> None:
+        """Deletes the push notification configuration for a task from memory."""
+        async with self.lock:
+            if task_id in self._push_notification_infos:
+                del self._push_notification_infos[task_id]
+
+    async def send_notification(self, task: Task) -> None:
+        """Sends a push notification for a task if configuration exists."""
+        push_info = await self.get_info(task.id)
+        if not push_info:
+            return
+        url = push_info.url
+
+        try:
+            response = await self._client.post(
+                url, json=task.model_dump(mode='json', exclude_none=True)
+            )
+            response.raise_for_status()
+            logger.info(f'Push-notification sent for URL: {url}')
+        except Exception as e:
+            logger.error(f'Error sending push-notification: {e}')