Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .env.example
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
ANTHROPIC_API_KEY="your-anthropic-api-key" # Needed if proxying *to* Anthropic
OPENAI_API_KEY="sk-..."
GEMINI_API_KEY="your-google-ai-studio-key"
GITHUB_TOKEN="your-github-token" # Needed if proxying *to* GitHub Models

# Optional: Provider Preference and Model Mapping
# Controls which provider (google or openai) is preferred for mapping haiku/sonnet.
Expand Down
53 changes: 48 additions & 5 deletions server.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,12 +15,14 @@
from datetime import datetime
import sys

litellm._turn_on_debug()

# Load environment variables from .env file
load_dotenv()

# Configure logging
logging.basicConfig(
level=logging.WARN, # Change to INFO level to show more details
level=logging.INFO, # Change to INFO level to show more details
format='%(asctime)s - %(levelname)s - %(message)s',
)
logger = logging.getLogger(__name__)
Expand Down Expand Up @@ -79,6 +81,7 @@ def format(self, record):

# Get API keys from environment
ANTHROPIC_API_KEY = os.environ.get("ANTHROPIC_API_KEY")
GITHUB_TOKEN = os.environ.get("GITHUB_TOKEN")
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")

Expand Down Expand Up @@ -112,6 +115,13 @@ def format(self, record):
"gemini-2.0-flash"
]

GITHUB_MODELS = [
"gpt-4.1-mini",
"gpt-4.1",
"gpt-4o",
"gpt-4o-mini",
]

# Helper function to clean schema for Gemini
def clean_gemini_schema(schema: Any) -> Any:
"""Recursively removes unsupported fields from a JSON schema for Gemini."""
Expand Down Expand Up @@ -202,6 +212,10 @@ def validate_model_field(cls, v, info): # Renamed to avoid conflict
clean_v = clean_v[7:]
elif clean_v.startswith('gemini/'):
clean_v = clean_v[7:]
elif clean_v.startswith('github/'):
clean_v = clean_v[7:]

logger.debug(f"📌 original_model: '{original_model}', clean_v: '{clean_v}'")

# --- Mapping Logic --- START ---
mapped = False
Expand All @@ -210,6 +224,9 @@ def validate_model_field(cls, v, info): # Renamed to avoid conflict
if PREFERRED_PROVIDER == "google" and SMALL_MODEL in GEMINI_MODELS:
new_model = f"gemini/{SMALL_MODEL}"
mapped = True
elif PREFERRED_PROVIDER == "github" and SMALL_MODEL in GITHUB_MODELS:
new_model = f"github/{SMALL_MODEL}"
mapped = True
else:
new_model = f"openai/{SMALL_MODEL}"
mapped = True
Expand All @@ -219,6 +236,9 @@ def validate_model_field(cls, v, info): # Renamed to avoid conflict
if PREFERRED_PROVIDER == "google" and BIG_MODEL in GEMINI_MODELS:
new_model = f"gemini/{BIG_MODEL}"
mapped = True
elif PREFERRED_PROVIDER == "github" and BIG_MODEL in GITHUB_MODELS:
new_model = f"github/{BIG_MODEL}"
mapped = True
else:
new_model = f"openai/{BIG_MODEL}"
mapped = True
Expand All @@ -231,13 +251,16 @@ def validate_model_field(cls, v, info): # Renamed to avoid conflict
elif clean_v in OPENAI_MODELS and not v.startswith('openai/'):
new_model = f"openai/{clean_v}"
mapped = True # Technically mapped to add prefix
elif clean_v in GITHUB_MODELS and not v.startswith('github/'):
new_model = f"github/{clean_v}"
mapped = True # Technically mapped to add prefix
# --- Mapping Logic --- END ---

if mapped:
logger.debug(f"📌 MODEL MAPPING: '{original_model}' ➡️ '{new_model}'")
else:
# If no mapping occurred and no prefix exists, log warning or decide default
if not v.startswith(('openai/', 'gemini/', 'anthropic/')):
if not v.startswith(('openai/', 'gemini/', 'anthropic/', 'github/')):
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

Expand Down Expand Up @@ -275,6 +298,8 @@ def validate_model_token_count(cls, v, info): # Renamed to avoid conflict
clean_v = clean_v[7:]
elif clean_v.startswith('gemini/'):
clean_v = clean_v[7:]
elif clean_v.startswith('github/'):
clean_v = clean_v[7:]

# --- Mapping Logic --- START ---
mapped = False
Expand All @@ -283,6 +308,9 @@ def validate_model_token_count(cls, v, info): # Renamed to avoid conflict
if PREFERRED_PROVIDER == "google" and SMALL_MODEL in GEMINI_MODELS:
new_model = f"gemini/{SMALL_MODEL}"
mapped = True
elif PREFERRED_PROVIDER == "github" and SMALL_MODEL in GITHUB_MODELS:
new_model = f"github/{SMALL_MODEL}"
mapped = True
else:
new_model = f"openai/{SMALL_MODEL}"
mapped = True
Expand All @@ -292,6 +320,9 @@ def validate_model_token_count(cls, v, info): # Renamed to avoid conflict
if PREFERRED_PROVIDER == "google" and BIG_MODEL in GEMINI_MODELS:
new_model = f"gemini/{BIG_MODEL}"
mapped = True
elif PREFERRED_PROVIDER == "github" and BIG_MODEL in GITHUB_MODELS:
new_model = f"github/{BIG_MODEL}"
mapped = True
else:
new_model = f"openai/{BIG_MODEL}"
mapped = True
Expand All @@ -304,12 +335,15 @@ def validate_model_token_count(cls, v, info): # Renamed to avoid conflict
elif clean_v in OPENAI_MODELS and not v.startswith('openai/'):
new_model = f"openai/{clean_v}"
mapped = True # Technically mapped to add prefix
elif clean_v in GITHUB_MODELS and not v.startswith('github/'):
new_model = f"github/{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:
if not v.startswith(('openai/', 'gemini/', 'anthropic/')):
if not v.startswith(('openai/', 'gemini/', 'anthropic/', 'github/')):
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

Expand Down Expand Up @@ -533,7 +567,7 @@ 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 anthropic_request.model.startswith("gemini/"):
if anthropic_request.model.startswith("openai/") or anthropic_request.model.startswith("github/") or anthropic_request.model.startswith("gemini/"):
max_tokens = min(max_tokens, 16384)
logger.debug(f"Capping max_tokens to 16384 for OpenAI/Gemini model (original value: {anthropic_request.max_tokens})")

Expand Down Expand Up @@ -628,6 +662,8 @@ def convert_litellm_to_anthropic(litellm_response: Union[Dict[str, Any], Any],
clean_model = clean_model[len("anthropic/"):]
elif clean_model.startswith("openai/"):
clean_model = clean_model[len("openai/"):]
elif clean_model.startswith("github/"):
clean_model = clean_model[len("github/"):]

# Check if this is a Claude model (which supports content blocks)
is_claude_model = clean_model.startswith("claude-")
Expand Down Expand Up @@ -1097,6 +1133,8 @@ async def create_message(
clean_model = clean_model[len("anthropic/"):]
elif clean_model.startswith("openai/"):
clean_model = clean_model[len("openai/"):]
elif clean_model.startswith("github/"):
clean_model = clean_model[len("github/"):]

logger.debug(f"📊 PROCESSING REQUEST: Model={request.model}, Stream={request.stream}")

Expand All @@ -1110,12 +1148,15 @@ async def create_message(
elif request.model.startswith("gemini/"):
litellm_request["api_key"] = GEMINI_API_KEY
logger.debug(f"Using Gemini API key for model: {request.model}")
elif request.model.startswith("github/"):
litellm_request["api_key"] = GITHUB_TOKEN
logger.debug(f"Using GITHUB_TOKEN for model: {request.model}")
else:
litellm_request["api_key"] = ANTHROPIC_API_KEY
logger.debug(f"Using Anthropic API key for model: {request.model}")

# For OpenAI models - modify request format to work with limitations
if "openai" in litellm_request["model"] and "messages" in litellm_request:
if ("openai" in litellm_request["model"] or "github" in litellm_request["model"]) and "messages" in litellm_request:
logger.debug(f"Processing OpenAI model request: {litellm_request['model']}")

# For OpenAI models, we need to convert content blocks to simple strings
Expand Down Expand Up @@ -1354,6 +1395,8 @@ async def count_tokens(
clean_model = clean_model[len("anthropic/"):]
elif clean_model.startswith("openai/"):
clean_model = clean_model[len("openai/"):]
elif clean_model.startswith("github/"):
clean_model = clean_model[len("github/"):]

# Convert the messages to a format LiteLLM can understand
converted_request = convert_anthropic_to_litellm(
Expand Down
2 changes: 1 addition & 1 deletion tests.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@
ANTHROPIC_API_URL = "https://api.anthropic.com/v1/messages"
PROXY_API_URL = "http://localhost:8082/v1/messages"
ANTHROPIC_VERSION = "2023-06-01"
MODEL = "claude-3-sonnet-20240229" # Change to your preferred model
MODEL = "claude-3-7-sonnet-20250219" # Change to your preferred model

# Headers
anthropic_headers = {
Expand Down