Add A2A server (#1537)

Co-authored-by: David Montague <[email protected]>

Author: Kludex

.pre-commit-config.yaml

@@ -60,6 +60,6 @@ repos:
     rev: v2.3.0
     hooks:
     - id: codespell
-      args: ['--skip', 'tests/models/cassettes/*']
+      args: ['--skip', 'tests/models/cassettes/*,docs/a2a/fasta2a.md']
       additional_dependencies:
       - tomli

docs/a2a.md

@@ -0,0 +1,96 @@
+# Agent2Agent (A2A) Protocol
+
+The [Agent2Agent (A2A) Protocol](https://google.github.io/A2A/) is an open standard introduced by Google that enables
+communication and interoperability between AI agents, regardless of the framework or vendor they are built on.
+
+At Pydantic, we built the [FastA2A](#fasta2a) library to make it easier to implement the A2A protocol in Python.
+
+We also built a convenience method that expose PydanticAI agents as A2A servers - let's have a quick look at how to use it:
+
+```py {title="agent_to_a2a.py" hl_lines="4"}
+from pydantic_ai import Agent
+
+agent = Agent('openai:gpt-4.1', instructions='Be fun!')
+app = agent.to_a2a()
+```
+
+_You can run the example with `uvicorn agent_to_a2a:app --host 0.0.0.0 --port 8000`_
+
+This will expose the agent as an A2A server, and you can start sending requests to it.
+
+See more about [exposing PydanticAI agents as A2A servers](#pydanticai-agent-to-a2a-server).
+
+## FastA2A
+
+**FastA2A** is an agentic framework agnostic implementation of the A2A protocol in Python.
+The library is designed to be used with any agentic framework, and is **not exclusive to PydanticAI**.
+
+### Design
+
+**FastA2A** is built on top of [Starlette](https://starlette.io), which means it's fully compatible with any ASGI server.
+
+Given the nature of the A2A protocol, it's important to understand the design before using it, as a developer
+you'll need to provide some components:
+
+- [`Storage`][fasta2a.Storage]: to save and load tasks
+- [`Broker`][fasta2a.Broker]: to schedule tasks
+- [`Worker`][fasta2a.Worker]: to execute tasks
+
+Let's have a look at how those components fit together:
+
+```mermaid
+flowchart TB
+    Server["HTTP Server"] <--> |Sends Requests/<br>Receives Results| TM
+
+    subgraph CC[Core Components]
+        direction RL
+        TM["TaskManager<br>(coordinates)"] --> |Schedules Tasks| Broker
+        TM <--> Storage
+        Broker["Broker<br>(queues & schedules)"] <--> Storage["Storage<br>(persistence)"]
+        Broker --> |Delegates Execution| Worker
+    end
+
+    Worker["Worker<br>(implementation)"]
+```
+
+FastA2A allows you to bring your own [`Storage`][fasta2a.Storage], [`Broker`][fasta2a.Broker] and [`Worker`][fasta2a.Worker].
+
+
+### Installation
+
+FastA2A is available on PyPI as [`fasta2a`](https://pypi.org/project/fasta2a/) so installation is as simple as:
+
+```bash
+pip/uv-add fasta2a
+```
+
+The only dependencies are:
+
+- [starlette](https://starlette.io): to expose the A2A server as an [ASGI application](https://asgi.readthedocs.io/en/latest/)
+- [pydantic](https://pydantic.dev): to validate the request/response messages
+- [opentelemetry-api](https://opentelemetry-python.readthedocs.io/en/latest): to provide tracing capabilities
+
+You can install PydanticAI with the `a2a` extra to include **FastA2A**:
+
+```bash
+pip/uv-add 'pydantic-ai[a2a]'
+```
+
+### PydanticAI Agent to A2A Server
+
+To expose a PydanticAI agent as an A2A server, you can use the `to_a2a` method:
+
+```python {title="agent_to_a2a.py"}
+from pydantic_ai import Agent
+
+agent = Agent('openai:gpt-4.1', instructions='Be fun!')
+app = agent.to_a2a()
+```
+
+Since `app` is an ASGI application, it can be used with any ASGI server.
+
+```bash
+uvicorn agent_to_a2a:app --host 0.0.0.0 --port 8000
+```
+
+Since the goal of `to_a2a` is to be a convenience method, it accepts the same arguments as the [`FastA2A`][fasta2a.FastA2A] constructor.

docs/api/fasta2a.md

@@ -0,0 +1,7 @@
+# `fasta2a`
+
+::: fasta2a
+
+::: fasta2a.schema
+
+::: fasta2a.client

fasta2a/README.md

@@ -0,0 +1,12 @@
+# FastA2A
+
+[![CI](https://github.com/pydantic/pydantic-ai/actions/workflows/ci.yml/badge.svg?event=push)](https://github.com/pydantic/pydantic-ai/actions/workflows/ci.yml?query=branch%3Amain)
+[![Coverage](https://coverage-badge.samuelcolvin.workers.dev/pydantic/pydantic-ai.svg)](https://coverage-badge.samuelcolvin.workers.dev/redirect/pydantic/pydantic-ai)
+[![PyPI](https://img.shields.io/pypi/v/fasta2a.svg)](https://pypi.python.org/pypi/fasta2a)
+[![python versions](https://img.shields.io/pypi/pyversions/fasta2a.svg)](https://github.com/pydantic/pydantic-ai)
+[![license](https://img.shields.io/github/license/pydantic/pydantic-ai.svg)](https://github.com/pydantic/pydantic-ai/blob/main/LICENSE)
+
+To make it easier to implement A2A servers, we've implemented FastA2A,
+a library built on top of Starlette and Pydantic to bring A2A to Python.
+
+See [the docs](https://ai.pydantic.dev/a2a/) for more information.

fasta2a/fasta2a/__init__.py

@@ -0,0 +1,7 @@
+from .applications import FastA2A
+from .broker import Broker
+from .schema import Skill
+from .storage import Storage
+from .worker import Worker
+
+__all__ = ['FastA2A', 'Skill', 'Storage', 'Broker', 'Worker']

fasta2a/fasta2a/applications.py

@@ -0,0 +1,124 @@
+from __future__ import annotations as _annotations
+
+from collections.abc import AsyncIterator, Sequence
+from contextlib import asynccontextmanager
+from typing import Any
+
+from starlette.applications import Starlette
+from starlette.middleware import Middleware
+from starlette.requests import Request
+from starlette.responses import Response
+from starlette.routing import Route
+from starlette.types import ExceptionHandler, Lifespan, Receive, Scope, Send
+
+from .broker import Broker
+from .schema import AgentCard, Provider, Skill, a2a_request_ta, a2a_response_ta, agent_card_ta
+from .storage import Storage
+from .task_manager import TaskManager
+
+
+class FastA2A(Starlette):
+    """The main class for the FastA2A library."""
+
+    def __init__(
+        self,
+        *,
+        storage: Storage,
+        broker: Broker,
+        # Agent card
+        name: str | None = None,
+        url: str = 'http://localhost:8000',
+        version: str = '1.0.0',
+        description: str | None = None,
+        provider: Provider | None = None,
+        skills: list[Skill] | None = None,
+        # Starlette
+        debug: bool = False,
+        routes: Sequence[Route] | None = None,
+        middleware: Sequence[Middleware] | None = None,
+        exception_handlers: dict[Any, ExceptionHandler] | None = None,
+        lifespan: Lifespan[FastA2A] | None = None,
+    ):
+        if lifespan is None:
+            lifespan = _default_lifespan
+
+        super().__init__(
+            debug=debug,
+            routes=routes,
+            middleware=middleware,
+            exception_handlers=exception_handlers,
+            lifespan=lifespan,
+        )
+
+        self.name = name or 'Agent'
+        self.url = url
+        self.version = version
+        self.description = description
+        self.provider = provider
+        self.skills = skills or []
+        # NOTE: For now, I don't think there's any reason to support any other input/output modes.
+        self.default_input_modes = ['application/json']
+        self.default_output_modes = ['application/json']
+
+        self.task_manager = TaskManager(broker=broker, storage=storage)
+
+        # Setup
+        self._agent_card_json_schema: bytes | None = None
+        self.router.add_route('/.well-known/agent.json', self._agent_card_endpoint, methods=['HEAD', 'GET', 'OPTIONS'])
+        self.router.add_route('/', self._agent_run_endpoint, methods=['POST'])
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope['type'] == 'http' and not self.task_manager.is_running:
+            raise RuntimeError('TaskManager was not properly initialized.')
+        await super().__call__(scope, receive, send)
+
+    async def _agent_card_endpoint(self, request: Request) -> Response:
+        if self._agent_card_json_schema is None:
+            agent_card = AgentCard(
+                name=self.name,
+                url=self.url,
+                version=self.version,
+                skills=self.skills,
+                default_input_modes=self.default_input_modes,
+                default_output_modes=self.default_output_modes,
+            )
+            if self.description is not None:
+                agent_card['description'] = self.description
+            if self.provider is not None:
+                agent_card['provider'] = self.provider
+            self._agent_card_json_schema = agent_card_ta.dump_json(agent_card)
+        return Response(content=self._agent_card_json_schema, media_type='application/json')
+
+    async def _agent_run_endpoint(self, request: Request) -> Response:
+        """This is the main endpoint for the A2A server.
+
+        Although the specification allows freedom of choice and implementation, I'm pretty sure about some decisions.
+
+        1. The server will always either send a "submitted" or a "failed" on `tasks/send`.
+            Never a "completed" on the first message.
+        2. There are three possible ends for the task:
+            2.1. The task was "completed" successfully.
+            2.2. The task was "canceled".
+            2.3. The task "failed".
+        3. The server will send a "working" on the first chunk on `tasks/pushNotification/get`.
+        """
+        data = await request.body()
+        a2a_request = a2a_request_ta.validate_json(data)
+
+        if a2a_request['method'] == 'tasks/send':
+            jsonrpc_response = await self.task_manager.send_task(a2a_request)
+        elif a2a_request['method'] == 'tasks/get':
+            jsonrpc_response = await self.task_manager.get_task(a2a_request)
+        elif a2a_request['method'] == 'tasks/cancel':
+            jsonrpc_response = await self.task_manager.cancel_task(a2a_request)
+        else:
+            raise NotImplementedError(f'Method {a2a_request["method"]} not implemented.')
+        return Response(
+            content=a2a_response_ta.dump_json(jsonrpc_response, by_alias=True), media_type='application/json'
+        )
+
+
+@asynccontextmanager
+async def _default_lifespan(app: FastA2A) -> AsyncIterator[None]:
+    async with app.task_manager:
+        yield

fasta2a/fasta2a/broker.py

@@ -0,0 +1,98 @@
+from __future__ import annotations as _annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import AsyncIterator
+from contextlib import AsyncExitStack
+from dataclasses import dataclass
+from typing import Annotated, Any, Generic, Literal, TypeVar
+
+import anyio
+from opentelemetry.trace import Span, get_current_span, get_tracer
+from pydantic import Discriminator
+from typing_extensions import Self, TypedDict
+
+from .schema import TaskIdParams, TaskSendParams
+
+tracer = get_tracer(__name__)
+
+
+@dataclass
+class Broker(ABC):
+    """The broker class is in charge of scheduling the tasks.
+
+    The HTTP server uses the broker to schedule tasks.
+
+    The simple implementation is the `InMemoryBroker`, which is the broker that
+    runs the tasks in the same process as the HTTP server. That said, this class can be
+    extended to support remote workers.
+    """
+
+    @abstractmethod
+    async def run_task(self, params: TaskSendParams) -> None:
+        """Send a task to be executed by the worker."""
+        raise NotImplementedError('send_run_task is not implemented yet.')
+
+    @abstractmethod
+    async def cancel_task(self, params: TaskIdParams) -> None:
+        """Cancel a task."""
+        raise NotImplementedError('send_cancel_task is not implemented yet.')
+
+    @abstractmethod
+    async def __aenter__(self) -> Self: ...
+
+    @abstractmethod
+    async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any): ...
+
+    @abstractmethod
+    def receive_task_operations(self) -> AsyncIterator[TaskOperation]:
+        """Receive task operations from the broker.
+
+        On a multi-worker setup, the broker will need to round-robin the task operations
+        between the workers.
+        """
+
+
+OperationT = TypeVar('OperationT')
+ParamsT = TypeVar('ParamsT')
+
+
+class _TaskOperation(TypedDict, Generic[OperationT, ParamsT]):
+    """A task operation."""
+
+    operation: OperationT
+    params: ParamsT
+    _current_span: Span
+
+
+_RunTask = _TaskOperation[Literal['run'], TaskSendParams]
+_CancelTask = _TaskOperation[Literal['cancel'], TaskIdParams]
+
+TaskOperation = Annotated['_RunTask | _CancelTask', Discriminator('operation')]
+
+
+class InMemoryBroker(Broker):
+    """A broker that schedules tasks in memory."""
+
+    async def __aenter__(self):
+        self.aexit_stack = AsyncExitStack()
+        await self.aexit_stack.__aenter__()
+
+        self._write_stream, self._read_stream = anyio.create_memory_object_stream[TaskOperation]()
+        await self.aexit_stack.enter_async_context(self._read_stream)
+        await self.aexit_stack.enter_async_context(self._write_stream)
+
+        return self
+
+    async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any):
+        await self.aexit_stack.__aexit__(exc_type, exc_value, traceback)
+
+    async def run_task(self, params: TaskSendParams) -> None:
+        await self._write_stream.send(_RunTask(operation='run', params=params, _current_span=get_current_span()))
+
+    async def cancel_task(self, params: TaskIdParams) -> None:
+        await self._write_stream.send(_CancelTask(operation='cancel', params=params, _current_span=get_current_span()))
+
+    async def receive_task_operations(self) -> AsyncIterator[TaskOperation]:
+        """Receive task operations from the broker."""
+        async for task_operation in self._read_stream:
+            yield task_operation