Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 4 additions & 3 deletions src/a2a/server/apps/jsonrpc/fastapi_app.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import logging

from collections.abc import Callable
from collections.abc import Awaitable, Callable
from typing import TYPE_CHECKING, Any


Expand Down Expand Up @@ -66,38 +66,39 @@
(SSE).
"""

def __init__( # noqa: PLR0913
self,
agent_card: AgentCard,
http_handler: RequestHandler,
extended_agent_card: AgentCard | None = None,
context_builder: CallContextBuilder | None = None,
card_modifier: Callable[[AgentCard], AgentCard] | None = None,
card_modifier: Callable[[AgentCard], Awaitable[AgentCard] | AgentCard]
| None = None,
extended_card_modifier: Callable[
[AgentCard, ServerCallContext], AgentCard
[AgentCard, ServerCallContext], Awaitable[AgentCard] | AgentCard
]
| None = None,
max_content_length: int | None = 10 * 1024 * 1024, # 10MB
) -> None:
"""Initializes the A2AFastAPIApplication.

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.
card_modifier: An optional callback to dynamically modify the public
agent card before it is served.
extended_card_modifier: An optional callback to dynamically modify
the extended agent card before it is served. It receives the
call context.
max_content_length: The maximum allowed content length for incoming
requests. Defaults to 10MB. Set to None for unbounded maximum.
"""

Check notice on line 101 in src/a2a/server/apps/jsonrpc/fastapi_app.py

View workflow job for this annotation

GitHub Actions / Lint Code Base

Copy/pasted code

see src/a2a/server/apps/rest/rest_adapter.py (57-83)
if not _package_fastapi_installed:
raise ImportError(
'The `fastapi` package is required to use the `A2AFastAPIApplication`.'
Expand Down
14 changes: 9 additions & 5 deletions src/a2a/server/apps/jsonrpc/jsonrpc_app.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
import traceback

from abc import ABC, abstractmethod
from collections.abc import AsyncGenerator, Callable
from collections.abc import AsyncGenerator, Awaitable, Callable
from typing import TYPE_CHECKING, Any

from pydantic import ValidationError
Expand Down Expand Up @@ -51,6 +51,7 @@
PREV_AGENT_CARD_WELL_KNOWN_PATH,
)
from a2a.utils.errors import MethodNotImplementedError
from a2a.utils.helpers import maybe_await


logger = logging.getLogger(__name__)
Expand Down Expand Up @@ -172,38 +173,39 @@
for model in A2ARequestModel.__args__
}

def __init__( # noqa: PLR0913
self,
agent_card: AgentCard,
http_handler: RequestHandler,
extended_agent_card: AgentCard | None = None,
context_builder: CallContextBuilder | None = None,
card_modifier: Callable[[AgentCard], AgentCard] | None = None,
card_modifier: Callable[[AgentCard], Awaitable[AgentCard] | AgentCard]
| None = None,
extended_card_modifier: Callable[
[AgentCard, ServerCallContext], AgentCard
[AgentCard, ServerCallContext], Awaitable[AgentCard] | AgentCard
]
| None = None,
max_content_length: int | None = 10 * 1024 * 1024, # 10MB
) -> None:
"""Initializes the JSONRPCApplication.

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.
card_modifier: An optional callback to dynamically modify the public
agent card before it is served.
extended_card_modifier: An optional callback to dynamically modify
the extended agent card before it is served. It receives the
call context.
max_content_length: The maximum allowed content length for incoming
requests. Defaults to 10MB. Set to None for unbounded maximum.
"""

Check notice on line 208 in src/a2a/server/apps/jsonrpc/jsonrpc_app.py

View workflow job for this annotation

GitHub Actions / Lint Code Base

Copy/pasted code

see src/a2a/server/apps/rest/rest_adapter.py (57-83)
if not _package_starlette_installed:
raise ImportError(
'Packages `starlette` and `sse-starlette` are required to use the'
Expand Down Expand Up @@ -576,7 +578,7 @@

card_to_serve = self.agent_card
if self.card_modifier:
card_to_serve = self.card_modifier(card_to_serve)
card_to_serve = await maybe_await(self.card_modifier(card_to_serve))

return JSONResponse(
card_to_serve.model_dump(
Expand Down Expand Up @@ -605,7 +607,9 @@
context = self._context_builder.build(request)
# If no base extended card is provided, pass the public card to the modifier
base_card = card_to_serve if card_to_serve else self.agent_card
card_to_serve = self.extended_card_modifier(base_card, context)
card_to_serve = await maybe_await(
self.extended_card_modifier(base_card, context)
)

if card_to_serve:
return JSONResponse(
Expand Down
7 changes: 4 additions & 3 deletions src/a2a/server/apps/jsonrpc/starlette_app.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import logging

from collections.abc import Callable
from collections.abc import Awaitable, Callable
from typing import TYPE_CHECKING, Any


Expand Down Expand Up @@ -48,18 +48,19 @@
(SSE).
"""

def __init__( # noqa: PLR0913
self,
agent_card: AgentCard,
http_handler: RequestHandler,
extended_agent_card: AgentCard | None = None,
context_builder: CallContextBuilder | None = None,
card_modifier: Callable[[AgentCard], AgentCard] | None = None,
card_modifier: Callable[[AgentCard], Awaitable[AgentCard] | AgentCard]
| None = None,
extended_card_modifier: Callable[
[AgentCard, ServerCallContext], AgentCard
[AgentCard, ServerCallContext], Awaitable[AgentCard] | AgentCard
]
| None = None,
max_content_length: int | None = 10 * 1024 * 1024, # 10MB

Check notice on line 63 in src/a2a/server/apps/jsonrpc/starlette_app.py

View workflow job for this annotation

GitHub Actions / Lint Code Base

Copy/pasted code

see src/a2a/server/apps/rest/rest_adapter.py (57-69)
) -> None:
"""Initializes the A2AStarletteApplication.

Expand Down
7 changes: 4 additions & 3 deletions src/a2a/server/apps/rest/fastapi_app.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import logging

from collections.abc import Callable
from collections.abc import Awaitable, Callable
from typing import TYPE_CHECKING, Any


Expand Down Expand Up @@ -43,35 +43,36 @@
(SSE).
"""

def __init__( # noqa: PLR0913
self,
agent_card: AgentCard,
http_handler: RequestHandler,
extended_agent_card: AgentCard | None = None,
context_builder: CallContextBuilder | None = None,
card_modifier: Callable[[AgentCard], AgentCard] | None = None,
card_modifier: Callable[[AgentCard], Awaitable[AgentCard] | AgentCard]
| None = None,
extended_card_modifier: Callable[
[AgentCard, ServerCallContext], AgentCard
[AgentCard, ServerCallContext], Awaitable[AgentCard] | AgentCard
]
| None = None,
):
"""Initializes the A2ARESTFastAPIApplication.

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.
card_modifier: An optional callback to dynamically modify the public
agent card before it is served.
extended_card_modifier: An optional callback to dynamically modify
the extended agent card before it is served. It receives the
call context.
"""

Check notice on line 75 in src/a2a/server/apps/rest/fastapi_app.py

View workflow job for this annotation

GitHub Actions / Lint Code Base

Copy/pasted code

see src/a2a/server/apps/rest/rest_adapter.py (57-86)
if not _package_fastapi_installed:
raise ImportError(
'The `fastapi` package is required to use the'
Expand Down
15 changes: 10 additions & 5 deletions src/a2a/server/apps/rest/rest_adapter.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@
from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable
from typing import TYPE_CHECKING, Any

from a2a.utils.helpers import maybe_await


if TYPE_CHECKING:
from sse_starlette.sse import EventSourceResponse
Expand Down Expand Up @@ -52,35 +54,36 @@
manages response generation including Server-Sent Events (SSE).
"""

def __init__( # noqa: PLR0913
self,
agent_card: AgentCard,
http_handler: RequestHandler,
extended_agent_card: AgentCard | None = None,
context_builder: CallContextBuilder | None = None,
card_modifier: Callable[[AgentCard], AgentCard] | None = None,
card_modifier: Callable[[AgentCard], Awaitable[AgentCard] | AgentCard]
| None = None,
extended_card_modifier: Callable[
[AgentCard, ServerCallContext], AgentCard
[AgentCard, ServerCallContext], Awaitable[AgentCard] | AgentCard
]
| None = None,
):

Check notice on line 69 in src/a2a/server/apps/rest/rest_adapter.py

View workflow job for this annotation

GitHub Actions / Lint Code Base

Copy/pasted code

see src/a2a/server/apps/jsonrpc/starlette_app.py (51-63)
"""Initializes the RESTApplication.

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.
card_modifier: An optional callback to dynamically modify the public
agent card before it is served.
extended_card_modifier: An optional callback to dynamically modify

Check notice on line 83 in src/a2a/server/apps/rest/rest_adapter.py

View workflow job for this annotation

GitHub Actions / Lint Code Base

Copy/pasted code

see src/a2a/server/apps/jsonrpc/fastapi_app.py (69-101)

Check notice on line 83 in src/a2a/server/apps/rest/rest_adapter.py

View workflow job for this annotation

GitHub Actions / Lint Code Base

Copy/pasted code

see src/a2a/server/apps/jsonrpc/jsonrpc_app.py (176-208)
the extended agent card before it is served. It receives the
call context.
"""

Check notice on line 86 in src/a2a/server/apps/rest/rest_adapter.py

View workflow job for this annotation

GitHub Actions / Lint Code Base

Copy/pasted code

see src/a2a/server/apps/rest/fastapi_app.py (46-75)
if not _package_starlette_installed:
raise ImportError(
'Packages `starlette` and `sse-starlette` are required to use'
Expand Down Expand Up @@ -150,7 +153,7 @@
"""
card_to_serve = self.agent_card
if self.card_modifier:
card_to_serve = self.card_modifier(card_to_serve)
card_to_serve = await maybe_await(self.card_modifier(card_to_serve))

return card_to_serve.model_dump(mode='json', exclude_none=True)

Expand Down Expand Up @@ -182,9 +185,11 @@

if self.extended_card_modifier:
context = self._context_builder.build(request)
card_to_serve = self.extended_card_modifier(card_to_serve, context)
card_to_serve = await maybe_await(
self.extended_card_modifier(card_to_serve, context)
)
elif self.card_modifier:
card_to_serve = self.card_modifier(card_to_serve)
card_to_serve = await maybe_await(self.card_modifier(card_to_serve))

return card_to_serve.model_dump(mode='json', exclude_none=True)

Expand Down
9 changes: 5 additions & 4 deletions src/a2a/server/request_handlers/grpc_handler.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
import logging

from abc import ABC, abstractmethod
from collections.abc import AsyncIterable, Sequence
from collections.abc import AsyncIterable, Awaitable, Sequence


try:
Expand Down Expand Up @@ -34,7 +34,7 @@
from a2a.types import AgentCard, TaskNotFoundError
from a2a.utils import proto_utils
from a2a.utils.errors import ServerError
from a2a.utils.helpers import validate, validate_async_generator
from a2a.utils.helpers import maybe_await, validate, validate_async_generator


logger = logging.getLogger(__name__)
Expand Down Expand Up @@ -89,7 +89,8 @@ def __init__(
agent_card: AgentCard,
request_handler: RequestHandler,
context_builder: CallContextBuilder | None = None,
card_modifier: Callable[[AgentCard], AgentCard] | None = None,
card_modifier: Callable[[AgentCard], Awaitable[AgentCard] | AgentCard]
| None = None,
):
"""Initializes the GrpcHandler.

Expand Down Expand Up @@ -339,7 +340,7 @@ async def GetAgentCard(
"""Get the agent card for the agent served."""
card_to_serve = self.agent_card
if self.card_modifier:
card_to_serve = self.card_modifier(card_to_serve)
card_to_serve = await maybe_await(self.card_modifier(card_to_serve))
return proto_utils.ToProto.agent_card(card_to_serve)

async def abort_context(
Expand Down
15 changes: 9 additions & 6 deletions src/a2a/server/request_handlers/jsonrpc_handler.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import logging

from collections.abc import AsyncIterable, Callable
from collections.abc import AsyncIterable, Awaitable, Callable

from a2a.server.context import ServerCallContext
from a2a.server.request_handlers.request_handler import RequestHandler
Expand Down Expand Up @@ -46,7 +46,7 @@
TaskStatusUpdateEvent,
)
from a2a.utils.errors import ServerError
from a2a.utils.helpers import validate
from a2a.utils.helpers import maybe_await, validate
from a2a.utils.telemetry import SpanKind, trace_class


Expand All @@ -63,10 +63,11 @@ def __init__(
request_handler: RequestHandler,
extended_agent_card: AgentCard | None = None,
extended_card_modifier: Callable[
[AgentCard, ServerCallContext], AgentCard
[AgentCard, ServerCallContext], Awaitable[AgentCard] | AgentCard
]
| None = None,
card_modifier: Callable[[AgentCard], AgentCard] | None = None,
card_modifier: Callable[[AgentCard], Awaitable[AgentCard] | AgentCard]
| None = None,
):
"""Initializes the JSONRPCHandler.

Expand Down Expand Up @@ -450,9 +451,11 @@ async def get_authenticated_extended_card(

card_to_serve = base_card
if self.extended_card_modifier and context:
card_to_serve = self.extended_card_modifier(base_card, context)
card_to_serve = await maybe_await(
self.extended_card_modifier(base_card, context)
)
elif self.card_modifier:
card_to_serve = self.card_modifier(base_card)
card_to_serve = await maybe_await(self.card_modifier(base_card))

return GetAuthenticatedExtendedCardResponse(
root=GetAuthenticatedExtendedCardSuccessResponse(
Expand Down
14 changes: 12 additions & 2 deletions src/a2a/utils/helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,8 @@
import json
import logging

from collections.abc import Callable
from typing import Any
from collections.abc import Awaitable, Callable
from typing import Any, TypeVar
from uuid import uuid4

from a2a.types import (
Expand All @@ -24,6 +24,9 @@
from a2a.utils.telemetry import trace_function


T = TypeVar('T')


logger = logging.getLogger(__name__)


Expand Down Expand Up @@ -368,3 +371,10 @@ def canonicalize_agent_card(agent_card: AgentCard) -> str:
# Recursively remove empty values
cleaned_dict = _clean_empty(card_dict)
return json.dumps(cleaned_dict, separators=(',', ':'), sort_keys=True)


async def maybe_await(value: T | Awaitable[T]) -> T:
"""Awaits a value if it's awaitable, otherwise simply provides it back."""
if inspect.isawaitable(value):
return await value
return value
28 changes: 28 additions & 0 deletions tests/server/request_handlers/test_grpc_handler.py
Original file line number Diff line number Diff line change
Expand Up @@ -209,6 +209,34 @@ async def test_get_agent_card_with_modifier(
) -> None:
"""Test GetAgentCard call with a card_modifier."""

async def modifier(card: types.AgentCard) -> types.AgentCard:
modified_card = card.model_copy(deep=True)
modified_card.name = 'Modified gRPC Agent'
return modified_card

grpc_handler_modified = GrpcHandler(
agent_card=sample_agent_card,
request_handler=mock_request_handler,
card_modifier=modifier,
)

request_proto = a2a_pb2.GetAgentCardRequest()
response = await grpc_handler_modified.GetAgentCard(
request_proto, mock_grpc_context
)

assert response.name == 'Modified gRPC Agent'
assert response.version == sample_agent_card.version


@pytest.mark.asyncio
async def test_get_agent_card_with_modifier_sync(
mock_request_handler: AsyncMock,
sample_agent_card: types.AgentCard,
mock_grpc_context: AsyncMock,
) -> None:
"""Test GetAgentCard call with a synchronous card_modifier."""

def modifier(card: types.AgentCard) -> types.AgentCard:
modified_card = card.model_copy(deep=True)
modified_card.name = 'Modified gRPC Agent'
Expand Down
51 changes: 51 additions & 0 deletions tests/server/request_handlers/test_jsonrpc_handler.py
Original file line number Diff line number Diff line change
Expand Up @@ -1295,6 +1295,57 @@ async def test_get_authenticated_extended_card_with_modifier(self) -> None:
skills=[],
)

async def modifier(
card: AgentCard, context: ServerCallContext
) -> AgentCard:
modified_card = card.model_copy(deep=True)
modified_card.name = 'Modified Card'
modified_card.description = (
f'Modified for context: {context.state.get("foo")}'
)
return modified_card

handler = JSONRPCHandler(
self.mock_agent_card,
mock_request_handler,
extended_agent_card=mock_base_card,
extended_card_modifier=modifier,
)
request = GetAuthenticatedExtendedCardRequest(id='ext-card-req-mod')
call_context = ServerCallContext(state={'foo': 'bar'})

# Act
response: GetAuthenticatedExtendedCardResponse = (
await handler.get_authenticated_extended_card(request, call_context)
)

# Assert
self.assertIsInstance(
response.root, GetAuthenticatedExtendedCardSuccessResponse
)
self.assertEqual(response.root.id, 'ext-card-req-mod')
modified_card = response.root.result
self.assertEqual(modified_card.name, 'Modified Card')
self.assertEqual(modified_card.description, 'Modified for context: bar')
self.assertEqual(modified_card.version, '1.0')

async def test_get_authenticated_extended_card_with_modifier_sync(
self,
) -> None:
"""Test successful retrieval of a synchronously dynamically modified extended agent card."""
# Arrange
mock_request_handler = AsyncMock(spec=DefaultRequestHandler)
mock_base_card = AgentCard(
name='Base Card',
description='Base details',
url='http://agent.example.com/api',
version='1.0',
capabilities=AgentCapabilities(),
default_input_modes=['text/plain'],
default_output_modes=['application/json'],
skills=[],
)

def modifier(card: AgentCard, context: ServerCallContext) -> AgentCard:
modified_card = card.model_copy(deep=True)
modified_card.name = 'Modified Card'
Expand Down
Loading
Loading