pygeai 0.1.51b3__py3-none-any.whl → 0.6.0b15__py3-none-any.whl

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

@@ -0,0 +1,51 @@
+import asyncio
+import logging
+
+from a2a.server.tasks.task_store import TaskStore
+from a2a.types import Task
+
+
+logger = logging.getLogger(__name__)
+
+
+class InMemoryTaskStore(TaskStore):
+    """In-memory implementation of TaskStore.
+
+    Stores task objects in a dictionary in memory. Task data is lost when the
+    server process stops.
+    """
+
+    def __init__(self) -> None:
+        """Initializes the InMemoryTaskStore."""
+        logger.debug('Initializing InMemoryTaskStore')
+        self.tasks: dict[str, Task] = {}
+        self.lock = asyncio.Lock()
+
+    async def save(self, task: Task) -> None:
+        """Saves or updates a task in the in-memory store."""
+        async with self.lock:
+            self.tasks[task.id] = task
+            logger.debug('Task %s saved successfully.', task.id)
+
+    async def get(self, task_id: str) -> Task | None:
+        """Retrieves a task from the in-memory store by ID."""
+        async with self.lock:
+            logger.debug('Attempting to get task with id: %s', task_id)
+            task = self.tasks.get(task_id)
+            if task:
+                logger.debug('Task %s retrieved successfully.', task_id)
+            else:
+                logger.debug('Task %s not found in store.', task_id)
+            return task
+
+    async def delete(self, task_id: str) -> None:
+        """Deletes a task from the in-memory store by ID."""
+        async with self.lock:
+            logger.debug('Attempting to delete task with id: %s', task_id)
+            if task_id in self.tasks:
+                del self.tasks[task_id]
+                logger.debug('Task %s deleted successfully.', task_id)
+            else:
+                logger.warning(
+                    'Attempted to delete nonexistent task with id: %s', task_id
+                )

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

@@ -0,0 +1,25 @@
+from abc import ABC, abstractmethod
+
+from a2a.types import PushNotificationConfig, Task
+
+
+class PushNotifier(ABC):
+    """PushNotifier interface to store, retrieve push notification for tasks and send push notifications."""
+
+    @abstractmethod
+    async def set_info(
+        self, task_id: str, notification_config: PushNotificationConfig
+    ) -> None:
+        """Sets or updates the push notification configuration for a task."""
+
+    @abstractmethod
+    async def get_info(self, task_id: str) -> PushNotificationConfig | None:
+        """Retrieves the push notification configuration for a task."""
+
+    @abstractmethod
+    async def delete_info(self, task_id: str) -> None:
+        """Deletes the push notification configuration for a task."""
+
+    @abstractmethod
+    async def send_notification(self, task: Task) -> None:
+        """Sends a push notification containing the latest task state."""

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

@@ -0,0 +1,151 @@
+import asyncio
+import logging
+
+from collections.abc import AsyncGenerator, AsyncIterator
+
+from a2a.server.events import Event, EventConsumer
+from a2a.server.tasks.task_manager import TaskManager
+from a2a.types import Message, Task, TaskState, TaskStatusUpdateEvent
+
+
+logger = logging.getLogger(__name__)
+
+
+class ResultAggregator:
+    """ResultAggregator is used to process the event streams from an AgentExecutor.
+
+    There are three main ways to use the ResultAggregator:
+    1) As part of a processing pipe. consume_and_emit will construct the updated
+       task as the events arrive, and re-emit those events for another consumer
+    2) As part of a blocking call. consume_all will process the entire stream and
+       return the final Task or Message object
+    3) As part of a push solution where the latest Task is emitted after processing an event.
+       consume_and_emit_task will consume the Event stream, process the events to the current
+       Task object and emit that Task object.
+    """
+
+    def __init__(self, task_manager: TaskManager):
+        """Initializes the ResultAggregator.
+
+        Args:
+            task_manager: The `TaskManager` instance to use for processing events
+                          and managing the task state.
+        """
+        self.task_manager = task_manager
+        self._message: Message | None = None
+
+    @property
+    async def current_result(self) -> Task | Message | None:
+        """Returns the current aggregated result (Task or Message).
+
+        This is the latest state processed from the event stream.
+
+        Returns:
+            The current `Task` object managed by the `TaskManager`, or the final
+            `Message` if one was received, or `None` if no result has been produced yet.
+        """
+        if self._message:
+            return self._message
+        return await self.task_manager.get_task()
+
+    async def consume_and_emit(
+        self, consumer: EventConsumer
+    ) -> AsyncGenerator[Event]:
+        """Processes the event stream from the consumer, updates the task state, and re-emits the same events.
+
+        Useful for streaming scenarios where the server needs to observe and
+        process events (e.g., save task state, send push notifications) while
+        forwarding them to the client.
+
+        Args:
+            consumer: The `EventConsumer` to read events from.
+
+        Yields:
+            The `Event` objects consumed from the `EventConsumer`.
+        """
+        async for event in consumer.consume_all():
+            await self.task_manager.process(event)
+            yield event
+
+    async def consume_all(
+        self, consumer: EventConsumer
+    ) -> Task | Message | None:
+        """Processes the entire event stream from the consumer and returns the final result.
+
+        Blocks until the event stream ends (queue is closed after final event or exception).
+
+        Args:
+            consumer: The `EventConsumer` to read events from.
+
+        Returns:
+            The final `Task` object or `Message` object after the stream is exhausted.
+            Returns `None` if the stream ends without producing a final result.
+
+        Raises:
+            BaseException: If the `EventConsumer` raises an exception during consumption.
+        """
+        async for event in consumer.consume_all():
+            if isinstance(event, Message):
+                self._message = event
+                return event
+            await self.task_manager.process(event)
+        return await self.task_manager.get_task()
+
+    async def consume_and_break_on_interrupt(
+        self, consumer: EventConsumer
+    ) -> tuple[Task | Message | None, bool]:
+        """Processes the event stream until completion or an interruptable state is encountered.
+
+        Interruptable states currently include `TaskState.auth_required`.
+        If interrupted, consumption continues in a background task.
+
+        Args:
+            consumer: The `EventConsumer` to read events from.
+
+        Returns:
+            A tuple containing:
+            - The current aggregated result (`Task` or `Message`) at the point of completion or interruption.
+            - A boolean indicating whether the consumption was interrupted (`True`) or completed naturally (`False`).
+
+        Raises:
+            BaseException: If the `EventConsumer` raises an exception during consumption.
+        """
+        event_stream = consumer.consume_all()
+        interrupted = False
+        async for event in event_stream:
+            if isinstance(event, Message):
+                self._message = event
+                return event, False
+            await self.task_manager.process(event)
+            if (
+                isinstance(event, Task | TaskStatusUpdateEvent)
+                and event.status.state == TaskState.auth_required
+            ):
+                # auth-required is a special state: the message should be
+                # escalated back to the caller, but the agent is expected to
+                # continue producing events once the authorization is received
+                # out-of-band. This is in contrast to input-required, where a
+                # new request is expected in order for the agent to make progress,
+                # so the agent should exit.
+                logger.debug(
+                    'Encountered an auth-required task: breaking synchronous message/send flow.'
+                )
+                # TODO: We should track all outstanding tasks to ensure they eventually complete.
+                asyncio.create_task(self._continue_consuming(event_stream))  # noqa: RUF006
+                interrupted = True
+                break
+        return await self.task_manager.get_task(), interrupted
+
+    async def _continue_consuming(
+        self, event_stream: AsyncIterator[Event]
+    ) -> None:
+        """Continues processing an event stream in a background task.
+
+        Used after an interruptable state (like auth_required) is encountered
+        in the synchronous consumption flow.
+
+        Args:
+            event_stream: The remaining `AsyncIterator` of events from the consumer.
+        """
+        async for event in event_stream:
+            await self.task_manager.process(event)

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

@@ -0,0 +1,253 @@
+import logging
+
+from a2a.server.events.event_queue import Event
+from a2a.server.tasks.task_store import TaskStore
+from a2a.types import (
+    InvalidParamsError,
+    Message,
+    Task,
+    TaskArtifactUpdateEvent,
+    TaskState,
+    TaskStatus,
+    TaskStatusUpdateEvent,
+)
+from a2a.utils import append_artifact_to_task
+from a2a.utils.errors import ServerError
+
+
+logger = logging.getLogger(__name__)
+
+
+class TaskManager:
+    """Helps manage a task's lifecycle during execution of a request.
+
+    Responsible for retrieving, saving, and updating the `Task` object based on
+    events received from the agent.
+    """
+
+    def __init__(
+        self,
+        task_id: str | None,
+        context_id: str | None,
+        task_store: TaskStore,
+        initial_message: Message | None,
+    ):
+        """Initializes the TaskManager.
+
+        Args:
+            task_id: The ID of the task, if known from the request.
+            context_id: The ID of the context, if known from the request.
+            task_store: The `TaskStore` instance for persistence.
+            initial_message: The `Message` that initiated the task, if any.
+                             Used when creating a new task object.
+        """
+        self.task_id = task_id
+        self.context_id = context_id
+        self.task_store = task_store
+        self._initial_message = initial_message
+        self._current_task: Task | None = None
+        logger.debug(
+            'TaskManager initialized with task_id: %s, context_id: %s',
+            task_id,
+            context_id,
+        )
+
+    async def get_task(self) -> Task | None:
+        """Retrieves the current task object, either from memory or the store.
+
+        If `task_id` is set, it first checks the in-memory `_current_task`,
+        then attempts to load it from the `task_store`.
+
+        Returns:
+            The `Task` object if found, otherwise `None`.
+        """
+        if not self.task_id:
+            logger.debug('task_id is not set, cannot get task.')
+            return None
+
+        if self._current_task:
+            return self._current_task
+
+        logger.debug(
+            'Attempting to get task from store with id: %s', self.task_id
+        )
+        self._current_task = await self.task_store.get(self.task_id)
+        if self._current_task:
+            logger.debug('Task %s retrieved successfully.', self.task_id)
+        else:
+            logger.debug('Task %s not found.', self.task_id)
+        return self._current_task
+
+    async def save_task_event(
+        self, event: Task | TaskStatusUpdateEvent | TaskArtifactUpdateEvent
+    ) -> Task | None:
+        """Processes a task-related event (Task, Status, Artifact) and saves the updated task state.
+
+        Ensures task and context IDs match or are set from the event.
+
+        Args:
+            event: The task-related event (`Task`, `TaskStatusUpdateEvent`, or `TaskArtifactUpdateEvent`).
+
+        Returns:
+            The updated `Task` object after processing the event.
+
+        Raises:
+            ServerError: If the task ID in the event conflicts with the TaskManager's ID
+                         when the TaskManager's ID is already set.
+        """
+        task_id_from_event = (
+            event.id if isinstance(event, Task) else event.taskId
+        )
+        # If task id is known, make sure it is matched
+        if self.task_id and self.task_id != task_id_from_event:
+            raise ServerError(
+                error=InvalidParamsError(
+                    message=f"Task in event doesn't match TaskManager {self.task_id} : {task_id_from_event}"
+                )
+            )
+        if not self.task_id:
+            self.task_id = task_id_from_event
+        if not self.context_id and self.context_id != event.contextId:
+            self.context_id = event.contextId
+
+        logger.debug(
+            'Processing save of task event of type %s for task_id: %s',
+            type(event).__name__,
+            task_id_from_event,
+        )
+        if isinstance(event, Task):
+            await self._save_task(event)
+            return event
+
+        task: Task = await self.ensure_task(event)
+
+        if isinstance(event, TaskStatusUpdateEvent):
+            logger.debug(
+                'Updating task %s status to: %s', task.id, event.status.state
+            )
+            if task.status.message:
+                if not task.history:
+                    task.history = [task.status.message]
+                else:
+                    task.history.append(task.status.message)
+
+            task.status = event.status
+        else:
+            logger.debug('Appending artifact to task %s', task.id)
+            append_artifact_to_task(task, event)
+
+        await self._save_task(task)
+        return task
+
+    async def ensure_task(
+        self, event: TaskStatusUpdateEvent | TaskArtifactUpdateEvent
+    ) -> Task:
+        """Ensures a Task object exists in memory, loading from store or creating new if needed.
+
+        Args:
+            event: The task-related event triggering the need for a Task object.
+
+        Returns:
+            An existing or newly created `Task` object.
+        """
+        task: Task | None = self._current_task
+        if not task and self.task_id:
+            logger.debug(
+                'Attempting to retrieve existing task with id: %s', self.task_id
+            )
+            task = await self.task_store.get(self.task_id)
+
+        if not task:
+            logger.info(
+                'Task not found or task_id not set. Creating new task for event (task_id: %s, context_id: %s).',
+                event.taskId,
+                event.contextId,
+            )
+            # streaming agent did not previously stream task object.
+            # Create a task object with the available information and persist the event
+            task = self._init_task_obj(event.taskId, event.contextId)
+            await self._save_task(task)
+
+        return task
+
+    async def process(self, event: Event) -> Event:
+        """Processes an event, updates the task state if applicable, stores it, and returns the event.
+
+        If the event is task-related (`Task`, `TaskStatusUpdateEvent`, `TaskArtifactUpdateEvent`),
+        the internal task state is updated and persisted.
+
+        Args:
+            event: The event object received from the agent.
+
+        Returns:
+            The same event object that was processed.
+        """
+        if isinstance(
+            event, Task | TaskStatusUpdateEvent | TaskArtifactUpdateEvent
+        ):
+            await self.save_task_event(event)
+
+        return event
+
+    def _init_task_obj(self, task_id: str, context_id: str) -> Task:
+        """Initializes a new task object in memory.
+
+        Args:
+            task_id: The ID for the new task.
+            context_id: The context ID for the new task.
+
+        Returns:
+            A new `Task` object with initial status and potentially the initial message in history.
+        """
+        logger.debug(
+            'Initializing new Task object with task_id: %s, context_id: %s',
+            task_id,
+            context_id,
+        )
+        history = [self._initial_message] if self._initial_message else []
+        return Task(
+            id=task_id,
+            contextId=context_id,
+            status=TaskStatus(state=TaskState.submitted),
+            history=history,
+        )
+
+    async def _save_task(self, task: Task) -> None:
+        """Saves the given task to the task store and updates the in-memory `_current_task`.
+
+        Args:
+            task: The `Task` object to save.
+        """
+        logger.debug('Saving task with id: %s', task.id)
+        await self.task_store.save(task)
+        self._current_task = task
+        if not self.task_id:
+            logger.info('New task created with id: %s', task.id)
+            self.task_id = task.id
+            self.context_id = task.contextId
+
+    def update_with_message(self, message: Message, task: Task) -> Task:
+        """Updates a task object in memory by adding a new message to its history.
+
+        If the task has a message in its current status, that message is moved
+        to the history first.
+
+        Args:
+            message: The new `Message` to add to the history.
+            task: The `Task` object to update.
+
+        Returns:
+            The updated `Task` object (updated in-place).
+        """
+        if task.status.message:
+            if task.history:
+                task.history.append(task.status.message)
+            else:
+                task.history = [task.status.message]
+            task.status.message = None
+        if task.history:
+            task.history.append(message)
+        else:
+            task.history = [message]
+        self._current_task = task
+        return task

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

@@ -0,0 +1,22 @@
+from abc import ABC, abstractmethod
+
+from a2a.types import Task
+
+
+class TaskStore(ABC):
+    """Agent Task Store interface.
+
+    Defines the methods for persisting and retrieving `Task` objects.
+    """
+
+    @abstractmethod
+    async def save(self, task: Task) -> None:
+        """Saves or updates a task in the store."""
+
+    @abstractmethod
+    async def get(self, task_id: str) -> Task | None:
+        """Retrieves a task from the store by ID."""
+
+    @abstractmethod
+    async def delete(self, task_id: str) -> None:
+        """Deletes a task from the store by ID."""

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

@@ -0,0 +1,155 @@
+import uuid
+
+from datetime import datetime, timezone
+from typing import Any
+
+from a2a.server.events import EventQueue
+from a2a.types import (
+    Artifact,
+    Message,
+    Part,
+    Role,
+    TaskArtifactUpdateEvent,
+    TaskState,
+    TaskStatus,
+    TaskStatusUpdateEvent,
+)
+
+
+class TaskUpdater:
+    """Helper class for agents to publish updates to a task's event queue.
+
+    Simplifies the process of creating and enqueueing standard task events.
+    """
+
+    def __init__(self, event_queue: EventQueue, task_id: str, context_id: str):
+        """Initializes the TaskUpdater.
+
+        Args:
+            event_queue: The `EventQueue` associated with the task.
+            task_id: The ID of the task.
+            context_id: The context ID of the task.
+        """
+        self.event_queue = event_queue
+        self.task_id = task_id
+        self.context_id = context_id
+
+    async def update_status(
+        self,
+        state: TaskState,
+        message: Message | None = None,
+        final: bool = False,
+        timestamp: str | None = None,
+    ) -> None:
+        """Updates the status of the task and publishes a `TaskStatusUpdateEvent`.
+
+        Args:
+            state: The new state of the task.
+            message: An optional message associated with the status update.
+            final: If True, indicates this is the final status update for the task.
+            timestamp: Optional ISO 8601 datetime string. Defaults to current time.
+        """
+        current_timestamp = (
+            timestamp if timestamp else datetime.now(timezone.utc).isoformat()
+        )
+        await self.event_queue.enqueue_event(
+            TaskStatusUpdateEvent(
+                taskId=self.task_id,
+                contextId=self.context_id,
+                final=final,
+                status=TaskStatus(
+                    state=state,
+                    message=message,
+                    timestamp=current_timestamp,
+                ),
+            )
+        )
+
+    async def add_artifact(
+        self,
+        parts: list[Part],
+        artifact_id: str | None = None,
+        name: str | None = None,
+        metadata: dict[str, Any] | None = None,
+    ) -> None:
+        """Adds an artifact chunk to the task and publishes a `TaskArtifactUpdateEvent`.
+
+        Args:
+            parts: A list of `Part` objects forming the artifact chunk.
+            artifact_id: The ID of the artifact. A new UUID is generated if not provided.
+            name: Optional name for the artifact.
+            metadata: Optional metadata for the artifact.
+        """
+        if not artifact_id:
+            artifact_id = str(uuid.uuid4())
+
+        await self.event_queue.enqueue_event(
+            TaskArtifactUpdateEvent(
+                taskId=self.task_id,
+                contextId=self.context_id,
+                artifact=Artifact(
+                    artifactId=artifact_id,
+                    name=name,
+                    parts=parts,
+                    metadata=metadata,
+                ),
+            )
+        )
+
+    async def complete(self, message: Message | None = None) -> None:
+        """Marks the task as completed and publishes a final status update."""
+        await self.update_status(
+            TaskState.completed,
+            message=message,
+            final=True,
+        )
+
+    async def failed(self, message: Message | None = None) -> None:
+        """Marks the task as failed and publishes a final status update."""
+        await self.update_status(TaskState.failed, message=message, final=True)
+
+    async def reject(self, message: Message | None = None) -> None:
+        """Marks the task as rejected and publishes a final status update."""
+        await self.update_status(
+            TaskState.rejected, message=message, final=True
+        )
+
+    async def submit(self, message: Message | None = None) -> None:
+        """Marks the task as submitted and publishes a status update."""
+        await self.update_status(
+            TaskState.submitted,
+            message=message,
+        )
+
+    async def start_work(self, message: Message | None = None) -> None:
+        """Marks the task as working and publishes a status update."""
+        await self.update_status(
+            TaskState.working,
+            message=message,
+        )
+
+    def new_agent_message(
+        self,
+        parts: list[Part],
+        metadata: dict[str, Any] | None = None,
+    ) -> Message:
+        """Creates a new message object sent by the agent for this task/context.
+
+        Note: This method only *creates* the message object. It does not
+              automatically enqueue it.
+
+        Args:
+            parts: A list of `Part` objects for the message content.
+            metadata: Optional metadata for the message.
+
+        Returns:
+            A new `Message` object.
+        """
+        return Message(
+            role=Role.agent,
+            taskId=self.task_id,
+            contextId=self.context_id,
+            messageId=str(uuid.uuid4()),
+            metadata=metadata,
+            parts=parts,
+        )