Skip to content

fix: support polymorphic AgentCard serialization in A2aAgent #6815

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: support polymorphic AgentCard serialization in A2aAgent #6815

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 9 additions & 3 deletions agentplatform/_genai/_agent_engines_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -665,9 +665,15 @@ 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)
elif isinstance(card, str):
class_method[_A2A_AGENT_CARD] = card
else:
class_method[_A2A_AGENT_CARD] = json.dumps(card)
class_methods_spec.append(class_method)

return class_methods_spec
Expand Down
12 changes: 9 additions & 3 deletions agentplatform/agent_engines/_agent_engines.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2005,9 +2005,15 @@ 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)
elif isinstance(card, str):
class_method[_A2A_AGENT_CARD] = card
else:
class_method[_A2A_AGENT_CARD] = json.dumps(card)
class_methods_spec.append(class_method)

return class_methods_spec
Expand Down
78 changes: 78 additions & 0 deletions tests/unit/agentplatform/genai/test_agent_engines.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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] == json.dumps(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] == 'null'

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]
78 changes: 78 additions & 0 deletions tests/unit/vertex_adk/test_agent_engine_templates_adk.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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] == json.dumps(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] == 'null'

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]
12 changes: 9 additions & 3 deletions vertexai/_genai/_agent_engines_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -633,9 +633,15 @@ 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)
elif isinstance(card, str):
class_method[_A2A_AGENT_CARD] = card
else:
class_method[_A2A_AGENT_CARD] = json.dumps(card)
class_methods_spec.append(class_method)

return class_methods_spec
Expand Down
12 changes: 9 additions & 3 deletions vertexai/agent_engines/_agent_engines.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1999,9 +1999,15 @@ 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)
elif isinstance(card, str):
class_method[_A2A_AGENT_CARD] = card
else:
class_method[_A2A_AGENT_CARD] = json.dumps(card)
class_methods_spec.append(class_method)

return class_methods_spec
Expand Down
Loading