Update main.py
Browse files
main.py
CHANGED
@@ -9,7 +9,6 @@ from typing import AsyncGenerator, Dict, List, Any
|
|
9 |
|
10 |
import aiohttp
|
11 |
import brotli
|
12 |
-
import importlib
|
13 |
from fastapi import FastAPI, Request, HTTPException
|
14 |
from fastapi.middleware.cors import CORSMiddleware
|
15 |
from fastapi.responses import StreamingResponse
|
@@ -67,17 +66,24 @@ HTTP_SESSION: aiohttp.ClientSession = None
|
|
67 |
RETRYABLE_STATUSES = {400, 429, 500, 502, 503, 504}
|
68 |
_ascii = string.ascii_letters + string.digits
|
69 |
|
70 |
-
# βββ
|
71 |
-
def _rand(n, pool=_ascii): return
|
72 |
def random_email(): return _rand(12) + "@gmail.com"
|
73 |
def random_id(): return _rand(21, string.digits)
|
74 |
def random_customer_id(): return "cus_" + _rand(12)
|
|
|
75 |
|
|
|
76 |
def build_payload(messages: List[Dict[str, Any]]) -> Dict[str, Any]:
|
|
|
|
|
|
|
|
|
|
|
77 |
return {
|
78 |
"messages": messages,
|
79 |
"agentMode": {},
|
80 |
-
"id":
|
81 |
"previewToken": None,
|
82 |
"userId": None,
|
83 |
"codeModelMode": True,
|
@@ -124,8 +130,8 @@ def build_payload(messages: List[Dict[str, Any]]) -> Dict[str, Any]:
|
|
124 |
"subscriptionCache": {
|
125 |
"status": "PREMIUM",
|
126 |
"customerId": random_customer_id(),
|
127 |
-
"expiryTimestamp":
|
128 |
-
"lastChecked":
|
129 |
"isTrialSubscription": False
|
130 |
},
|
131 |
"beastMode": False,
|
@@ -133,7 +139,7 @@ def build_payload(messages: List[Dict[str, Any]]) -> Dict[str, Any]:
|
|
133 |
"designerMode": False
|
134 |
}
|
135 |
|
136 |
-
# βββ Retry βββ
|
137 |
class RetryableStatusError(Exception):
|
138 |
def __init__(self, status: int, text: str):
|
139 |
super().__init__(f"status={status} body={text[:100]}...")
|
@@ -153,7 +159,6 @@ def log_retry(retry_state):
|
|
153 |
async def get_blackbox_response(*, data, stream: bool, request_id: str) -> AsyncGenerator[str, None]:
|
154 |
global HTTP_SESSION
|
155 |
if not HTTP_SESSION:
|
156 |
-
importlib.import_module("brotli")
|
157 |
HTTP_SESSION = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=REQUEST_TIMEOUT))
|
158 |
|
159 |
async with HTTP_SESSION.post(BLACKBOX_URL, json=data, headers=HEADERS, timeout=REQUEST_TIMEOUT) as resp:
|
|
|
9 |
|
10 |
import aiohttp
|
11 |
import brotli
|
|
|
12 |
from fastapi import FastAPI, Request, HTTPException
|
13 |
from fastapi.middleware.cors import CORSMiddleware
|
14 |
from fastapi.responses import StreamingResponse
|
|
|
66 |
RETRYABLE_STATUSES = {400, 429, 500, 502, 503, 504}
|
67 |
_ascii = string.ascii_letters + string.digits
|
68 |
|
69 |
+
# βββ Helpers βββ
|
70 |
+
def _rand(n, pool=_ascii): return ''.join(random.choice(pool) for _ in range(n))
|
71 |
def random_email(): return _rand(12) + "@gmail.com"
|
72 |
def random_id(): return _rand(21, string.digits)
|
73 |
def random_customer_id(): return "cus_" + _rand(12)
|
74 |
+
def generate_7char_id(): return _rand(7)
|
75 |
|
76 |
+
# βββ Payload Builder βββ
|
77 |
def build_payload(messages: List[Dict[str, Any]]) -> Dict[str, Any]:
|
78 |
+
msg_id = generate_7char_id()
|
79 |
+
if messages:
|
80 |
+
messages[-1]["id"] = msg_id
|
81 |
+
|
82 |
+
now = int(time.time())
|
83 |
return {
|
84 |
"messages": messages,
|
85 |
"agentMode": {},
|
86 |
+
"id": msg_id,
|
87 |
"previewToken": None,
|
88 |
"userId": None,
|
89 |
"codeModelMode": True,
|
|
|
130 |
"subscriptionCache": {
|
131 |
"status": "PREMIUM",
|
132 |
"customerId": random_customer_id(),
|
133 |
+
"expiryTimestamp": now + 60 * 86400, # 60 days
|
134 |
+
"lastChecked": int(time.time() * 1000), # milliseconds
|
135 |
"isTrialSubscription": False
|
136 |
},
|
137 |
"beastMode": False,
|
|
|
139 |
"designerMode": False
|
140 |
}
|
141 |
|
142 |
+
# βββ Retry Logic βββ
|
143 |
class RetryableStatusError(Exception):
|
144 |
def __init__(self, status: int, text: str):
|
145 |
super().__init__(f"status={status} body={text[:100]}...")
|
|
|
159 |
async def get_blackbox_response(*, data, stream: bool, request_id: str) -> AsyncGenerator[str, None]:
|
160 |
global HTTP_SESSION
|
161 |
if not HTTP_SESSION:
|
|
|
162 |
HTTP_SESSION = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=REQUEST_TIMEOUT))
|
163 |
|
164 |
async with HTTP_SESSION.post(BLACKBOX_URL, json=data, headers=HEADERS, timeout=REQUEST_TIMEOUT) as resp:
|