Update main.py

main.py

@@ -4,22 +4,28 @@ from fastapi import FastAPI, HTTPException, Request
 from fastapi.responses import StreamingResponse, HTMLResponse
 from pydantic import BaseModel
 import httpx
-from functools import lru_cache
 
 load_dotenv()
 
 app = FastAPI()
 
+# Get API keys and secret endpoint from environment variables
 api_keys_str = os.getenv('API_KEYS')
 valid_api_keys = api_keys_str.split(',') if api_keys_str else []
 secret_api_endpoint = os.getenv('SECRET_API_ENDPOINT')
+secret_api_endpoint_2 = os.getenv('SECRET_API_ENDPOINT_2')
 
-if not secret_api_endpoint:
-    raise HTTPException(status_code=500, detail="API endpoint is not configured in environment variables.")
+# Validate if the main secret API endpoint is set
+if not secret_api_endpoint or not secret_api_endpoint_2:
+    raise HTTPException(status_code=500, detail="API endpoint(s) are not configured in environment variables.")
+
+# Define models that should use the secondary endpoint
+alternate_models = {"gpt-4o-mini", "claude-3-haiku", "llama-3.1-70b", "mixtral-8x7b"}
 
 class Payload(BaseModel):
     model: str
     messages: list
+    stream: bool
 
 @app.get("/", response_class=HTMLResponse)
 async def root():
@@ -65,15 +71,20 @@ async def fetch_models():
 async def get_completion(payload: Payload, request: Request):
     api_key = request.headers.get("Authorization")
     
+    # Validate API key
     if api_key not in valid_api_keys:
         raise HTTPException(status_code=403, detail="Forbidden: Invalid API key. Join dsc.gg/chadgang and DM @mr_leaderyt on Discord for a free API key :)")
     
-    payload_dict = {**payload.dict(), "stream": True}
+    # Determine which endpoint to use based on the model
+    endpoint = secret_api_endpoint_2 if payload.model in alternate_models else secret_api_endpoint
+
+    # Use the payload directly as it includes stream and other user data
+    payload_dict = payload.dict()
 
     async def stream_generator(payload_dict):
         async with httpx.AsyncClient() as client:
             try:
-                async with client.stream("POST", f"{secret_api_endpoint}/v1/chat/completions", json=payload_dict, timeout=10) as response:
+                async with client.stream("POST", f"{endpoint}/v1/chat/completions", json=payload_dict, timeout=10) as response:
                     response.raise_for_status()
                     async for line in response.aiter_lines():
                         if line: