Merge pull request #59 from geojaz/feat/vertexai-auth

Feat: Enables authentication to google via VertexAI and application default creds

.env.example

@@ -20,8 +20,14 @@ OPENAI_BASE_URL="https://api.openai.com/v1"
 
 # Example Google mapping:
 # PREFERRED_PROVIDER="google"
-# BIG_MODEL="gemini-2.5-pro-preview-03-25"
+# BIG_MODEL="gemini-2.5-pro"
-# SMALL_MODEL="gemini-2.0-flash"
+# SMALL_MODEL="gemini-2.5-flash"
+
+# Example Google with vertex AI auth via ADC:
+# PREFERRED_PROVIDER="google"
+# USE_VERTEX_AUTH=true
+# BIG_MODEL="gemini-2.5-pro"
+# SMALL_MODEL="gemini-2.5-flash"
 
 # Example "just an Anthropic proxy" mode:
 # PREFERRED_PROVIDER="anthropic"

README.md

@@ -13,6 +13,7 @@ A proxy server that lets you use Anthropic clients with Gemini, OpenAI, or Anthr
 
 - OpenAI API key 🔑
 - Google AI Studio (Gemini) API key (if using Google provider) 🔑
+- Google Cloud Project with Vertex AI API enabled (if using Application Default Credentials for Gemini) ☁️
 - [uv](https://github.com/astral-sh/uv) installed.
 
 ### Setup 🛠️
@@ -40,7 +41,10 @@ A proxy server that lets you use Anthropic clients with Gemini, OpenAI, or Anthr
 
    *   `ANTHROPIC_API_KEY`: (Optional) Needed only if proxying *to* Anthropic models.
    *   `OPENAI_API_KEY`: Your OpenAI API key (Required if using the default OpenAI preference or as fallback).
-   *   `GEMINI_API_KEY`: Your Google AI Studio (Gemini) API key (Required if PREFERRED_PROVIDER=google).
+   *   `GEMINI_API_KEY`: Your Google AI Studio (Gemini) API key (Required if `PREFERRED_PROVIDER=google` and `USE_VERTEX_AUTH=true`).
+   *   `USE_VERTEX_AUTH` (Optional): Set to `true` to use Application Default Credentials (ADC) will be used (no static API key required). Note: when USE_VERTEX_AUTH=true, you must configure `VERTEX_PROJECT` and `VERTEX_LOCATION`.
+   *   `VERTEX_PROJECT` (Optional): Your Google Cloud Project ID (Required if `PREFERRED_PROVIDER=google` and `USE_VERTEX_AUTH=true`).
+   *   `VERTEX_LOCATION` (Optional): The Google Cloud region for Vertex AI (e.g., `us-central1`) (Required if `PREFERRED_PROVIDER=google` and `USE_VERTEX_AUTH=true`).
    *   `PREFERRED_PROVIDER` (Optional): Set to `openai` (default), `google`, or `anthropic`. This determines the primary backend for mapping `haiku`/`sonnet`.
    *   `BIG_MODEL` (Optional): The model to map `sonnet` requests to. Defaults to `gpt-4.1` (if `PREFERRED_PROVIDER=openai`) or `gemini-2.5-pro-preview-03-25`. Ignored when `PREFERRED_PROVIDER=anthropic`.
    *   `SMALL_MODEL` (Optional): The model to map `haiku` requests to. Defaults to `gpt-4.1-mini` (if `PREFERRED_PROVIDER=openai`) or `gemini-2.0-flash`. Ignored when `PREFERRED_PROVIDER=anthropic`.
@@ -151,13 +155,24 @@ GEMINI_API_KEY="your-google-key" # Needed if PREFERRED_PROVIDER=google
 # SMALL_MODEL="gpt-4.1-mini" # Optional, it's the default
 ```
 
-**Example 2: Prefer Google**
+**Example 2a: Prefer Google (using GEMINI_API_KEY)**
 ```dotenv
 GEMINI_API_KEY="your-google-key"
 OPENAI_API_KEY="your-openai-key" # Needed for fallback
 PREFERRED_PROVIDER="google"
-# BIG_MODEL="gemini-2.5-pro-preview-03-25" # Optional, it's the default for Google pref
+# BIG_MODEL="gemini-2.5-pro" # Optional, it's the default for Google pref
-# SMALL_MODEL="gemini-2.0-flash" # Optional, it's the default for Google pref
+# SMALL_MODEL="gemini-2.5-flash" # Optional, it's the default for Google pref
+```
+
+**Example 2b: Prefer Google (using Vertex AI with Application Default Credentials)**
+```dotenv
+OPENAI_API_KEY="your-openai-key" # Needed for fallback
+PREFERRED_PROVIDER="google"
+VERTEX_PROJECT="your-gcp-project-id"
+VERTEX_LOCATION="us-central1"
+USE_VERTEX_AUTH=true
+# BIG_MODEL="gemini-2.5-pro" # Optional, it's the default for Google pref
+# SMALL_MODEL="gemini-2.5-flash" # Optional, it's the default for Google pref
 ```
 
 **Example 3: Use Direct Anthropic ("Just an Anthropic Proxy" Mode)**

pyproject.toml

@@ -9,7 +9,9 @@ dependencies = [
     "uvicorn>=0.34.0",
     "httpx>=0.25.0",
     "pydantic>=2.0.0",
-    "litellm>=1.40.14",
+    "litellm>=1.77.7",
     "python-dotenv>=1.0.0",
+    "google-auth>=2.41.1",
+    "google-cloud-aiplatform>=1.120.0",
 ]
 

server.py

@@ -82,6 +82,13 @@ ANTHROPIC_API_KEY = os.environ.get("ANTHROPIC_API_KEY")
 OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
 GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
 
+# Get Vertex AI project and location from environment (if set)
+VERTEX_PROJECT = os.environ.get("VERTEX_PROJECT", "unset")
+VERTEX_LOCATION = os.environ.get("VERTEX_LOCATION", "unset")
+
+# Option to use Gemini API key instead of ADC for Vertex AI
+USE_VERTEX_AUTH = os.environ.get("USE_VERTEX_AUTH", "False").lower() == "true"
+
 # Get OpenAI base URL from environment (if set)
 OPENAI_BASE_URL = os.environ.get("OPENAI_BASE_URL")
 
@@ -1125,8 +1132,14 @@ async def create_message(
             else:
                 logger.debug(f"Using OpenAI API key for model: {request.model}")
         elif request.model.startswith("gemini/"):
-            litellm_request["api_key"] = GEMINI_API_KEY
+            if USE_VERTEX_AUTH:
-            logger.debug(f"Using Gemini API key for model: {request.model}")
+                litellm_request["vertex_project"] = VERTEX_PROJECT
+                litellm_request["vertex_location"] = VERTEX_LOCATION
+                litellm_request["custom_llm_provider"] = "vertex_ai"
+                logger.debug(f"Using Gemini ADC with project={VERTEX_PROJECT}, location={VERTEX_LOCATION} and model: {request.model}")
+            else:
+                litellm_request["api_key"] = GEMINI_API_KEY
+                logger.debug(f"Using Gemini API key for model: {request.model}")
         else:
             litellm_request["api_key"] = ANTHROPIC_API_KEY
             logger.debug(f"Using Anthropic API key for model: {request.model}")