feat: A2A + Agent Engine integration

PiperOrigin-RevId: 791846986

Author: vertex-sdk-bot

setup.py

@@ -142,6 +142,10 @@
     "google-adk >= 1.0.0, < 2.0.0",
 ]
 
+a2a_extra_require = [
+    "a2a-sdk >= 0.3.0",
+]
+
 reasoning_engine_extra_require = [
     "cloudpickle >= 3.0, < 4.0",
     "google-cloud-trace < 2",
@@ -325,6 +329,7 @@
         "ray": ray_extra_require,
         "ray_testing": ray_testing_extra_require,
         "adk": adk_extra_require,
+        "a2a": a2a_extra_require,
         "reasoningengine": reasoning_engine_extra_require,
         "agent_engines": agent_engines_extra_require,
         "evaluation": evaluation_extra_require,

tests/unit/vertexai/genai/test_agent_engines.py

@@ -1872,7 +1872,7 @@ def setup_method(self):
                     "register the API methods: "
                     "https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/develop/custom#custom-methods. "
                     "Error: {Unsupported api mode: `UNKNOWN_API_MODE`, "
-                    "Supported modes are: ``, `async`, `async_stream`, `stream`.}"
+                    "Supported modes are: ``, `a2a_extension`, `async`, `async_stream`, `stream`.}"
                 ),
             ),
         ],

vertexai/_genai/_agent_engines_utils.py

@@ -42,6 +42,9 @@
     Union,
 )
 
+import json
+import httpx
+
 import proto
 
 from google.api_core import exceptions
@@ -102,6 +105,27 @@
     Session = Any
 
 
+try:
+    from a2a.types import (
+        AgentCard,
+        TransportProtocol,
+        Message,
+        TaskIdParams,
+    )
+    from a2a.client import ClientConfig, ClientFactory
+
+    AgentCard = AgentCard
+    TransportProtocol = TransportProtocol
+    ClientConfig = ClientConfig
+    ClientFactory = ClientFactory
+    TaskIdParams = TaskIdParams
+except (ImportError, AttributeError):
+    AgentCard = None
+    TransportProtocol = None
+    ClientConfig = None
+    ClientFactory = None
+    TaskIdParams = None
+
 _ACTIONS_KEY = "actions"
 _ACTION_APPEND = "append"
 _AGENT_FRAMEWORK_ATTR = "agent_framework"
@@ -144,6 +168,8 @@
 _REQUIREMENTS_FILE = "requirements.txt"
 _STANDARD_API_MODE = ""
 _STREAM_API_MODE = "stream"
+_A2A_EXTENSION_MODE = "a2a_extension"
+_A2A_AGENT_CARD = "a2a_agent_card"
 _WARNINGS_KEY = "warnings"
 _WARNING_MISSING = "missing"
 _WARNING_INCOMPATIBLE = "incompatible"
@@ -454,11 +480,32 @@ def _generate_class_methods_spec_or_raise(
 
             class_method = _to_proto(schema_dict)
             class_method[_MODE_KEY_IN_SCHEMA] = mode
+            if hasattr(agent_engine, "agent_card"):
+                class_method[_A2A_AGENT_CARD] = getattr(
+                    agent_engine, "agent_card"
+                ).model_dump_json()
             class_methods_spec.append(class_method)
 
     return class_methods_spec
 
 
+def _is_pydantic_serializable(param: inspect.Parameter) -> bool:
+    """Checks if the parameter is pydantic serializable."""
+
+    if param.annotation == inspect.Parameter.empty:
+        return True
+
+    if isinstance(param.annotation, str):
+        return False
+
+    pydantic = _import_pydantic_or_raise()
+    try:
+        pydantic.TypeAdapter(param.annotation)
+        return True
+    except Exception:
+        return False
+
+
 def _generate_schema(
     f: Callable[..., Any],
     *,
@@ -518,6 +565,7 @@ def _generate_schema(
             inspect.Parameter.KEYWORD_ONLY,
             inspect.Parameter.POSITIONAL_ONLY,
         )
+        and _is_pydantic_serializable(param)
     }
     parameters = pydantic.create_model(f.__name__, **fields_dict).schema()
     # Postprocessing
@@ -827,6 +875,7 @@ def _register_api_methods_or_raise(
             _ASYNC_API_MODE: _wrap_async_query_operation,
             _STREAM_API_MODE: _wrap_stream_query_operation,
             _ASYNC_STREAM_API_MODE: _wrap_async_stream_query_operation,
+            _A2A_EXTENSION_MODE: _wrap_a2a_operation,
         }
         if isinstance(wrap_operation_fn, dict) and api_mode in wrap_operation_fn:
             # Override the default function with user-specified function if it exists.
@@ -1438,6 +1487,62 @@ async def _method(self: genai_types.AgentEngine, **kwargs) -> AsyncIterator[Any]
     return _method
 
 
+def _wrap_a2a_operation(
+    method_name: str, agent_card: str
+) -> Callable[..., list]:
+    async def _method(self, **kwargs) -> list:
+        """Wraps an Agent Engine method, creating a callable for A2A API."""
+        if not self.api_client:
+            raise ValueError("api_client is not initialized.")
+        if not self.api_resource:
+            raise ValueError("api_resource is not initialized.")
+        a2a_agent_card = AgentCard(**json.loads(agent_card))
+        a2a_agent_card.preferred_transport = TransportProtocol.http_json
+        # AE cannot support streaming yet. Turn off streaming for now.
+        a2a_agent_card.capabilities.streaming = False
+        # agent_card is set on the class_methods before set_up is invoked.
+        # Ensure that the agent_card url is set correctly before the client is created.
+        a2a_agent_card.url = f"{self.api_resource.name}/a2a"
+
+        # Using a2a client, inject the auth token from the global config.
+        config = ClientConfig(
+            supported_transports=[
+                TransportProtocol.http_json,
+            ],
+            use_client_preference=True,
+            httpx_client=httpx.AsyncClient(
+                headers={
+                    "Authorization": (
+                        f"Bearer { self.api_client.credentials.token}"
+                    )
+                }
+            ),
+        )
+        factory = ClientFactory(config)
+        client = factory.create(a2a_agent_card)
+
+        match method_name:
+            case "on_message_send":
+                response = client.send_message(Message(**kwargs))
+            case "on_get_task":
+                response = client.get_task(TaskIdParams(**kwargs))
+            case "on_cancel_task":
+                response = client.cancel_task(TaskIdParams(**kwargs))
+            case "handle_authenticated_agent_card":
+                response = await client.get_card()
+
+        if inspect.isasyncgen(response):
+            # Response is an async generator, collect the chunks.
+            chunks = []
+            async for chunk in response:
+                chunks.append(chunk)
+            return chunks
+        else:
+            return response
+
+    return _method
+
+
 def _yield_parsed_json(http_response: google_genai_types.HttpResponse) -> Iterator[Any]:
     """Converts the body of the HTTP Response message to JSON format.
 
@@ -1531,3 +1636,4 @@ def _validate_resource_limits_or_raise(resource_limits: dict[str, str]) -> None:
             f"Memory size of {memory_str} requires at least {min_cpu} CPUs."
             f" Got {cpu}"
         )
+

vertexai/_genai/agent_engines.py

@@ -3669,6 +3669,7 @@ def _register_api_methods(self, *, agent: types.AgentEngine) -> types.AgentEngin
                     "async": _agent_engines_utils._wrap_async_query_operation,
                     "stream": _agent_engines_utils._wrap_stream_query_operation,
                     "async_stream": _agent_engines_utils._wrap_async_stream_query_operation,
+                    "a2a_extension": _agent_engines_utils._wrap_a2a_operation,
                 },
             )
         except Exception as e:

vertexai/agent_engines/_agent_engines.py

@@ -16,28 +16,14 @@
 import abc
 import inspect
 import io
+import json
 import logging
 import os
 import sys
 import tarfile
 import types
 import typing
-from typing import (
-    Any,
-    AsyncIterable,
-    Callable,
-    Coroutine,
-    Dict,
-    Iterable,
-    List,
-    Optional,
-    Protocol,
-    Sequence,
-    Tuple,
-    Union,
-)
-
-import proto
+from typing import Any, AsyncIterable, Callable, Coroutine, Dict, Iterable, List, Optional, Protocol, Sequence, Tuple, Union
 
 from google.api_core import exceptions
 from google.cloud import storage
@@ -47,6 +33,9 @@
 from google.cloud.aiplatform_v1 import types as aip_types
 from google.cloud.aiplatform_v1.types import reasoning_engine_service
 from vertexai.agent_engines import _utils
+import httpx
+import proto
+
 from google.protobuf import field_mask_pb2
 
 
@@ -60,6 +49,8 @@
 _ASYNC_API_MODE = "async"
 _STREAM_API_MODE = "stream"
 _ASYNC_STREAM_API_MODE = "async_stream"
+_A2A_EXTENSION_MODE = "a2a_extension"
+_A2A_AGENT_CARD = "a2a_agent_card"
 _MODE_KEY_IN_SCHEMA = "api_mode"
 _METHOD_NAME_KEY_IN_SCHEMA = "name"
 _DEFAULT_METHOD_NAME = "query"
@@ -111,6 +102,27 @@
 except (ImportError, AttributeError):
     ADKAgent = None
 
+try:
+    from a2a.types import (
+        AgentCard,
+        TransportProtocol,
+        Message,
+        TaskIdParams,
+    )
+    from a2a.client import ClientConfig, ClientFactory
+
+    AgentCard = AgentCard
+    TransportProtocol = TransportProtocol
+    ClientConfig = ClientConfig
+    ClientFactory = ClientFactory
+    TaskIdParams = TaskIdParams
+except (ImportError, AttributeError):
+    AgentCard = None
+    TransportProtocol = None
+    ClientConfig = None
+    ClientFactory = None
+    TaskIdParams = None
+
 
 @typing.runtime_checkable
 class Queryable(Protocol):
@@ -1498,6 +1510,58 @@ async def _method(self, **kwargs) -> AsyncIterable[Any]:
     return _method
 
 
+def _wrap_a2a_operation(
+    method_name: str, agent_card: str
+) -> Callable[..., list]:
+    async def _method(self, **kwargs) -> list:
+        """Wraps an Agent Engine method, creating a callable for A2A API."""
+        a2a_agent_card = AgentCard(**json.loads(agent_card))
+        a2a_agent_card.preferred_transport = TransportProtocol.http_json
+        # AE cannot support streaming yet. Turn off streaming for now.
+        a2a_agent_card.capabilities.streaming = False
+        # agent_card is set on the class_methods before set_up is invoked.
+        # Ensure that the agent_card url is set correctly before the client is created.
+        a2a_agent_card.url = f"https://{initializer.global_config.api_endpoint}/v1/{self.resource_name}/a2a"
+
+        # Using a2a client, inject the auth token from the global config.
+        config = ClientConfig(
+            supported_transports=[
+                TransportProtocol.http_json,
+            ],
+            use_client_preference=True,
+            httpx_client=httpx.AsyncClient(
+                headers={
+                    "Authorization": (
+                        f"Bearer {initializer.global_config.credentials.token}"
+                    )
+                }
+            ),
+        )
+        factory = ClientFactory(config)
+        client = factory.create(a2a_agent_card)
+
+        match method_name:
+            case "on_message_send":
+                response = client.send_message(Message(**kwargs))
+            case "on_get_task":
+                response = client.get_task(TaskIdParams(**kwargs))
+            case "on_cancel_task":
+                response = client.cancel_task(TaskIdParams(**kwargs))
+            case "handle_authenticated_agent_card":
+                response = await client.get_card()
+
+        if inspect.isasyncgen(response):
+            # Response is an async generator, collect the chunks.
+            chunks = []
+            async for chunk in response:
+                chunks.append(chunk)
+            return chunks
+        else:
+            return response
+
+    return _method
+
+
 def _unregister_api_methods(
     obj: "AgentEngine", operation_schemas: Sequence[_utils.JsonDict]
 ):
@@ -1573,6 +1637,7 @@ def _register_api_methods_or_raise(
             _ASYNC_API_MODE: _wrap_async_query_operation,
             _STREAM_API_MODE: _wrap_stream_query_operation,
             _ASYNC_STREAM_API_MODE: _wrap_async_stream_query_operation,
+            _A2A_EXTENSION_MODE: _wrap_a2a_operation,
         }
         if isinstance(wrap_operation_fn, dict) and api_mode in wrap_operation_fn:
             # Override the default function with user-specified function if it exists.
@@ -1661,6 +1726,11 @@ def _generate_class_methods_spec_or_raise(
 
             class_method = _utils.to_proto(schema_dict)
             class_method[_MODE_KEY_IN_SCHEMA] = mode
+            # A2A agent card is a special case, when running in A2A mode,
+            if hasattr(agent_engine, "agent_card"):
+                class_method[_A2A_AGENT_CARD] = getattr(
+                    agent_engine, "agent_card"
+                ).model_dump_json()
             class_methods_spec.append(class_method)
 
     return class_methods_spec

vertexai/agent_engines/_utils.py

@@ -501,6 +501,22 @@ def scan_requirements(
     return {module: importlib_metadata.version(module) for module in modules_found}
 
 
+def _is_pydantic_serializable(param: inspect.Parameter) -> bool:
+    """Checks if the parameter is pydantic serializable."""
+
+    if param.annotation == inspect.Parameter.empty:
+        return True
+
+    if isinstance(param.annotation, str):
+        return False
+    pydantic = _import_pydantic_or_raise()
+    try:
+        pydantic.TypeAdapter(param.annotation)
+        return True
+    except Exception:
+        return False
+
+
 def generate_schema(
     f: Callable[..., Any],
     *,
@@ -560,6 +576,7 @@ def generate_schema(
             inspect.Parameter.KEYWORD_ONLY,
             inspect.Parameter.POSITIONAL_ONLY,
         )
+        and _is_pydantic_serializable(param)
     }
     parameters = pydantic.create_model(f.__name__, **fields_dict).schema()
     # Postprocessing