|
from fastapi import FastAPI, HTTPException, Header |
|
from fastapi.middleware.cors import CORSMiddleware |
|
from fastapi.responses import StreamingResponse |
|
from pydantic import BaseModel |
|
import openai |
|
from typing import List, Optional, Union |
|
import logging |
|
import httpx |
|
import uuid |
|
import time |
|
import json |
|
from datetime import datetime, timezone |
|
import requests |
|
import uvicorn |
|
import random |
|
|
|
logging.basicConfig( |
|
level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" |
|
) |
|
logger = logging.getLogger(__name__) |
|
|
|
app = FastAPI() |
|
|
|
app.add_middleware( |
|
CORSMiddleware, |
|
allow_origins=["*"], |
|
allow_credentials=True, |
|
allow_methods=["*"], |
|
allow_headers=["*"], |
|
) |
|
|
|
MAX_RETRIES = 3 |
|
|
|
class ChatRequest(BaseModel): |
|
messages: List[dict] |
|
model: str |
|
temperature: Optional[float] = 0.7 |
|
stream: Optional[bool] = False |
|
tools: Optional[List[dict]] = [] |
|
tool_choice: Optional[str] = "auto" |
|
|
|
class EmbeddingRequest(BaseModel): |
|
input: Union[str, List[str]] |
|
model: str |
|
encoding_format: Optional[str] = "float" |
|
|
|
async def verify_authorization(authorization: str = Header(None)): |
|
print("Authorization header:", authorization) |
|
if not authorization: |
|
logger.error("Missing Authorization header") |
|
raise HTTPException(status_code=401, detail="Missing Authorization header") |
|
if not authorization.startswith("Bearer "): |
|
logger.error("Invalid Authorization header format") |
|
raise HTTPException( |
|
status_code=401, detail="Invalid Authorization header format" |
|
) |
|
token = authorization.replace("Bearer ", "") |
|
return token |
|
|
|
def get_openai_models(api_keys): |
|
api_key = random.choice(api_keys) |
|
try: |
|
client = openai.OpenAI(api_key=api_key) |
|
models = client.models.list() |
|
return models.model_dump() |
|
except Exception as e: |
|
logger.error(f"Error getting models from OpenAI with key {api_key}: {e}") |
|
return {"error": str(e)} |
|
|
|
def get_gemini_models(api_keys): |
|
api_key = random.choice(api_keys) |
|
base_url = "https://generativelanguage.googleapis.com/v1beta" |
|
url = f"{base_url}/models?key={api_key}" |
|
|
|
try: |
|
response = requests.get(url) |
|
if response.status_code == 200: |
|
gemini_models = response.json() |
|
return convert_to_openai_models_format(gemini_models) |
|
else: |
|
logger.error(f"Error getting models from Gemini with key {api_key}: {response.status_code} - {response.text}") |
|
return {"error": f"Gemini API error: {response.status_code} - {response.text}"} |
|
|
|
except requests.RequestException as e: |
|
logger.error(f"Request failed: {e}") |
|
return {"error": f"Request failed: {e}"} |
|
|
|
def convert_to_openai_models_format(gemini_models): |
|
openai_format = {"object": "list", "data": []} |
|
|
|
for model in gemini_models.get("models", []): |
|
openai_model = { |
|
"id": model["name"].split("/")[-1], |
|
"object": "model", |
|
"created": int(datetime.now(timezone.utc).timestamp()), |
|
"owned_by": "google", |
|
"permission": [], |
|
"root": model["name"], |
|
"parent": None, |
|
} |
|
openai_format["data"].append(openai_model) |
|
|
|
return openai_format |
|
|
|
def convert_messages_to_gemini_format(messages): |
|
gemini_messages = [] |
|
for msg in messages: |
|
role = "user" if msg["role"] == "user" else "model" |
|
parts = [] |
|
if isinstance(msg["content"], str): |
|
parts.append({"text": msg["content"]}) |
|
elif isinstance(msg["content"], list): |
|
for content in msg["content"]: |
|
if isinstance(content, str): |
|
parts.append({"text": content}) |
|
elif isinstance(content, dict) and content["type"] == "text": |
|
parts.append({"text": content["text"]}) |
|
elif isinstance(content, dict) and content["type"] == "image_url": |
|
image_url = content["image_url"]["url"] |
|
if image_url.startswith("data:image"): |
|
parts.append( |
|
{ |
|
"inline_data": { |
|
"mime_type": "image/jpeg", |
|
"data": image_url.split(",")[1], |
|
} |
|
} |
|
) |
|
else: |
|
parts.append( |
|
{ |
|
"image_url": { |
|
"url": image_url, |
|
} |
|
} |
|
) |
|
gemini_messages.append({"role": role, "parts": parts}) |
|
return gemini_messages |
|
|
|
async def convert_gemini_response_to_openai(response, model, stream=False): |
|
if stream: |
|
chunk = response |
|
if not chunk["candidates"]: |
|
return None |
|
|
|
return { |
|
"id": "chatcmpl-" + str(uuid.uuid4()), |
|
"object": "chat.completion.chunk", |
|
"created": int(time.time()), |
|
"model": model, |
|
"choices": [ |
|
{ |
|
"index": 0, |
|
"delta": { |
|
"content": chunk["candidates"][0]["content"]["parts"][0]["text"] |
|
}, |
|
"finish_reason": None, |
|
} |
|
], |
|
} |
|
else: |
|
content = response["candidates"][0]["content"]["parts"][0]["text"] |
|
return { |
|
"id": "chatcmpl-" + str(uuid.uuid4()), |
|
"object": "chat.completion", |
|
"created": int(time.time()), |
|
"model": model, |
|
"choices": [ |
|
{ |
|
"index": 0, |
|
"message": { |
|
"role": "assistant", |
|
"content": content, |
|
}, |
|
"finish_reason": "stop", |
|
} |
|
], |
|
"usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}, |
|
} |
|
|
|
@app.get("/v1/models") |
|
@app.get("/hf/v1/models") |
|
async def list_models(authorization: str = Header(None)): |
|
token = await verify_authorization(authorization) |
|
api_keys = [key.strip() for key in token.split(',')] |
|
|
|
all_models = [] |
|
error_messages = [] |
|
|
|
for api_key in api_keys: |
|
if api_key.startswith("sk-"): |
|
response = get_openai_models([api_key]) |
|
else: |
|
response = get_gemini_models([api_key]) |
|
|
|
if "error" in response: |
|
error_messages.append(response["error"]) |
|
else: |
|
if isinstance(response, dict) and 'data' in response: |
|
all_models.extend(response['data']) |
|
else: |
|
logger.warning(f"Unexpected response format from model list API for key {api_key}: {response}") |
|
|
|
if error_messages and not all_models: |
|
raise HTTPException(status_code=500, detail=f"Errors encountered: {', '.join(error_messages)}") |
|
|
|
return {"data": all_models, "object": "list"} |
|
|
|
@app.post("/v1/chat/completions") |
|
@app.post("/hf/v1/chat/completions") |
|
async def chat_completion(request: ChatRequest, authorization: str = Header(None)): |
|
token = await verify_authorization(authorization) |
|
api_keys = [key.strip() for key in token.split(',')] |
|
logger.info(f"Chat completion request - Model: {request.model}") |
|
|
|
retries = 0 |
|
|
|
while retries < MAX_RETRIES: |
|
api_key = random.choice(api_keys) |
|
try: |
|
logger.info(f"Attempt {retries + 1} with API key: {api_key}") |
|
|
|
if api_key.startswith("sk-"): |
|
client = openai.OpenAI(api_key=api_key) |
|
|
|
if request.stream: |
|
logger.info("Streaming response enabled") |
|
|
|
async def generate(): |
|
try: |
|
stream_response = client.chat.completions.create( |
|
model=request.model, |
|
messages=request.messages, |
|
temperature=request.temperature, |
|
stream=True, |
|
) |
|
|
|
for chunk in stream_response: |
|
chunk_json = chunk.model_dump_json() |
|
yield f"data: {chunk_json}\n\n" |
|
yield "data: [DONE]\n\n" |
|
except Exception as e: |
|
logger.error(f"Stream error: {str(e)}") |
|
raise |
|
|
|
return StreamingResponse(content=generate(), media_type="text/event-stream") |
|
|
|
else: |
|
response = client.chat.completions.create( |
|
model=request.model, |
|
messages=request.messages, |
|
temperature=request.temperature, |
|
) |
|
logger.info("Chat completion successful") |
|
return response.model_dump() |
|
else: |
|
gemini_messages = convert_messages_to_gemini_format(request.messages) |
|
payload = { |
|
"contents": gemini_messages, |
|
"generationConfig": { |
|
"temperature": request.temperature, |
|
} |
|
} |
|
|
|
if request.stream: |
|
logger.info("Streaming response enabled") |
|
|
|
async def generate(): |
|
nonlocal api_key, retries, api_keys |
|
|
|
while retries < MAX_RETRIES: |
|
try: |
|
async with httpx.AsyncClient() as client: |
|
stream_url = f"https://generativelanguage.googleapis.com/v1beta/models/{request.model}:streamGenerateContent?alt=sse&key={api_key}" |
|
async with client.stream("POST", stream_url, json=payload, timeout=60.0) as response: |
|
if response.status_code == 429: |
|
logger.warning(f"Rate limit reached for key: {api_key}") |
|
retries += 1 |
|
if retries >= MAX_RETRIES: |
|
yield f"data: {json.dumps({'error': 'Max retries reached'})}\n\n" |
|
break |
|
|
|
api_keys.remove(api_key) |
|
if not api_keys: |
|
yield f"data: {json.dumps({'error': 'All API keys exhausted'})}\n\n" |
|
break |
|
|
|
api_key = random.choice(api_keys) |
|
logger.info(f"Retrying with a new API key: {api_key}") |
|
continue |
|
|
|
if response.status_code != 200: |
|
logger.error(f"Error in streaming response with key {api_key}: {response.status_code} - {response.text}") |
|
|
|
retries += 1 |
|
if retries >= MAX_RETRIES: |
|
yield f"data: {json.dumps({'error': 'Max retries reached'})}\n\n" |
|
break |
|
|
|
api_keys.remove(api_key) |
|
if not api_keys: |
|
yield f"data: {json.dumps({'error': 'All API keys exhausted'})}\n\n" |
|
break |
|
|
|
api_key = random.choice(api_keys) |
|
logger.info(f"Retrying with a new API key: {api_key}") |
|
continue |
|
|
|
async for line in response.aiter_lines(): |
|
if line.startswith("data: "): |
|
try: |
|
chunk = json.loads(line[6:]) |
|
if not chunk.get("candidates"): |
|
continue |
|
|
|
content = chunk["candidates"][0]["content"]["parts"][0]["text"] |
|
|
|
new_chunk = { |
|
"id": "chatcmpl-" + str(uuid.uuid4()), |
|
"object": "chat.completion.chunk", |
|
"created": int(time.time()), |
|
"model": request.model, |
|
"choices": [ |
|
{ |
|
"index": 0, |
|
"delta": { |
|
"content": content |
|
}, |
|
"finish_reason": None, |
|
} |
|
], |
|
} |
|
yield f"data: {json.dumps(new_chunk)}\n\n" |
|
|
|
except json.JSONDecodeError: |
|
continue |
|
yield "data: [DONE]\n\n" |
|
return |
|
except Exception as e: |
|
logger.error(f"Stream error: {str(e)}") |
|
retries += 1 |
|
if retries >= MAX_RETRIES: |
|
yield f"data: {json.dumps({'error': 'Max retries reached'})}\n\n" |
|
break |
|
|
|
api_keys.remove(api_key) |
|
if not api_keys: |
|
yield f"data: {json.dumps({'error': 'All API keys exhausted'})}\n\n" |
|
break |
|
|
|
api_key = random.choice(api_keys) |
|
logger.info(f"Retrying with a new API key: {api_key}") |
|
continue |
|
|
|
return StreamingResponse(content=generate(), media_type="text/event-stream") |
|
else: |
|
async with httpx.AsyncClient() as client: |
|
non_stream_url = f"https://generativelanguage.googleapis.com/v1beta/models/{request.model}:generateContent?key={api_key}" |
|
response = await client.post(non_stream_url, json=payload) |
|
|
|
if response.status_code != 200: |
|
logger.error(f"Error in non-streaming response with key {api_key}: {response.status_code} - {response.text}") |
|
|
|
retries += 1 |
|
if retries >= MAX_RETRIES: |
|
raise HTTPException(status_code=500, detail="Max retries reached") |
|
|
|
api_keys.remove(api_key) |
|
if not api_keys: |
|
raise HTTPException(status_code=500, detail="All API keys exhausted") |
|
|
|
api_key = random.choice(api_keys) |
|
logger.info(f"Retrying with a new API key: {api_key}") |
|
continue |
|
|
|
gemini_response = response.json() |
|
logger.info("Chat completion successful") |
|
return await convert_gemini_response_to_openai(gemini_response, request.model) |
|
|
|
except Exception as e: |
|
logger.error(f"Error in chat completion: {str(e)}") |
|
if isinstance(e, HTTPException): |
|
raise e |
|
|
|
retries += 1 |
|
if retries >= MAX_RETRIES: |
|
logger.error("Max retries reached, giving up") |
|
raise HTTPException(status_code=500, detail="Max retries reached") |
|
|
|
api_keys.remove(api_key) |
|
if not api_keys: |
|
raise HTTPException(status_code=500, detail="All API keys exhausted") |
|
|
|
api_key = random.choice(api_keys) |
|
logger.info(f"Retrying with a new API key: {api_key}") |
|
continue |
|
|
|
raise HTTPException(status_code=500, detail="Unexpected error in chat completion") |
|
|
|
|
|
@app.get("/health") |
|
@app.get("/") |
|
async def health_check(): |
|
logger.info("Health check endpoint called") |
|
return {"status": "healthy"} |
|
|
|
if __name__ == "__main__": |
|
uvicorn.run(app, host="0.0.0.0", port=8080) |
