Skip to content

Commit dac71d6

Browse files
vertex-sdk-botcopybara-github
vertex-sdk-bot
authored and
copybara-github
committed
fix: support polymorphic AgentCard serialization in A2aAgent
PiperOrigin-RevId: 918921714
1 parent a99f340 commit dac71d6

7 files changed

Lines changed: 202 additions & 22 deletions

File tree

‎agentplatform/_genai/_agent_engines_utils.py‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -665,9 +665,17 @@ 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 card is not None:
670+
if hasattr(card, "model_dump_json"):
671+
class_method[_A2A_AGENT_CARD] = card.model_dump_json()
672+
elif hasattr(card, "DESCRIPTOR"):
673+
class_method[_A2A_AGENT_CARD] = json_format.MessageToJson(card)
674+
else:
675+
raise TypeError(
676+
f"Unsupported AgentCard type: {type(card)}. "
677+
"Only Pydantic models and Protobuf Messages are supported."
678+
)
671679
class_methods_spec.append(class_method)
672680

673681
return class_methods_spec

‎agentplatform/_genai/agent_engines.py‎

Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2498,13 +2498,23 @@ def _create_config(
24982498

24992499
if hasattr(agent, "agent_card"):
25002500
agent_card = getattr(agent, "agent_card")
2501-
if agent_card:
2501+
if agent_card is not None:
25022502
try:
2503-
from google.protobuf import json_format
2504-
2505-
agent_engine_spec["agent_card"] = json_format.MessageToDict(
2506-
agent_card
2507-
)
2503+
if hasattr(agent_card, "model_dump"):
2504+
agent_engine_spec["agent_card"] = agent_card.model_dump(
2505+
exclude_none=True
2506+
)
2507+
elif hasattr(agent_card, "DESCRIPTOR"):
2508+
from google.protobuf import json_format
2509+
2510+
agent_engine_spec["agent_card"] = json_format.MessageToDict(
2511+
agent_card
2512+
)
2513+
else:
2514+
raise TypeError(
2515+
f"Unsupported AgentCard type: {type(agent_card)}. "
2516+
"Only Pydantic models and Protobuf Messages are supported."
2517+
)
25082518
except Exception as e:
25092519
raise ValueError(
25102520
f"Failed to convert agent card to dict (serialization error): {e}"

‎agentplatform/agent_engines/_agent_engines.py‎

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2003,11 +2003,19 @@ def _generate_class_methods_spec_or_raise(
20032003
class_method[_MODE_KEY_IN_SCHEMA] = mode
20042004
# A2A agent card is a special case, when running in A2A mode,
20052005
if hasattr(agent_engine, "agent_card"):
2006-
from google.protobuf import json_format
2007-
2008-
class_method[_A2A_AGENT_CARD] = json_format.MessageToJson(
2009-
getattr(agent_engine, "agent_card")
2010-
)
2006+
card = getattr(agent_engine, "agent_card")
2007+
if card is not None:
2008+
if hasattr(card, "model_dump_json"):
2009+
class_method[_A2A_AGENT_CARD] = card.model_dump_json()
2010+
elif hasattr(card, "DESCRIPTOR"):
2011+
from google.protobuf import json_format
2012+
2013+
class_method[_A2A_AGENT_CARD] = json_format.MessageToJson(card)
2014+
else:
2015+
raise TypeError(
2016+
f"Unsupported AgentCard type: {type(card)}. "
2017+
"Only Pydantic models and Protobuf Messages are supported."
2018+
)
20112019
class_methods_spec.append(class_method)
20122020

20132021
return class_methods_spec

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

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4090,3 +4090,72 @@ 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+
4101+
class DummyAgentEngine:
4102+
def __init__(self, card=None, has_card=True):
4103+
if has_card:
4104+
self.agent_card = card
4105+
4106+
def set_up(self):
4107+
pass
4108+
4109+
def query(self, query: str) -> str:
4110+
return query
4111+
4112+
4113+
class TestAgentEngineGenerateClassMethodsSpec:
4114+
"""Tests Pydantic, Protobuf, and No Card AgentCard serialization in _generate_class_methods_spec_or_raise."""
4115+
4116+
def test_pydantic_card_serialization(self):
4117+
agent_engine = DummyAgentEngine(DummyPydanticCard())
4118+
specs = _agent_engines_utils._generate_class_methods_spec_or_raise(
4119+
agent=agent_engine,
4120+
operations={"standard": ["query"]},
4121+
)
4122+
assert len(specs) == 1
4123+
assert specs[0][_agent_engines_utils._A2A_AGENT_CARD] == '{"name": "pydantic_card"}'
4124+
4125+
def test_protobuf_card_serialization(self):
4126+
from google.protobuf import struct_pb2
4127+
card = struct_pb2.Struct()
4128+
agent_engine = DummyAgentEngine(card)
4129+
specs = _agent_engines_utils._generate_class_methods_spec_or_raise(
4130+
agent=agent_engine,
4131+
operations={"standard": ["query"]},
4132+
)
4133+
assert len(specs) == 1
4134+
assert specs[0][_agent_engines_utils._A2A_AGENT_CARD] == "{}"
4135+
4136+
def test_no_card_serialization(self):
4137+
agent_engine = DummyAgentEngine(has_card=False)
4138+
specs = _agent_engines_utils._generate_class_methods_spec_or_raise(
4139+
agent=agent_engine,
4140+
operations={"standard": ["query"]},
4141+
)
4142+
assert len(specs) == 1
4143+
assert _agent_engines_utils._A2A_AGENT_CARD not in specs[0]
4144+
4145+
def test_none_card_serialization(self):
4146+
agent_engine = DummyAgentEngine(None)
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 _agent_engines_utils._A2A_AGENT_CARD not in specs[0]
4153+
4154+
def test_unsupported_card_serialization_raises_type_error(self):
4155+
agent_engine = DummyAgentEngine({"unsupported": "type"})
4156+
with pytest.raises(TypeError) as excinfo:
4157+
_agent_engines_utils._generate_class_methods_spec_or_raise(
4158+
agent=agent_engine,
4159+
operations={"standard": ["query"]},
4160+
)
4161+
assert "Unsupported AgentCard type" in str(excinfo.value)

‎tests/unit/vertex_adk/test_agent_engine_templates_adk.py‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1426,3 +1426,72 @@ 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+
1437+
class DummyAgentEngine:
1438+
def __init__(self, card=None, has_card=True):
1439+
if has_card:
1440+
self.agent_card = card
1441+
1442+
def set_up(self):
1443+
pass
1444+
1445+
def query(self, query: str) -> str:
1446+
return query
1447+
1448+
1449+
class TestAgentEngineGenerateClassMethodsSpec:
1450+
"""Tests Pydantic, Protobuf, and No Card AgentCard serialization in _generate_class_methods_spec_or_raise."""
1451+
1452+
def test_pydantic_card_serialization(self):
1453+
agent_engine = DummyAgentEngine(DummyPydanticCard())
1454+
specs = _agent_engines._generate_class_methods_spec_or_raise(
1455+
agent_engine=agent_engine,
1456+
operations={"standard": ["query"]},
1457+
)
1458+
assert len(specs) == 1
1459+
assert specs[0][_agent_engines._A2A_AGENT_CARD] == '{"name": "pydantic_card"}'
1460+
1461+
def test_protobuf_card_serialization(self):
1462+
from google.protobuf import struct_pb2
1463+
card = struct_pb2.Struct()
1464+
agent_engine = DummyAgentEngine(card)
1465+
specs = _agent_engines._generate_class_methods_spec_or_raise(
1466+
agent_engine=agent_engine,
1467+
operations={"standard": ["query"]},
1468+
)
1469+
assert len(specs) == 1
1470+
assert specs[0][_agent_engines._A2A_AGENT_CARD] == "{}"
1471+
1472+
def test_no_card_serialization(self):
1473+
agent_engine = DummyAgentEngine(has_card=False)
1474+
specs = _agent_engines._generate_class_methods_spec_or_raise(
1475+
agent_engine=agent_engine,
1476+
operations={"standard": ["query"]},
1477+
)
1478+
assert len(specs) == 1
1479+
assert _agent_engines._A2A_AGENT_CARD not in specs[0]
1480+
1481+
def test_none_card_serialization(self):
1482+
agent_engine = DummyAgentEngine(None)
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 _agent_engines._A2A_AGENT_CARD not in specs[0]
1489+
1490+
def test_unsupported_card_serialization_raises_type_error(self):
1491+
agent_engine = DummyAgentEngine({"unsupported": "type"})
1492+
with pytest.raises(TypeError) as excinfo:
1493+
_agent_engines._generate_class_methods_spec_or_raise(
1494+
agent_engine=agent_engine,
1495+
operations={"standard": ["query"]},
1496+
)
1497+
assert "Unsupported AgentCard type" in str(excinfo.value)

‎vertexai/_genai/_agent_engines_utils.py‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -633,9 +633,17 @@ 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 card is not None:
638+
if hasattr(card, "model_dump_json"):
639+
class_method[_A2A_AGENT_CARD] = card.model_dump_json()
640+
elif hasattr(card, "DESCRIPTOR"):
641+
class_method[_A2A_AGENT_CARD] = json_format.MessageToJson(card)
642+
else:
643+
raise TypeError(
644+
f"Unsupported AgentCard type: {type(card)}. "
645+
"Only Pydantic models and Protobuf Messages are supported."
646+
)
639647
class_methods_spec.append(class_method)
640648

641649
return class_methods_spec

‎vertexai/agent_engines/_agent_engines.py‎

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1997,11 +1997,19 @@ def _generate_class_methods_spec_or_raise(
19971997
class_method[_MODE_KEY_IN_SCHEMA] = mode
19981998
# A2A agent card is a special case, when running in A2A mode,
19991999
if hasattr(agent_engine, "agent_card"):
2000-
from google.protobuf import json_format
2001-
2002-
class_method[_A2A_AGENT_CARD] = json_format.MessageToJson(
2003-
getattr(agent_engine, "agent_card")
2004-
)
2000+
card = getattr(agent_engine, "agent_card")
2001+
if card is not None:
2002+
if hasattr(card, "model_dump_json"):
2003+
class_method[_A2A_AGENT_CARD] = card.model_dump_json()
2004+
elif hasattr(card, "DESCRIPTOR"):
2005+
from google.protobuf import json_format
2006+
2007+
class_method[_A2A_AGENT_CARD] = json_format.MessageToJson(card)
2008+
else:
2009+
raise TypeError(
2010+
f"Unsupported AgentCard type: {type(card)}. "
2011+
"Only Pydantic models and Protobuf Messages are supported."
2012+
)
20052013
class_methods_spec.append(class_method)
20062014

20072015
return class_methods_spec

0 commit comments

Comments
 (0)