diff --git a/agentplatform/_genai/_agent_engines_utils.py b/agentplatform/_genai/_agent_engines_utils.py index 124123f0e1..7a7d51a0c0 100644 --- a/agentplatform/_genai/_agent_engines_utils.py +++ b/agentplatform/_genai/_agent_engines_utils.py @@ -652,10 +652,9 @@ 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", None) + if card is not None: + class_method[_A2A_AGENT_CARD] = _serialize_agent_card_to_json(card) class_methods_spec.append(class_method) return class_methods_spec @@ -2002,7 +2001,7 @@ def _is_adk_agent(agent_engine: _AgentEngineInterface) -> bool: def _add_telemetry_enablement_env( - env_vars: Optional[Dict[str, Union[str, Any]]] + env_vars: Optional[Dict[str, Union[str, Any]]], ) -> Optional[Dict[str, Union[str, Any]]]: """Adds telemetry enablement env var to the env vars. @@ -2263,3 +2262,55 @@ def _import_autogen_tools_or_warn() -> Optional[types.ModuleType]: "call `pip install google-cloud-aiplatform[ag2]`." ) return None + + +def _serialize_agent_card_to_dict(card: Any) -> Optional[Dict[str, Any]]: + """Validates and serializes an AgentCard to a dictionary representation. + + Args: + card: The AgentCard instance (Protobuf Message). + + Returns: + The serialized card as a dictionary. + + Raises: + TypeError: If the card type is not supported. + """ + if card is None: + return None + + if hasattr(card, "DESCRIPTOR"): + from google.protobuf import json_format + + return typing.cast(dict[str, Any], json_format.MessageToDict(card)) + else: + raise TypeError( + f"Unsupported AgentCard type: {type(card)}. " + "Only Protobuf Messages are supported in agentplatform." + ) + + +def _serialize_agent_card_to_json(card: Any) -> Optional[str]: + """Validates and serializes an AgentCard to a JSON string representation. + + Args: + card: The AgentCard instance (Protobuf Message). + + Returns: + The serialized card as a JSON string. + + Raises: + TypeError: If the card type is not supported. + """ + if card is None: + return None + + if hasattr(card, "DESCRIPTOR"): + from google.protobuf import json_format + + return typing.cast(str, json_format.MessageToJson(card)) + else: + raise TypeError( + f"Unsupported AgentCard type: {type(card)}. " + "Only Protobuf Messages are supported in agentplatform." + ) diff --git a/agentplatform/_genai/agent_engines.py b/agentplatform/_genai/agent_engines.py index f160d7327b..6d5d174318 100644 --- a/agentplatform/_genai/agent_engines.py +++ b/agentplatform/_genai/agent_engines.py @@ -2498,12 +2498,12 @@ 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 + agent_engine_spec["agent_card"] = ( + _agent_engines_utils._serialize_agent_card_to_dict( + agent_card + ) ) except Exception as e: raise ValueError( diff --git a/agentplatform/agent_engines/_agent_engines.py b/agentplatform/agent_engines/_agent_engines.py index 663b2843e6..27618ec3eb 100644 --- a/agentplatform/agent_engines/_agent_engines.py +++ b/agentplatform/agent_engines/_agent_engines.py @@ -2005,11 +2005,11 @@ 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: + class_method[_A2A_AGENT_CARD] = ( + _agent_engines_utils._serialize_agent_card_to_json(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 a643967f01..32048c51da 100644 --- a/tests/unit/agentplatform/genai/test_agent_engines.py +++ b/tests/unit/agentplatform/genai/test_agent_engines.py @@ -4170,3 +4170,362 @@ 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()) + 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) + + 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, + ) + 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) + + @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"): + 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"): + 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"): + 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, + ) + 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) + + @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/tests/unit/vertex_langchain/test_agent_engines.py b/tests/unit/vertex_langchain/test_agent_engines.py index 8c3610577e..c15b741b8c 100644 --- a/tests/unit/vertex_langchain/test_agent_engines.py +++ b/tests/unit/vertex_langchain/test_agent_engines.py @@ -1061,6 +1061,53 @@ def test_create_agent_engine( retry=_TEST_RETRY, ) + def test_create_agent_engine_with_agent_card( + self, + create_agent_engine_mock, + cloud_storage_create_bucket_mock, + tarfile_open_mock, + cloudpickle_dump_mock, + cloudpickle_load_mock, + importlib_metadata_version_mock, + get_agent_engine_mock, + get_gca_resource_mock, + ): + class CapitalizeEngineWithCard(CapitalizeEngine): + def __init__(self, card): + self.agent_card = card + + from google.protobuf import struct_pb2 + + card = struct_pb2.Struct() + card["name"] = "test_agent_card" + agent = CapitalizeEngineWithCard(card) + + agent_engines.create( + agent, + display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, + requirements=_TEST_AGENT_ENGINE_REQUIREMENTS, + extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH], + ) + + expected_reasoning_engine = types.ReasoningEngine( + display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, + spec=types.ReasoningEngineSpec( + package_spec=_TEST_AGENT_ENGINE_PACKAGE_SPEC, + agent_framework=_agent_engines._DEFAULT_AGENT_FRAMEWORK, + ), + ) + from google.protobuf import json_format + + expected_class_method = struct_pb2.Struct() + expected_class_method.CopyFrom(_TEST_AGENT_ENGINE_QUERY_SCHEMA) + expected_class_method["a2a_agent_card"] = json_format.MessageToJson(card) + expected_reasoning_engine.spec.class_methods.append(expected_class_method) + + create_agent_engine_mock.assert_called_with( + parent=_TEST_PARENT, + reasoning_engine=expected_reasoning_engine, + ) + def test_create_agent_engine_requirements_from_file( self, create_agent_engine_mock, diff --git a/vertexai/_genai/_agent_engines_utils.py b/vertexai/_genai/_agent_engines_utils.py index f876238401..dc741bab3c 100644 --- a/vertexai/_genai/_agent_engines_utils.py +++ b/vertexai/_genai/_agent_engines_utils.py @@ -652,10 +652,9 @@ 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", None) + if card is not None: + class_method[_A2A_AGENT_CARD] = _serialize_agent_card_to_json(card) class_methods_spec.append(class_method) return class_methods_spec @@ -2119,7 +2118,7 @@ def _is_adk_agent(agent_engine: _AgentEngineInterface) -> bool: def _add_telemetry_enablement_env( - env_vars: Optional[Dict[str, Union[str, Any]]] + env_vars: Optional[Dict[str, Union[str, Any]]], ) -> Optional[Dict[str, Union[str, Any]]]: """Adds telemetry enablement env var to the env vars. @@ -2148,3 +2147,59 @@ def _add_telemetry_enablement_env( return env_vars return env_vars | env_to_add + + +def _serialize_agent_card_to_dict(card: Any) -> Optional[Dict[str, Any]]: + """Validates and serializes an AgentCard to a dictionary representation. + + Args: + card: The AgentCard instance (Pydantic model or Protobuf Message). + + Returns: + The serialized card as a dictionary. + + Raises: + TypeError: If the card type is not supported. + """ + if card is None: + return None + + if hasattr(card, "model_dump"): + return typing.cast(dict[str, Any], card.model_dump(exclude_none=True)) + elif hasattr(card, "DESCRIPTOR"): + from google.protobuf import json_format + + return typing.cast(dict[str, Any], json_format.MessageToDict(card)) + else: + raise TypeError( + f"Unsupported AgentCard type: {type(card)}. " + "Only Pydantic models and Protobuf Messages are supported." + ) + + +def _serialize_agent_card_to_json(card: Any) -> Optional[str]: + """Validates and serializes an AgentCard to a JSON string representation. + + Args: + card: The AgentCard instance (Pydantic model or Protobuf Message). + + Returns: + The serialized card as a JSON string. + + Raises: + TypeError: If the card type is not supported. + """ + if card is None: + return None + + if hasattr(card, "model_dump_json"): + return typing.cast(str, card.model_dump_json()) + elif hasattr(card, "DESCRIPTOR"): + from google.protobuf import json_format + + return typing.cast(str, json_format.MessageToJson(card)) + else: + raise TypeError( + f"Unsupported AgentCard type: {type(card)}. " + "Only Pydantic models and Protobuf Messages are supported." + ) diff --git a/vertexai/_genai/agent_engines.py b/vertexai/_genai/agent_engines.py index 8e1c1b1aa1..f642e7da51 100644 --- a/vertexai/_genai/agent_engines.py +++ b/vertexai/_genai/agent_engines.py @@ -2498,12 +2498,12 @@ 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 + agent_engine_spec["agent_card"] = ( + _agent_engines_utils._serialize_agent_card_to_dict( + agent_card + ) ) except Exception as e: raise ValueError( diff --git a/vertexai/agent_engines/_agent_engines.py b/vertexai/agent_engines/_agent_engines.py index c191e78dd5..5a87dfebd1 100644 --- a/vertexai/agent_engines/_agent_engines.py +++ b/vertexai/agent_engines/_agent_engines.py @@ -46,6 +46,7 @@ from google.cloud.aiplatform_v1 import types as aip_types from google.cloud.aiplatform_v1.types import reasoning_engine_service from vertexai.agent_engines import _utils +from vertexai._genai import _agent_engines_utils import httpx import proto @@ -1997,11 +1998,11 @@ 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: + class_method[_A2A_AGENT_CARD] = ( + _agent_engines_utils._serialize_agent_card_to_json(card) + ) class_methods_spec.append(class_method) return class_methods_spec