Better error handling and retry

app.py

@@ -2,55 +2,97 @@ import os
 import httpx
 from fastapi import FastAPI, Request, HTTPException
 from fastapi.responses import JSONResponse, Response
+import time
 
 app = FastAPI()
 
-# Environment variables
 BASE_URL = os.getenv("BASE_URL", "https://ollama.com")
 MASTER_API_KEY = os.getenv("MASTER_API_KEY", "dummy")
 
-# Load Ollama keys from env
+# Load Ollama keys
 OLLAMA_KEYS = []
 for i in range(1, 10):
     key = os.getenv(f"OLLAMA_KEY_{i}")
     if key:
         OLLAMA_KEYS.append(key)
 
+# Track key failures
+key_failures = {k: 0 for k in OLLAMA_KEYS}
+last_used_key = {}
+
 @app.get("/")
 def root():
     return {
         "status": "ok",
         "master_key": MASTER_API_KEY[:10] + "...",
-        "ollama_keys": len(OLLAMA_KEYS)
+        "ollama_keys": len(OLLAMA_KEYS),
+        "key_failures": sum(key_failures.values())
     }
 
+def get_key():
+    """Get least failed key"""
+    for k, v in sorted(key_failures.items(), key=lambda x: x[1]):
+        if v < 3:
+            return k
+    # Reset if all failed
+    for k in key_failures:
+        key_failures[k] = 0
+    return OLLAMA_KEYS[0]
+
 @app.post("/v1/chat/completions")
 async def chat(req: Request):
-    # Auth
     auth_key = req.headers.get("Authorization", "").replace("Bearer ", "")
     if auth_key != MASTER_API_KEY:
         return JSONResponse({"error": "Unauthorized"}, status_code=401)
     
     body = await req.json()
+    max_retries = 3
     
-    # Try each key
-    for key in OLLAMA_KEYS:
+    for attempt in range(max_retries):
+        key = get_key()
+        
         try:
-            async with httpx.AsyncClient(timeout=60) as client:
+            async with httpx.AsyncClient(timeout=120.0) as client:
                 resp = await client.post(
                     f"{BASE_URL}/v1/chat/completions",
                     json=body,
-                    headers={"Authorization": f"Bearer {key}"}
+                    headers={
+                        "Authorization": f"Bearer {key}",
+                        "Content-Type": "application/json"
+                    }
                 )
+                
                 if resp.status_code == 200:
+                    key_failures[key] = 0  # Reset on success
                     return Response(resp.content, status_code=200)
+                    
+                elif resp.status_code == 429:
+                    # Rate limited - try another key
+                    key_failures[key] = key_failures.get(key, 0) + 1
+                    continue
+                    
+                elif resp.status_code >= 500:
+                    key_failures[key] = key_failures.get(key, 0) + 1
+                    continue
+                    
+                else:
+                    # Other error - return as is
+                    return Response(resp.content, status_code=resp.status_code)
+                    
+        except httpx.TimeoutException:
+            key_failures[key] = key_failures.get(key, 0) + 1
+            continue
         except Exception as e:
+            key_failures[key] = key_failures.get(key, 0) + 1
             continue
     
-    return JSONResponse({"error": "all keys failed"}, status_code=500)
+    return JSONResponse({
+        "error": "all keys failed",
+        "key_failures": {k[:10]: v for k, v in key_failures.items()}
+    }, status_code=500)
 
 @app.get("/v1/models")
-def models():
+def models(req: Request):
     auth_key = req.headers.get("Authorization", "").replace("Bearer ", "")
     if auth_key != MASTER_API_KEY:
         return JSONResponse({"error": "Unauthorized"}, status_code=401)
@@ -67,4 +109,4 @@ def models():
         except:
             pass
     
-    return JSONResponse({"error": "error"}, status_code=500)
+    return JSONResponse({"error": "no keys available"}, status_code=500)