Upload app.py

app.py

@@ -1,17 +1,26 @@
-from fastapi import FastAPI, Request, Header
-from fastapi.responses import StreamingResponse
+from fastapi import FastAPI, HTTPException, Header
 from fastapi.middleware.cors import CORSMiddleware
-import httpx
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel
+import openai
+from typing import List, Optional, Union
 import logging
-import os
+import httpx
+import uuid
+import time
+import json
+from datetime import datetime, timezone
+import requests
+import uvicorn
+import random
 
-# Configure logging
-logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
+)
 logger = logging.getLogger(__name__)
 
 app = FastAPI()
 
-# Add CORS middleware to allow all origins
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -20,53 +29,374 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-TARGET_HOST = "https://generativelanguage.googleapis.com"
+MAX_RETRIES = 3
 
-# Use a persistent client for connection pooling
-client = httpx.AsyncClient(base_url=TARGET_HOST)
+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"
 
-@app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS"])
-async def reverse_proxy(request: Request, path: str, authorization: str = Header(None)):
-    """
-    A reverse proxy that forwards requests to the Google Generative Language API.
-    It extracts the API key from the Authorization header.
-    """
-    if not authorization or not authorization.startswith("Bearer "):
-        return {"error": "Authorization header with Bearer token is required."}, 401
-    
-    api_key = authorization.replace("Bearer ", "").strip()
+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
 
-    # Build the target URL
-    url = httpx.URL(path=f"/{path}", query=request.url.query.encode("utf-8"))
+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(',')]
     
-    # Add the API key to the query parameters
-    # httpx URL query parameters are immutable, so we create a new one
-    params = list(url.query)
-    params.append(("key", api_key))
-    url = url.copy_with(query=b'&'.join([f"{k}={v}".encode() for k, v in params]))
-
-    logger.info(f"Forwarding request to: {url}")
-
-    # Prepare the request to be forwarded
-    rp_req = client.build_request(
-        method=request.method,
-        url=url,
-        headers=request.headers,
-        content=await request.body(),
-    )
+    all_models = []
+    error_messages = []
     
-    # Stream the response back to the client
-    try:
-        rp_resp = await client.send(rp_req, stream=True)
-    except httpx.RequestError as e:
-        logger.error(f"An error occurred while requesting {url}: {e}")
-        return {"error": "Failed to connect to the upstream server."}, 502
+    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")
 
-    return StreamingResponse(
-        rp_resp.aiter_raw(),
-        status_code=rp_resp.status_code,
-        headers=rp_resp.headers,
-    )
 
 @app.get("/health")
 @app.get("/")
@@ -75,6 +405,4 @@ async def health_check():
     return {"status": "healthy"}
 
 if __name__ == "__main__":
-    import uvicorn
-    port = int(os.environ.get("PORT", 8080))
-    uvicorn.run(app, host="0.0.0.0", port=port)
+    uvicorn.run(app, host="0.0.0.0", port=8080)