diff --git a/README.md b/README.md index 0e8cb9e..84edaf8 100644 --- a/README.md +++ b/README.md @@ -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. diff --git a/src/gete/runtime/__init__.py b/src/gete/runtime/__init__.py index b340019..0c298cf 100644 --- a/src/gete/runtime/__init__.py +++ b/src/gete/runtime/__init__.py @@ -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 @@ -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) diff --git a/src/gete/runtime/model.py b/src/gete/runtime/model.py new file mode 100644 index 0000000..76034a4 --- /dev/null +++ b/src/gete/runtime/model.py @@ -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"]}, + ) diff --git a/src/gete/schema/agent.json b/src/gete/schema/agent.json index c396cfd..08159c5 100644 --- a/src/gete/schema/agent.json +++ b/src/gete/schema/agent.json @@ -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.", diff --git a/src/gete/validate.py b/src/gete/validate.py index 67b6f3c..aafa488 100644 --- a/src/gete/validate.py +++ b/src/gete/validate.py @@ -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}" diff --git a/tests/test_runtime_build.py b/tests/test_runtime_build.py index 34a6e82..58ed998 100644 --- a/tests/test_runtime_build.py +++ b/tests/test_runtime_build.py @@ -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: diff --git a/tests/test_schema.py b/tests/test_schema.py index 4d14f88..dc53d23 100644 --- a/tests/test_schema.py +++ b/tests/test_schema.py @@ -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 @@ -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( diff --git a/tests/test_validate.py b/tests/test_validate.py index 5eb0d05..7370f89 100644 --- a/tests/test_validate.py +++ b/tests/test_validate.py @@ -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(