Update main.py

main.py

@@ -3,11 +3,21 @@ from fastapi import FastAPI, Request, HTTPException
 from starlette.responses import StreamingResponse
 from starlette.background import BackgroundTask
 import os
+import random
 from contextlib import asynccontextmanager
 
 # --- Configuration ---
 # The target URL is configurable via an environment variable.
 TARGET_URL = os.getenv("TARGET_URL", "https://api.gmi-serving.com/v1/chat")
+# Number of retries for specific error codes.
+MAX_RETRIES = 5
+# HTTP status codes that will trigger a retry.
+RETRY_STATUS_CODES = {429, 502, 503, 504}
+
+# --- Helper Function ---
+def generate_random_ip():
+    """Generates a random, valid-looking IPv4 address."""
+    return ".".join(str(random.randint(1, 254)) for _ in range(4))
 
 # --- HTTPX Client Lifecycle Management ---
 @asynccontextmanager
@@ -29,8 +39,8 @@ app = FastAPI(docs_url=None, redoc_url=None, lifespan=lifespan)
 # --- Reverse Proxy Logic ---
 async def _reverse_proxy(request: Request):
     """
-    Forwards a request specifically for the /completions endpoint to the target URL.
-    It injects required headers and allows for a user-provided Authorization header.
+    Forwards a request to the target URL with retry logic and spoofed IP headers.
+    It allows for a user-provided Authorization header.
     """
     client: httpx.AsyncClient = request.app.state.http_client
 
@@ -48,7 +58,10 @@ async def _reverse_proxy(request: Request):
     # 2. Get the user's authorization key from the incoming request.
     authorization_header = request.headers.get("authorization")
 
-    # 3. Set the specific, required headers for the target API.
+    # 3. Generate a random IP for spoofing headers.
+    random_ip = generate_random_ip()
+
+    # 4. Set the specific, required headers for the target API.
     #    This will overwrite any conflicting headers from the original request.
     specific_headers = {
         "accept": "application/json, text/plain, */*",
@@ -64,35 +77,63 @@ async def _reverse_proxy(request: Request):
         "sec-fetch-mode": "cors",
         "sec-fetch-site": "same-origin",
         "user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/139.0.0.0 Safari/537.36",
+        "x-forwarded-for": random_ip,
+        "x-real-ip": random_ip,
     }
     request_headers.update(specific_headers)
 
-    # 4. Add the user's authorization key to the headers if it exists.
+    # 5. Add the user's authorization key to the headers if it exists.
     if authorization_header:
         request_headers["authorization"] = authorization_header
 
-    # Build the final request to the target service.
-    rp_req = client.build_request(
-        method=request.method,
-        url=url,
-        headers=request_headers,
-        content=await request.body(),
+    # Read the request body once, as it will be reused in case of retries.
+    body = await request.body()
+    
+    # --- Retry Logic ---
+    last_exception = None
+    for attempt in range(MAX_RETRIES):
+        try:
+            # Build the request for each attempt to ensure the content stream is fresh.
+            rp_req = client.build_request(
+                method=request.method,
+                url=url,
+                headers=request_headers,
+                content=body,
+            )
+            # Send the request and get a streaming response.
+            rp_resp = await client.send(rp_req, stream=True)
+
+            # If the status code is not in our retry list, we can return the response.
+            if rp_resp.status_code not in RETRY_STATUS_CODES:
+                return StreamingResponse(
+                    rp_resp.aiter_raw(),
+                    status_code=rp_resp.status_code,
+                    headers=rp_resp.headers,
+                    background=BackgroundTask(rp_resp.aclose),
+                )
+            
+            # If it's a retryable error but this is the last attempt, return the error response.
+            if attempt == MAX_RETRIES - 1:
+                 return StreamingResponse(
+                    rp_resp.aiter_raw(),
+                    status_code=rp_resp.status_code,
+                    headers=rp_resp.headers,
+                    background=BackgroundTask(rp_resp.aclose),
+                )
+
+            # Otherwise, close the unsuccessful response before trying again.
+            await rp_resp.aclose()
+
+        except httpx.ConnectError as e:
+            last_exception = e
+            # Continue to the next attempt if a connection error occurs.
+
+    # If all retry attempts failed with a connection error, raise a final exception.
+    raise HTTPException(
+        status_code=502,
+        detail=f"Bad Gateway: Cannot connect to target service after {MAX_RETRIES} attempts. {last_exception}"
     )
 
-    try:
-        # Send the request and get a streaming response.
-        rp_resp = await client.send(rp_req, stream=True)
-    except httpx.ConnectError as e:
-        # This error occurs if the target service is down or unreachable.
-        raise HTTPException(status_code=502, detail=f"Bad Gateway: Cannot connect to target service. {e}")
-
-    # Stream the response from the target service back to the original client.
-    return StreamingResponse(
-        rp_resp.aiter_raw(),
-        status_code=rp_resp.status_code,
-        headers=rp_resp.headers,
-        background=BackgroundTask(rp_resp.aclose),
-    )
 
 # --- API Endpoint ---
 @app.api_route(