Skip to content
Merged
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
20 changes: 20 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,26 @@ registration:
engine: my-app_1234567890 # from the console URL; omit to deploy without listing
```

A plain `model` is called in the region the agent runs in. A model offered
only in another Vertex AI location, such as a preview model that lives only
in `global`, names that location:

```yaml
model:
id: gemini-3-flash-preview
location: global
```

Agent Engine sets `GOOGLE_CLOUD_LOCATION` to the agent's region, and a full
resource path in `model` still goes to that region's endpoint, so the
location is declared here instead.

The location is handed to the Gemini client, so `validate` refuses it for a
model ADK does not serve as Gemini. Such a model is always called on Vertex
AI, including under `gete run`: a `GOOGLE_API_KEY` for the Gemini Developer
API is not enough there, and the project comes from `GOOGLE_CLOUD_PROJECT`
and application default credentials.

Tools that do not say `effect: read` count as writes, which is what the
`has_write_tools` policies key on.

Expand Down
3 changes: 2 additions & 1 deletion src/gete/runtime/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
from gete.policies import applicable, compose_instruction
from gete.redact import RedactRules
from gete.runtime.callbacks import bind_tool_call, redact_results, safe_tool_error
from gete.runtime.model import build_model
from gete.runtime.tools import build_tools
from gete.shared_credentials import SHARED_CREDENTIALS

Expand Down Expand Up @@ -49,7 +50,7 @@ def _build(resolved: Resolved) -> Any:
return LlmAgent(
# ADK wants an identifier. The display name lives in the registration.
name=agent.name.replace("-", "_"),
model=agent.data["model"],
model=build_model(agent.data["model"]),
description=agent.data["description"],
instruction=compose_instruction(
resolved.policies, resolved.data, _with_shared_credential_rules(agent)
Expand Down
49 changes: 49 additions & 0 deletions src/gete/runtime/model.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,49 @@
"""The model an agent declares, as ADK's LlmAgent takes it."""

from collections.abc import Mapping
from typing import Any

from gete.errors import DeclarationError


def gemini_class(model_id: str) -> type[Any] | None:
"""ADK's class for the model when it is a Gemini, or None.

The class comes from ADK's own registry, the one LlmAgent resolves a
plain name through, so a model ADK serves with another class - or does
not know - is not mistaken for a Gemini.
"""
from google.adk.models import Gemini
from google.adk.models.registry import LLMRegistry

try:
resolved = LLMRegistry.resolve(model_id)
except ValueError:
return None
return resolved if issubclass(resolved, Gemini) else None


def build_model(declared: str | Mapping[str, str]) -> Any:
"""A plain name as it is, or a Gemini bound to the declared location.

A plain name is called in the region the agent runs in, which Agent
Engine sets in GOOGLE_CLOUD_LOCATION. A full resource path naming another
location does not help - the request still goes to the regional endpoint -
so the location has to reach the client itself.
"""
if isinstance(declared, str):
return declared
model_id = declared["id"]
cls = gemini_class(model_id)
if cls is None:
raise DeclarationError(
f"model.id: {model_id!r} is not a Gemini model; "
"only a Gemini model's client can be pointed at a location"
)
# client_kwargs rather than a subclass overriding api_client: the field is
# what ADK reads when it builds both its generate and its live clients,
# so one setting covers both.
return cls(
model=model_id,
client_kwargs={"enterprise": True, "location": declared["location"]},
)
24 changes: 22 additions & 2 deletions src/gete/schema/agent.json
Original file line number Diff line number Diff line change
Expand Up @@ -28,8 +28,28 @@
"minLength": 1
},
"model": {
"type": "string",
"minLength": 1
"description": "A model called in the region the agent runs in, as a plain string, or an object naming another location to call it in. Some models are offered only in global, and Agent Engine sets the region itself. The two shapes share one schema rather than a oneOf, so a malformed object is reported by the key at fault, not as the string it is not.",
"type": [
"string",
"object"
],
"minLength": 1,
"additionalProperties": false,
"required": [
"id",
"location"
],
"properties": {
"id": {
"type": "string",
"minLength": 1
},
"location": {
"description": "Vertex AI location to call the model in, such as global. Becomes part of the endpoint's host name and URL path, so the shape is part of a URL.",
"type": "string",
"pattern": "^[a-z]+(-[a-z0-9]+)*$"
}
}
},
"instruction": {
"description": "Path relative to agent.yaml, or the instruction text itself.",
Expand Down
10 changes: 10 additions & 0 deletions src/gete/validate.py
Original file line number Diff line number Diff line change
Expand Up @@ -243,6 +243,16 @@ def _agent_problems(
requirements = agent.requirements
if requirements is not None and not requirements.is_file():
found.append(f"requirements: {project.display(requirements)} does not exist")
model = agent.data["model"]
if not isinstance(model, str):
from gete.runtime.model import gemini_class

# Agent Engine would take the archive and fail on the first turn.
if gemini_class(model["id"]) is None:
found.append(
f"model.id: {model['id']!r} is not a Gemini model; only a "
"Gemini model's client can be pointed at a location"
)
for index, tool in enumerate(agent.tools):
found.extend(
f"tools[{index}]: {message}"
Expand Down
47 changes: 47 additions & 0 deletions tests/test_runtime_build.py
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,53 @@ def test_agent_carries_name_model_description_and_policy_prefixed_instruction(
assert agent.instruction == "Never approve.\n\nYou sort mail."


def test_a_model_with_a_location_is_called_there_on_vertex_ai(
project: ProjectBuilder,
) -> None:
from google.adk.models import Gemini

agent = build(
resolved_path(
project,
"mail-triage",
{"model": {"id": "gemini-3-flash-preview", "location": "global"}},
)
)
assert isinstance(agent.model, Gemini)
assert agent.model.model == "gemini-3-flash-preview"


def test_a_model_location_wins_over_the_region_agent_engine_sets(
project: ProjectBuilder, monkeypatch: pytest.MonkeyPatch
) -> None:
"""Agent Engine sets GOOGLE_CLOUD_LOCATION, and a declaration may not."""
monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "example-project")
monkeypatch.setenv("GOOGLE_CLOUD_LOCATION", "us-central1")
agent = build(
resolved_path(
project,
"mail-triage",
{"model": {"id": "gemini-3-flash-preview", "location": "global"}},
)
)
client = agent.model.api_client
assert client.vertexai
assert client._api_client.location == "global"


def test_a_location_is_refused_for_a_model_that_is_not_gemini(
project: ProjectBuilder,
) -> None:
"""A deployment validate never saw must not call the wrong publisher's model."""
path = resolved_path(
project,
"mail-triage",
{"model": {"id": "claude-sonnet-4@20250514", "location": "global"}},
)
with pytest.raises(GeteError, match="is not a Gemini model"):
build(path)


def test_write_policy_applies_only_to_agents_with_write_tools(
project: ProjectBuilder,
) -> None:
Expand Down
43 changes: 43 additions & 0 deletions tests/test_schema.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
"""Shape of the declarations. The JSON Schemas decide what a document may say."""

import re
from pathlib import Path
from typing import Any

Expand Down Expand Up @@ -193,6 +194,48 @@ def test_agent_name_must_be_an_rfc1034_label(name: str) -> None:
validate_document("agent", {**AGENT, "name": name}, source="agent.yaml")


def test_a_model_may_name_the_location_it_is_called_in() -> None:
validate_document(
"agent",
{**AGENT, "model": {"id": "gemini-3-flash-preview", "location": "global"}},
source="agent.yaml",
)


@pytest.mark.parametrize(
("model", "says"),
[
({"id": "gemini-3-flash-preview"}, "'location' is a required property"),
({"location": "global"}, "'id' is a required property"),
({"id": "", "location": "global"}, "model.id: '' should be non-empty"),
(
{"id": "gemini-3-flash-preview", "location": "global", "region": "us"},
"'region' was unexpected",
),
],
)
def test_a_model_object_needs_an_id_and_a_location(
model: dict[str, Any], says: str
) -> None:
"""A model called where the agent runs is written as the plain string instead.

The message names the key at fault, not the string the object is not.
"""
with pytest.raises(DeclarationError, match=re.escape(says)):
validate_document("agent", {**AGENT, "model": model}, source="agent.yaml")


@pytest.mark.parametrize("location", ["Global", "us central1", "../global", "evil.com"])
def test_a_model_location_must_look_like_a_location(location: str) -> None:
"""It becomes part of the endpoint's host name and URL path."""
with pytest.raises(DeclarationError, match="model"):
validate_document(
"agent",
{**AGENT, "model": {"id": "gemini-3-flash-preview", "location": location}},
source="agent.yaml",
)


def test_connections_must_be_unique() -> None:
with pytest.raises(DeclarationError, match="connections"):
validate_document(
Expand Down
21 changes: 21 additions & 0 deletions tests/test_validate.py
Original file line number Diff line number Diff line change
Expand Up @@ -288,6 +288,27 @@ def test_known_builtin_tool_passes(project: ProjectBuilder) -> None:
assert problems(project) == []


def test_a_gemini_model_may_name_its_location(project: ProjectBuilder) -> None:
project.write_agent(
"mail-triage",
{"model": {"id": "gemini-3-flash-preview", "location": "global"}},
)
assert problems(project) == []


@pytest.mark.parametrize("model_id", ["claude-sonnet-4@20250514", "no-such-model"])
def test_a_location_is_refused_for_a_model_that_is_not_gemini(
project: ProjectBuilder, model_id: str
) -> None:
"""The location reaches the model through Gemini's client and nowhere else."""
project.write_agent(
"mail-triage", {"model": {"id": model_id, "location": "global"}}
)
assert any(
f"model.id: {model_id!r} is not a Gemini model" in p for p in problems(project)
)


def test_mcp_host_must_be_allowed_by_its_connection(project: ProjectBuilder) -> None:
"""Otherwise the token would be sent to a host the connection never declared."""
project.write_agent(
Expand Down
Loading