Update main.py

main.py

@@ -1,26 +1,29 @@
 import httpx
-from fastapi import FastAPI, Request
+from fastapi import FastAPI, Request, HTTPException
 from fastapi.responses import StreamingResponse
 import json
+import asyncio
 
 app = FastAPI()
 
+# List of API URLs, with the primary first and fallbacks after
+API_URLS = [
+    "https://api.deepinfra.com/v1/openai/chat/completions",
+    "https://stage.api.deepinfra.com/v1/openai/chat/completions",
+]
+
 @app.post("/v1/openai/chat/completions")
 async def proxy_deepinfra(request: Request):
     """
-    Proxies chat completion requests to the DeepInfra API.
+    Proxies chat completion requests to the DeepInfra API with fallback support.
     """
-    client = httpx.AsyncClient()
-    url = "https://api.deepinfra.com/v1/openai/chat/completions"
-
-    # Extract the raw body from the request
-    body = await request.body()
-    # Decode the body to a string and then load it as JSON
     try:
-        data = json.loads(body.decode('utf-8'))
+        body = await request.json()
     except json.JSONDecodeError:
-        return {"error": "Invalid JSON in request body"}, 400
+        raise HTTPException(status_code=400, detail="Invalid JSON in request body")
 
+    # The original curl command headers can be used as a base
+    # and overridden by headers from the incoming request if needed.
     headers = {
         'sec-ch-ua-platform': request.headers.get('sec-ch-ua-platform', '"Windows"'),
         'Referer': request.headers.get('Referer', 'https://deepinfra.com/'),
@@ -32,12 +35,28 @@ async def proxy_deepinfra(request: Request):
         'Content-Type': request.headers.get('Content-Type', 'application/json'),
     }
 
-    async def stream_response():
-        async with client.stream("POST", url, headers=headers, json=data, timeout=None) as response:
-            async for chunk in response.aiter_bytes():
-                yield chunk
+    async def stream_generator():
+        last_error = None
+        for url in API_URLS:
+            try:
+                async with httpx.AsyncClient() as client:
+                    async with client.stream("POST", url, headers=headers, json=body, timeout=30.0) as response:
+                        # Raise an exception for bad status codes to trigger the fallback
+                        response.raise_for_status()
+                        async for chunk in response.aiter_bytes():
+                            yield chunk
+                        # If we successfully stream the response, we're done.
+                        return
+            except (httpx.RequestError, httpx.HTTPStatusError) as e:
+                # Store the last error and try the next URL
+                last_error = e
+                print(f"Error with URL {url}: {e}. Trying next fallback.")
+                continue
+        # If all URLs failed, raise an exception based on the last error.
+        if last_error:
+            raise HTTPException(status_code=502, detail=f"All API endpoints failed. Last error: {last_error}")
 
-    return StreamingResponse(stream_response(), media_type="text/event-stream")
+    return StreamingResponse(stream_generator(), media_type="text/event-stream")
 
 if __name__ == "__main__":
     import uvicorn