Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
124 changes: 124 additions & 0 deletions tests/unit/vertex_langchain/test_agent_engines.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
63 changes: 59 additions & 4 deletions vertexai/_genai/_agent_engines_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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."
)
10 changes: 5 additions & 5 deletions vertexai/_genai/agent_engines.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
11 changes: 6 additions & 5 deletions vertexai/agent_engines/_agent_engines.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down
Loading