Update main.py

main.py

@@ -13,13 +13,7 @@ import importlib
 from fastapi import FastAPI, Request, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import StreamingResponse
-from tenacity import (
-    retry,
-    stop_after_attempt,
-    wait_exponential,
-    retry_if_exception_type,
-    RetryError,
-)
+from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type, RetryError
 
 # ─── Logging ───
 logging.basicConfig(
@@ -70,17 +64,10 @@ app.add_middleware(
 )
 
 HTTP_SESSION: aiohttp.ClientSession = None
-
-class RetryableStatusError(Exception):
-    def __init__(self, status: int, text: str):
-        self.status = status
-        self.text = text
-        super().__init__(f"status={status} body={text[:100]}...")
-
 RETRYABLE_STATUSES = {400, 429, 500, 502, 503, 504}
-
 _ascii = string.ascii_letters + string.digits
 
+# ─── Utility ───
 def _rand(n, pool=_ascii): return "".join(random.choice(pool) for _ in range(n))
 def random_email(): return _rand(12) + "@gmail.com"
 def random_id(): return _rand(21, string.digits)
@@ -146,12 +133,14 @@ def build_payload(messages: List[Dict[str, Any]]) -> Dict[str, Any]:
         "designerMode": False
     }
 
-# ─── Retry Logic ───
+# ─── Retry ───
+class RetryableStatusError(Exception):
+    def __init__(self, status: int, text: str):
+        super().__init__(f"status={status} body={text[:100]}...")
+
 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/3 due to %s", rid, attempt, err)
+    logger.warning("[%s] retry %s/3 due to %s", rid, retry_state.attempt_number, retry_state.outcome.exception())
 
 @retry(
     stop=stop_after_attempt(3),
@@ -164,30 +153,23 @@ 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")
+        importlib.import_module("brotli")
         HTTP_SESSION = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=REQUEST_TIMEOUT))
 
-    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 Exception as e:
-        logger.exception("[%s] Unexpected upstream error", request_id)
-        raise HTTPException(status_code=502, detail="Failed to contact upstream") from e
+    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()
 
 # ─── Middleware ───
 @app.middleware("http")
@@ -231,13 +213,11 @@ async def chat_completions(request: Request):
                 "object": "chat.completion",
                 "created": int(time.time()),
                 "model": "DeepResearch",
-                "choices": [
-                    {
-                        "index": 0,
-                        "message": {"role": "assistant", "content": answer},
-                        "finish_reason": "stop",
-                    }
-                ],
+                "choices": [{
+                    "index": 0,
+                    "message": {"role": "assistant", "content": answer},
+                    "finish_reason": "stop",
+                }],
             }
 
         async def event_stream():
@@ -249,12 +229,8 @@ async def chat_completions(request: Request):
                     "model": "DeepResearch",
                     "choices": [{"index": 0, "delta": {"content": chunk}}],
                 }
-                yield f"data: {json.dumps(msg)}
-
-"
-            yield "data: [DONE]
-
-"
+                yield f"data: {json.dumps(msg)}\n\n"
+            yield "data: [DONE]\n\n"
 
         return StreamingResponse(event_stream(), media_type="text/event-stream")
 
@@ -265,4 +241,4 @@ async def chat_completions(request: Request):
         raise HTTPException(status_code=502, detail="Blackbox upstream failed")
     except Exception as e:
         logger.exception("[%s] error", rid)
-        raise HTTPException(status_code=500, detail="Internal proxy error") 
+        raise HTTPException(status_code=500, detail="Internal proxy error")