Create main.py

main.py

@@ -0,0 +1,291 @@
+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
+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,
+)
+
+# ─── Logging ───
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s %(levelname)s %(message)s",
+    datefmt="%H:%M:%S",
+)
+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
+
+# ─── Static Headers ───
+HEADERS = {
+    "accept": "*/*",
+    "content-type": "application/json",
+    "origin": "https://www.blackbox.ai",
+    "referer": "https://www.blackbox.ai/",
+    "user-agent": (
+        "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
+        "AppleWebKit/537.36 (KHTML, like Gecko) "
+        "Chrome/136.0.0.0 Safari/537.36"
+    ),
+}
+
+# ─── FastAPI Setup ───
+app = FastAPI()
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+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
+        self.text = text
+        super().__init__(f"status={status} body={text[:100]}...")
+
+RETRYABLE_STATUSES = {400, 429, 500, 502, 503, 504}
+
+# ─── Random Data ───
+_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]:
+    return {
+        "messages": messages,
+        "id": _rand(8),
+        "agentMode": {},
+        "codeModelMode": True,
+        "trendingAgentMode": {},
+        "isMicMode": False,
+        "userSystemPrompt": None,
+        "maxTokens": 1024,
+        "playgroundTopP": None,
+        "playgroundTemperature": None,
+        "isChromeExt": False,
+        "githubToken": "",
+        "clickedAnswer2": False,
+        "clickedAnswer3": False,
+        "clickedForceWebSearch": False,
+        "visitFromDelta": False,
+        "isMemoryEnabled": False,
+        "mobileClient": False,
+        "userSelectedModel": None,
+        "validated": str(uuid.uuid4()),
+        "imageGenerationMode": False,
+        "webSearchModePrompt": False,
+        "deepSearchMode": True,
+        "domains": None,
+        "vscodeClient": False,
+        "codeInterpreterMode": False,
+        "customProfile": {
+            "name": "",
+            "occupation": "",
+            "traits": [],
+            "additionalInfo": "",
+            "enableNewChats": False
+        },
+        "session": {
+            "user": {
+                "name": "S.C gaming",
+                "email": random_email(),
+                "image": "https://lh3.googleusercontent.com/a/default",
+                "id": random_id()
+            },
+            "expires": "2025-12-31T23:59:59Z",
+            "isNewUser": False
+        },
+        "isPremium": True,
+        "subscriptionCache": {
+            "status": "PREMIUM",
+            "customerId": random_customer_id(),
+            "expiryTimestamp": int(time.time()) + 30 * 86400,
+            "lastChecked": int(time.time()),
+            "isTrialSubscription": False
+        },
+        "beastMode": False,
+        "reasoningMode": False,
+        "designerMode": False
+    }
+
+# ─── Retry Wrapper ───
+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)
+
+@retry(
+    stop=stop_after_attempt(10),
+    wait=wait_exponential(min=1, max=10),
+    retry=retry_if_exception_type(
+        (aiohttp.ClientConnectionError, aiohttp.ClientResponseError, asyncio.TimeoutError, RetryableStatusError)
+    ),
+    before_sleep=log_retry
+)
+async def get_blackbox_response(*, data, stream: bool, request_id: str) -> AsyncGenerator[str, None]:
+    assert HTTP_SESSION
+    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()
+
+# ─── 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())
+    logger.info("[%s] %s %s", rid, request.method, request.url.path)
+    start = time.perf_counter()
+    resp = await call_next(request)
+    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"}
+
+@app.get("/health")
+async def health():
+    return {"status": "ok"}
+
+# ─── Chat Completion ───
+@app.post("/v1/chat/completions")
+async def chat_completions(request: Request):
+    rid = request.state.request_id
+    try:
+        body = await request.json()
+        messages = body.get("messages", [])
+        if not messages:
+            raise HTTPException(status_code=400, detail="Missing 'messages'")
+        stream = body.get("stream", False)
+
+        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:"):
+                    raise HTTPException(status_code=502, detail=part)
+                chunks.append(part)
+            answer = "".join(chunks) or "No response."
+            return {
+                "id": str(uuid.uuid4()),
+                "object": "chat.completion",
+                "created": int(time.time()),
+                "model": "DeepResearch",
+                "choices": [
+                    {
+                        "index": 0,
+                        "message": {"role": "assistant", "content": answer},
+                        "finish_reason": "stop",
+                    }
+                ],
+            }
+
+        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")
+
+        return StreamingResponse(event_stream(), media_type="text/event-stream")
+
+    except json.JSONDecodeError:
+        raise HTTPException(status_code=400, detail="Invalid JSON")
+    except RetryError as re:
+        logger.error("[%s] retries failed: %s", rid, re)
+        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")
+
+# ─── 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")