diff --git a/agentplatform/_genai/runtimes.py b/agentplatform/_genai/runtimes.py index 2cc5bbd676..479badf162 100644 --- a/agentplatform/_genai/runtimes.py +++ b/agentplatform/_genai/runtimes.py @@ -1911,6 +1911,7 @@ def create( container_spec=config.container_spec, keep_alive_probe=keep_alive_probe, build_config=config.build_config, + agent_card=config.agent_card, ) operation = self._create(config=api_config) logger.info(f"Create Agent Runtime backing LRO: {operation.name}") @@ -2230,6 +2231,7 @@ def _create_config( keep_alive_probe: Optional[dict[str, Any]] = None, traffic_config: Optional[types.ReasoningEngineTrafficConfigDict] = None, build_config: Optional[types.ReasoningEngineSpecBuildSpecDict] = None, + agent_card: Optional[dict[str, Any]] = None, ) -> types.UpdateRuntimeConfigDict: import sys @@ -2404,14 +2406,15 @@ def _create_config( agent=agent, ) - if hasattr(agent, "agent_card"): - agent_card = getattr(agent, "agent_card") - if agent_card: + # An agent_card given in the config takes precedence; it is added below. + if agent_card is None and hasattr(agent, "agent_card"): + agent_object_card = getattr(agent, "agent_card") + if agent_object_card: try: from google.protobuf import json_format runtime_spec["agent_card"] = json_format.MessageToDict( - agent_card + agent_object_card ) except Exception as e: raise ValueError( @@ -2452,6 +2455,12 @@ def _create_config( if build_spec: runtime_spec["build_spec"] = build_spec + if agent_card is not None: + if runtime_spec is None: + runtime_spec = {} + runtime_spec["agent_card"] = agent_card + update_masks.append("spec.agent_card") + if runtime_spec is not None: config["spec"] = runtime_spec @@ -2724,6 +2733,7 @@ def update( keep_alive_probe=keep_alive_probe, traffic_config=traffic_config, build_config=config.build_config, + agent_card=config.agent_card, ) operation = self._update(name=name, config=api_config) logger.info(f"Update Agent Runtime backing LRO: {operation.name}") diff --git a/agentplatform/_genai/types/common.py b/agentplatform/_genai/types/common.py index 34eec20ffd..9cac27746a 100644 --- a/agentplatform/_genai/types/common.py +++ b/agentplatform/_genai/types/common.py @@ -32104,6 +32104,16 @@ class AgentRuntimeConfig(_common.BaseModel): default=None, description="""The build config for the Agent Runtime. Allows bringing your own Cloud Build private worker pool (BYOBP) and, optionally, a build-time service account for the container build. Supported keys: `worker_pool` (the resource name of the Cloud Build WorkerPool to use for the build) and `service_account` (the service account that Cloud Build uses to run the build; only applicable when `worker_pool` is specified).""", ) + agent_card: Optional[dict[str, Any]] = Field( + default=None, + description="""The A2A agent card for the Agent Runtime, as a dictionary with + the A2A JSON field names (camelCase, such as `preferredTransport`). + Use it when there is no `agent` object to read the card from, for + example with `container_spec` or `source_packages`. If the `agent` also + has an `agent_card`, this value is used instead. On `update()`, it + replaces the deployed card; pass `{}` to remove it. The service sets + the card's `url` to the runtime's A2A endpoint.""", + ) class AgentRuntimeConfigDict(TypedDict, total=False): @@ -32294,6 +32304,15 @@ class AgentRuntimeConfigDict(TypedDict, total=False): build_config: Optional[ReasoningEngineSpecBuildSpecDict] """The build config for the Agent Runtime. Allows bringing your own Cloud Build private worker pool (BYOBP) and, optionally, a build-time service account for the container build. Supported keys: `worker_pool` (the resource name of the Cloud Build WorkerPool to use for the build) and `service_account` (the service account that Cloud Build uses to run the build; only applicable when `worker_pool` is specified).""" + agent_card: Optional[dict[str, Any]] + """The A2A agent card for the Agent Runtime, as a dictionary with + the A2A JSON field names (camelCase, such as `preferredTransport`). + Use it when there is no `agent` object to read the card from, for + example with `container_spec` or `source_packages`. If the `agent` also + has an `agent_card`, this value is used instead. On `update()`, it + replaces the deployed card; pass `{}` to remove it. The service sets + the card's `url` to the runtime's A2A endpoint.""" + AgentRuntimeConfigOrDict = Union[AgentRuntimeConfig, AgentRuntimeConfigDict] diff --git a/tests/unit/agentplatform/genai/test_agent_engines.py b/tests/unit/agentplatform/genai/test_agent_engines.py index d8d0ab90e1..5e52b86139 100644 --- a/tests/unit/agentplatform/genai/test_agent_engines.py +++ b/tests/unit/agentplatform/genai/test_agent_engines.py @@ -578,6 +578,7 @@ def register_operations(self) -> Dict[str, List[str]]: }, "max_seconds": 60, } +_TEST_AGENT_CARD = {"name": "test-agent", "defaultInputModes": ["text"]} _TEST_AGENT_ENGINE_SPEC = _genai_types.ReasoningEngineSpecDict( agent_framework=_TEST_AGENT_ENGINE_FRAMEWORK, class_methods=[_TEST_AGENT_ENGINE_CLASS_METHOD_1], @@ -1434,6 +1435,18 @@ def test_create_runtime_config_with_container_spec(self): ) assert "keep_alive_probe" not in config["spec"].get("deployment_spec", {}) + def test_create_runtime_config_with_container_spec_and_agent_card(self): + container_spec = {"image_uri": "gcr.io/test-project/test-image"} + config = self.client.runtimes._create_config( + mode="create", + container_spec=container_spec, + class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS, + agent_card=_TEST_AGENT_CARD, + ) + assert config["spec"]["container_spec"] == container_spec + assert config["spec"]["agent_card"] == _TEST_AGENT_CARD + assert "update_mask" not in config + def test_create_runtime_config_with_container_spec_and_keep_alive_probe( self, ): @@ -1644,6 +1657,32 @@ def test_update_runtime_config_without_agent_card_omits_mask( assert "agent_card" not in config["spec"] assert "spec.agent_card" not in config["update_mask"].split(",") + @pytest.mark.parametrize("agent_card", [_TEST_AGENT_CARD, {}]) + def test_update_runtime_config_with_only_agent_card(self, agent_card): + config = self.client.runtimes._create_config( + mode="update", + agent_card=agent_card, + ) + assert config["spec"] == {"agent_card": agent_card} + assert config["update_mask"] == "spec.agent_card" + + @mock.patch.object(_runtimes_utils, "_prepare") + def test_update_runtime_config_agent_card_overrides_agent(self, mock_prepare): + from google.protobuf import struct_pb2 + + agent_object_card = struct_pb2.Struct() + agent_object_card["version"] = "1.0.0" + config_card = {"name": "test-agent", "version": "2.0.0"} + config = self.client.runtimes._create_config( + mode="update", + agent=CapitalizeEngineWithAgentCard(agent_card=agent_object_card), + staging_bucket=_TEST_STAGING_BUCKET, + requirements=_TEST_AGENT_ENGINE_REQUIREMENTS, + agent_card=config_card, + ) + assert config["spec"]["agent_card"] == config_card + assert config["update_mask"].split(",").count("spec.agent_card") == 1 + @mock.patch.object(_runtimes_utils, "_prepare") def test_update_runtime_clear_service_account(self, mock_prepare): config = self.client.runtimes._create_config( @@ -2394,6 +2433,7 @@ def test_create_runtime_with_env_vars_dict( container_spec=None, keep_alive_probe=None, build_config=None, + agent_card=None, ) request_mock.assert_called_with( "post", @@ -2501,6 +2541,7 @@ def test_create_runtime_with_custom_service_account( container_spec=None, keep_alive_probe=None, build_config=None, + agent_card=None, ) request_mock.assert_called_with( "post", @@ -2607,6 +2648,7 @@ def test_create_runtime_with_experimental_mode( container_spec=None, keep_alive_probe=None, build_config=None, + agent_card=None, ) request_mock.assert_called_with( "post", @@ -2782,6 +2824,7 @@ def test_create_runtime_with_class_methods( container_spec=None, keep_alive_probe=None, build_config=None, + agent_card=None, ) request_mock.assert_called_with( "post", @@ -2883,6 +2926,7 @@ def test_create_runtime_with_agent_framework( container_spec=None, keep_alive_probe=None, build_config=None, + agent_card=None, ) request_mock.assert_called_with( "post", @@ -3246,6 +3290,35 @@ def test_update_runtime_description(self, mock_await_operation): None, ) + @mock.patch.object(_runtimes_utils, "_await_operation") + def test_update_runtime_agent_card(self, mock_await_operation): + mock_await_operation.return_value = _genai_types.RuntimeOperation( + response=_genai_types.ReasoningEngine( + name=_TEST_AGENT_ENGINE_RESOURCE_NAME, + spec=_TEST_AGENT_ENGINE_SPEC, + ) + ) + with mock.patch.object( + self.client.runtimes._api_client, "request" + ) as request_mock: + request_mock.return_value = genai_types.HttpResponse(body="") + self.client.runtimes.update( + name=_TEST_AGENT_ENGINE_RESOURCE_NAME, + config=_genai_types.AgentRuntimeConfig( + agent_card=_TEST_AGENT_CARD, + ), + ) + request_mock.assert_called_with( + "patch", + f"{_TEST_AGENT_ENGINE_RESOURCE_NAME}?updateMask=spec.agent_card", + { + "_url": {"name": _TEST_AGENT_ENGINE_RESOURCE_NAME}, + "spec": {"agent_card": _TEST_AGENT_CARD}, + "_query": {"updateMask": "spec.agent_card"}, + }, + None, + ) + def test_delete_runtime(self): with mock.patch.object( self.client.runtimes._api_client, "request"