Skip to content

feat(agentplatform): add agent_card to AgentRuntimeConfig for runtimes.create and update #7183

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat(agentplatform): add agent_card to AgentRuntimeConfig for runtimes.create and update #7183

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
18 changes: 14 additions & 4 deletions agentplatform/_genai/runtimes.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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}")
Expand Down Expand Up @@ -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

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

Expand Down Expand Up @@ -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}")
Expand Down
19 changes: 19 additions & 0 deletions agentplatform/_genai/types/common.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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]

Expand Down
73 changes: 73 additions & 0 deletions tests/unit/agentplatform/genai/test_agent_engines.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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],
Expand Down Expand Up @@ -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,
):
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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"
Expand Down
Loading