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

pygeai/vendor/a2a/server/agent_execution/context.py

@@ -0,0 +1,155 @@
+import uuid
+
+from a2a.server.context import ServerCallContext
+from a2a.types import (
+    InvalidParamsError,
+    Message,
+    MessageSendConfiguration,
+    MessageSendParams,
+    Task,
+)
+from a2a.utils import get_message_text
+from a2a.utils.errors import ServerError
+
+
+class RequestContext:
+    """Request Context.
+
+    Holds information about the current request being processed by the server,
+    including the incoming message, task and context identifiers, and related
+    tasks.
+    """
+
+    def __init__(  # noqa: PLR0913
+        self,
+        request: MessageSendParams | None = None,
+        task_id: str | None = None,
+        context_id: str | None = None,
+        task: Task | None = None,
+        related_tasks: list[Task] | None = None,
+        call_context: ServerCallContext | None = None,
+    ):
+        """Initializes the RequestContext.
+
+        Args:
+            request: The incoming `MessageSendParams` request payload.
+            task_id: The ID of the task explicitly provided in the request or path.
+            context_id: The ID of the context explicitly provided in the request or path.
+            task: The existing `Task` object retrieved from the store, if any.
+            related_tasks: A list of other tasks related to the current request (e.g., for tool use).
+            call_context: The server call context associated with this request.
+        """
+        if related_tasks is None:
+            related_tasks = []
+        self._params = request
+        self._task_id = task_id
+        self._context_id = context_id
+        self._current_task = task
+        self._related_tasks = related_tasks
+        self._call_context = call_context
+        # If the task id and context id were provided, make sure they
+        # match the request. Otherwise, create them
+        if self._params:
+            if task_id:
+                self._params.message.taskId = task_id
+                if task and task.id != task_id:
+                    raise ServerError(InvalidParamsError(message='bad task id'))
+            else:
+                self._check_or_generate_task_id()
+            if context_id:
+                self._params.message.contextId = context_id
+                if task and task.contextId != context_id:
+                    raise ServerError(
+                        InvalidParamsError(message='bad context id')
+                    )
+            else:
+                self._check_or_generate_context_id()
+
+    def get_user_input(self, delimiter: str = '\n') -> str:
+        """Extracts text content from the user's message parts.
+
+        Args:
+            delimiter: The string to use when joining multiple text parts.
+
+        Returns:
+            A single string containing all text content from the user message,
+            joined by the specified delimiter. Returns an empty string if no
+            user message is present or if it contains no text parts.
+        """
+        if not self._params:
+            return ''
+
+        return get_message_text(self._params.message, delimiter)
+
+    def attach_related_task(self, task: Task) -> None:
+        """Attaches a related task to the context.
+
+        This is useful for scenarios like tool execution where a new task
+        might be spawned.
+
+        Args:
+            task: The `Task` object to attach.
+        """
+        self._related_tasks.append(task)
+
+    @property
+    def message(self) -> Message | None:
+        """The incoming `Message` object from the request, if available."""
+        return self._params.message if self._params else None
+
+    @property
+    def related_tasks(self) -> list[Task]:
+        """A list of tasks related to the current request."""
+        return self._related_tasks
+
+    @property
+    def current_task(self) -> Task | None:
+        """The current `Task` object being processed."""
+        return self._current_task
+
+    @current_task.setter
+    def current_task(self, task: Task) -> None:
+        """Sets the current task object."""
+        self._current_task = task
+
+    @property
+    def task_id(self) -> str | None:
+        """The ID of the task associated with this context."""
+        return self._task_id
+
+    @property
+    def context_id(self) -> str | None:
+        """The ID of the conversation context associated with this task."""
+        return self._context_id
+
+    @property
+    def configuration(self) -> MessageSendConfiguration | None:
+        """The `MessageSendConfiguration` from the request, if available."""
+        if not self._params:
+            return None
+        return self._params.configuration
+
+    @property
+    def call_context(self) -> ServerCallContext | None:
+        """The server call context associated with this request."""
+        return self._call_context
+
+    def _check_or_generate_task_id(self) -> None:
+        """Ensures a task ID is present, generating one if necessary."""
+        if not self._params:
+            return
+
+        if not self._task_id and not self._params.message.taskId:
+            self._params.message.taskId = str(uuid.uuid4())
+        if self._params.message.taskId:
+            self._task_id = self._params.message.taskId
+
+    def _check_or_generate_context_id(self) -> None:
+        """Ensures a context ID is present, generating one if necessary."""
+        if not self._params:
+            return
+
+        if not self._context_id and not self._params.message.contextId:
+            self._params.message.contextId = str(uuid.uuid4())
+        if self._params.message.contextId:
+            self._context_id = self._params.message.contextId

pygeai/vendor/a2a/server/agent_execution/request_context_builder.py

@@ -0,0 +1,20 @@
+from abc import ABC, abstractmethod
+
+from a2a.server.agent_execution import RequestContext
+from a2a.server.context import ServerCallContext
+from a2a.types import MessageSendParams, Task
+
+
+class RequestContextBuilder(ABC):
+    """Builds request context to be supplied to agent executor."""
+
+    @abstractmethod
+    async def build(
+        self,
+        params: MessageSendParams | None = None,
+        task_id: str | None = None,
+        context_id: str | None = None,
+        task: Task | None = None,
+        context: ServerCallContext | None = None,
+    ) -> RequestContext:
+        pass

pygeai/vendor/a2a/server/agent_execution/simple_request_context_builder.py

@@ -0,0 +1,77 @@
+import asyncio
+
+from a2a.server.agent_execution import RequestContext, RequestContextBuilder
+from a2a.server.context import ServerCallContext
+from a2a.server.tasks import TaskStore
+from a2a.types import MessageSendParams, Task
+
+
+class SimpleRequestContextBuilder(RequestContextBuilder):
+    """Builds request context and populates referred tasks."""
+
+    def __init__(
+        self,
+        should_populate_referred_tasks: bool = False,
+        task_store: TaskStore | None = None,
+    ) -> None:
+        """Initializes the SimpleRequestContextBuilder.
+
+        Args:
+            should_populate_referred_tasks: If True, the builder will fetch tasks
+                referenced in `params.message.referenceTaskIds` and populate the
+                `related_tasks` field in the RequestContext. Defaults to False.
+            task_store: The TaskStore instance to use for fetching referred tasks.
+                Required if `should_populate_referred_tasks` is True.
+        """
+        self._task_store = task_store
+        self._should_populate_referred_tasks = should_populate_referred_tasks
+
+    async def build(
+        self,
+        params: MessageSendParams | None = None,
+        task_id: str | None = None,
+        context_id: str | None = None,
+        task: Task | None = None,
+        context: ServerCallContext | None = None,
+    ) -> RequestContext:
+        """Builds the request context for an agent execution.
+
+        This method assembles the RequestContext object. If the builder was
+        initialized with `should_populate_referred_tasks=True`, it fetches all tasks
+        referenced in `params.message.referenceTaskIds` from the `task_store`.
+
+        Args:
+            params: The parameters of the incoming message send request.
+            task_id: The ID of the task being executed.
+            context_id: The ID of the current execution context.
+            task: The primary task object associated with the request.
+            context: The server call context, containing metadata about the call.
+
+        Returns:
+            An instance of RequestContext populated with the provided information
+            and potentially a list of related tasks.
+        """
+        related_tasks: list[Task] | None = None
+
+        if (
+            self._task_store
+            and self._should_populate_referred_tasks
+            and params
+            and params.message.referenceTaskIds
+        ):
+            tasks = await asyncio.gather(
+                *[
+                    self._task_store.get(task_id)
+                    for task_id in params.message.referenceTaskIds
+                ]
+            )
+            related_tasks = [x for x in tasks if x is not None]
+
+        return RequestContext(
+            request=params,
+            task_id=task_id,
+            context_id=context_id,
+            task=task,
+            related_tasks=related_tasks,
+            call_context=context,
+        )

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

@@ -0,0 +1,16 @@
+"""HTTP application components for the A2A server."""
+
+from a2a.server.apps.jsonrpc import (
+    A2AFastAPIApplication,
+    A2AStarletteApplication,
+    CallContextBuilder,
+    JSONRPCApplication,
+)
+
+
+__all__ = [
+    'A2AFastAPIApplication',
+    'A2AStarletteApplication',
+    'CallContextBuilder',
+    'JSONRPCApplication',
+]

pygeai/vendor/a2a/server/apps/jsonrpc/__init__.py

@@ -0,0 +1,16 @@
+"""A2A JSON-RPC Applications."""
+
+from a2a.server.apps.jsonrpc.fastapi_app import A2AFastAPIApplication
+from a2a.server.apps.jsonrpc.jsonrpc_app import (
+    CallContextBuilder,
+    JSONRPCApplication,
+)
+from a2a.server.apps.jsonrpc.starlette_app import A2AStarletteApplication
+
+
+__all__ = [
+    'A2AFastAPIApplication',
+    'A2AStarletteApplication',
+    'CallContextBuilder',
+    'JSONRPCApplication',
+]

pygeai/vendor/a2a/server/apps/jsonrpc/fastapi_app.py

@@ -0,0 +1,88 @@
+import logging
+
+from typing import Any
+
+from fastapi import FastAPI, Request, Response
+
+from a2a.server.apps.jsonrpc.jsonrpc_app import (
+    CallContextBuilder,
+    JSONRPCApplication,
+)
+from a2a.server.request_handlers.jsonrpc_handler import RequestHandler
+from a2a.types import AgentCard
+
+
+logger = logging.getLogger(__name__)
+
+
+class A2AFastAPIApplication(JSONRPCApplication):
+    """A FastAPI application implementing the A2A protocol server endpoints.
+
+    Handles incoming JSON-RPC requests, routes them to the appropriate
+    handler methods, and manages response generation including Server-Sent Events
+    (SSE).
+    """
+
+    def __init__(
+        self,
+        agent_card: AgentCard,
+        http_handler: RequestHandler,
+        extended_agent_card: AgentCard | None = None,
+        context_builder: CallContextBuilder | None = None,
+    ):
+        """Initializes the A2AStarletteApplication.
+
+        Args:
+            agent_card: The AgentCard describing the agent's capabilities.
+            http_handler: The handler instance responsible for processing A2A
+              requests via http.
+            extended_agent_card: An optional, distinct AgentCard to be served
+              at the authenticated extended card endpoint.
+            context_builder: The CallContextBuilder used to construct the
+              ServerCallContext passed to the http_handler. If None, no
+              ServerCallContext is passed.
+        """
+        super().__init__(
+            agent_card=agent_card,
+            http_handler=http_handler,
+            extended_agent_card=extended_agent_card,
+            context_builder=context_builder,
+        )
+
+    def build(
+        self,
+        agent_card_url: str = '/.well-known/agent.json',
+        rpc_url: str = '/',
+        extended_agent_card_url: str = '/agent/authenticatedExtendedCard',
+        **kwargs: Any,
+    ) -> FastAPI:
+        """Builds and returns the FastAPI application instance.
+
+        Args:
+            agent_card_url: The URL for the agent card endpoint.
+            rpc_url: The URL for the A2A JSON-RPC endpoint.
+            extended_agent_card_url: The URL for the authenticated extended agent card endpoint.
+            **kwargs: Additional keyword arguments to pass to the FastAPI constructor.
+
+        Returns:
+            A configured FastAPI application instance.
+        """
+        app = FastAPI(**kwargs)
+
+        @app.post(rpc_url)
+        async def handle_a2a_request(request: Request) -> Response:
+            return await self._handle_requests(request)
+
+        @app.get(agent_card_url)
+        async def get_agent_card(request: Request) -> Response:
+            return await self._handle_get_agent_card(request)
+
+        if self.agent_card.supportsAuthenticatedExtendedCard:
+
+            @app.get(extended_agent_card_url)
+            async def get_extended_agent_card(request: Request) -> Response:
+                return await self._handle_get_authenticated_extended_agent_card(
+                    request
+                )
+
+        return app