Update main.py

main.py

@@ -8,6 +8,9 @@ import uuid
 from typing import AsyncGenerator, Dict, List, Any
 
 import aiohttp
+import brotli
+import sys
+import importlib
 from fastapi import FastAPI, Request, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import StreamingResponse
@@ -30,14 +33,9 @@ logger = logging.getLogger("proxy")
 # ─── Config ───
 BLACKBOX_URL = "https://www.blackbox.ai/api/chat"
 REQUEST_TIMEOUT = 300
-WORKER_COUNT = 10
 
 # ─── Headers ───
 HEADERS = {
-    "authority": "www.blackbox.ai",
-    "method": "POST",
-    "path": "/api/chat",
-    "scheme": "https",
     "accept": "*/*",
     "accept-encoding": "gzip, deflate, br, zstd",
     "accept-language": "en-US,en;q=0.9",
@@ -69,8 +67,6 @@ app.add_middleware(
 )
 
 HTTP_SESSION: aiohttp.ClientSession = None
-REQUEST_QUEUE: asyncio.Queue = asyncio.Queue()
-WORKER_TASKS: List[asyncio.Task] = []
 
 class RetryableStatusError(Exception):
     def __init__(self, status: int, text: str):
@@ -155,7 +151,7 @@ def log_retry(retry_state):
     rid = retry_state.kwargs.get("request_id", "unknown")
     attempt = retry_state.attempt_number
     err = retry_state.outcome.exception()
-    logger.warning("[%s] retry %s/10 due to %s", rid, attempt, err)
+    logger.warning("[%s] retry %s/3 due to %s", rid, attempt, err)
 
 @retry(
     stop=stop_after_attempt(3),
@@ -168,20 +164,30 @@ def log_retry(retry_state):
 async def get_blackbox_response(*, data, stream: bool, request_id: str) -> AsyncGenerator[str, None]:
     global HTTP_SESSION
     if not HTTP_SESSION:
+        try:
+            importlib.import_module("brotli")
+        except ImportError:
+            logger.error("Missing Brotli module. Install with: pip install brotli")
+            raise HTTPException(status_code=502, detail="Missing Brotli module on server")
         HTTP_SESSION = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=REQUEST_TIMEOUT))
-    async with HTTP_SESSION.post(BLACKBOX_URL, json=data, headers=HEADERS, timeout=REQUEST_TIMEOUT) as resp:
-        if resp.status != 200:
-            body = await resp.text()
-            logger.error("[%s] Upstream %s error: %s", request_id, BLACKBOX_URL, resp.status)
-            if resp.status in RETRYABLE_STATUSES:
-                raise RetryableStatusError(resp.status, body)
-            raise HTTPException(status_code=502, detail=f"Upstream error {resp.status}")
-        if stream:
-            async for chunk in resp.content.iter_any():
-                if chunk:
-                    yield chunk.decode("utf-8", "ignore")
-        else:
-            yield await resp.text()
+
+    try:
+        async with HTTP_SESSION.post(BLACKBOX_URL, json=data, headers=HEADERS, timeout=REQUEST_TIMEOUT) as resp:
+            if resp.status != 200:
+                body = await resp.text()
+                logger.error("[%s] Upstream %s error: %s", request_id, BLACKBOX_URL, resp.status)
+                if resp.status in RETRYABLE_STATUSES:
+                    raise RetryableStatusError(resp.status, body)
+                raise HTTPException(status_code=502, detail=f"Upstream error {resp.status}")
+            if stream:
+                async for chunk in resp.content.iter_any():
+                    if chunk:
+                        yield chunk.decode("utf-8", "ignore")
+            else:
+                yield await resp.text()
+    except aiohttp.ContentEncodingError as e:
+        logger.error("[%s] Brotli decode failed: %s", request_id, e)
+        raise HTTPException(status_code=502, detail="Install Brotli for br-encoded content") from e
 
 @app.middleware("http")
 async def add_request_id(request: Request, call_next):