From b64350df69206a5b8b998e64090db222364e726a Mon Sep 17 00:00:00 2001 From: Rahul Sengottuvelu Date: Wed, 9 Apr 2025 16:23:48 -0400 Subject: [PATCH] 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 | 21 ++++ README.md | 66 ++++++++++-- server.py | 286 +++++++++++++++++++++++++++++++++------------------ 3 files changed, 265 insertions(+), 108 deletions(-) create mode 100644 .env.example diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..0e8c44f --- /dev/null +++ b/.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" \ No newline at end of file diff --git a/README.md b/README.md index 6113839..5816939 100644 --- a/README.md +++ b/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 | -|--------------|--------------| -| haiku | gpt-4o-mini (default) | -| sonnet | gpt-4o (default) | +| Claude Model | Default Mapping | When BIG_MODEL/SMALL_MODEL is a Gemini model | +|--------------|--------------|---------------------------| +| haiku | openai/gpt-4o-mini | gemini/[model-name] | +| 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") -- `SMALL_MODEL`: The OpenAI model to use for Claude Haiku models (default: "gpt-4o-mini") +- `BIG_MODEL`: The model to use for Claude Sonnet models (default: "gpt-4o") +- `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 🧩 diff --git a/server.py b/server.py index d4ba467..d745b59 100644 --- a/server.py +++ b/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") -SMALL_MODEL = os.environ.get("SMALL_MODEL", "gpt-4o-mini") +# Default to latest Gemini models if not set +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): - # Store the original model name + def validate_model_field(cls, v, info): # Renamed to avoid conflict original_model = v - - # Check if we're using OpenAI models and need to swap - if USE_OPENAI_MODELS: - # Remove anthropic/ prefix if it exists - if v.startswith('anthropic/'): - v = v[10:] # Remove 'anthropic/' prefix - - # Swap Haiku with small model (default: gpt-4o-mini) - if 'haiku' in v.lower(): + new_model = v # Default to original value + + logger.debug(f"πŸ“‹ MODEL VALIDATION: Original='{original_model}', Preferred='{PREFERRED_PROVIDER}', BIG='{BIG_MODEL}', SMALL='{SMALL_MODEL}'") + + # Remove provider prefixes for easier matching + clean_v = v + if clean_v.startswith('anthropic/'): + clean_v = clean_v[10:] + 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}") - v = new_model - - # Swap any Sonnet model with big model (default: gpt-4o) - elif 'sonnet' in v.lower(): + mapped = True + + # Map Sonnet to BIG_MODEL based on provider preference + elif 'sonnet' in clean_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}") - v = new_model - - # Keep the model as is but add openai/ prefix if not already present - elif not v.startswith('openai/'): - new_model = f"openai/{v}" - logger.debug(f"πŸ“Œ MODEL MAPPING: {original_model} ➑️ {new_model}") - v = new_model - - # Store the original model in the values dictionary - # This will be accessible as request.original_model - values = info.data - if isinstance(values, dict): - values['original_model'] = original_model - - return v + mapped = True + + # Add prefixes to non-mapped models if they match known lists + elif not mapped: + if clean_v in GEMINI_MODELS and not v.startswith('gemini/'): + new_model = f"gemini/{clean_v}" + mapped = True # Technically mapped to add prefix + elif clean_v in OPENAI_MODELS and not v.startswith('openai/'): + new_model = f"openai/{clean_v}" + mapped = True # Technically mapped to add prefix + # --- Mapping Logic --- END --- + + if mapped: + logger.debug(f"πŸ“Œ MODEL MAPPING: '{original_model}' ➑️ '{new_model}'") else: - # Original behavior - ensure anthropic/ prefix - original_model = v - if not v.startswith('anthropic/'): - new_model = f"anthropic/{v}" - logger.debug(f"πŸ“Œ MODEL MAPPING: {original_model} ➑️ {new_model}") - - # Store original model - values = info.data - if isinstance(values, dict): - values['original_model'] = original_model - - return new_model - return v + # If no mapping occurred and no prefix exists, log warning or decide default + if not v.startswith(('openai/', 'gemini/', 'anthropic/')): + logger.warning(f"⚠️ No prefix or mapping rule for model: '{original_model}'. Using as is.") + new_model = v # Ensure we return the original if no rule applied + + # Store the original model in the values dictionary + values = info.data + if isinstance(values, dict): + values['original_model'] = original_model + + return new_model 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): - # Store the original model name + def validate_model_token_count(cls, v, info): # Renamed to avoid conflict + # 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 - - # Same validation as MessagesRequest - if USE_OPENAI_MODELS: - # Remove anthropic/ prefix if it exists - if v.startswith('anthropic/'): - v = v[10:] - - # Swap Haiku with small model (default: gpt-4o-mini) - if 'haiku' in v.lower(): + new_model = v # Default to original value + + logger.debug(f"πŸ“‹ TOKEN COUNT VALIDATION: Original='{original_model}', Preferred='{PREFERRED_PROVIDER}', BIG='{BIG_MODEL}', SMALL='{SMALL_MODEL}'") + + # Remove provider prefixes for easier matching + clean_v = v + if clean_v.startswith('anthropic/'): + clean_v = clean_v[10:] + 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}") - v = new_model - - # Swap any Sonnet model with big model (default: gpt-4o) - elif 'sonnet' in v.lower(): + mapped = True + + # Map Sonnet to BIG_MODEL based on provider preference + elif 'sonnet' in clean_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}") - v = new_model - - # Keep the model as is but add openai/ prefix if not already present - elif not v.startswith('openai/'): - new_model = f"openai/{v}" - logger.debug(f"πŸ“Œ MODEL MAPPING: {original_model} ➑️ {new_model}") - v = new_model - - # Store the original model in the values dictionary - values = info.data - if isinstance(values, dict): - values['original_model'] = original_model - - return v + mapped = True + + # Add prefixes to non-mapped models if they match known lists + elif not mapped: + if clean_v in GEMINI_MODELS and not v.startswith('gemini/'): + new_model = f"gemini/{clean_v}" + mapped = True # Technically mapped to add prefix + elif clean_v in OPENAI_MODELS and not v.startswith('openai/'): + new_model = f"openai/{clean_v}" + mapped = True # Technically mapped to add prefix + # --- Mapping Logic --- END --- + + if mapped: + logger.debug(f"πŸ“Œ TOKEN COUNT MAPPING: '{original_model}' ➑️ '{new_model}'") else: - # Original behavior - ensure anthropic/ prefix - if not v.startswith('anthropic/'): - new_model = f"anthropic/{v}" - logger.debug(f"πŸ“Œ MODEL MAPPING: {original_model} ➑️ {new_model}") - - # Store original model - values = info.data - if isinstance(values, dict): - values['original_model'] = original_model - - return new_model - return v + if not v.startswith(('openai/', 'gemini/', 'anthropic/')): + logger.warning(f"⚠️ No prefix or mapping rule for token count model: '{original_model}'. Using as is.") + new_model = v # Ensure we return the original if no rule applied + + # Store the original model in the values dictionary + values = info.data + if isinstance(values, dict): + values['original_model'] = original_model + + return new_model 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}")