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

pygeai/vendor/a2a/server/events/event_consumer.py

@@ -0,0 +1,149 @@
+import asyncio
+import logging
+import sys
+
+from collections.abc import AsyncGenerator
+
+from a2a.server.events.event_queue import Event, EventQueue
+from a2a.types import (
+    InternalError,
+    Message,
+    Task,
+    TaskState,
+    TaskStatusUpdateEvent,
+)
+from a2a.utils.errors import ServerError
+from a2a.utils.telemetry import SpanKind, trace_class
+
+
+# This is an alias to the exception for closed queue
+QueueClosed: type[Exception] = asyncio.QueueEmpty
+
+# When using python 3.13 or higher, the closed queue signal is QueueShutdown
+if sys.version_info >= (3, 13):
+    QueueClosed = asyncio.QueueShutDown
+
+logger = logging.getLogger(__name__)
+
+
+@trace_class(kind=SpanKind.SERVER)
+class EventConsumer:
+    """Consumer to read events from the agent event queue."""
+
+    def __init__(self, queue: EventQueue):
+        """Initializes the EventConsumer.
+
+        Args:
+            queue: The `EventQueue` instance to consume events from.
+        """
+        self.queue = queue
+        self._timeout = 0.5
+        self._exception: BaseException | None = None
+        logger.debug('EventConsumer initialized')
+
+    async def consume_one(self) -> Event:
+        """Consume one event from the agent event queue non-blocking.
+
+        Returns:
+            The next event from the queue.
+
+        Raises:
+            ServerError: If the queue is empty when attempting to dequeue
+                immediately.
+        """
+        logger.debug('Attempting to consume one event.')
+        try:
+            event = await self.queue.dequeue_event(no_wait=True)
+        except asyncio.QueueEmpty as e:
+            logger.warning('Event queue was empty in consume_one.')
+            raise ServerError(
+                InternalError(message='Agent did not return any response')
+            ) from e
+
+        logger.debug(f'Dequeued event of type: {type(event)} in consume_one.')
+
+        self.queue.task_done()
+
+        return event
+
+    async def consume_all(self) -> AsyncGenerator[Event]:
+        """Consume all the generated streaming events from the agent.
+
+        This method yields events as they become available from the queue
+        until a final event is received or the queue is closed. It also
+        monitors for exceptions set by the `agent_task_callback`.
+
+        Yields:
+            Events dequeued from the queue.
+
+        Raises:
+            BaseException: If an exception was set by the `agent_task_callback`.
+        """
+        logger.debug('Starting to consume all events from the queue.')
+        while True:
+            if self._exception:
+                raise self._exception
+            try:
+                # We use a timeout when waiting for an event from the queue.
+                # This is required because it allows the loop to check if
+                # `self._exception` has been set by the `agent_task_callback`.
+                # Without the timeout, loop might hang indefinitely if no events are
+                # enqueued by the agent and the agent simply threw an exception
+                event = await asyncio.wait_for(
+                    self.queue.dequeue_event(), timeout=self._timeout
+                )
+                logger.debug(
+                    f'Dequeued event of type: {type(event)} in consume_all.'
+                )
+                self.queue.task_done()
+                logger.debug(
+                    'Marked task as done in event queue in consume_all'
+                )
+
+                is_final_event = (
+                    (isinstance(event, TaskStatusUpdateEvent) and event.final)
+                    or isinstance(event, Message)
+                    or (
+                        isinstance(event, Task)
+                        and event.status.state
+                        in (
+                            TaskState.completed,
+                            TaskState.canceled,
+                            TaskState.failed,
+                            TaskState.rejected,
+                            TaskState.unknown,
+                        )
+                    )
+                )
+
+                # Make sure the yield is after the close events, otherwise
+                # the caller may end up in a blocked state where this
+                # generator isn't called again to close things out and the
+                # other part is waiting for an event or a closed queue.
+                if is_final_event:
+                    logger.debug('Stopping event consumption in consume_all.')
+                    await self.queue.close()
+                    yield event
+                    break
+                yield event
+            except TimeoutError:
+                # continue polling until there is a final event
+                continue
+            except QueueClosed:
+                # Confirm that the queue is closed, e.g. we aren't on
+                # python 3.12 and get a queue empty error on an open queue
+                if self.queue.is_closed():
+                    break
+
+    def agent_task_callback(self, agent_task: asyncio.Task[None]) -> None:
+        """Callback to handle exceptions from the agent's execution task.
+
+        If the agent's asyncio task raises an exception, this callback is
+        invoked, and the exception is stored to be re-raised by the consumer loop.
+
+        Args:
+            agent_task: The asyncio.Task that completed.
+        """
+        logger.debug('Agent task callback triggered.')
+        if agent_task.exception() is not None:
+            self._exception = agent_task.exception()

pygeai/vendor/a2a/server/events/event_queue.py

@@ -0,0 +1,156 @@
+import asyncio
+import logging
+import sys
+
+from a2a.types import (
+    Message,
+    Task,
+    TaskArtifactUpdateEvent,
+    TaskStatusUpdateEvent,
+)
+from a2a.utils.telemetry import SpanKind, trace_class
+
+
+logger = logging.getLogger(__name__)
+
+
+Event = Message | Task | TaskStatusUpdateEvent | TaskArtifactUpdateEvent
+"""Type alias for events that can be enqueued."""
+
+DEFAULT_MAX_QUEUE_SIZE = 1024
+
+
+@trace_class(kind=SpanKind.SERVER)
+class EventQueue:
+    """Event queue for A2A responses from agent.
+
+    Acts as a buffer between the agent's asynchronous execution and the
+    server's response handling (e.g., streaming via SSE). Supports tapping
+    to create child queues that receive the same events.
+    """
+
+    def __init__(self, max_queue_size: int = DEFAULT_MAX_QUEUE_SIZE) -> None:
+        """Initializes the EventQueue."""
+        # Make sure the `asyncio.Queue` is bounded.
+        # If it's unbounded (maxsize=0), then `queue.put()` never needs to wait,
+        # and so the streaming won't work correctly.
+        if max_queue_size <= 0:
+            raise ValueError('max_queue_size must be greater than 0')
+
+        self.queue: asyncio.Queue[Event] = asyncio.Queue(maxsize=max_queue_size)
+        self._children: list[EventQueue] = []
+        self._is_closed = False
+        self._lock = asyncio.Lock()
+        logger.debug('EventQueue initialized.')
+
+    async def enqueue_event(self, event: Event) -> None:
+        """Enqueues an event to this queue and all its children.
+
+        Args:
+            event: The event object to enqueue.
+        """
+        async with self._lock:
+            if self._is_closed:
+                logger.warning('Queue is closed. Event will not be enqueued.')
+                return
+
+        logger.debug(f'Enqueuing event of type: {type(event)}')
+
+        # Make sure to use put instead of put_nowait to avoid blocking the event loop.
+        await self.queue.put(event)
+        for child in self._children:
+            await child.enqueue_event(event)
+
+    async def dequeue_event(self, no_wait: bool = False) -> Event:
+        """Dequeues an event from the queue.
+
+        This implementation expects that dequeue to raise an exception when
+        the queue has been closed. In python 3.13+ this is naturally provided
+        by the QueueShutDown exception generated when the queue has closed and
+        the user is awaiting the queue.get method. Python<=3.12 this needs to
+        manage this lifecycle itself. The current implementation can lead to
+        blocking if the dequeue_event is called before the EventQueue has been
+        closed but when there are no events on the queue. Two ways to avoid this
+        are to call this with no_wait = True which won't block, but is the
+        callers responsibility to retry as appropriate. Alternatively, one can
+        use a async Task management solution to cancel the get task if the queue
+        has closed or some other condition is met. The implementation of the
+        EventConsumer uses an async.wait with a timeout to abort the
+        dequeue_event call and retry, when it will return with a closed error.
+
+        Args:
+            no_wait: If True, retrieve an event immediately or raise `asyncio.QueueEmpty`.
+                     If False (default), wait until an event is available.
+
+        Returns:
+            The next event from the queue.
+
+        Raises:
+            asyncio.QueueEmpty: If `no_wait` is True and the queue is empty.
+            asyncio.QueueShutDown: If the queue has been closed and is empty.
+        """
+        async with self._lock:
+            if self._is_closed and self.queue.empty():
+                logger.warning('Queue is closed. Event will not be dequeued.')
+                raise asyncio.QueueEmpty('Queue is closed.')
+
+        if no_wait:
+            logger.debug('Attempting to dequeue event (no_wait=True).')
+            event = self.queue.get_nowait()
+            logger.debug(
+                f'Dequeued event (no_wait=True) of type: {type(event)}'
+            )
+            return event
+
+        logger.debug('Attempting to dequeue event (waiting).')
+        event = await self.queue.get()
+        logger.debug(f'Dequeued event (waited) of type: {type(event)}')
+        return event
+
+    def task_done(self) -> None:
+        """Signals that a formerly enqueued task is complete.
+
+        Used in conjunction with `dequeue_event` to track processed items.
+        """
+        logger.debug('Marking task as done in EventQueue.')
+        self.queue.task_done()
+
+    def tap(self) -> 'EventQueue':
+        """Taps the event queue to create a new child queue that receives all future events.
+
+        Returns:
+            A new `EventQueue` instance that will receive all events enqueued
+            to this parent queue from this point forward.
+        """
+        logger.debug('Tapping EventQueue to create a child queue.')
+        queue = EventQueue()
+        self._children.append(queue)
+        return queue
+
+    async def close(self) -> None:
+        """Closes the queue for future push events.
+
+        Once closed, `dequeue_event` will eventually raise `asyncio.QueueShutDown`
+        when the queue is empty. Also closes all child queues.
+        """
+        logger.debug('Closing EventQueue.')
+        async with self._lock:
+            # If already closed, just return.
+            if self._is_closed:
+                return
+            self._is_closed = True
+        # If using python 3.13 or higher, use the shutdown method
+        if sys.version_info >= (3, 13):
+            self.queue.shutdown()
+            for child in self._children:
+                await child.close()
+        # Otherwise, join the queue
+        else:
+            tasks = [asyncio.create_task(self.queue.join())]
+            for child in self._children:
+                tasks.append(asyncio.create_task(child.close()))
+            await asyncio.wait(tasks, return_when=asyncio.ALL_COMPLETED)
+
+    def is_closed(self) -> bool:
+        """Checks if the queue is closed."""
+        return self._is_closed

pygeai/vendor/a2a/server/events/in_memory_queue_manager.py

@@ -0,0 +1,85 @@
+import asyncio
+
+from a2a.server.events.event_queue import EventQueue
+from a2a.server.events.queue_manager import (
+    NoTaskQueue,
+    QueueManager,
+    TaskQueueExists,
+)
+from a2a.utils.telemetry import SpanKind, trace_class
+
+
+@trace_class(kind=SpanKind.SERVER)
+class InMemoryQueueManager(QueueManager):
+    """InMemoryQueueManager is used for a single binary management.
+
+    This implements the `QueueManager` interface using in-memory storage for event
+    queues. It requires all incoming interactions for a given task ID to hit the
+    same binary instance.
+
+    This implementation is suitable for single-instance deployments but needs
+    a distributed approach for scalable deployments.
+    """
+
+    def __init__(self) -> None:
+        """Initializes the InMemoryQueueManager."""
+        self._task_queue: dict[str, EventQueue] = {}
+        self._lock = asyncio.Lock()
+
+    async def add(self, task_id: str, queue: EventQueue) -> None:
+        """Adds a new event queue for a task ID.
+
+        Raises:
+            TaskQueueExists: If a queue for the given `task_id` already exists.
+        """
+        async with self._lock:
+            if task_id in self._task_queue:
+                raise TaskQueueExists()
+            self._task_queue[task_id] = queue
+
+    async def get(self, task_id: str) -> EventQueue | None:
+        """Retrieves the event queue for a task ID.
+
+        Returns:
+            The `EventQueue` instance for the `task_id`, or `None` if not found.
+        """
+        async with self._lock:
+            if task_id not in self._task_queue:
+                return None
+            return self._task_queue[task_id]
+
+    async def tap(self, task_id: str) -> EventQueue | None:
+        """Taps the event queue for a task ID to create a child queue.
+
+        Returns:
+            A new child `EventQueue` instance, or `None` if the task ID is not found.
+        """
+        async with self._lock:
+            if task_id not in self._task_queue:
+                return None
+            return self._task_queue[task_id].tap()
+
+    async def close(self, task_id: str) -> None:
+        """Closes and removes the event queue for a task ID.
+
+        Raises:
+            NoTaskQueue: If no queue exists for the given `task_id`.
+        """
+        async with self._lock:
+            if task_id not in self._task_queue:
+                raise NoTaskQueue()
+            queue = self._task_queue.pop(task_id)
+            await queue.close()
+
+    async def create_or_tap(self, task_id: str) -> EventQueue:
+        """Creates a new event queue for a task ID if one doesn't exist, otherwise taps the existing one.
+
+        Returns:
+            A new or child `EventQueue` instance for the `task_id`.
+        """
+        async with self._lock:
+            if task_id not in self._task_queue:
+                queue = EventQueue()
+                self._task_queue[task_id] = queue
+                return queue
+            return self._task_queue[task_id].tap()

pygeai/vendor/a2a/server/events/queue_manager.py

@@ -0,0 +1,35 @@
+from abc import ABC, abstractmethod
+
+from a2a.server.events.event_queue import EventQueue
+
+
+class QueueManager(ABC):
+    """Interface for managing the event queue lifecycles per task."""
+
+    @abstractmethod
+    async def add(self, task_id: str, queue: EventQueue) -> None:
+        """Adds a new event queue associated with a task ID."""
+
+    @abstractmethod
+    async def get(self, task_id: str) -> EventQueue | None:
+        """Retrieves the event queue for a task ID."""
+
+    @abstractmethod
+    async def tap(self, task_id: str) -> EventQueue | None:
+        """Creates a child event queue (tap) for an existing task ID."""
+
+    @abstractmethod
+    async def close(self, task_id: str) -> None:
+        """Closes and removes the event queue for a task ID."""
+
+    @abstractmethod
+    async def create_or_tap(self, task_id: str) -> EventQueue:
+        """Creates a queue if one doesn't exist, otherwise taps the existing one."""
+
+
+class TaskQueueExists(Exception):  # noqa: N818
+    """Exception raised when attempting to add a queue for a task ID that already exists."""
+
+
+class NoTaskQueue(Exception):  # noqa: N818
+    """Exception raised when attempting to access or close a queue for a task ID that does not exist."""

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

@@ -0,0 +1,20 @@
+"""Request handler components for the A2A server."""
+
+from a2a.server.request_handlers.default_request_handler import (
+    DefaultRequestHandler,
+)
+from a2a.server.request_handlers.jsonrpc_handler import JSONRPCHandler
+from a2a.server.request_handlers.request_handler import RequestHandler
+from a2a.server.request_handlers.response_helpers import (
+    build_error_response,
+    prepare_response_object,
+)
+
+
+__all__ = [
+    'DefaultRequestHandler',
+    'JSONRPCHandler',
+    'RequestHandler',
+    'build_error_response',
+    'prepare_response_object',
+]