fix: update agent engine utils to support python-a2a sdk 1.0

PiperOrigin-RevId: 913139425

Author: vertex-sdk-bot

vertexai/_genai/_agent_engines_utils.py

@@ -110,6 +110,21 @@
 
 
 try:
+    from a2a.utils.constants import TransportProtocol as _TpTest
+    _A2A_SDK_VERSION = "v1"
+except ImportError:
+    try:
+        from a2a.types import TransportProtocol as _TpTest
+        _A2A_SDK_VERSION = "v0"
+    except ImportError:
+        _A2A_SDK_VERSION = None
+
+if _A2A_SDK_VERSION == "v1":
+    from a2a.types import AgentCard, Message
+    from a2a.client import ClientConfig, ClientFactory
+    from a2a.utils.constants import TransportProtocol
+    from a2a.compat.v0_3.types import TaskIdParams, TaskQueryParams
+elif _A2A_SDK_VERSION == "v0":
     from a2a.types import (
         AgentCard,
         TransportProtocol,
@@ -118,15 +133,7 @@
         TaskQueryParams,
     )
     from a2a.client import ClientConfig, ClientFactory
-
-    AgentCard = AgentCard
-    TransportProtocol = TransportProtocol
-    Message = Message
-    ClientConfig = ClientConfig
-    ClientFactory = ClientFactory
-    TaskIdParams = TaskIdParams
-    TaskQueryParams = TaskQueryParams
-except (ImportError, AttributeError):
+else:
     AgentCard = None
     TransportProtocol = None
     Message = None
@@ -1737,7 +1744,7 @@ async def _method(self: genai_types.AgentEngine, **kwargs) -> AsyncIterator[Any]
     return _method
 
 
-def _wrap_a2a_operation(method_name: str, agent_card: str) -> Callable[..., list[Any]]:
+def _wrap_a2a_operation_v03(method_name: str, agent_card: str) -> Callable[..., list[Any]]:
     """Wraps an Agent Engine method, creating a callable for A2A API.
 
     Args:
@@ -1854,6 +1861,119 @@ async def _method(self, **kwargs) -> Any:  # type: ignore[no-untyped-def]
     return _method  # type: ignore[return-value]
 
 
+def _wrap_a2a_operation(method_name: str, agent_card: str) -> Callable[..., list[Any]]:
+    """Wraps an Agent Engine method, creating a callable for A2A API (v1.0.0+).
+
+    Args:
+        method_name: The name of the Agent Engine method to call.
+        agent_card: The agent card to use for the A2A API call.
+            Example:
+                {'additionalInterfaces': None,
+                'capabilities': {'extensions': None,
+                'pushNotifications': None,
+                'stateTransitionHistory': None,
+                'streaming': False},
+                'defaultInputModes': ['text'],
+                'defaultOutputModes': ['text'],
+                'description': (
+                    'A helpful assistant agent that can answer questions.'
+                ),
+                'documentationUrl': None,
+                'iconUrl': None,
+                'name': 'Q&A Agent',
+                'preferredTransport': 'JSONRPC',
+                'protocolVersion': '0.3.0',
+                'provider': None,
+                'security': None,
+                'securitySchemes': None,
+                'signatures': None,
+                'skills': [{
+                    'description': (
+                        'A helpful assistant agent that can answer questions.'
+                    ),
+                    'examples': ['Who is leading 2025 F1 Standings?',
+                        'Where can i find an active volcano?'],
+                    'id': 'question_answer',
+                'inputModes': None,
+                'name': 'Q&A Agent',
+                'outputModes': None,
+                'security': None,
+                'tags': ['Question-Answer']}],
+                'supportsAuthenticatedExtendedCard': True,
+                'url': 'http://localhost:8080/',
+                'version': '1.0.0'}
+    Returns:
+        A callable object that executes the method on the Agent Engine via
+        the A2A API.
+    """
+
+    async def _method(self, **kwargs) -> Any:  # type: ignore[no-untyped-def]
+        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_format.ParseDict(json.loads(agent_card), a2a_agent_card, ignore_unknown_fields=True)
+
+        if a2a_agent_card.supported_interfaces:
+            interface = a2a_agent_card.supported_interfaces[0]
+            if interface.protocol_binding != TransportProtocol.HTTP_JSON:
+                raise ValueError("Only HTTP+JSON is supported for preferred transport on agent card")
+        else:
+            raise ValueError("Agent card does not define any supported interfaces.")
+
+        base_url = self.api_client._api_client._http_options.base_url.rstrip("/")
+        api_version = self.api_client._api_client._http_options.api_version
+        a2a_agent_card.supported_interfaces[0].url = f"{base_url}/{api_version}/{self.api_resource.name}/a2a"
+
+        config = ClientConfig(
+            supported_protocol_bindings=[
+                TransportProtocol.HTTP_JSON,
+            ],
+            use_client_preference=True,
+            httpx_client=httpx.AsyncClient(
+                headers={
+                    "Authorization": (
+                        f"Bearer {self.api_client._api_client._credentials.token}"
+                    )
+                },
+                timeout=(
+                    self.api_client._api_client._http_options.timeout / 1000.0
+                    if self.api_client._api_client._http_options.timeout
+                    else None
+                ),
+            ),
+        )
+        factory = ClientFactory(config)
+        client = factory.create(a2a_agent_card)
+
+        if method_name == "on_message_send":
+            response = client.send_message(Message(**kwargs))
+            chunks = []
+            async for chunk in response:
+                chunks.append(chunk)
+            return chunks
+        elif method_name == "on_get_task":
+            return await client.get_task(TaskQueryParams(**kwargs))
+        elif method_name == "on_cancel_task":
+            return await client.cancel_task(TaskIdParams(**kwargs))
+        elif method_name == "handle_authenticated_agent_card":
+            return await client.get_card()
+        elif method_name == "on_get_extended_agent_card":
+            from a2a.types import GetExtendedAgentCardRequest
+            req = GetExtendedAgentCardRequest()
+            return await client.get_extended_agent_card(req)
+        else:
+            raise ValueError(f"Unknown method name: {method_name}")
+
+    return _method  # type: ignore[return-value]
+
+
+if _A2A_SDK_VERSION != "v1":
+    _wrap_a2a_operation = _wrap_a2a_operation_v03
+
+
 def _yield_parsed_json(http_response: google_genai_types.HttpResponse) -> Iterator[Any]:
     """Converts the body of the HTTP Response message to JSON format.