Skip to content

Commit 1156006

Browse files
vertex-sdk-botcopybara-github
vertex-sdk-bot
authored and
copybara-github
committed
fix: support both Pydantic and Protobuf AgentCard during A2aAgent serialization
PiperOrigin-RevId: 918370966
1 parent f1a6a5e commit 1156006

6 files changed

Lines changed: 192 additions & 12 deletions

File tree

‎agentplatform/_genai/_agent_engines_utils.py‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -665,9 +665,15 @@ def _generate_class_methods_spec_or_raise(
665665
class_method = _to_proto(schema_dict)
666666
class_method[_MODE_KEY_IN_SCHEMA] = mode
667667
if hasattr(agent, "agent_card"):
668-
class_method[_A2A_AGENT_CARD] = json_format.MessageToJson(
669-
getattr(agent, "agent_card")
670-
)
668+
card = getattr(agent, "agent_card")
669+
if hasattr(card, "model_dump_json"):
670+
class_method[_A2A_AGENT_CARD] = card.model_dump_json()
671+
elif hasattr(card, "DESCRIPTOR"):
672+
class_method[_A2A_AGENT_CARD] = json_format.MessageToJson(card)
673+
elif isinstance(card, str):
674+
class_method[_A2A_AGENT_CARD] = card
675+
else:
676+
class_method[_A2A_AGENT_CARD] = json.dumps(card)
671677
class_methods_spec.append(class_method)
672678

673679
return class_methods_spec

‎agentplatform/agent_engines/_agent_engines.py‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2005,9 +2005,15 @@ def _generate_class_methods_spec_or_raise(
20052005
if hasattr(agent_engine, "agent_card"):
20062006
from google.protobuf import json_format
20072007

2008-
class_method[_A2A_AGENT_CARD] = json_format.MessageToJson(
2009-
getattr(agent_engine, "agent_card")
2010-
)
2008+
card = getattr(agent_engine, "agent_card")
2009+
if hasattr(card, "model_dump_json"):
2010+
class_method[_A2A_AGENT_CARD] = card.model_dump_json()
2011+
elif hasattr(card, "DESCRIPTOR"):
2012+
class_method[_A2A_AGENT_CARD] = json_format.MessageToJson(card)
2013+
elif isinstance(card, str):
2014+
class_method[_A2A_AGENT_CARD] = card
2015+
else:
2016+
class_method[_A2A_AGENT_CARD] = json.dumps(card)
20112017
class_methods_spec.append(class_method)
20122018

20132019
return class_methods_spec

‎tests/unit/agentplatform/genai/test_agent_engines.py‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4090,3 +4090,81 @@ def test_delete_agent_engine_force(self):
40904090
{"_url": {"name": _TEST_AGENT_ENGINE_RESOURCE_NAME}, "force": True},
40914091
None,
40924092
)
4093+
4094+
4095+
class DummyPydanticCard:
4096+
def model_dump_json(self):
4097+
return '{"name": "pydantic_card"}'
4098+
4099+
4100+
class DummyProtoCard:
4101+
DESCRIPTOR = mock.Mock()
4102+
4103+
4104+
class DummyFallbackCard(dict):
4105+
def __init__(self):
4106+
super().__init__({"fallback": "yes"})
4107+
4108+
4109+
class DummyAgentEngine:
4110+
def __init__(self, card=None, has_card=True):
4111+
if has_card:
4112+
self.agent_card = card
4113+
4114+
def set_up(self):
4115+
pass
4116+
4117+
def query(self, query: str) -> str:
4118+
return query
4119+
4120+
4121+
class TestAgentEngineGenerateClassMethodsSpec:
4122+
"""Tests Pydantic, Protobuf, None, No Card, and Fallback AgentCard serialization in _generate_class_methods_spec_or_raise."""
4123+
4124+
def test_pydantic_card_serialization(self):
4125+
agent_engine = DummyAgentEngine(DummyPydanticCard())
4126+
specs = _agent_engines_utils._generate_class_methods_spec_or_raise(
4127+
agent=agent_engine,
4128+
operations={"standard": ["query"]},
4129+
)
4130+
assert len(specs) == 1
4131+
assert specs[0][_agent_engines_utils._A2A_AGENT_CARD] == '{"name": "pydantic_card"}'
4132+
4133+
@mock.patch("google3.net.proto2.python.public.json_format.MessageToJson")
4134+
def test_protobuf_card_serialization(self, mock_message_to_json):
4135+
mock_message_to_json.return_value = '{"name": "proto_card"}'
4136+
agent_engine = DummyAgentEngine(DummyProtoCard())
4137+
specs = _agent_engines_utils._generate_class_methods_spec_or_raise(
4138+
agent=agent_engine,
4139+
operations={"standard": ["query"]},
4140+
)
4141+
assert len(specs) == 1
4142+
assert specs[0][_agent_engines_utils._A2A_AGENT_CARD] == '{"name": "proto_card"}'
4143+
4144+
def test_fallback_card_serialization(self):
4145+
card = DummyFallbackCard()
4146+
agent_engine = DummyAgentEngine(card)
4147+
specs = _agent_engines_utils._generate_class_methods_spec_or_raise(
4148+
agent=agent_engine,
4149+
operations={"standard": ["query"]},
4150+
)
4151+
assert len(specs) == 1
4152+
assert specs[0][_agent_engines_utils._A2A_AGENT_CARD] == json.dumps(card)
4153+
4154+
def test_none_card_serialization(self):
4155+
agent_engine = DummyAgentEngine(None)
4156+
specs = _agent_engines_utils._generate_class_methods_spec_or_raise(
4157+
agent=agent_engine,
4158+
operations={"standard": ["query"]},
4159+
)
4160+
assert len(specs) == 1
4161+
assert specs[0][_agent_engines_utils._A2A_AGENT_CARD] == 'null'
4162+
4163+
def test_no_card_serialization(self):
4164+
agent_engine = DummyAgentEngine(has_card=False)
4165+
specs = _agent_engines_utils._generate_class_methods_spec_or_raise(
4166+
agent=agent_engine,
4167+
operations={"standard": ["query"]},
4168+
)
4169+
assert len(specs) == 1
4170+
assert _agent_engines_utils._A2A_AGENT_CARD not in specs[0]

‎tests/unit/vertex_adk/test_agent_engine_templates_adk.py‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1426,3 +1426,81 @@ def test_default_instrumentor_builder_mtls_no_cert_source(
14261426
mock_exporter.call_args.kwargs["endpoint"]
14271427
== adk_template._DEFAULT_TELEMETRY_ENDPOINT
14281428
)
1429+
1430+
1431+
class DummyPydanticCard:
1432+
def model_dump_json(self):
1433+
return '{"name": "pydantic_card"}'
1434+
1435+
1436+
class DummyProtoCard:
1437+
DESCRIPTOR = mock.Mock()
1438+
1439+
1440+
class DummyFallbackCard(dict):
1441+
def __init__(self):
1442+
super().__init__({"fallback": "yes"})
1443+
1444+
1445+
class DummyAgentEngine:
1446+
def __init__(self, card=None, has_card=True):
1447+
if has_card:
1448+
self.agent_card = card
1449+
1450+
def set_up(self):
1451+
pass
1452+
1453+
def query(self, query: str) -> str:
1454+
return query
1455+
1456+
1457+
class TestAgentEngineGenerateClassMethodsSpec:
1458+
"""Tests Pydantic, Protobuf, None, No Card, and Fallback AgentCard serialization in _generate_class_methods_spec_or_raise."""
1459+
1460+
def test_pydantic_card_serialization(self):
1461+
agent_engine = DummyAgentEngine(DummyPydanticCard())
1462+
specs = _agent_engines._generate_class_methods_spec_or_raise(
1463+
agent_engine=agent_engine,
1464+
operations={"standard": ["query"]},
1465+
)
1466+
assert len(specs) == 1
1467+
assert specs[0][_agent_engines._A2A_AGENT_CARD] == '{"name": "pydantic_card"}'
1468+
1469+
@mock.patch("google3.net.proto2.python.public.json_format.MessageToJson")
1470+
def test_protobuf_card_serialization(self, mock_message_to_json):
1471+
mock_message_to_json.return_value = '{"name": "proto_card"}'
1472+
agent_engine = DummyAgentEngine(DummyProtoCard())
1473+
specs = _agent_engines._generate_class_methods_spec_or_raise(
1474+
agent_engine=agent_engine,
1475+
operations={"standard": ["query"]},
1476+
)
1477+
assert len(specs) == 1
1478+
assert specs[0][_agent_engines._A2A_AGENT_CARD] == '{"name": "proto_card"}'
1479+
1480+
def test_fallback_card_serialization(self):
1481+
card = DummyFallbackCard()
1482+
agent_engine = DummyAgentEngine(card)
1483+
specs = _agent_engines._generate_class_methods_spec_or_raise(
1484+
agent_engine=agent_engine,
1485+
operations={"standard": ["query"]},
1486+
)
1487+
assert len(specs) == 1
1488+
assert specs[0][_agent_engines._A2A_AGENT_CARD] == json.dumps(card)
1489+
1490+
def test_none_card_serialization(self):
1491+
agent_engine = DummyAgentEngine(None)
1492+
specs = _agent_engines._generate_class_methods_spec_or_raise(
1493+
agent_engine=agent_engine,
1494+
operations={"standard": ["query"]},
1495+
)
1496+
assert len(specs) == 1
1497+
assert specs[0][_agent_engines._A2A_AGENT_CARD] == 'null'
1498+
1499+
def test_no_card_serialization(self):
1500+
agent_engine = DummyAgentEngine(has_card=False)
1501+
specs = _agent_engines._generate_class_methods_spec_or_raise(
1502+
agent_engine=agent_engine,
1503+
operations={"standard": ["query"]},
1504+
)
1505+
assert len(specs) == 1
1506+
assert _agent_engines._A2A_AGENT_CARD not in specs[0]

‎vertexai/_genai/_agent_engines_utils.py‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -633,9 +633,15 @@ def _generate_class_methods_spec_or_raise(
633633
class_method = _to_proto(schema_dict)
634634
class_method[_MODE_KEY_IN_SCHEMA] = mode
635635
if hasattr(agent, "agent_card"):
636-
class_method[_A2A_AGENT_CARD] = json_format.MessageToJson(
637-
getattr(agent, "agent_card")
638-
)
636+
card = getattr(agent, "agent_card")
637+
if hasattr(card, "model_dump_json"):
638+
class_method[_A2A_AGENT_CARD] = card.model_dump_json()
639+
elif hasattr(card, "DESCRIPTOR"):
640+
class_method[_A2A_AGENT_CARD] = json_format.MessageToJson(card)
641+
elif isinstance(card, str):
642+
class_method[_A2A_AGENT_CARD] = card
643+
else:
644+
class_method[_A2A_AGENT_CARD] = json.dumps(card)
639645
class_methods_spec.append(class_method)
640646

641647
return class_methods_spec

‎vertexai/agent_engines/_agent_engines.py‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1999,9 +1999,15 @@ def _generate_class_methods_spec_or_raise(
19991999
if hasattr(agent_engine, "agent_card"):
20002000
from google.protobuf import json_format
20012001

2002-
class_method[_A2A_AGENT_CARD] = json_format.MessageToJson(
2003-
getattr(agent_engine, "agent_card")
2004-
)
2002+
card = getattr(agent_engine, "agent_card")
2003+
if hasattr(card, "model_dump_json"):
2004+
class_method[_A2A_AGENT_CARD] = card.model_dump_json()
2005+
elif hasattr(card, "DESCRIPTOR"):
2006+
class_method[_A2A_AGENT_CARD] = json_format.MessageToJson(card)
2007+
elif isinstance(card, str):
2008+
class_method[_A2A_AGENT_CARD] = card
2009+
else:
2010+
class_method[_A2A_AGENT_CARD] = json.dumps(card)
20052011
class_methods_spec.append(class_method)
20062012

20072013
return class_methods_spec

0 commit comments

Comments
 (0)