Update main.py

main.py

@@ -10,12 +10,18 @@ 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
@@ -65,15 +71,18 @@ 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 :)")
     
+    # Determine which endpoint to use based on the model
+    endpoint = secret_api_endpoint_2 if payload.model in alternate_models else secret_api_endpoint
     payload_dict = {**payload.dict(), "stream": True}
 
     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:
@@ -96,4 +105,4 @@ async def startup_event():
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    uvicorn.run(app, host="0.0.0.0", port=8000)