diff --git a/tests/unit/vertex_langchain/test_agent_engines.py b/tests/unit/vertex_langchain/test_agent_engines.py index 8c3610577e..02c9bfa9ed 100644 --- a/tests/unit/vertex_langchain/test_agent_engines.py +++ b/tests/unit/vertex_langchain/test_agent_engines.py @@ -1061,6 +1061,130 @@ def test_create_agent_engine( retry=_TEST_RETRY, ) + def test_create_agent_engine_with_protobuf_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_with_pydantic_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 + + import pydantic + + class DummyPydanticCard(pydantic.BaseModel): + name: str = "test_pydantic_card" + + card = DummyPydanticCard() + 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 struct_pb2 + + expected_class_method = struct_pb2.Struct() + expected_class_method.CopyFrom(_TEST_AGENT_ENGINE_QUERY_SCHEMA) + expected_class_method["a2a_agent_card"] = card.model_dump_json() + 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_with_invalid_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 + + agent = CapitalizeEngineWithCard(card="invalid_card_type_string") + + with pytest.raises( + TypeError, + match="Unsupported AgentCard type", + ): + agent_engines.create( + agent, + display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, + requirements=_TEST_AGENT_ENGINE_REQUIREMENTS, + extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH], + ) + 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..e3a2fb8b68 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 @@ -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