diff --git a/.env.example b/.env.example index 4df12e2..32d4533 100644 --- a/.env.example +++ b/.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" -# SMALL_MODEL="gemini-2.0-flash" +# BIG_MODEL="gemini-2.5-pro" +# 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" diff --git a/README.md b/README.md index f4c46d2..e8bc6e0 100644 --- a/README.md +++ b/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 -# SMALL_MODEL="gemini-2.0-flash" # 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.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)** diff --git a/pyproject.toml b/pyproject.toml index b025174..f64c095 100644 --- a/pyproject.toml +++ b/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", ] diff --git a/server.py b/server.py index 2adaa88..bf2ac6c 100644 --- a/server.py +++ b/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 - logger.debug(f"Using Gemini API key for model: {request.model}") + if USE_VERTEX_AUTH: + 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}")