Skip to content

Commit 151d45e

Browse files
authored
improvement: google vertex support, make google key optional (#6537)
This makes the Google API Key optional so you can use hit Google Vertex using environment variabes and application default credentials (ADC)
1 parent 36549d6 commit 151d45e

File tree

5 files changed

+109
-49
lines changed

5 files changed

+109
-49
lines changed

docs/guides/configuration/llm_providers.md

Lines changed: 50 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,7 @@ Below we describe how to connect marimo to your AI provider.
7272

7373
**Requirements**
7474

75-
* `pip install openai`
75+
* `pip install openai` or `uv add openai`
7676

7777
**Configuration**
7878

@@ -142,20 +142,43 @@ Use `profile_name` for a non-default named profile, or rely on env vars/standard
142142

143143
**Requirements**
144144

145-
* API key from [Google AI Studio](https://aistudio.google.com/app/apikey)
146145
* `pip install google-genai`
147146

148-
**Configuration**
147+
You can use Google AI via two backends: **Google AI Studio** (API key) or **Google Vertex AI** (no API key required).
148+
149+
#### Using Google AI Studio (API key)
150+
151+
1. Sign up at [Google AI Studio](https://aistudio.google.com/app/apikey) and obtain your API key.
152+
2. Configure `marimo.toml` (or set these in the editor Settings):
149153

150154
```toml title="marimo.toml"
151155
[ai.models]
152-
chat_model = "google/gemini-2.5-pro" # also see gemini-2.0-flash, etc.
153-
# Model list: https://ai.google.dev/gemini-api/docs/models/gemini
156+
chat_model = "google/gemini-2.5-pro"
157+
# or any model from https://ai.google.dev/gemini-api/docs/models/gemini
154158

155159
[ai.google]
156160
api_key = "AI..."
157161
```
158162

163+
#### Using Google Vertex AI (no API key required)
164+
165+
1. Ensure you have access to a Google Cloud project with Vertex AI enabled.
166+
2. Set the following environment variables before starting marimo:
167+
168+
```bash
169+
export GOOGLE_GENAI_USE_VERTEXAI=true
170+
export GOOGLE_CLOUD_PROJECT='your-project-id'
171+
export GOOGLE_CLOUD_LOCATION='us-central1'
172+
```
173+
174+
* `GOOGLE_GENAI_USE_VERTEXAI=true` tells the client to use Vertex AI.
175+
* `GOOGLE_CLOUD_PROJECT` is your GCP project ID.
176+
* `GOOGLE_CLOUD_LOCATION` is your region (e.g., `us-central1`).
177+
178+
3. No API key is needed in your `marimo.toml` for Vertex AI.
179+
180+
For details and advanced configuration, see the `google-genai` Python client docs: `https://googleapis.github.io/python-genai/#create-a-client`.
181+
159182
### GitHub Copilot
160183

161184
Use Copilot for code refactoring or the chat panel (Copilot subscription required).
@@ -193,7 +216,7 @@ Route to many providers through OpenRouter with a single API.
193216
**Requirements**
194217

195218
* Create an API key: [OpenRouter Dashboard](https://openrouter.ai/)
196-
* `pip install openai` (OpenRouter is OpenAI‑compatible)
219+
* `pip install openai` or `uv add openai` (OpenRouter is OpenAI‑compatible)
197220

198221
**Configuration**
199222

@@ -219,7 +242,11 @@ Run open-source LLMs locally and connect via an OpenAI‑compatible API.
219242
**Requirements**
220243

221244
* Install [Ollama](https://ollama.com/)
222-
* Pull a model
245+
* `pip install openai` or `uv add openai`
246+
247+
**Setup**
248+
249+
1. Pull a model
223250

224251
```bash
225252
# View available models at https://ollama.com/library
@@ -230,26 +257,27 @@ Run open-source LLMs locally and connect via an OpenAI‑compatible API.
230257
ollama ls
231258
```
232259

233-
3. Start the Ollama server:
260+
2. Start the Ollama server:
234261

235262
```bash
236263
ollama serve
237264
# In another terminal, run a model (optional)
238265
ollama run codellama
239266
```
240267

241-
4. Visit <http://127.0.0.1:11434> to confirm that the server is running.
268+
3. Visit <http://127.0.0.1:11434> to confirm that the server is running.
242269

243270
!!! note "Port already in use"
244271
If you get a "port already in use" error, you may need to close an existing Ollama instance. On Windows, click the up arrow in the taskbar, find the Ollama icon, and select "Quit". This is a known issue (see [Ollama Issue #3575](https://github.com/ollama/ollama/issues/3575)). Once you've closed the existing Ollama instance, you should be able to run `ollama serve` successfully.
245272

246-
5. Install the OpenAI client (`pip install openai` or `uv add openai`).
247-
248-
6. Start marimo:
273+
**Configuration**
249274

250-
```bash
251-
marimo edit notebook.py
252-
```
275+
```toml title="marimo.toml"
276+
[ai.models]
277+
chat_model = "ollama/llama3.1:latest"
278+
edit_model = "ollama/codellama"
279+
autocomplete_model = "ollama/codellama" # or another model from `ollama ls`
280+
```
253281

254282
??? warning "Important: Use the `/v1` endpoint"
255283

@@ -270,15 +298,6 @@ Run open-source LLMs locally and connect via an OpenAI‑compatible API.
270298
curl http://127.0.0.1:11434/v1/models
271299
```
272300

273-
7. Configure `marimo.toml` (or use Settings):
274-
275-
```toml title="marimo.toml"
276-
[ai.models]
277-
chat_model = "ollama/llama3.1:latest"
278-
edit_model = "ollama/codellama"
279-
autocomplete_model = "ollama/codellama" # or another model from `ollama ls`
280-
```
281-
282301
### OpenAI-compatible providers
283302

284303
Many providers expose OpenAI-compatible endpoints. Point `base_url` at the provider and use their models.
@@ -288,7 +307,7 @@ Common examples include [GROQ](https://console.groq.com/docs/openai), DeepSeek,
288307

289308
* Provider API key
290309
* Provider OpenAI-compatible `base_url`
291-
* `pip install openai`
310+
* `pip install openai` or `uv add openai`
292311

293312
**Configuration**
294313

@@ -334,7 +353,7 @@ Use DeepSeek via its OpenAI‑compatible API.
334353
**Requirements**
335354

336355
* DeepSeek API key
337-
* `pip install openai`
356+
* `pip install openai` or `uv add openai`
338357

339358
**Configuration**
340359

@@ -354,7 +373,7 @@ Use Grok models via xAI's OpenAI‑compatible API.
354373
**Requirements**
355374

356375
* xAI API key
357-
* `pip install openai`
376+
* `pip install openai` or `uv add openai`
358377

359378
**Configuration**
360379

@@ -374,7 +393,7 @@ Connect to a local model served by LM Studio's OpenAI‑compatible endpoint.
374393
**Requirements**
375394

376395
* Install LM Studio and start its server
377-
* `pip install openai`
396+
* `pip install openai` or `uv add openai`
378397

379398
**Configuration**
380399

@@ -393,7 +412,7 @@ Use Mistral via its OpenAI‑compatible API.
393412
**Requirements**
394413

395414
* Mistral API key
396-
* `pip install openai`
415+
* `pip install openai` or `uv add openai`
397416

398417
**Configuration**
399418

@@ -413,7 +432,7 @@ Access multiple hosted models via Together AI's OpenAI‑compatible API.
413432
**Requirements**
414433

415434
* Together AI API key
416-
* `pip install openai`
435+
* `pip install openai` or `uv add openai`
417436

418437
**Configuration**
419438

@@ -433,7 +452,7 @@ Use Vercel's v0 OpenAI‑compatible models for app-oriented generation.
433452
**Requirements**
434453

435454
* v0 API key
436-
* `pip install openai`
455+
* `pip install openai` or `uv add openai`
437456

438457
**Configuration**
439458

marimo/_server/ai/config.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -162,11 +162,12 @@ def for_google(cls, config: AiConfig) -> AnyProviderConfig:
162162
ai_config,
163163
"Google AI",
164164
fallback_key=fallback_key,
165-
require_key=True,
165+
require_key=False,
166166
)
167167
return cls(
168168
base_url=_get_base_url(ai_config),
169169
api_key=key,
170+
ssl_verify=True,
170171
tools=_get_tools(config.get("mode", "manual")),
171172
)
172173

marimo/_server/ai/providers.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -777,6 +777,23 @@ def get_client(self, config: AnyProviderConfig) -> GoogleClient:
777777
)
778778
from google import genai # type: ignore
779779

780+
# If no API key is provided, try to use environment variables and ADC
781+
# This supports Google Vertex AI usage without explicit API keys
782+
if not config.api_key:
783+
# Check if GOOGLE_GENAI_USE_VERTEXAI is set to enable Vertex AI mode
784+
use_vertex = (
785+
os.getenv("GOOGLE_GENAI_USE_VERTEXAI", "").lower() == "true"
786+
)
787+
if use_vertex:
788+
project = os.getenv("GOOGLE_CLOUD_PROJECT")
789+
location = os.getenv("GOOGLE_CLOUD_LOCATION", "us-central1")
790+
return genai.Client(
791+
vertexai=True, project=project, location=location
792+
).aio
793+
else:
794+
# Try default initialization which may work with environment variables
795+
return genai.Client().aio
796+
780797
return genai.Client(api_key=config.api_key).aio
781798

782799
async def stream_completion(

tests/_server/ai/test_ai_config.py

Lines changed: 14 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -506,14 +506,16 @@ def test_for_google_config_key_takes_precedence(self) -> None:
506506

507507
@patch.dict(os.environ, {}, clear=True)
508508
def test_for_google_no_fallback_available(self) -> None:
509-
"""Test Google config fails when no config key and no env vars."""
509+
"""Test Google config succeeds with empty key when no env vars."""
510510
config: AiConfig = {"google": {}}
511511

512-
with pytest.raises(HTTPException) as exc_info:
513-
AnyProviderConfig.for_google(config)
512+
provider_config = AnyProviderConfig.for_google(config)
514513

515-
assert exc_info.value.status_code == HTTPStatus.BAD_REQUEST
516-
assert "Google AI API key not configured" in str(exc_info.value.detail)
514+
assert provider_config == AnyProviderConfig(
515+
base_url=None,
516+
api_key="",
517+
ssl_verify=True,
518+
)
517519

518520
@patch.dict(os.environ, {"GITHUB_TOKEN": "env-github-token"})
519521
def test_for_github_with_fallback_key(self) -> None:
@@ -996,14 +998,15 @@ def test_anthropic_config_missing(self):
996998
assert "Anthropic API key not configured" in str(exc_info.value.detail)
997999

9981000
def test_google_config_missing(self):
999-
"""Test error when Google config is missing."""
1001+
"""Test Google config defaults to empty key when config is missing."""
10001002
config: AiConfig = {}
10011003

1002-
with pytest.raises(HTTPException) as exc_info:
1003-
AnyProviderConfig.for_google(config)
1004-
1005-
assert exc_info.value.status_code == HTTPStatus.BAD_REQUEST
1006-
assert "Google AI API key not configured" in str(exc_info.value.detail)
1004+
provider_config = AnyProviderConfig.for_google(config)
1005+
assert provider_config == AnyProviderConfig(
1006+
base_url=None,
1007+
api_key="",
1008+
ssl_verify=True,
1009+
)
10071010

10081011
def test_bedrock_config_missing(self):
10091012
"""Test when Bedrock config is missing, should not error since could use environment variables."""

tests/_server/api/endpoints/test_ai.py

Lines changed: 26 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -591,13 +591,29 @@ async def mock_stream():
591591

592592
@staticmethod
593593
@with_session(SESSION_ID)
594-
@patch("google.genai.client.AsyncClient")
594+
@patch("google.genai.Client")
595595
def test_google_ai_completion_without_token(
596596
client: TestClient, google_ai_mock: Any
597597
) -> None:
598-
del google_ai_mock
599598
user_config_manager = get_session_config_manager(client)
600599

600+
# Mock the google client and its aio attribute
601+
google_client_mock = MagicMock()
602+
google_aio_mock = MagicMock()
603+
google_client_mock.aio = google_aio_mock
604+
google_ai_mock.return_value = google_client_mock
605+
606+
# Mock async stream
607+
async def mock_stream():
608+
yield MagicMock(
609+
text="import pandas as pd",
610+
thought=None,
611+
)
612+
613+
google_aio_mock.models.generate_content_stream = AsyncMock(
614+
side_effect=lambda **kwargs: mock_stream() # noqa: ARG005
615+
)
616+
601617
config = {
602618
"ai": {
603619
"open_ai": {"model": "gemini-1.5-pro"},
@@ -617,10 +633,14 @@ def test_google_ai_completion_without_token(
617633
"code": "",
618634
},
619635
)
620-
assert response.status_code == 400, response.text
621-
assert response.json() == {
622-
"detail": "Google AI API key not configured. Go to Settings > AI to configure."
623-
}
636+
637+
assert response.status_code == 200, response.text
638+
prompt = (
639+
google_aio_mock.models.generate_content_stream.call_args.kwargs[
640+
"contents"
641+
]
642+
)
643+
assert prompt[0]["parts"][0]["text"] == "Help me create a dataframe"
624644

625645
@staticmethod
626646
@with_session(SESSION_ID)

0 commit comments

Comments
 (0)