From 04781ac581947bd4508dca47f3f5a91495226374 Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Thu, 21 May 2026 02:14:00 -0700 Subject: [PATCH] fix: support polymorphic AgentCard serialization in A2aAgent PiperOrigin-RevId: 918921714 --- agentplatform/_genai/_agent_engines_utils.py | 10 ++- agentplatform/agent_engines/_agent_engines.py | 10 ++- .../agentplatform/genai/test_agent_engines.py | 78 +++++++++++++++++++ .../test_agent_engine_templates_adk.py | 78 +++++++++++++++++++ vertexai/_genai/_agent_engines_utils.py | 10 ++- vertexai/agent_engines/_agent_engines.py | 10 ++- 6 files changed, 184 insertions(+), 12 deletions(-) diff --git a/agentplatform/_genai/_agent_engines_utils.py b/agentplatform/_genai/_agent_engines_utils.py index dfa4f52cf5..f0004834d6 100644 --- a/agentplatform/_genai/_agent_engines_utils.py +++ b/agentplatform/_genai/_agent_engines_utils.py @@ -665,9 +665,13 @@ def _generate_class_methods_spec_or_raise( class_method = _to_proto(schema_dict) class_method[_MODE_KEY_IN_SCHEMA] = mode if hasattr(agent, "agent_card"): - class_method[_A2A_AGENT_CARD] = json_format.MessageToJson( - getattr(agent, "agent_card") - ) + card = getattr(agent, "agent_card") + if hasattr(card, "model_dump_json"): + class_method[_A2A_AGENT_CARD] = card.model_dump_json() + elif hasattr(card, "DESCRIPTOR"): + class_method[_A2A_AGENT_CARD] = json_format.MessageToJson(card) + else: + class_method[_A2A_AGENT_CARD] = card class_methods_spec.append(class_method) return class_methods_spec diff --git a/agentplatform/agent_engines/_agent_engines.py b/agentplatform/agent_engines/_agent_engines.py index 8f929b8b47..da19663d0b 100644 --- a/agentplatform/agent_engines/_agent_engines.py +++ b/agentplatform/agent_engines/_agent_engines.py @@ -2005,9 +2005,13 @@ def _generate_class_methods_spec_or_raise( if hasattr(agent_engine, "agent_card"): from google.protobuf import json_format - class_method[_A2A_AGENT_CARD] = json_format.MessageToJson( - getattr(agent_engine, "agent_card") - ) + card = getattr(agent_engine, "agent_card") + if hasattr(card, "model_dump_json"): + class_method[_A2A_AGENT_CARD] = card.model_dump_json() + elif hasattr(card, "DESCRIPTOR"): + class_method[_A2A_AGENT_CARD] = json_format.MessageToJson(card) + else: + class_method[_A2A_AGENT_CARD] = card class_methods_spec.append(class_method) return class_methods_spec diff --git a/tests/unit/agentplatform/genai/test_agent_engines.py b/tests/unit/agentplatform/genai/test_agent_engines.py index 94db652c5e..d09fb16a3a 100644 --- a/tests/unit/agentplatform/genai/test_agent_engines.py +++ b/tests/unit/agentplatform/genai/test_agent_engines.py @@ -4090,3 +4090,81 @@ def test_delete_agent_engine_force(self): {"_url": {"name": _TEST_AGENT_ENGINE_RESOURCE_NAME}, "force": True}, None, ) + + +class DummyPydanticCard: + def model_dump_json(self): + return '{"name": "pydantic_card"}' + + +class DummyProtoCard: + DESCRIPTOR = mock.Mock() + + +class DummyFallbackCard(dict): + def __init__(self): + super().__init__({"fallback": "yes"}) + + +class DummyAgentEngine: + def __init__(self, card=None, has_card=True): + if has_card: + self.agent_card = card + + def set_up(self): + pass + + def query(self, query: str) -> str: + return query + + +class TestAgentEngineGenerateClassMethodsSpec: + """Tests Pydantic, Protobuf, None, No Card, and Fallback AgentCard serialization in _generate_class_methods_spec_or_raise.""" + + def test_pydantic_card_serialization(self): + agent_engine = DummyAgentEngine(DummyPydanticCard()) + specs = _agent_engines_utils._generate_class_methods_spec_or_raise( + agent=agent_engine, + operations={"standard": ["query"]}, + ) + assert len(specs) == 1 + assert specs[0][_agent_engines_utils._A2A_AGENT_CARD] == '{"name": "pydantic_card"}' + + @mock.patch("google3.net.proto2.python.public.json_format.MessageToJson") + def test_protobuf_card_serialization(self, mock_message_to_json): + mock_message_to_json.return_value = '{"name": "proto_card"}' + agent_engine = DummyAgentEngine(DummyProtoCard()) + specs = _agent_engines_utils._generate_class_methods_spec_or_raise( + agent=agent_engine, + operations={"standard": ["query"]}, + ) + assert len(specs) == 1 + assert specs[0][_agent_engines_utils._A2A_AGENT_CARD] == '{"name": "proto_card"}' + + def test_fallback_card_serialization(self): + card = DummyFallbackCard() + agent_engine = DummyAgentEngine(card) + specs = _agent_engines_utils._generate_class_methods_spec_or_raise( + agent=agent_engine, + operations={"standard": ["query"]}, + ) + assert len(specs) == 1 + assert specs[0][_agent_engines_utils._A2A_AGENT_CARD] == card + + def test_none_card_serialization(self): + agent_engine = DummyAgentEngine(None) + specs = _agent_engines_utils._generate_class_methods_spec_or_raise( + agent=agent_engine, + operations={"standard": ["query"]}, + ) + assert len(specs) == 1 + assert specs[0][_agent_engines_utils._A2A_AGENT_CARD] is None + + def test_no_card_serialization(self): + agent_engine = DummyAgentEngine(has_card=False) + specs = _agent_engines_utils._generate_class_methods_spec_or_raise( + agent=agent_engine, + operations={"standard": ["query"]}, + ) + assert len(specs) == 1 + assert _agent_engines_utils._A2A_AGENT_CARD not in specs[0] diff --git a/tests/unit/vertex_adk/test_agent_engine_templates_adk.py b/tests/unit/vertex_adk/test_agent_engine_templates_adk.py index ca4503a581..60b274fd1f 100644 --- a/tests/unit/vertex_adk/test_agent_engine_templates_adk.py +++ b/tests/unit/vertex_adk/test_agent_engine_templates_adk.py @@ -1426,3 +1426,81 @@ def test_default_instrumentor_builder_mtls_no_cert_source( mock_exporter.call_args.kwargs["endpoint"] == adk_template._DEFAULT_TELEMETRY_ENDPOINT ) + + +class DummyPydanticCard: + def model_dump_json(self): + return '{"name": "pydantic_card"}' + + +class DummyProtoCard: + DESCRIPTOR = mock.Mock() + + +class DummyFallbackCard(dict): + def __init__(self): + super().__init__({"fallback": "yes"}) + + +class DummyAgentEngine: + def __init__(self, card=None, has_card=True): + if has_card: + self.agent_card = card + + def set_up(self): + pass + + def query(self, query: str) -> str: + return query + + +class TestAgentEngineGenerateClassMethodsSpec: + """Tests Pydantic, Protobuf, None, No Card, and Fallback AgentCard serialization in _generate_class_methods_spec_or_raise.""" + + def test_pydantic_card_serialization(self): + agent_engine = DummyAgentEngine(DummyPydanticCard()) + specs = _agent_engines._generate_class_methods_spec_or_raise( + agent_engine=agent_engine, + operations={"standard": ["query"]}, + ) + assert len(specs) == 1 + assert specs[0][_agent_engines._A2A_AGENT_CARD] == '{"name": "pydantic_card"}' + + @mock.patch("google3.net.proto2.python.public.json_format.MessageToJson") + def test_protobuf_card_serialization(self, mock_message_to_json): + mock_message_to_json.return_value = '{"name": "proto_card"}' + agent_engine = DummyAgentEngine(DummyProtoCard()) + specs = _agent_engines._generate_class_methods_spec_or_raise( + agent_engine=agent_engine, + operations={"standard": ["query"]}, + ) + assert len(specs) == 1 + assert specs[0][_agent_engines._A2A_AGENT_CARD] == '{"name": "proto_card"}' + + def test_fallback_card_serialization(self): + card = DummyFallbackCard() + agent_engine = DummyAgentEngine(card) + specs = _agent_engines._generate_class_methods_spec_or_raise( + agent_engine=agent_engine, + operations={"standard": ["query"]}, + ) + assert len(specs) == 1 + assert specs[0][_agent_engines._A2A_AGENT_CARD] == card + + def test_none_card_serialization(self): + agent_engine = DummyAgentEngine(None) + specs = _agent_engines._generate_class_methods_spec_or_raise( + agent_engine=agent_engine, + operations={"standard": ["query"]}, + ) + assert len(specs) == 1 + assert specs[0][_agent_engines._A2A_AGENT_CARD] is None + + def test_no_card_serialization(self): + agent_engine = DummyAgentEngine(has_card=False) + specs = _agent_engines._generate_class_methods_spec_or_raise( + agent_engine=agent_engine, + operations={"standard": ["query"]}, + ) + assert len(specs) == 1 + assert _agent_engines._A2A_AGENT_CARD not in specs[0] diff --git a/vertexai/_genai/_agent_engines_utils.py b/vertexai/_genai/_agent_engines_utils.py index 92d91addbb..4da35f4c88 100644 --- a/vertexai/_genai/_agent_engines_utils.py +++ b/vertexai/_genai/_agent_engines_utils.py @@ -633,9 +633,13 @@ def _generate_class_methods_spec_or_raise( class_method = _to_proto(schema_dict) class_method[_MODE_KEY_IN_SCHEMA] = mode if hasattr(agent, "agent_card"): - class_method[_A2A_AGENT_CARD] = json_format.MessageToJson( - getattr(agent, "agent_card") - ) + card = getattr(agent, "agent_card") + if hasattr(card, "model_dump_json"): + class_method[_A2A_AGENT_CARD] = card.model_dump_json() + elif hasattr(card, "DESCRIPTOR"): + class_method[_A2A_AGENT_CARD] = json_format.MessageToJson(card) + else: + class_method[_A2A_AGENT_CARD] = card class_methods_spec.append(class_method) return class_methods_spec diff --git a/vertexai/agent_engines/_agent_engines.py b/vertexai/agent_engines/_agent_engines.py index c191e78dd5..75b0e26095 100644 --- a/vertexai/agent_engines/_agent_engines.py +++ b/vertexai/agent_engines/_agent_engines.py @@ -1999,9 +1999,13 @@ def _generate_class_methods_spec_or_raise( if hasattr(agent_engine, "agent_card"): from google.protobuf import json_format - class_method[_A2A_AGENT_CARD] = json_format.MessageToJson( - getattr(agent_engine, "agent_card") - ) + card = getattr(agent_engine, "agent_card") + if hasattr(card, "model_dump_json"): + class_method[_A2A_AGENT_CARD] = card.model_dump_json() + elif hasattr(card, "DESCRIPTOR"): + class_method[_A2A_AGENT_CARD] = json_format.MessageToJson(card) + else: + class_method[_A2A_AGENT_CARD] = card class_methods_spec.append(class_method) return class_methods_spec