feat: add provider preference and update model mapping. Introduce PREFERRED_PROVIDER env var (default: google). Update default BIG/SMALL models to Gemini latest. Refactor model mapping logic to respect provider preference. Add .env.example and update README.

.env.example

@@ -0,0 +1,21 @@
+# Required API Keys
+ANTHROPIC_API_KEY="your-anthropic-api-key" # Needed if proxying *to* Anthropic
+OPENAI_API_KEY="sk-..."
+GEMINI_API_KEY="your-google-ai-studio-key"
+
+# Optional: Provider Preference and Model Mapping
+# Controls which provider (google or openai) is preferred for mapping haiku/sonnet.
+# Defaults to google if not set.
+PREFERRED_PROVIDER="google"
+
+# Optional: Specify the exact models to map haiku/sonnet to.
+# If PREFERRED_PROVIDER=google, these MUST be valid Gemini model names known to the server.
+# Defaults to gemini-1.5-pro-latest and gemini-1.5-flash-latest if PREFERRED_PROVIDER=google.
+# Defaults to gpt-4o and gpt-4o-mini if PREFERRED_PROVIDER=openai.
+# BIG_MODEL="gemini-1.5-pro-latest"
+# SMALL_MODEL="gemini-1.5-flash-latest"
+
+# Example OpenAI mapping:
+# PREFERRED_PROVIDER="openai"
+# BIG_MODEL="gpt-4o"
+# SMALL_MODEL="gpt-4o-mini" 

README.md

@@ -31,8 +31,13 @@ A proxy server that lets you use Claude Code with OpenAI models like GPT-4o / gp
    ```
    OPENAI_API_KEY=your-openai-key
    # Optional: customize which models are used
+   # For OpenAI models (default)
    # BIG_MODEL=gpt-4o
    # SMALL_MODEL=gpt-4o-mini
+   
+   # For Gemini models
+   # BIG_MODEL=gemini-2.5-pro-preview-03-25
+   # SMALL_MODEL=gemini-2.0-flash
    ```
 
 4. **Start the proxy server**:
@@ -56,30 +61,75 @@ A proxy server that lets you use Claude Code with OpenAI models like GPT-4o / gp
 
 ## Model Mapping 🗺️
 
-The proxy automatically maps Claude models to OpenAI models:
+The proxy automatically maps Claude models to either OpenAI or Gemini models based on the configured model:
 
-| Claude Model | OpenAI Model |
+| Claude Model | Default Mapping | When BIG_MODEL/SMALL_MODEL is a Gemini model |
-|--------------|--------------|
+|--------------|--------------|---------------------------|
-| haiku | gpt-4o-mini (default) |
+| haiku | openai/gpt-4o-mini | gemini/[model-name] |
-| sonnet | gpt-4o (default) |
+| sonnet | openai/gpt-4o | gemini/[model-name] |
+
+### Supported Models
+
+#### OpenAI Models
+The following OpenAI models are supported with automatic `openai/` prefix handling:
+- o3-mini
+- o1
+- o1-mini
+- o1-pro
+- gpt-4.5-preview
+- gpt-4o
+- gpt-4o-audio-preview
+- chatgpt-4o-latest
+- gpt-4o-mini
+- gpt-4o-mini-audio-preview
+
+#### Gemini Models
+The following Gemini models are supported with automatic `gemini/` prefix handling:
+- gemini-2.5-pro-preview-03-25
+- gemini-2.0-flash
+
+### Model Prefix Handling
+The proxy automatically adds the appropriate prefix to model names:
+- OpenAI models get the `openai/` prefix 
+- Gemini models get the `gemini/` prefix
+- The BIG_MODEL and SMALL_MODEL will get the appropriate prefix based on whether they're in the OpenAI or Gemini model lists
+
+For example:
+- `gpt-4o` becomes `openai/gpt-4o`
+- `gemini-2.5-pro-preview-03-25` becomes `gemini/gemini-2.5-pro-preview-03-25`
+- When BIG_MODEL is set to a Gemini model, Claude Sonnet will map to `gemini/[model-name]`
 
 ### Customizing Model Mapping
 
-You can customize which OpenAI models are used via environment variables:
+You can customize which models are used via environment variables:
 
-- `BIG_MODEL`: The OpenAI model to use for Claude Sonnet models (default: "gpt-4o")
+- `BIG_MODEL`: The model to use for Claude Sonnet models (default: "gpt-4o")
-- `SMALL_MODEL`: The OpenAI model to use for Claude Haiku models (default: "gpt-4o-mini")
+- `SMALL_MODEL`: The model to use for Claude Haiku models (default: "gpt-4o-mini")
 
 Add these to your `.env` file to customize:
 ```
 OPENAI_API_KEY=your-openai-key
+# For OpenAI models (default)
 BIG_MODEL=gpt-4o
 SMALL_MODEL=gpt-4o-mini
+
+# For Gemini models
+# BIG_MODEL=gemini-2.5-pro-preview-03-25
+# SMALL_MODEL=gemini-2.0-flash
 ```
 
 Or set them directly when running the server:
 ```bash
+# Using OpenAI models
 BIG_MODEL=gpt-4o SMALL_MODEL=gpt-4o-mini uv run uvicorn server:app --host 0.0.0.0 --port 8082
+
+# Using Gemini models
+BIG_MODEL=gemini-2.5-pro-preview-03-25 SMALL_MODEL=gemini-2.0-flash uv run uvicorn server:app --host 0.0.0.0 --port 8082
+```
+
+To use a mix of OpenAI and Gemini models:
+```bash
+BIG_MODEL=gemini-2.5-pro-preview-03-25 SMALL_MODEL=gpt-4o-mini uv run uvicorn server:app --host 0.0.0.0 --port 8082
 ```
 
 ## How It Works 🧩

server.py

@@ -75,19 +75,65 @@ for handler in logger.handlers:
     if isinstance(handler, logging.StreamHandler):
         handler.setFormatter(ColorizedFormatter('%(asctime)s - %(levelname)s - %(message)s'))
 
-# Flag to enable model swapping between Anthropic and OpenAI
-# Always use OpenAI models
-USE_OPENAI_MODELS = True
-
 app = FastAPI()
 
 # Get API keys from environment
 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 preferred provider (default to google)
+PREFERRED_PROVIDER = os.environ.get("PREFERRED_PROVIDER", "google").lower()
 
 # Get model mapping configuration from environment
-BIG_MODEL = os.environ.get("BIG_MODEL", "gpt-4o")
+# Default to latest Gemini models if not set
-SMALL_MODEL = os.environ.get("SMALL_MODEL", "gpt-4o-mini")
+BIG_MODEL = os.environ.get("BIG_MODEL", "gemini-1.5-pro-latest")
+SMALL_MODEL = os.environ.get("SMALL_MODEL", "gemini-1.5-flash-latest")
+
+# List of OpenAI models
+OPENAI_MODELS = [
+    "o3-mini",
+    "o1",
+    "o1-mini",
+    "o1-pro",
+    "gpt-4.5-preview",
+    "gpt-4o",
+    "gpt-4o-audio-preview",
+    "chatgpt-4o-latest",
+    "gpt-4o-mini",
+    "gpt-4o-mini-audio-preview"
+]
+
+# List of Gemini models
+GEMINI_MODELS = [
+    "gemini-2.5-pro-preview-03-25",
+    "gemini-2.0-flash",
+    "gemini-1.5-pro-latest", # Added default big model
+    "gemini-1.5-flash-latest" # Added default small model
+]
+
+# Helper function to clean schema for Gemini
+def clean_gemini_schema(schema: Any) -> Any:
+    """Recursively removes unsupported fields from a JSON schema for Gemini."""
+    if isinstance(schema, dict):
+        # Remove specific keys unsupported by Gemini tool parameters
+        schema.pop("additionalProperties", None)
+        schema.pop("default", None)
+
+        # Check for unsupported 'format' in string types
+        if schema.get("type") == "string" and "format" in schema:
+            allowed_formats = {"enum", "date-time"}
+            if schema["format"] not in allowed_formats:
+                logger.debug(f"Removing unsupported format '{schema['format']}' for string type in Gemini schema.")
+                schema.pop("format")
+
+        # Recursively clean nested schemas (properties, items, etc.)
+        for key, value in list(schema.items()): # Use list() to allow modification during iteration
+            schema[key] = clean_gemini_schema(value)
+    elif isinstance(schema, list):
+        # Recursively clean items in a list
+        return [clean_gemini_schema(item) for item in schema]
+    return schema
 
 # Models for Anthropic API requests
 class ContentBlockText(BaseModel):
@@ -142,55 +188,65 @@ class MessagesRequest(BaseModel):
     original_model: Optional[str] = None  # Will store the original model name
     
     @field_validator('model')
-    def validate_model(cls, v, info):
+    def validate_model_field(cls, v, info): # Renamed to avoid conflict
-        # Store the original model name
         original_model = v
-        
+        new_model = v # Default to original value
-        # Check if we're using OpenAI models and need to swap
+
-        if USE_OPENAI_MODELS:
+        logger.debug(f"📋 MODEL VALIDATION: Original='{original_model}', Preferred='{PREFERRED_PROVIDER}', BIG='{BIG_MODEL}', SMALL='{SMALL_MODEL}'")
-            # Remove anthropic/ prefix if it exists
+
-            if v.startswith('anthropic/'):
+        # Remove provider prefixes for easier matching
-                v = v[10:]  # Remove 'anthropic/' prefix
+        clean_v = v
-            
+        if clean_v.startswith('anthropic/'):
-            # Swap Haiku with small model (default: gpt-4o-mini)
+            clean_v = clean_v[10:]
-            if 'haiku' in v.lower():
+        elif clean_v.startswith('openai/'):
+            clean_v = clean_v[7:]
+        elif clean_v.startswith('gemini/'):
+            clean_v = clean_v[7:]
+
+        # --- Mapping Logic --- START ---
+        mapped = False
+        # Map Haiku to SMALL_MODEL based on provider preference
+        if 'haiku' in clean_v.lower():
+            if PREFERRED_PROVIDER == "google" and SMALL_MODEL in GEMINI_MODELS:
+                new_model = f"gemini/{SMALL_MODEL}"
+                mapped = True
+            else:
                 new_model = f"openai/{SMALL_MODEL}"
-                logger.debug(f"📌 MODEL MAPPING: {original_model} ➡️ {new_model}")
+                mapped = True
-                v = new_model
+
-            
+        # Map Sonnet to BIG_MODEL based on provider preference
-            # Swap any Sonnet model with big model (default: gpt-4o)
+        elif 'sonnet' in clean_v.lower():
-            elif 'sonnet' in v.lower():
+            if PREFERRED_PROVIDER == "google" and BIG_MODEL in GEMINI_MODELS:
+                new_model = f"gemini/{BIG_MODEL}"
+                mapped = True
+            else:
                 new_model = f"openai/{BIG_MODEL}"
-                logger.debug(f"📌 MODEL MAPPING: {original_model} ➡️ {new_model}")
+                mapped = True
-                v = new_model
+
-            
+        # Add prefixes to non-mapped models if they match known lists
-            # Keep the model as is but add openai/ prefix if not already present
+        elif not mapped:
-            elif not v.startswith('openai/'):
+            if clean_v in GEMINI_MODELS and not v.startswith('gemini/'):
-                new_model = f"openai/{v}"
+                new_model = f"gemini/{clean_v}"
-                logger.debug(f"📌 MODEL MAPPING: {original_model} ➡️ {new_model}")
+                mapped = True # Technically mapped to add prefix
-                v = new_model
+            elif clean_v in OPENAI_MODELS and not v.startswith('openai/'):
-                
+                new_model = f"openai/{clean_v}"
-            # Store the original model in the values dictionary
+                mapped = True # Technically mapped to add prefix
-            # This will be accessible as request.original_model
+        # --- Mapping Logic --- END ---
-            values = info.data
+
-            if isinstance(values, dict):
+        if mapped:
-                values['original_model'] = original_model
+            logger.debug(f"📌 MODEL MAPPING: '{original_model}' ➡️ '{new_model}'")
-                
-            return v
         else:
-            # Original behavior - ensure anthropic/ prefix
+             # If no mapping occurred and no prefix exists, log warning or decide default
-            original_model = v
+             if not v.startswith(('openai/', 'gemini/', 'anthropic/')):
-            if not v.startswith('anthropic/'):
+                 logger.warning(f"⚠️ No prefix or mapping rule for model: '{original_model}'. Using as is.")
-                new_model = f"anthropic/{v}"
+             new_model = v # Ensure we return the original if no rule applied
-                logger.debug(f"📌 MODEL MAPPING: {original_model} ➡️ {new_model}")
+
-                
+        # Store the original model in the values dictionary
-                # Store original model
+        values = info.data
-                values = info.data
+        if isinstance(values, dict):
-                if isinstance(values, dict):
+            values['original_model'] = original_model
-                    values['original_model'] = original_model
+
-                    
+        return new_model
-                return new_model
-            return v
 
 class TokenCountRequest(BaseModel):
     model: str
@@ -202,53 +258,67 @@ class TokenCountRequest(BaseModel):
     original_model: Optional[str] = None  # Will store the original model name
     
     @field_validator('model')
-    def validate_model(cls, v, info):
+    def validate_model_token_count(cls, v, info): # Renamed to avoid conflict
-        # Store the original model name
+        # Use the same logic as MessagesRequest validator
+        # NOTE: Pydantic validators might not share state easily if not class methods
+        # Re-implementing the logic here for clarity, could be refactored
         original_model = v
-        
+        new_model = v # Default to original value
-        # Same validation as MessagesRequest
+
-        if USE_OPENAI_MODELS:
+        logger.debug(f"📋 TOKEN COUNT VALIDATION: Original='{original_model}', Preferred='{PREFERRED_PROVIDER}', BIG='{BIG_MODEL}', SMALL='{SMALL_MODEL}'")
-            # Remove anthropic/ prefix if it exists
+
-            if v.startswith('anthropic/'):
+        # Remove provider prefixes for easier matching
-                v = v[10:]  
+        clean_v = v
-            
+        if clean_v.startswith('anthropic/'):
-            # Swap Haiku with small model (default: gpt-4o-mini)
+            clean_v = clean_v[10:]
-            if 'haiku' in v.lower():
+        elif clean_v.startswith('openai/'):
+            clean_v = clean_v[7:]
+        elif clean_v.startswith('gemini/'):
+            clean_v = clean_v[7:]
+
+        # --- Mapping Logic --- START ---
+        mapped = False
+        # Map Haiku to SMALL_MODEL based on provider preference
+        if 'haiku' in clean_v.lower():
+            if PREFERRED_PROVIDER == "google" and SMALL_MODEL in GEMINI_MODELS:
+                new_model = f"gemini/{SMALL_MODEL}"
+                mapped = True
+            else:
                 new_model = f"openai/{SMALL_MODEL}"
-                logger.debug(f"📌 MODEL MAPPING: {original_model} ➡️ {new_model}")
+                mapped = True
-                v = new_model
+
-            
+        # Map Sonnet to BIG_MODEL based on provider preference
-            # Swap any Sonnet model with big model (default: gpt-4o)
+        elif 'sonnet' in clean_v.lower():
-            elif 'sonnet' in v.lower():
+            if PREFERRED_PROVIDER == "google" and BIG_MODEL in GEMINI_MODELS:
+                new_model = f"gemini/{BIG_MODEL}"
+                mapped = True
+            else:
                 new_model = f"openai/{BIG_MODEL}"
-                logger.debug(f"📌 MODEL MAPPING: {original_model} ➡️ {new_model}")
+                mapped = True
-                v = new_model
+
-            
+        # Add prefixes to non-mapped models if they match known lists
-            # Keep the model as is but add openai/ prefix if not already present
+        elif not mapped:
-            elif not v.startswith('openai/'):
+            if clean_v in GEMINI_MODELS and not v.startswith('gemini/'):
-                new_model = f"openai/{v}"
+                new_model = f"gemini/{clean_v}"
-                logger.debug(f"📌 MODEL MAPPING: {original_model} ➡️ {new_model}")
+                mapped = True # Technically mapped to add prefix
-                v = new_model
+            elif clean_v in OPENAI_MODELS and not v.startswith('openai/'):
-            
+                new_model = f"openai/{clean_v}"
-            # Store the original model in the values dictionary
+                mapped = True # Technically mapped to add prefix
-            values = info.data
+        # --- Mapping Logic --- END ---
-            if isinstance(values, dict):
+
-                values['original_model'] = original_model
+        if mapped:
-                
+            logger.debug(f"📌 TOKEN COUNT MAPPING: '{original_model}' ➡️ '{new_model}'")
-            return v
         else:
-            # Original behavior - ensure anthropic/ prefix
+             if not v.startswith(('openai/', 'gemini/', 'anthropic/')):
-            if not v.startswith('anthropic/'):
+                 logger.warning(f"⚠️ No prefix or mapping rule for token count model: '{original_model}'. Using as is.")
-                new_model = f"anthropic/{v}"
+             new_model = v # Ensure we return the original if no rule applied
-                logger.debug(f"📌 MODEL MAPPING: {original_model} ➡️ {new_model}")
+
-                
+        # Store the original model in the values dictionary
-                # Store original model
+        values = info.data
-                values = info.data
+        if isinstance(values, dict):
-                if isinstance(values, dict):
+            values['original_model'] = original_model
-                    values['original_model'] = original_model
+
-                    
+        return new_model
-                return new_model
-            return v
 
 class TokenCountResponse(BaseModel):
     input_tokens: int
@@ -463,9 +533,9 @@ def convert_anthropic_to_litellm(anthropic_request: MessagesRequest) -> Dict[str
     
     # Cap max_tokens for OpenAI models to their limit of 16384
     max_tokens = anthropic_request.max_tokens
-    if anthropic_request.model.startswith("openai/") or USE_OPENAI_MODELS:
+    if anthropic_request.model.startswith("openai/") or anthropic_request.model.startswith("gemini/"):
         max_tokens = min(max_tokens, 16384)
-        logger.debug(f"Capping max_tokens to 16384 for OpenAI model (original value: {anthropic_request.max_tokens})")
+        logger.debug(f"Capping max_tokens to 16384 for OpenAI/Gemini model (original value: {anthropic_request.max_tokens})")
     
     # Create LiteLLM request dict
     litellm_request = {
@@ -489,24 +559,37 @@ def convert_anthropic_to_litellm(anthropic_request: MessagesRequest) -> Dict[str
     # Convert tools to OpenAI format
     if anthropic_request.tools:
         openai_tools = []
+        is_gemini_model = anthropic_request.model.startswith("gemini/")
+
         for tool in anthropic_request.tools:
             # Convert to dict if it's a pydantic model
             if hasattr(tool, 'dict'):
                 tool_dict = tool.dict()
             else:
-                tool_dict = tool
+                # Ensure tool_dict is a dictionary, handle potential errors if 'tool' isn't dict-like
-                
+                try:
+                    tool_dict = dict(tool) if not isinstance(tool, dict) else tool
+                except (TypeError, ValueError):
+                     logger.error(f"Could not convert tool to dict: {tool}")
+                     continue # Skip this tool if conversion fails
+
+            # Clean the schema if targeting a Gemini model
+            input_schema = tool_dict.get("input_schema", {})
+            if is_gemini_model:
+                 logger.debug(f"Cleaning schema for Gemini tool: {tool_dict.get('name')}")
+                 input_schema = clean_gemini_schema(input_schema)
+
             # Create OpenAI-compatible function tool
             openai_tool = {
                 "type": "function",
                 "function": {
                     "name": tool_dict["name"],
                     "description": tool_dict.get("description", ""),
-                    "parameters": tool_dict["input_schema"]
+                    "parameters": input_schema # Use potentially cleaned schema
                 }
             }
             openai_tools.append(openai_tool)
-            
+
         litellm_request["tools"] = openai_tools
     
     # Convert tool_choice to OpenAI format if present
@@ -1024,6 +1107,9 @@ async def create_message(
         if request.model.startswith("openai/"):
             litellm_request["api_key"] = OPENAI_API_KEY
             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}")
         else:
             litellm_request["api_key"] = ANTHROPIC_API_KEY
             logger.debug(f"Using Anthropic API key for model: {request.model}")