Update main.py

main.py

@@ -1,12 +1,10 @@
 import asyncio
 import json
 import logging
-import os
 import random
 import string
 import time
 import uuid
-from http import HTTPStatus
 from typing import AsyncGenerator, Dict, List, Any
 
 import aiohttp
@@ -31,12 +29,10 @@ logger = logging.getLogger("proxy")
 
 # ─── Config ───
 BLACKBOX_URL = "https://www.blackbox.ai/api/chat"
-CONNECTION_LIMIT = 200
-CONNECTION_LIMIT_PER_HOST = 50
 REQUEST_TIMEOUT = 300
 WORKER_COUNT = 10
 
-# ─── Required Headers ───
+# ─── Headers ───
 HEADERS = {
     "authority": "www.blackbox.ai",
     "method": "POST",
@@ -62,7 +58,7 @@ HEADERS = {
     "content-type": "application/json",
 }
 
-# ─── FastAPI Setup ───
+# ─── FastAPI ───
 app = FastAPI()
 app.add_middleware(
     CORSMiddleware,
@@ -76,7 +72,6 @@ HTTP_SESSION: aiohttp.ClientSession = None
 REQUEST_QUEUE: asyncio.Queue = asyncio.Queue()
 WORKER_TASKS: List[asyncio.Task] = []
 
-# ─── Retryable Error ───
 class RetryableStatusError(Exception):
     def __init__(self, status: int, text: str):
         self.status = status
@@ -86,18 +81,21 @@ class RetryableStatusError(Exception):
 RETRYABLE_STATUSES = {400, 429, 500, 502, 503, 504}
 
 _ascii = string.ascii_letters + string.digits
+
 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)
 def random_customer_id(): return "cus_" + _rand(12)
 
-# ─── Payload Generator ───
 def build_payload(messages: List[Dict[str, Any]]) -> Dict[str, Any]:
     unique_id = str(uuid.uuid4())
+    msg_id = _rand(8)
+    if messages and isinstance(messages[-1], dict):
+        msg_id = messages[-1].get("id", msg_id)
     return {
         "messages": messages,
         "agentMode": {},
-        "id": messages[-1]["id"] if messages else _rand(8),
+        "id": msg_id,
         "previewToken": None,
         "userId": None,
         "codeModelMode": True,
@@ -153,7 +151,6 @@ def build_payload(messages: List[Dict[str, Any]]) -> Dict[str, Any]:
         "designerMode": False
     }
 
-# ─── Retry Wrapper ───
 def log_retry(retry_state):
     rid = retry_state.kwargs.get("request_id", "unknown")
     attempt = retry_state.attempt_number
@@ -161,7 +158,7 @@ def log_retry(retry_state):
     logger.warning("[%s] retry %s/10 due to %s", rid, attempt, err)
 
 @retry(
-    stop=stop_after_attempt(10),
+    stop=stop_after_attempt(3),
     wait=wait_exponential(min=1, max=10),
     retry=retry_if_exception_type(
         (aiohttp.ClientConnectionError, aiohttp.ClientResponseError, asyncio.TimeoutError, RetryableStatusError)
@@ -169,7 +166,9 @@ def log_retry(retry_state):
     before_sleep=log_retry
 )
 async def get_blackbox_response(*, data, stream: bool, request_id: str) -> AsyncGenerator[str, None]:
-    assert HTTP_SESSION
+    global HTTP_SESSION
+    if not HTTP_SESSION:
+        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()
@@ -184,23 +183,6 @@ async def get_blackbox_response(*, data, stream: bool, request_id: str) -> Async
         else:
             yield await resp.text()
 
-# ─── Worker Thread ───
-async def _worker():
-    while True:
-        try:
-            data, request_id, out_q = await REQUEST_QUEUE.get()
-            try:
-                async for piece in get_blackbox_response(data=data, stream=False, request_id=request_id):
-                    await out_q.put(piece)
-            except Exception as e:
-                await out_q.put(f"Error:{e}")
-            finally:
-                await out_q.put(None)
-                REQUEST_QUEUE.task_done()
-        except asyncio.CancelledError:
-            break
-
-# ─── Middleware ───
 @app.middleware("http")
 async def add_request_id(request: Request, call_next):
     request.state.request_id = rid = str(uuid.uuid4())
@@ -210,7 +192,6 @@ async def add_request_id(request: Request, call_next):
     logger.info("[%s] finished in %.2fs", rid, time.perf_counter() - start)
     return resp
 
-# ─── Root & Health ───
 @app.get("/")
 async def root():
     return {"message": "API is running"}
@@ -219,7 +200,6 @@ async def root():
 async def health():
     return {"status": "ok"}
 
-# ─── Chat Completion ───
 @app.post("/v1/chat/completions")
 async def chat_completions(request: Request):
     rid = request.state.request_id
@@ -232,14 +212,9 @@ async def chat_completions(request: Request):
         payload = build_payload(messages)
 
         if not stream:
-            q: asyncio.Queue = asyncio.Queue()
-            await REQUEST_QUEUE.put((payload, rid, q))
             chunks: List[str] = []
-            while True:
-                part = await q.get()
-                if part is None:
-                    break
-                if isinstance(part, str) and part.startswith("Error:"):
+            async for part in get_blackbox_response(data=payload, stream=False, request_id=rid):
+                if part.startswith("Error:"):
                     raise HTTPException(status_code=502, detail=part)
                 chunks.append(part)
             answer = "".join(chunks) or "No response."
@@ -258,20 +233,16 @@ async def chat_completions(request: Request):
             }
 
         async def event_stream():
-            try:
-                async for chunk in get_blackbox_response(data=payload, stream=True, request_id=rid):
-                    msg = {
-                        "id": str(uuid.uuid4()),
-                        "object": "chat.completion.chunk",
-                        "created": int(time.time()),
-                        "model": "DeepResearch",
-                        "choices": [{"index": 0, "delta": {"content": chunk}}],
-                    }
-                    yield f"data: {json.dumps(msg)}\n\n"
-                yield "data: [DONE]\n\n"
-            except Exception as e:
-                logger.error("[%s] stream error: %s", rid, e)
-                raise HTTPException(status_code=500, detail="streaming error")
+            async for chunk in get_blackbox_response(data=payload, stream=True, request_id=rid):
+                msg = {
+                    "id": str(uuid.uuid4()),
+                    "object": "chat.completion.chunk",
+                    "created": int(time.time()),
+                    "model": "DeepResearch",
+                    "choices": [{"index": 0, "delta": {"content": chunk}}],
+                }
+                yield f"data: {json.dumps(msg)}\n\n"
+            yield "data: [DONE]\n\n"
 
         return StreamingResponse(event_stream(), media_type="text/event-stream")
 
@@ -283,23 +254,3 @@ async def chat_completions(request: Request):
     except Exception as e:
         logger.exception("[%s] error", rid)
         raise HTTPException(status_code=500, detail="Internal proxy error")
-
-# ─── Startup & Shutdown ───
-@app.on_event("startup")
-async def startup():
-    global HTTP_SESSION, WORKER_TASKS
-    HTTP_SESSION = aiohttp.ClientSession(
-        connector=aiohttp.TCPConnector(limit=CONNECTION_LIMIT, limit_per_host=CONNECTION_LIMIT_PER_HOST),
-        timeout=aiohttp.ClientTimeout(total=REQUEST_TIMEOUT),
-    )
-    WORKER_TASKS = [asyncio.create_task(_worker()) for _ in range(WORKER_COUNT)]
-    logger.info("Started %d workers", WORKER_COUNT)
-
-@app.on_event("shutdown")
-async def shutdown():
-    for t in WORKER_TASKS:
-        t.cancel()
-    await asyncio.gather(*WORKER_TASKS, return_exceptions=True)
-    if HTTP_SESSION:
-        await HTTP_SESSION.close()
-    logger.info("Shutdown complete")