Update main.py

main.py

@@ -7,21 +7,24 @@ from typing import Any, Dict, List, Optional
 
 import httpx
 import uvicorn
-from fastapi import FastAPI, HTTPException, Depends
+from fastapi import FastAPI, HTTPException, Depends, Request, status
 from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from pydantic import BaseModel
 from starlette.middleware.cors import CORSMiddleware
 from starlette.responses import StreamingResponse, Response
 
-# Configure Logging
+from dotenv import load_dotenv
+
+# Retry Mechanism Libraries
+from tenacity import retry, wait_exponential, stop_after_attempt, retry_if_exception_type
+
+# Initialize Logging
 logging.basicConfig(
     level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
 )
 logger = logging.getLogger(__name__)
 
 # Load Environment Variables for Sensitive Information
-from dotenv import load_dotenv
-
 load_dotenv()
 
 app = FastAPI()
@@ -55,21 +58,23 @@ OUTGOING_HEADERS = {
     'sec-fetch-dest': 'empty',
     'sec-fetch-mode': 'cors',
     'sec-fetch-site': 'none',
-    'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/130.0.0.0 Safari/537.36',
+    'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) '
+                  'AppleWebKit/537.36 (KHTML, like Gecko) '
+                  'Chrome/130.0.0.0 Safari/537.36',
 }
 
 # Updated ALLOWED_MODELS List (No Duplicates)
 ALLOWED_MODELS = [
-    {"id": "claude-3.5-sonnet", "name": "claude-3.5-sonnet"},
-    {"id": "sider", "name": "sider"},
-    {"id": "gpt-4o-mini", "name": "gpt-4o-mini"},
-    {"id": "claude-3-haiku", "name": "claude-3-haiku"},
-    {"id": "claude-3.5-haiku", "name": "claude-3.5-haiku"},
-    {"id": "gemini-1.5-flash", "name": "gemini-1.5-flash"},
-    {"id": "llama-3", "name": "llama-3"},
-    {"id": "gpt-4o", "name": "gpt-4o"},
-    {"id": "gemini-1.5-pro", "name": "gemini-1.5-pro"},
-    {"id": "llama-3.1-405b", "name": "llama-3.1-405b"},
+    {"id": "claude-3.5-sonnet", "name": "Claude 3.5 Sonnet"},
+    {"id": "sider", "name": "Sider"},
+    {"id": "gpt-4o-mini", "name": "GPT-4o Mini"},
+    {"id": "claude-3-haiku", "name": "Claude 3 Haiku"},
+    {"id": "claude-3.5-haiku", "name": "Claude 3.5 Haiku"},
+    {"id": "gemini-1.5-flash", "name": "Gemini 1.5 Flash"},
+    {"id": "llama-3", "name": "Llama 3"},
+    {"id": "gpt-4o", "name": "GPT-4o"},
+    {"id": "gemini-1.5-pro", "name": "Gemini 1.5 Pro"},
+    {"id": "llama-3.1-405b", "name": "Llama 3.1 405b"},
 ]
 
 # Configure CORS
@@ -84,19 +89,16 @@ app.add_middleware(
 # Security Dependency
 security = HTTPBearer()
 
-
 # Pydantic Models
 class Message(BaseModel):
     role: str
     content: str
 
-
 class ChatRequest(BaseModel):
     model: str
     messages: List[Message]
     stream: Optional[bool] = False
 
-
 # Utility Functions
 def create_chat_completion_data(content: str, model: str, finish_reason: Optional[str] = None) -> Dict[str, Any]:
     return {
@@ -114,7 +116,6 @@ def create_chat_completion_data(content: str, model: str, finish_reason: Optiona
         "usage": None,
     }
 
-
 def verify_app_secret(credentials: HTTPAuthorizationCredentials = Depends(security)):
     if credentials.credentials != APP_SECRET:
         logger.warning(f"Invalid APP_SECRET provided: {credentials.credentials}")
@@ -122,6 +123,26 @@ def verify_app_secret(credentials: HTTPAuthorizationCredentials = Depends(securi
     logger.info("APP_SECRET verified successfully.")
     return credentials.credentials
 
+# Retry Configuration using Tenacity
+def is_retryable_exception(exception):
+    return isinstance(exception, httpx.HTTPStatusError) and exception.response.status_code == 429
+
+@retry(
+    retry=retry_if_exception_type(httpx.HTTPStatusError) & is_retryable_exception,
+    wait=wait_exponential(multiplier=1, min=2, max=10),
+    stop=stop_after_attempt(5),
+    reraise=True
+)
+async def send_request_with_retry(json_data: Dict[str, Any]) -> httpx.Response:
+    async with httpx.AsyncClient() as client:
+        response = await client.post(
+            'https://sider.ai/api/v3/completion/text',  # Updated endpoint
+            headers=OUTGOING_HEADERS,
+            json=json_data,
+            timeout=120.0
+        )
+        response.raise_for_status()
+        return response
 
 # CORS Preflight Options Endpoint
 @app.options("/hf/v1/chat/completions")
@@ -135,19 +156,18 @@ async def chat_completions_options():
         },
     )
 
-
 # List Available Models
 @app.get("/hf/v1/models")
 async def list_models():
     return {"object": "list", "data": ALLOWED_MODELS}
 
-
 # Chat Completions Endpoint
 @app.post("/hf/v1/chat/completions")
 async def chat_completions(
-    request: ChatRequest, app_secret: str = Depends(verify_app_secret)
+    request: ChatRequest, app_secret: str = Depends(verify_app_secret), req: Request = None
 ):
-    logger.info(f"Received chat completion request for model: {request.model}")
+    client_ip = req.client.host if req else "unknown"
+    logger.info(f"Received chat completion request from {client_ip} for model: {request.model}")
 
     # Validate Selected Model
     if request.model not in [model['id'] for model in ALLOWED_MODELS]:
@@ -207,39 +227,41 @@ async def chat_completions(
 
     logger.debug(f"JSON Data Sent to External API: {json.dumps(json_data, indent=2)}")
 
+    try:
+        response = await send_request_with_retry(json_data)
+    except httpx.HTTPStatusError as e:
+        status_code = e.response.status_code
+        if status_code == 429:
+            retry_after = e.response.headers.get("Retry-After", "60")
+            logger.warning(f"Rate limited by Sider AI. Retry after {retry_after} seconds.")
+            raise HTTPException(
+                status_code=429,
+                detail=f"Rate limited by external service. Please retry after {retry_after} seconds."
+            )
+        else:
+            logger.error(f"HTTP error occurred: {e} - Response: {e.response.text}")
+            raise HTTPException(status_code=status_code, detail=str(e))
+    except httpx.RequestError as e:
+        logger.error(f"An error occurred while requesting: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
     async def generate():
-        async with httpx.AsyncClient() as client:
-            try:
-                async with client.stream(
-                    'POST',
-                    'https://sider.ai/api/v2/completion/text',  # Updated endpoint
-                    headers=OUTGOING_HEADERS,
-                    json=json_data,
-                    timeout=120.0
-                ) as response:
-                    response.raise_for_status()
-                    async for line in response.aiter_lines():
-                        if line and ("[DONE]" not in line):
-                            # Assuming the line starts with 'data: ' followed by JSON
-                            if line.startswith("data: "):
-                                json_line = line[6:]
-                                if json_line.startswith("{"):
-                                    try:
-                                        data = json.loads(json_line)
-                                        content = data.get("data", {}).get("text", "")
-                                        logger.debug(f"Received content: {content}")
-                                        yield f"data: {json.dumps(create_chat_completion_data(content, request.model))}\n\n"
-                                    except json.JSONDecodeError as e:
-                                        logger.error(f"JSON decode error: {e} - Line: {json_line}")
-                    # Send the stop signal
-                    yield f"data: {json.dumps(create_chat_completion_data('', request.model, 'stop'))}\n\n"
-                    yield "data: [DONE]\n\n"
-            except httpx.HTTPStatusError as e:
-                logger.error(f"HTTP error occurred: {e} - Response: {e.response.text}")
-                raise HTTPException(status_code=e.response.status_code, detail=str(e))
-            except httpx.RequestError as e:
-                logger.error(f"An error occurred while requesting: {e}")
-                raise HTTPException(status_code=500, detail=str(e))
+        async for line in response.aiter_lines():
+            if line and ("[DONE]" not in line):
+                # Assuming the line starts with 'data: ' followed by JSON
+                if line.startswith("data: "):
+                    json_line = line[6:]
+                    if json_line.startswith("{"):
+                        try:
+                            data = json.loads(json_line)
+                            content = data.get("data", {}).get("text", "")
+                            logger.debug(f"Received content: {content}")
+                            yield f"data: {json.dumps(create_chat_completion_data(content, request.model))}\n\n"
+                        except json.JSONDecodeError as e:
+                            logger.error(f"JSON decode error: {e} - Line: {json_line}")
+        # Send the stop signal
+        yield f"data: {json.dumps(create_chat_completion_data('', request.model, 'stop'))}\n\n"
+        yield "data: [DONE]\n\n"
 
     if request.stream:
         logger.info("Streaming response initiated.")
@@ -273,7 +295,6 @@ async def chat_completions(
             "usage": None,
         }
 
-
 # Entry Point
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)