diff --git a/docs/guides/configuration/llm_providers.md b/docs/guides/configuration/llm_providers.md index 8a0763f9836..34f849b7f57 100644 --- a/docs/guides/configuration/llm_providers.md +++ b/docs/guides/configuration/llm_providers.md @@ -72,7 +72,7 @@ Below we describe how to connect marimo to your AI provider. **Requirements** -* `pip install openai` +* `pip install openai` or `uv add openai` **Configuration** @@ -142,20 +142,43 @@ Use `profile_name` for a non-default named profile, or rely on env vars/standard **Requirements** -* API key from [Google AI Studio](https://aistudio.google.com/app/apikey) * `pip install google-genai` -**Configuration** +You can use Google AI via two backends: **Google AI Studio** (API key) or **Google Vertex AI** (no API key required). + +#### Using Google AI Studio (API key) + +1. Sign up at [Google AI Studio](https://aistudio.google.com/app/apikey) and obtain your API key. +2. Configure `marimo.toml` (or set these in the editor Settings): ```toml title="marimo.toml" [ai.models] -chat_model = "google/gemini-2.5-pro" # also see gemini-2.0-flash, etc. -# Model list: https://ai.google.dev/gemini-api/docs/models/gemini +chat_model = "google/gemini-2.5-pro" +# or any model from https://ai.google.dev/gemini-api/docs/models/gemini [ai.google] api_key = "AI..." ``` +#### Using Google Vertex AI (no API key required) + +1. Ensure you have access to a Google Cloud project with Vertex AI enabled. +2. Set the following environment variables before starting marimo: + +```bash +export GOOGLE_GENAI_USE_VERTEXAI=true +export GOOGLE_CLOUD_PROJECT='your-project-id' +export GOOGLE_CLOUD_LOCATION='us-central1' +``` + +* `GOOGLE_GENAI_USE_VERTEXAI=true` tells the client to use Vertex AI. +* `GOOGLE_CLOUD_PROJECT` is your GCP project ID. +* `GOOGLE_CLOUD_LOCATION` is your region (e.g., `us-central1`). + +3. No API key is needed in your `marimo.toml` for Vertex AI. + +For details and advanced configuration, see the `google-genai` Python client docs: `https://googleapis.github.io/python-genai/#create-a-client`. + ### GitHub Copilot 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. **Requirements** * Create an API key: [OpenRouter Dashboard](https://openrouter.ai/) -* `pip install openai` (OpenRouter is OpenAI‑compatible) +* `pip install openai` or `uv add openai` (OpenRouter is OpenAI‑compatible) **Configuration** @@ -219,7 +242,11 @@ Run open-source LLMs locally and connect via an OpenAI‑compatible API. **Requirements** * Install [Ollama](https://ollama.com/) -* Pull a model +* `pip install openai` or `uv add openai` + +**Setup** + +1. Pull a model ```bash # View available models at https://ollama.com/library @@ -230,7 +257,7 @@ Run open-source LLMs locally and connect via an OpenAI‑compatible API. ollama ls ``` -3. Start the Ollama server: +2. Start the Ollama server: ```bash ollama serve @@ -238,18 +265,19 @@ Run open-source LLMs locally and connect via an OpenAI‑compatible API. ollama run codellama ``` -4. Visit to confirm that the server is running. +3. Visit to confirm that the server is running. !!! note "Port already in use" 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. -5. Install the OpenAI client (`pip install openai` or `uv add openai`). - -6. Start marimo: +**Configuration** - ```bash - marimo edit notebook.py - ``` +```toml title="marimo.toml" +[ai.models] +chat_model = "ollama/llama3.1:latest" +edit_model = "ollama/codellama" +autocomplete_model = "ollama/codellama" # or another model from `ollama ls` +``` ??? warning "Important: Use the `/v1` endpoint" @@ -270,15 +298,6 @@ Run open-source LLMs locally and connect via an OpenAI‑compatible API. curl http://127.0.0.1:11434/v1/models ``` -7. Configure `marimo.toml` (or use Settings): - -```toml title="marimo.toml" -[ai.models] -chat_model = "ollama/llama3.1:latest" -edit_model = "ollama/codellama" -autocomplete_model = "ollama/codellama" # or another model from `ollama ls` -``` - ### OpenAI-compatible providers 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, * Provider API key * Provider OpenAI-compatible `base_url` -* `pip install openai` +* `pip install openai` or `uv add openai` **Configuration** @@ -334,7 +353,7 @@ Use DeepSeek via its OpenAI‑compatible API. **Requirements** * DeepSeek API key -* `pip install openai` +* `pip install openai` or `uv add openai` **Configuration** @@ -354,7 +373,7 @@ Use Grok models via xAI's OpenAI‑compatible API. **Requirements** * xAI API key -* `pip install openai` +* `pip install openai` or `uv add openai` **Configuration** @@ -374,7 +393,7 @@ Connect to a local model served by LM Studio's OpenAI‑compatible endpoint. **Requirements** * Install LM Studio and start its server -* `pip install openai` +* `pip install openai` or `uv add openai` **Configuration** @@ -393,7 +412,7 @@ Use Mistral via its OpenAI‑compatible API. **Requirements** * Mistral API key -* `pip install openai` +* `pip install openai` or `uv add openai` **Configuration** @@ -413,7 +432,7 @@ Access multiple hosted models via Together AI's OpenAI‑compatible API. **Requirements** * Together AI API key -* `pip install openai` +* `pip install openai` or `uv add openai` **Configuration** @@ -433,7 +452,7 @@ Use Vercel's v0 OpenAI‑compatible models for app-oriented generation. **Requirements** * v0 API key -* `pip install openai` +* `pip install openai` or `uv add openai` **Configuration** diff --git a/marimo/_server/ai/config.py b/marimo/_server/ai/config.py index 0f23815b6fd..7a42876ba79 100644 --- a/marimo/_server/ai/config.py +++ b/marimo/_server/ai/config.py @@ -162,11 +162,12 @@ def for_google(cls, config: AiConfig) -> AnyProviderConfig: ai_config, "Google AI", fallback_key=fallback_key, - require_key=True, + require_key=False, ) return cls( base_url=_get_base_url(ai_config), api_key=key, + ssl_verify=True, tools=_get_tools(config.get("mode", "manual")), ) diff --git a/marimo/_server/ai/providers.py b/marimo/_server/ai/providers.py index 0b8f73d27d5..a7d3d57a615 100644 --- a/marimo/_server/ai/providers.py +++ b/marimo/_server/ai/providers.py @@ -777,6 +777,23 @@ def get_client(self, config: AnyProviderConfig) -> GoogleClient: ) from google import genai # type: ignore + # If no API key is provided, try to use environment variables and ADC + # This supports Google Vertex AI usage without explicit API keys + if not config.api_key: + # Check if GOOGLE_GENAI_USE_VERTEXAI is set to enable Vertex AI mode + use_vertex = ( + os.getenv("GOOGLE_GENAI_USE_VERTEXAI", "").lower() == "true" + ) + if use_vertex: + project = os.getenv("GOOGLE_CLOUD_PROJECT") + location = os.getenv("GOOGLE_CLOUD_LOCATION", "us-central1") + return genai.Client( + vertexai=True, project=project, location=location + ).aio + else: + # Try default initialization which may work with environment variables + return genai.Client().aio + return genai.Client(api_key=config.api_key).aio async def stream_completion( diff --git a/tests/_server/ai/test_ai_config.py b/tests/_server/ai/test_ai_config.py index 17a7df44ec4..7846b9106bf 100644 --- a/tests/_server/ai/test_ai_config.py +++ b/tests/_server/ai/test_ai_config.py @@ -506,14 +506,16 @@ def test_for_google_config_key_takes_precedence(self) -> None: @patch.dict(os.environ, {}, clear=True) def test_for_google_no_fallback_available(self) -> None: - """Test Google config fails when no config key and no env vars.""" + """Test Google config succeeds with empty key when no env vars.""" config: AiConfig = {"google": {}} - with pytest.raises(HTTPException) as exc_info: - AnyProviderConfig.for_google(config) + provider_config = AnyProviderConfig.for_google(config) - assert exc_info.value.status_code == HTTPStatus.BAD_REQUEST - assert "Google AI API key not configured" in str(exc_info.value.detail) + assert provider_config == AnyProviderConfig( + base_url=None, + api_key="", + ssl_verify=True, + ) @patch.dict(os.environ, {"GITHUB_TOKEN": "env-github-token"}) def test_for_github_with_fallback_key(self) -> None: @@ -996,14 +998,15 @@ def test_anthropic_config_missing(self): assert "Anthropic API key not configured" in str(exc_info.value.detail) def test_google_config_missing(self): - """Test error when Google config is missing.""" + """Test Google config defaults to empty key when config is missing.""" config: AiConfig = {} - with pytest.raises(HTTPException) as exc_info: - AnyProviderConfig.for_google(config) - - assert exc_info.value.status_code == HTTPStatus.BAD_REQUEST - assert "Google AI API key not configured" in str(exc_info.value.detail) + provider_config = AnyProviderConfig.for_google(config) + assert provider_config == AnyProviderConfig( + base_url=None, + api_key="", + ssl_verify=True, + ) def test_bedrock_config_missing(self): """Test when Bedrock config is missing, should not error since could use environment variables.""" diff --git a/tests/_server/api/endpoints/test_ai.py b/tests/_server/api/endpoints/test_ai.py index c49c711e222..188c2d90dd7 100644 --- a/tests/_server/api/endpoints/test_ai.py +++ b/tests/_server/api/endpoints/test_ai.py @@ -591,13 +591,29 @@ async def mock_stream(): @staticmethod @with_session(SESSION_ID) - @patch("google.genai.client.AsyncClient") + @patch("google.genai.Client") def test_google_ai_completion_without_token( client: TestClient, google_ai_mock: Any ) -> None: - del google_ai_mock user_config_manager = get_session_config_manager(client) + # Mock the google client and its aio attribute + google_client_mock = MagicMock() + google_aio_mock = MagicMock() + google_client_mock.aio = google_aio_mock + google_ai_mock.return_value = google_client_mock + + # Mock async stream + async def mock_stream(): + yield MagicMock( + text="import pandas as pd", + thought=None, + ) + + google_aio_mock.models.generate_content_stream = AsyncMock( + side_effect=lambda **kwargs: mock_stream() # noqa: ARG005 + ) + config = { "ai": { "open_ai": {"model": "gemini-1.5-pro"}, @@ -617,10 +633,14 @@ def test_google_ai_completion_without_token( "code": "", }, ) - assert response.status_code == 400, response.text - assert response.json() == { - "detail": "Google AI API key not configured. Go to Settings > AI to configure." - } + + assert response.status_code == 200, response.text + prompt = ( + google_aio_mock.models.generate_content_stream.call_args.kwargs[ + "contents" + ] + ) + assert prompt[0]["parts"][0]["text"] == "Help me create a dataframe" @staticmethod @with_session(SESSION_ID)