diff --git a/.env.example b/.env.example index 8b23d0a53..bba1116c4 100644 --- a/.env.example +++ b/.env.example @@ -12,6 +12,13 @@ ANTHROPIC_API_KEY= OPEN_ROUTER_API_KEY= XAI_API_KEY= +# Gemini on Vertex AI (Optional - instead of the Gemini API key) +# Runs every Gemini call, helper agents included, on Vertex with Application +# Default Credentials (`gcloud auth application-default login`). +# GOOGLE_GENAI_USE_VERTEXAI=true +# GOOGLE_CLOUD_PROJECT= +# GOOGLE_CLOUD_LOCATION=global + # Google Cloud Vision OCR (Optional - for advanced OCR processing) OCR_API_KEY= diff --git a/artemis/agents/explorer/explorer.py b/artemis/agents/explorer/explorer.py index f39dfab28..5f521a34d 100644 --- a/artemis/agents/explorer/explorer.py +++ b/artemis/agents/explorer/explorer.py @@ -49,7 +49,7 @@ from artemis.context import ArtemisContext from artemis.data_engine.trace import trace from artemis.graph.state import State -from artemis.llm.google import is_gemini_model, strip_provider_prefix +from artemis.llm.google import is_gemini_model, strip_provider_prefix, use_vertex_backend from artemis.utils.ocr_api import is_ocr_configured from third_party.mobile_use.utils.logger import get_logger @@ -130,6 +130,10 @@ def _detect_native_engine(self) -> bool: """ if not is_gemini_model(self.model_name): return False + # On Vertex the universal engine (LangChain, routed to Vertex) is used; + # the native client here is built from the developer API key. + if use_vertex_backend(): + return False if getattr(self.ctx, "_genai_client", None) is not None: return True api_key = getattr(settings, "GOOGLE_API_KEY", None) diff --git a/artemis/agents/video_analyzer/video_analyzer.py b/artemis/agents/video_analyzer/video_analyzer.py index 9704a5796..30a665411 100644 --- a/artemis/agents/video_analyzer/video_analyzer.py +++ b/artemis/agents/video_analyzer/video_analyzer.py @@ -72,6 +72,7 @@ is_gemini_model, resolve_video_processing, strip_provider_prefix, + use_vertex_backend, ) from artemis.llm.reliability import retry_policy_for from artemis.services.llm import _record_llm_event, _record_llm_retry, get_llm @@ -421,7 +422,8 @@ def _init_engine(self) -> None: llm_cfg.model if (llm_cfg and hasattr(llm_cfg, "model")) else "gemini-3.8-flash" ) - has_google_key = bool( + # The native path uploads through the Files API, which Vertex does not have. + has_google_key = not use_vertex_backend() and bool( settings.GOOGLE_API_KEY and settings.GOOGLE_API_KEY.get_secret_value() ) diff --git a/artemis/llm/google/__init__.py b/artemis/llm/google/__init__.py index 282719da7..e4c78fbae 100644 --- a/artemis/llm/google/__init__.py +++ b/artemis/llm/google/__init__.py @@ -37,6 +37,9 @@ strip_provider_prefix, supports_agentic_video, supports_thinking_level, + use_vertex_backend, + vertex_location, + vertex_project, ) from artemis.llm.google.usage import normalize_usage, usage_from_message @@ -54,4 +57,7 @@ "supports_agentic_video", "supports_thinking_level", "usage_from_message", + "use_vertex_backend", + "vertex_location", + "vertex_project", ] diff --git a/artemis/llm/google/provider.py b/artemis/llm/google/provider.py index c9f4379eb..759acb1e9 100644 --- a/artemis/llm/google/provider.py +++ b/artemis/llm/google/provider.py @@ -30,6 +30,7 @@ from __future__ import annotations +import os import re from typing import Any, Literal, NamedTuple @@ -78,6 +79,28 @@ def is_google_family_provider(value: Any) -> bool: return _normalize_provider(value) in _GOOGLE_FAMILY_PROVIDERS +def use_vertex_backend() -> bool: + """True when ``GOOGLE_GENAI_USE_VERTEXAI`` routes Gemini calls to Vertex AI. + + The flag is google-genai's own switch. With it set, every ``google`` + endpoint runs on Vertex through Application Default Credentials, so one + environment line moves all agents (helpers included) off the API key. + """ + return os.environ.get("GOOGLE_GENAI_USE_VERTEXAI", "").strip().lower() in {"1", "true", "yes"} + + +def vertex_project() -> str | None: + """Google Cloud project for Vertex AI; None lets google-auth use the ADC default.""" + return os.environ.get("VERTEX_AI_PROJECT") or os.environ.get("GOOGLE_CLOUD_PROJECT") + + +def vertex_location() -> str: + """Vertex AI location; Gemini 3.x models are served from ``global``.""" + return ( + os.environ.get("VERTEX_AI_LOCATION") or os.environ.get("GOOGLE_CLOUD_LOCATION") or "global" + ) + + def is_google_chat_model(model: Any) -> bool: """True when a LangChain chat model is backed by Google (by class name). diff --git a/artemis/llm/router.py b/artemis/llm/router.py index f2e6b42b6..192cce919 100644 --- a/artemis/llm/router.py +++ b/artemis/llm/router.py @@ -28,7 +28,12 @@ from pydantic import BaseModel, Field from artemis.config.settings import settings -from artemis.llm.google.provider import supports_thinking_level +from artemis.llm.google.provider import ( + supports_thinking_level, + use_vertex_backend, + vertex_location, + vertex_project, +) from third_party.mobile_use.utils.logger import get_logger logger = get_logger(__name__) @@ -229,22 +234,39 @@ def create_model(cls, endpoint: ModelEndpoint) -> BaseChatModel: HarmCategory, ) - api_key = ( - endpoint.api_key - or (settings.GOOGLE_API_KEY.get_secret_value() if settings.GOOGLE_API_KEY else None) - or os.environ.get("GOOGLE_API_KEY") - or os.environ.get("GEMINI_API_KEY") - ) # Gemini 1.x/2.x only understand thinking_budget. thinking_level = ( endpoint.thinking_level if supports_thinking_level(endpoint.model_name) else None ) + # GOOGLE_GENAI_USE_VERTEXAI moves every google endpoint, helper agents + # included, to Vertex AI with Application Default Credentials. The + # google-genai client keeps the developer-API message conversion. + if use_vertex_backend(): + auth: dict[str, Any] = { + "vertexai": True, + "project": vertex_project(), + "location": vertex_location(), + } + else: + auth = { + "api_key": ( + endpoint.api_key + or ( + settings.GOOGLE_API_KEY.get_secret_value() + if settings.GOOGLE_API_KEY + else None + ) + or os.environ.get("GOOGLE_API_KEY") + or os.environ.get("GEMINI_API_KEY") + ) + } + kwargs: dict[str, Any] = { "model": endpoint.model_name, "temperature": endpoint.temperature, "max_output_tokens": endpoint.max_tokens, - "api_key": api_key, + **auth, "timeout": endpoint.timeout_seconds, "thinking_budget": endpoint.thinking_budget, "thinking_level": thinking_level, diff --git a/tests/conftest.py b/tests/conftest.py index b99e8d4dd..dea6bedda 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -26,6 +26,12 @@ from artemis.drivers.mock.mock_driver import MockDeviceDriver +@pytest.fixture(autouse=True) +def _gemini_api_key_backend(monkeypatch): + """Keep tests on the API-key backend even when a local .env routes Gemini to Vertex.""" + monkeypatch.delenv("GOOGLE_GENAI_USE_VERTEXAI", raising=False) + + @pytest.fixture def mock_driver(): """Provide an isolated mock mobile driver.""" diff --git a/tests/unit/test_vertex_backend.py b/tests/unit/test_vertex_backend.py new file mode 100644 index 000000000..7f6e83c44 --- /dev/null +++ b/tests/unit/test_vertex_backend.py @@ -0,0 +1,119 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for routing Gemini endpoints to Vertex AI.""" + +from types import SimpleNamespace + +import pytest + +from artemis.agents.explorer.explorer import Explorer +from artemis.llm.google import use_vertex_backend, vertex_location, vertex_project +from artemis.llm.router import ModelEndpoint, ModelFactory, ModelProvider + + +class _RecordingChatModel: + """Stands in for ChatGoogleGenerativeAI and keeps the constructor kwargs.""" + + def __init__(self, **kwargs): + self.kwargs = kwargs + + +@pytest.fixture +def recording_chat_model(monkeypatch): + monkeypatch.setattr("langchain_google_genai.ChatGoogleGenerativeAI", _RecordingChatModel) + monkeypatch.setattr("artemis.llm.router._patch_langchain_google_genai", lambda: None) + + +@pytest.fixture +def clean_vertex_env(monkeypatch): + for name in ( + "GOOGLE_GENAI_USE_VERTEXAI", + "VERTEX_AI_PROJECT", + "VERTEX_AI_LOCATION", + "GOOGLE_CLOUD_PROJECT", + "GOOGLE_CLOUD_LOCATION", + ): + monkeypatch.delenv(name, raising=False) + return monkeypatch + + +@pytest.mark.parametrize("value", ["true", "True", "1", "yes"]) +def test_use_vertex_backend_reads_genai_flag(clean_vertex_env, value): + clean_vertex_env.setenv("GOOGLE_GENAI_USE_VERTEXAI", value) + assert use_vertex_backend() is True + + +@pytest.mark.parametrize("value", ["", "false", "0", "no"]) +def test_use_vertex_backend_off(clean_vertex_env, value): + clean_vertex_env.setenv("GOOGLE_GENAI_USE_VERTEXAI", value) + assert use_vertex_backend() is False + + +def test_vertex_location_defaults_to_global(clean_vertex_env): + assert vertex_location() == "global" + clean_vertex_env.setenv("GOOGLE_CLOUD_LOCATION", "us-central1") + assert vertex_location() == "us-central1" + clean_vertex_env.setenv("VERTEX_AI_LOCATION", "europe-west4") + assert vertex_location() == "europe-west4" + + +def test_vertex_project_prefers_artemis_variable(clean_vertex_env): + assert vertex_project() is None + clean_vertex_env.setenv("GOOGLE_CLOUD_PROJECT", "from-gcloud") + assert vertex_project() == "from-gcloud" + clean_vertex_env.setenv("VERTEX_AI_PROJECT", "from-artemis") + assert vertex_project() == "from-artemis" + + +def test_google_endpoint_uses_api_key_by_default(clean_vertex_env, recording_chat_model): + ep = ModelEndpoint( + provider=ModelProvider.GOOGLE, model_name="gemini-3.8-flash", api_key="test-key" + ) + model = ModelFactory.create_model(ep) + assert model.kwargs["api_key"] == "test-key" + assert "vertexai" not in model.kwargs + + +def test_google_endpoint_moves_to_vertex_with_flag(clean_vertex_env, recording_chat_model): + clean_vertex_env.setenv("GOOGLE_GENAI_USE_VERTEXAI", "true") + clean_vertex_env.setenv("GOOGLE_CLOUD_PROJECT", "test-project") + ep = ModelEndpoint( + provider=ModelProvider.GOOGLE, model_name="gemini-3.5-flash-lite", api_key="test-key" + ) + model = ModelFactory.create_model(ep) + assert model.kwargs["vertexai"] is True + assert model.kwargs["project"] == "test-project" + assert model.kwargs["location"] == "global" + assert "api_key" not in model.kwargs + assert model.kwargs["model"] == "gemini-3.5-flash-lite" + + +def test_flag_keeps_explicit_project_and_location(clean_vertex_env, recording_chat_model): + clean_vertex_env.setenv("GOOGLE_GENAI_USE_VERTEXAI", "1") + clean_vertex_env.setenv("VERTEX_AI_PROJECT", "artemis-project") + clean_vertex_env.setenv("VERTEX_AI_LOCATION", "us-central1") + ep = ModelEndpoint(provider=ModelProvider.GOOGLE, model_name="gemini-3.8-flash") + model = ModelFactory.create_model(ep) + assert model.kwargs["project"] == "artemis-project" + assert model.kwargs["location"] == "us-central1" + + +def test_explorer_skips_native_engine_on_vertex(clean_vertex_env, monkeypatch): + explorer = SimpleNamespace( + model_name="gemini-3.8-flash", ctx=SimpleNamespace(_genai_client=object()) + ) + assert Explorer._detect_native_engine(explorer) is True + monkeypatch.setenv("GOOGLE_GENAI_USE_VERTEXAI", "true") + assert Explorer._detect_native_engine(explorer) is False