diff --git a/agentplatform/_genai/_agent_engines_utils.py b/agentplatform/_genai/_agent_engines_utils.py index dfa4f52cf5..68c0606bdf 100644 --- a/agentplatform/_genai/_agent_engines_utils.py +++ b/agentplatform/_genai/_agent_engines_utils.py @@ -665,9 +665,17 @@ 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 card is not None: + 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: + raise TypeError( + f"Unsupported AgentCard type: {type(card)}. " + "Only Pydantic models and Protobuf Messages are supported." + ) class_methods_spec.append(class_method) return class_methods_spec diff --git a/agentplatform/_genai/agent_engines.py b/agentplatform/_genai/agent_engines.py index f160d7327b..68449ea100 100644 --- a/agentplatform/_genai/agent_engines.py +++ b/agentplatform/_genai/agent_engines.py @@ -2498,13 +2498,23 @@ def _create_config( if hasattr(agent, "agent_card"): agent_card = getattr(agent, "agent_card") - if agent_card: + if agent_card is not None: try: - from google.protobuf import json_format - - agent_engine_spec["agent_card"] = json_format.MessageToDict( - agent_card - ) + if hasattr(agent_card, "model_dump"): + agent_engine_spec["agent_card"] = agent_card.model_dump( + exclude_none=True + ) + elif hasattr(agent_card, "DESCRIPTOR"): + from google.protobuf import json_format + + agent_engine_spec["agent_card"] = json_format.MessageToDict( + agent_card + ) + else: + raise TypeError( + f"Unsupported AgentCard type: {type(agent_card)}. " + "Only Pydantic models and Protobuf Messages are supported." + ) except Exception as e: raise ValueError( f"Failed to convert agent card to dict (serialization error): {e}" diff --git a/agentplatform/agent_engines/_agent_engines.py b/agentplatform/agent_engines/_agent_engines.py index 8f929b8b47..224023ae2d 100644 --- a/agentplatform/agent_engines/_agent_engines.py +++ b/agentplatform/agent_engines/_agent_engines.py @@ -2003,11 +2003,19 @@ def _generate_class_methods_spec_or_raise( class_method[_MODE_KEY_IN_SCHEMA] = mode # A2A agent card is a special case, when running in A2A mode, 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 card is not None: + if hasattr(card, "model_dump_json"): + class_method[_A2A_AGENT_CARD] = card.model_dump_json() + elif hasattr(card, "DESCRIPTOR"): + from google.protobuf import json_format + + class_method[_A2A_AGENT_CARD] = json_format.MessageToJson(card) + else: + raise TypeError( + f"Unsupported AgentCard type: {type(card)}. " + "Only Pydantic models and Protobuf Messages are supported." + ) 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..42e3254d8c 100644 --- a/tests/unit/agentplatform/genai/test_agent_engines.py +++ b/tests/unit/agentplatform/genai/test_agent_engines.py @@ -4090,3 +4090,341 @@ def test_delete_agent_engine_force(self): {"_url": {"name": _TEST_AGENT_ENGINE_RESOURCE_NAME}, "force": True}, None, ) + + +class DummyPydanticCard: + def model_dump(self, exclude_none=True): + return {"name": "pydantic_card"} + def model_dump_json(self): + return '{"name": "pydantic_card"}' + + + +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, and No Card 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"}' + + def test_protobuf_card_serialization(self): + from google.protobuf import struct_pb2 + card = struct_pb2.Struct() + 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] == "{}" + + 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] + + 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 _agent_engines_utils._A2A_AGENT_CARD not in specs[0] + + def test_unsupported_card_serialization_raises_type_error(self): + agent_engine = DummyAgentEngine({"unsupported": "type"}) + with pytest.raises(TypeError) as excinfo: + _agent_engines_utils._generate_class_methods_spec_or_raise( + agent=agent_engine, + operations={"standard": ["query"]}, + ) + assert "Unsupported AgentCard type" in str(excinfo.value) + + +class TestVertexAIAgentEngineGenerateClassMethodsSpec: + """Tests Pydantic, Protobuf, and No Card AgentCard serialization in vertexai namespace _generate_class_methods_spec_or_raise.""" + + def test_pydantic_card_serialization(self): + from vertexai._genai import _agent_engines_utils as vertexai_utils + + agent_engine = DummyAgentEngine(DummyPydanticCard()) + specs = vertexai_utils._generate_class_methods_spec_or_raise( + agent=agent_engine, + operations={"standard": ["query"]}, + ) + assert len(specs) == 1 + assert specs[0][vertexai_utils._A2A_AGENT_CARD] == '{"name": "pydantic_card"}' + + def test_protobuf_card_serialization(self): + from vertexai._genai import _agent_engines_utils as vertexai_utils + from google.protobuf import struct_pb2 + + card = struct_pb2.Struct() + agent_engine = DummyAgentEngine(card) + specs = vertexai_utils._generate_class_methods_spec_or_raise( + agent=agent_engine, + operations={"standard": ["query"]}, + ) + assert len(specs) == 1 + assert specs[0][vertexai_utils._A2A_AGENT_CARD] == "{}" + + def test_no_card_serialization(self): + from vertexai._genai import _agent_engines_utils as vertexai_utils + + agent_engine = DummyAgentEngine(has_card=False) + specs = vertexai_utils._generate_class_methods_spec_or_raise( + agent=agent_engine, + operations={"standard": ["query"]}, + ) + assert len(specs) == 1 + assert vertexai_utils._A2A_AGENT_CARD not in specs[0] + + def test_none_card_serialization(self): + from vertexai._genai import _agent_engines_utils as vertexai_utils + + agent_engine = DummyAgentEngine(None) + specs = vertexai_utils._generate_class_methods_spec_or_raise( + agent=agent_engine, + operations={"standard": ["query"]}, + ) + assert len(specs) == 1 + assert vertexai_utils._A2A_AGENT_CARD not in specs[0] + + def test_unsupported_card_serialization_raises_type_error(self): + from vertexai._genai import _agent_engines_utils as vertexai_utils + + agent_engine = DummyAgentEngine({"unsupported": "type"}) + with pytest.raises(TypeError) as excinfo: + vertexai_utils._generate_class_methods_spec_or_raise( + agent=agent_engine, + operations={"standard": ["query"]}, + ) + assert "Unsupported AgentCard type" in str(excinfo.value) + + +class TestAgentEnginesCreateConfigAgentCard: + """Tests polymorphic AgentCard serialization in agentplatform _create_config.""" + + @mock.patch.object(_agent_engines_utils, "_prepare") + def test_create_config_with_pydantic_card(self, mock_prepare): + class DummyPydanticCard: + def model_dump(self, exclude_none=True): + return {"name": "pydantic_card"} + def model_dump_json(self): + return '{"name": "pydantic_card"}' + + class CapitalizeEngineWithCard(CapitalizeEngine): + def __init__(self, card): + self.agent_card = card + + agent = CapitalizeEngineWithCard(DummyPydanticCard()) + client = agentplatform.Client( + project=_TEST_PROJECT, + location=_TEST_LOCATION, + credentials=_TEST_CREDENTIALS, + ) + config = client.agent_engines._create_config( + mode="create", + agent=agent, + staging_bucket=_TEST_STAGING_BUCKET, + ) + assert config["spec"]["agent_card"] == {"name": "pydantic_card"} + + @mock.patch.object(_agent_engines_utils, "_prepare") + def test_create_config_with_protobuf_card(self, mock_prepare): + from google.protobuf import struct_pb2 + + class CapitalizeEngineWithCard(CapitalizeEngine): + def __init__(self, card): + self.agent_card = card + + card = struct_pb2.Struct() + card["key"] = "val" + agent = CapitalizeEngineWithCard(card) + client = agentplatform.Client( + project=_TEST_PROJECT, + location=_TEST_LOCATION, + credentials=_TEST_CREDENTIALS, + ) + config = client.agent_engines._create_config( + mode="create", + agent=agent, + staging_bucket=_TEST_STAGING_BUCKET, + ) + assert config["spec"]["agent_card"] == {"key": "val"} + + @mock.patch.object(_agent_engines_utils, "_prepare") + def test_create_config_with_unsupported_card_raises_type_error(self, mock_prepare): + class CapitalizeEngineWithCard(CapitalizeEngine): + def __init__(self, card): + self.agent_card = card + + agent = CapitalizeEngineWithCard("unsupported_string_type") + client = agentplatform.Client( + project=_TEST_PROJECT, + location=_TEST_LOCATION, + credentials=_TEST_CREDENTIALS, + ) + with pytest.raises(TypeError) as excinfo: + client.agent_engines._create_config( + mode="create", + agent=agent, + staging_bucket=_TEST_STAGING_BUCKET, + ) + assert "Unsupported AgentCard type" in str(excinfo.value) + + +class TestVertexAIAgentEnginesCreateConfigAgentCard: + """Tests polymorphic AgentCard serialization in vertexai namespace _create_config.""" + + @mock.patch.object(_agent_engines_utils, "_prepare") + def test_create_config_with_pydantic_card(self, mock_prepare): + import vertexai as vertexai_sdk + from vertexai._genai import _agent_engines_utils as vertexai_utils + + class DummyPydanticCard: + def model_dump(self, exclude_none=True): + return {"name": "pydantic_card"} + def model_dump_json(self): + return '{"name": "pydantic_card"}' + + class CapitalizeEngineWithCard(CapitalizeEngine): + def __init__(self, card): + self.agent_card = card + + with mock.patch.object(vertexai_utils, "_prepare") as mock_vertexai_prepare: + agent = CapitalizeEngineWithCard(DummyPydanticCard()) + client = vertexai_sdk.Client( + project=_TEST_PROJECT, + location=_TEST_LOCATION, + credentials=_TEST_CREDENTIALS, + ) + config = client.agent_engines._create_config( + mode="create", + agent=agent, + staging_bucket=_TEST_STAGING_BUCKET, + ) + assert config["spec"]["agent_card"] == {"name": "pydantic_card"} + + @mock.patch.object(_agent_engines_utils, "_prepare") + def test_create_config_with_protobuf_card(self, mock_prepare): + import vertexai as vertexai_sdk + from vertexai._genai import _agent_engines_utils as vertexai_utils + from google.protobuf import struct_pb2 + + class CapitalizeEngineWithCard(CapitalizeEngine): + def __init__(self, card): + self.agent_card = card + + with mock.patch.object(vertexai_utils, "_prepare") as mock_vertexai_prepare: + card = struct_pb2.Struct() + card["key"] = "val" + agent = CapitalizeEngineWithCard(card) + client = vertexai_sdk.Client( + project=_TEST_PROJECT, + location=_TEST_LOCATION, + credentials=_TEST_CREDENTIALS, + ) + config = client.agent_engines._create_config( + mode="create", + agent=agent, + staging_bucket=_TEST_STAGING_BUCKET, + ) + assert config["spec"]["agent_card"] == {"key": "val"} + + @mock.patch.object(_agent_engines_utils, "_prepare") + def test_create_config_with_unsupported_card_raises_type_error(self, mock_prepare): + import vertexai as vertexai_sdk + from vertexai._genai import _agent_engines_utils as vertexai_utils + + class CapitalizeEngineWithCard(CapitalizeEngine): + def __init__(self, card): + self.agent_card = card + + with mock.patch.object(vertexai_utils, "_prepare") as mock_vertexai_prepare: + agent = CapitalizeEngineWithCard("unsupported_string_type") + client = vertexai_sdk.Client( + project=_TEST_PROJECT, + location=_TEST_LOCATION, + credentials=_TEST_CREDENTIALS, + ) + with pytest.raises(TypeError) as excinfo: + client.agent_engines._create_config( + mode="create", + agent=agent, + staging_bucket=_TEST_STAGING_BUCKET, + ) + assert "Unsupported AgentCard type" in str(excinfo.value) + + +class TestAgentEnginesCreateConfigRealAgentCard: + """Tests polymorphic AgentCard serialization in agentplatform _create_config utilizing real Pydantic and Protobuf structures.""" + + @mock.patch.object(_agent_engines_utils, "_prepare") + def test_create_config_with_real_pydantic_card(self, mock_prepare): + import pydantic + + class RealPydanticCard(pydantic.BaseModel): + name: str = "real_pydantic_card_instance" + + class CapitalizeEngineWithCard(CapitalizeEngine): + def __init__(self, card): + self.agent_card = card + + agent = CapitalizeEngineWithCard(RealPydanticCard()) + client = agentplatform.Client( + project=_TEST_PROJECT, + location=_TEST_LOCATION, + credentials=_TEST_CREDENTIALS, + ) + config = client.agent_engines._create_config( + mode="create", + agent=agent, + staging_bucket=_TEST_STAGING_BUCKET, + ) + assert config["spec"]["agent_card"] == {"name": "real_pydantic_card_instance"} + + @mock.patch.object(_agent_engines_utils, "_prepare") + def test_create_config_with_real_protobuf_card(self, mock_prepare): + from google.protobuf import struct_pb2 + + class CapitalizeEngineWithCard(CapitalizeEngine): + def __init__(self, card): + self.agent_card = card + + card = struct_pb2.Struct() + card["name"] = "real_protobuf_card_instance" + agent = CapitalizeEngineWithCard(card) + client = agentplatform.Client( + project=_TEST_PROJECT, + location=_TEST_LOCATION, + credentials=_TEST_CREDENTIALS, + ) + config = client.agent_engines._create_config( + mode="create", + agent=agent, + staging_bucket=_TEST_STAGING_BUCKET, + ) + assert config["spec"]["agent_card"] == {"name": "real_protobuf_card_instance"} diff --git a/vertexai/_genai/_agent_engines_utils.py b/vertexai/_genai/_agent_engines_utils.py index 92d91addbb..ae2bf6080d 100644 --- a/vertexai/_genai/_agent_engines_utils.py +++ b/vertexai/_genai/_agent_engines_utils.py @@ -633,9 +633,17 @@ 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 card is not None: + 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: + raise TypeError( + f"Unsupported AgentCard type: {type(card)}. " + "Only Pydantic models and Protobuf Messages are supported." + ) class_methods_spec.append(class_method) return class_methods_spec diff --git a/vertexai/_genai/agent_engines.py b/vertexai/_genai/agent_engines.py index 8e1c1b1aa1..4a3fd14681 100644 --- a/vertexai/_genai/agent_engines.py +++ b/vertexai/_genai/agent_engines.py @@ -2498,13 +2498,23 @@ def _create_config( if hasattr(agent, "agent_card"): agent_card = getattr(agent, "agent_card") - if agent_card: + if agent_card is not None: try: - from google.protobuf import json_format - - agent_engine_spec["agent_card"] = json_format.MessageToDict( - agent_card - ) + if hasattr(agent_card, "model_dump"): + agent_engine_spec["agent_card"] = agent_card.model_dump( + exclude_none=True + ) + elif hasattr(agent_card, "DESCRIPTOR"): + from google.protobuf import json_format + + agent_engine_spec["agent_card"] = json_format.MessageToDict( + agent_card + ) + else: + raise TypeError( + f"Unsupported AgentCard type: {type(agent_card)}. " + "Only Pydantic models and Protobuf Messages are supported." + ) except Exception as e: raise ValueError( f"Failed to convert agent card to dict (serialization error): {e}" diff --git a/vertexai/agent_engines/_agent_engines.py b/vertexai/agent_engines/_agent_engines.py index c191e78dd5..7f7c8f186b 100644 --- a/vertexai/agent_engines/_agent_engines.py +++ b/vertexai/agent_engines/_agent_engines.py @@ -1997,11 +1997,19 @@ def _generate_class_methods_spec_or_raise( class_method[_MODE_KEY_IN_SCHEMA] = mode # A2A agent card is a special case, when running in A2A mode, 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 card is not None: + if hasattr(card, "model_dump_json"): + class_method[_A2A_AGENT_CARD] = card.model_dump_json() + elif hasattr(card, "DESCRIPTOR"): + from google.protobuf import json_format + + class_method[_A2A_AGENT_CARD] = json_format.MessageToJson(card) + else: + raise TypeError( + f"Unsupported AgentCard type: {type(card)}. " + "Only Pydantic models and Protobuf Messages are supported." + ) class_methods_spec.append(class_method) return class_methods_spec