danvics/claude-code-proxy
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: support custom endpoint with openai compatible apis

Browse source
This commit is contained in:
Zhizhen He 2025-08-05 22:10:41 +08:00
parent e9c8cf8de6
commit a595457352
No known key found for this signature in database
GPG key ID: B5A215D74FCFD6A7
2 changed files with 23 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
.env.example
View file

@ -7,6 +7,7 @@ GEMINI_API_KEY="your-google-ai-studio-key"
# Controls which provider (google or openai) is preferred for mapping haiku/sonnet.
# Defaults to openai if not set.
PREFERRED_PROVIDER="openai"
OPENAI_BASE_URL="https://api.openai.com/v1"
# 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.
@ -18,4 +19,4 @@ PREFERRED_PROVIDER="openai"
# Example Google mapping:
# PREFERRED_PROVIDER="google"
# BIG_MODEL="gemini-2.5-pro-preview-03-25"
# SMALL_MODEL="gemini-2.0-flash"
# SMALL_MODEL="gemini-2.0-flash"

27
server.py
View file

@ -82,6 +82,9 @@ 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 OpenAI base URL from environment (if set)
OPENAI_BASE_URL = os.environ.get("OPENAI_BASE_URL")
# Get preferred provider (default to openai)
PREFERRED_PROVIDER = os.environ.get("PREFERRED_PROVIDER", "openai").lower()
@ -1106,7 +1109,12 @@ async def create_message(
# Determine which API key to use based on the model
if request.model.startswith("openai/"):
litellm_request["api_key"] = OPENAI_API_KEY
logger.debug(f"Using OpenAI API key for model: {request.model}")
# Use custom OpenAI base URL if configured
if OPENAI_BASE_URL:
litellm_request["api_base"] = OPENAI_BASE_URL
logger.debug(f"Using OpenAI API key and custom base URL {OPENAI_BASE_URL} for model: {request.model}")
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}")
@ -1386,11 +1394,18 @@ async def count_tokens(
200 # Assuming success at this point
)
# Prepare token counter arguments
token_counter_args = {
"model": converted_request["model"],
"messages": converted_request["messages"],
}
# Add custom base URL for OpenAI models if configured
if request.model.startswith("openai/") and OPENAI_BASE_URL:
token_counter_args["api_base"] = OPENAI_BASE_URL
# Count tokens
token_count = token_counter(
model=converted_request["model"],
messages=converted_request["messages"],
)
token_count = token_counter(**token_counter_args)
# Return Anthropic-style response
return TokenCountResponse(input_tokens=token_count)
@ -1462,4 +1477,4 @@ if __name__ == "__main__":
sys.exit(0)
# Configure uvicorn to run with minimal logs
uvicorn.run(app, host="0.0.0.0", port=8082, log_level="error")
uvicorn.run(app, host="0.0.0.0", port=8082, log_level="error")