From a595457352a0b249ff95ca0b88a8836979bb8f2c Mon Sep 17 00:00:00 2001 From: Zhizhen He Date: Tue, 5 Aug 2025 22:10:41 +0800 Subject: [PATCH] feat: support custom endpoint with openai compatible apis --- .env.example | 3 ++- server.py | 27 +++++++++++++++++++++------ 2 files changed, 23 insertions(+), 7 deletions(-) diff --git a/.env.example b/.env.example index 798ad8a..22adc1b 100644 --- a/.env.example +++ b/.env.example @@ -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" \ No newline at end of file +# SMALL_MODEL="gemini-2.0-flash" diff --git a/server.py b/server.py index f4966b2..4c0e636 100644 --- a/server.py +++ b/server.py @@ -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") \ No newline at end of file + uvicorn.run(app, host="0.0.0.0", port=8082, log_level="error")