Update main.py
Browse files
main.py
CHANGED
@@ -1,607 +1,5 @@
|
|
1 |
-
import
|
2 |
-
import
|
3 |
-
import random
|
4 |
-
import string
|
5 |
-
import uuid
|
6 |
-
import json
|
7 |
-
import logging
|
8 |
-
import asyncio
|
9 |
-
import time
|
10 |
-
from collections import defaultdict
|
11 |
-
from typing import List, Dict, Any, Optional, AsyncGenerator, Union
|
12 |
-
from datetime import datetime
|
13 |
-
|
14 |
-
from aiohttp import ClientSession, ClientTimeout, ClientError
|
15 |
-
from fastapi import FastAPI, HTTPException, Request, Depends, Header
|
16 |
-
from fastapi.responses import StreamingResponse, JSONResponse, RedirectResponse
|
17 |
-
from pydantic import BaseModel
|
18 |
-
|
19 |
-
# Configure logging
|
20 |
-
logging.basicConfig(
|
21 |
-
level=logging.INFO,
|
22 |
-
format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
|
23 |
-
handlers=[logging.StreamHandler()]
|
24 |
-
)
|
25 |
-
logger = logging.getLogger(__name__)
|
26 |
-
|
27 |
-
# Load environment variables
|
28 |
-
API_KEYS = os.getenv('API_KEYS', '').split(',') # Comma-separated API keys
|
29 |
-
RATE_LIMIT = int(os.getenv('RATE_LIMIT', '60')) # Requests per minute
|
30 |
-
AVAILABLE_MODELS = os.getenv('AVAILABLE_MODELS', '') # Comma-separated available models
|
31 |
-
|
32 |
-
if not API_KEYS or API_KEYS == ['']:
|
33 |
-
logger.error("No API keys found. Please set the API_KEYS environment variable.")
|
34 |
-
raise Exception("API_KEYS environment variable not set.")
|
35 |
-
|
36 |
-
# Process available models
|
37 |
-
if AVAILABLE_MODELS:
|
38 |
-
AVAILABLE_MODELS = [model.strip() for model in AVAILABLE_MODELS.split(',') if model.strip()]
|
39 |
-
else:
|
40 |
-
AVAILABLE_MODELS = [] # If empty, all models are available
|
41 |
-
|
42 |
-
# Simple in-memory rate limiter based solely on IP addresses
|
43 |
-
rate_limit_store = defaultdict(lambda: {"count": 0, "timestamp": time.time()})
|
44 |
-
|
45 |
-
# Define cleanup interval and window
|
46 |
-
CLEANUP_INTERVAL = 60 # seconds
|
47 |
-
RATE_LIMIT_WINDOW = 60 # seconds
|
48 |
-
|
49 |
-
async def cleanup_rate_limit_stores():
|
50 |
-
"""
|
51 |
-
Periodically cleans up stale entries in the rate_limit_store to prevent memory bloat.
|
52 |
-
"""
|
53 |
-
while True:
|
54 |
-
current_time = time.time()
|
55 |
-
ips_to_delete = [ip for ip, value in rate_limit_store.items() if current_time - value["timestamp"] > RATE_LIMIT_WINDOW * 2]
|
56 |
-
for ip in ips_to_delete:
|
57 |
-
del rate_limit_store[ip]
|
58 |
-
logger.debug(f"Cleaned up rate_limit_store for IP: {ip}")
|
59 |
-
await asyncio.sleep(CLEANUP_INTERVAL)
|
60 |
-
|
61 |
-
async def rate_limiter_per_ip(request: Request):
|
62 |
-
"""
|
63 |
-
Rate limiter that enforces a limit based on the client's IP address.
|
64 |
-
"""
|
65 |
-
client_ip = request.client.host
|
66 |
-
current_time = time.time()
|
67 |
-
|
68 |
-
# Initialize or update the count and timestamp
|
69 |
-
if current_time - rate_limit_store[client_ip]["timestamp"] > RATE_LIMIT_WINDOW:
|
70 |
-
rate_limit_store[client_ip] = {"count": 1, "timestamp": current_time}
|
71 |
-
else:
|
72 |
-
if rate_limit_store[client_ip]["count"] >= RATE_LIMIT:
|
73 |
-
logger.warning(f"Rate limit exceeded for IP address: {client_ip}")
|
74 |
-
raise HTTPException(status_code=429, detail='Rate limit exceeded for IP address | NiansuhAI')
|
75 |
-
rate_limit_store[client_ip]["count"] += 1
|
76 |
-
|
77 |
-
async def get_api_key(request: Request, authorization: str = Header(None)) -> str:
|
78 |
-
"""
|
79 |
-
Dependency to extract and validate the API key from the Authorization header.
|
80 |
-
"""
|
81 |
-
client_ip = request.client.host
|
82 |
-
if authorization is None or not authorization.startswith('Bearer '):
|
83 |
-
logger.warning(f"Invalid or missing authorization header from IP: {client_ip}")
|
84 |
-
raise HTTPException(status_code=401, detail='Invalid authorization header format')
|
85 |
-
api_key = authorization[7:]
|
86 |
-
if api_key not in API_KEYS:
|
87 |
-
logger.warning(f"Invalid API key attempted: {api_key} from IP: {client_ip}")
|
88 |
-
raise HTTPException(status_code=401, detail='Invalid API key')
|
89 |
-
return api_key
|
90 |
-
|
91 |
-
# Custom exception for model not working
|
92 |
-
class ModelNotWorkingException(Exception):
|
93 |
-
def __init__(self, model: str):
|
94 |
-
self.model = model
|
95 |
-
self.message = f"The model '{model}' is currently not working. Please try another model or wait for it to be fixed."
|
96 |
-
super().__init__(self.message)
|
97 |
-
|
98 |
-
# Mock implementations for ImageResponse and to_data_uri
|
99 |
-
class ImageResponse:
|
100 |
-
def __init__(self, url: str, alt: str):
|
101 |
-
self.url = url
|
102 |
-
self.alt = alt
|
103 |
-
|
104 |
-
def to_data_uri(image: Any) -> str:
|
105 |
-
return "data:image/png;base64,..." # Replace with actual base64 data
|
106 |
-
|
107 |
-
class Blackbox:
|
108 |
-
url = "https://www.blackbox.ai"
|
109 |
-
api_endpoint = "https://www.blackbox.ai/api/chat"
|
110 |
-
working = True
|
111 |
-
supports_stream = True
|
112 |
-
supports_system_message = True
|
113 |
-
supports_message_history = True
|
114 |
-
|
115 |
-
default_model = 'blackboxai'
|
116 |
-
image_models = ['ImageGeneration']
|
117 |
-
models = [
|
118 |
-
default_model,
|
119 |
-
'blackboxai-pro',
|
120 |
-
"llama-3.1-8b",
|
121 |
-
'llama-3.1-70b',
|
122 |
-
'llama-3.1-405b',
|
123 |
-
'gpt-4o',
|
124 |
-
'gemini-pro',
|
125 |
-
'gemini-1.5-flash',
|
126 |
-
'claude-sonnet-3.5',
|
127 |
-
'PythonAgent',
|
128 |
-
'JavaAgent',
|
129 |
-
'JavaScriptAgent',
|
130 |
-
'HTMLAgent',
|
131 |
-
'GoogleCloudAgent',
|
132 |
-
'AndroidDeveloper',
|
133 |
-
'SwiftDeveloper',
|
134 |
-
'Next.jsAgent',
|
135 |
-
'MongoDBAgent',
|
136 |
-
'PyTorchAgent',
|
137 |
-
'ReactAgent',
|
138 |
-
'XcodeAgent',
|
139 |
-
'AngularJSAgent',
|
140 |
-
*image_models,
|
141 |
-
'Niansuh',
|
142 |
-
]
|
143 |
-
|
144 |
-
# Filter models based on AVAILABLE_MODELS
|
145 |
-
if AVAILABLE_MODELS:
|
146 |
-
models = [model for model in models if model in AVAILABLE_MODELS]
|
147 |
-
|
148 |
-
agentMode = {
|
149 |
-
'ImageGeneration': {'mode': True, 'id': "ImageGenerationLV45LJp", 'name': "Image Generation"},
|
150 |
-
'Niansuh': {'mode': True, 'id': "NiansuhAIk1HgESy", 'name': "Niansuh"},
|
151 |
-
}
|
152 |
-
trendingAgentMode = {
|
153 |
-
"blackboxai": {},
|
154 |
-
"gemini-1.5-flash": {'mode': True, 'id': 'Gemini'},
|
155 |
-
"llama-3.1-8b": {'mode': True, 'id': "llama-3.1-8b"},
|
156 |
-
'llama-3.1-70b': {'mode': True, 'id': "llama-3.1-70b"},
|
157 |
-
'llama-3.1-405b': {'mode': True, 'id': "llama-3.1-405b"},
|
158 |
-
'blackboxai-pro': {'mode': True, 'id': "BLACKBOXAI-PRO"},
|
159 |
-
'PythonAgent': {'mode': True, 'id': "Python Agent"},
|
160 |
-
'JavaAgent': {'mode': True, 'id': "Java Agent"},
|
161 |
-
'JavaScriptAgent': {'mode': True, 'id': "JavaScript Agent"},
|
162 |
-
'HTMLAgent': {'mode': True, 'id': "HTML Agent"},
|
163 |
-
'GoogleCloudAgent': {'mode': True, 'id': "Google Cloud Agent"},
|
164 |
-
'AndroidDeveloper': {'mode': True, 'id': "Android Developer"},
|
165 |
-
'SwiftDeveloper': {'mode': True, 'id': "Swift Developer"},
|
166 |
-
'Next.jsAgent': {'mode': True, 'id': "Next.js Agent"},
|
167 |
-
'MongoDBAgent': {'mode': True, 'id': "MongoDB Agent"},
|
168 |
-
'PyTorchAgent': {'mode': True, 'id': "PyTorch Agent"},
|
169 |
-
'ReactAgent': {'mode': True, 'id': "React Agent"},
|
170 |
-
'XcodeAgent': {'mode': True, 'id': "Xcode Agent"},
|
171 |
-
'AngularJSAgent': {'mode': True, 'id': "AngularJS Agent"},
|
172 |
-
}
|
173 |
-
|
174 |
-
userSelectedModel = {
|
175 |
-
"gpt-4o": "gpt-4o",
|
176 |
-
"gemini-pro": "gemini-pro",
|
177 |
-
'claude-sonnet-3.5': "claude-sonnet-3.5",
|
178 |
-
}
|
179 |
-
|
180 |
-
model_prefixes = {
|
181 |
-
'gpt-4o': '@GPT-4o',
|
182 |
-
'gemini-pro': '@Gemini-PRO',
|
183 |
-
'claude-sonnet-3.5': '@Claude-Sonnet-3.5',
|
184 |
-
'PythonAgent': '@Python Agent',
|
185 |
-
'JavaAgent': '@Java Agent',
|
186 |
-
'JavaScriptAgent': '@JavaScript Agent',
|
187 |
-
'HTMLAgent': '@HTML Agent',
|
188 |
-
'GoogleCloudAgent': '@Google Cloud Agent',
|
189 |
-
'AndroidDeveloper': '@Android Developer',
|
190 |
-
'SwiftDeveloper': '@Swift Developer',
|
191 |
-
'Next.jsAgent': '@Next.js Agent',
|
192 |
-
'MongoDBAgent': '@MongoDB Agent',
|
193 |
-
'PyTorchAgent': '@PyTorch Agent',
|
194 |
-
'ReactAgent': '@React Agent',
|
195 |
-
'XcodeAgent': '@Xcode Agent',
|
196 |
-
'AngularJSAgent': '@AngularJS Agent',
|
197 |
-
'blackboxai-pro': '@BLACKBOXAI-PRO',
|
198 |
-
'ImageGeneration': '@Image Generation',
|
199 |
-
'Niansuh': '@Niansuh',
|
200 |
-
}
|
201 |
-
|
202 |
-
model_referers = {
|
203 |
-
"blackboxai": f"{url}/?model=blackboxai",
|
204 |
-
"gpt-4o": f"{url}/?model=gpt-4o",
|
205 |
-
"gemini-pro": f"{url}/?model=gemini-pro",
|
206 |
-
"claude-sonnet-3.5": f"{url}/?model=claude-sonnet-3.5"
|
207 |
-
}
|
208 |
-
|
209 |
-
model_aliases = {
|
210 |
-
"gemini-flash": "gemini-1.5-flash",
|
211 |
-
"claude-3.5-sonnet": "claude-sonnet-3.5",
|
212 |
-
"flux": "ImageGeneration",
|
213 |
-
"niansuh": "Niansuh",
|
214 |
-
}
|
215 |
-
|
216 |
-
@classmethod
|
217 |
-
def get_model(cls, model: str) -> Optional[str]:
|
218 |
-
if model in cls.models:
|
219 |
-
return model
|
220 |
-
elif model in cls.userSelectedModel and cls.userSelectedModel[model] in cls.models:
|
221 |
-
return model
|
222 |
-
elif model in cls.model_aliases and cls.model_aliases[model] in cls.models:
|
223 |
-
return cls.model_aliases[model]
|
224 |
-
else:
|
225 |
-
return cls.default_model if cls.default_model in cls.models else None
|
226 |
-
|
227 |
-
@classmethod
|
228 |
-
async def create_async_generator(
|
229 |
-
cls,
|
230 |
-
model: str,
|
231 |
-
messages: List[Dict[str, str]],
|
232 |
-
proxy: Optional[str] = None,
|
233 |
-
image: Any = None,
|
234 |
-
image_name: Optional[str] = None,
|
235 |
-
webSearchMode: bool = False,
|
236 |
-
**kwargs
|
237 |
-
) -> AsyncGenerator[Any, None]:
|
238 |
-
model = cls.get_model(model)
|
239 |
-
if model is None:
|
240 |
-
logger.error(f"Model {model} is not available.")
|
241 |
-
raise ModelNotWorkingException(model)
|
242 |
-
|
243 |
-
logger.info(f"Selected model: {model}")
|
244 |
-
|
245 |
-
if not cls.working or model not in cls.models:
|
246 |
-
logger.error(f"Model {model} is not working or not supported.")
|
247 |
-
raise ModelNotWorkingException(model)
|
248 |
-
|
249 |
-
headers = {
|
250 |
-
"accept": "*/*",
|
251 |
-
"accept-language": "en-US,en;q=0.9",
|
252 |
-
"cache-control": "no-cache",
|
253 |
-
"content-type": "application/json",
|
254 |
-
"origin": cls.url,
|
255 |
-
"pragma": "no-cache",
|
256 |
-
"priority": "u=1, i",
|
257 |
-
"referer": cls.model_referers.get(model, cls.url),
|
258 |
-
"sec-ch-ua": '"Chromium";v="129", "Not=A?Brand";v="8"',
|
259 |
-
"sec-ch-ua-mobile": "?0",
|
260 |
-
"sec-ch-ua-platform": '"Linux"',
|
261 |
-
"sec-fetch-dest": "empty",
|
262 |
-
"sec-fetch-mode": "cors",
|
263 |
-
"sec-fetch-site": "same-origin",
|
264 |
-
"user-agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/129.0.0.0 Safari/537.36",
|
265 |
-
}
|
266 |
-
|
267 |
-
if model in cls.model_prefixes:
|
268 |
-
prefix = cls.model_prefixes[model]
|
269 |
-
if not messages[0]['content'].startswith(prefix):
|
270 |
-
logger.debug(f"Adding prefix '{prefix}' to the first message.")
|
271 |
-
messages[0]['content'] = f"{prefix} {messages[0]['content']}"
|
272 |
-
|
273 |
-
random_id = ''.join(random.choices(string.ascii_letters + string.digits, k=7))
|
274 |
-
messages[-1]['id'] = random_id
|
275 |
-
messages[-1]['role'] = 'user'
|
276 |
-
|
277 |
-
# Don't log the full message content for privacy
|
278 |
-
logger.debug(f"Generated message ID: {random_id} for model: {model}")
|
279 |
-
|
280 |
-
if image is not None:
|
281 |
-
messages[-1]['data'] = {
|
282 |
-
'fileText': '',
|
283 |
-
'imageBase64': to_data_uri(image),
|
284 |
-
'title': image_name
|
285 |
-
}
|
286 |
-
messages[-1]['content'] = 'FILE:BB\n$#$\n\n$#$\n' + messages[-1]['content']
|
287 |
-
logger.debug("Image data added to the message.")
|
288 |
-
|
289 |
-
data = {
|
290 |
-
"messages": messages,
|
291 |
-
"id": random_id,
|
292 |
-
"previewToken": None,
|
293 |
-
"userId": None,
|
294 |
-
"codeModelMode": True,
|
295 |
-
"agentMode": {},
|
296 |
-
"trendingAgentMode": {},
|
297 |
-
"isMicMode": False,
|
298 |
-
"userSystemPrompt": None,
|
299 |
-
"maxTokens": 99999999,
|
300 |
-
"playgroundTopP": 0.9,
|
301 |
-
"playgroundTemperature": 0.5,
|
302 |
-
"isChromeExt": False,
|
303 |
-
"githubToken": None,
|
304 |
-
"clickedAnswer2": False,
|
305 |
-
"clickedAnswer3": False,
|
306 |
-
"clickedForceWebSearch": False,
|
307 |
-
"visitFromDelta": False,
|
308 |
-
"mobileClient": False,
|
309 |
-
"userSelectedModel": None,
|
310 |
-
"webSearchMode": webSearchMode,
|
311 |
-
"validated": "00f37b34-a166-4efb-bce5-1312d87f2f94"
|
312 |
-
}
|
313 |
-
|
314 |
-
if model in cls.agentMode:
|
315 |
-
data["agentMode"] = cls.agentMode[model]
|
316 |
-
elif model in cls.trendingAgentMode:
|
317 |
-
data["trendingAgentMode"] = cls.trendingAgentMode[model]
|
318 |
-
elif model in cls.userSelectedModel:
|
319 |
-
data["userSelectedModel"] = cls.userSelectedModel[model]
|
320 |
-
logger.info(f"Sending request to {cls.api_endpoint} with data (excluding messages).")
|
321 |
-
|
322 |
-
timeout = ClientTimeout(total=60) # Set an appropriate timeout
|
323 |
-
retry_attempts = 10 # Set the number of retry attempts
|
324 |
-
|
325 |
-
for attempt in range(retry_attempts):
|
326 |
-
try:
|
327 |
-
async with ClientSession(headers=headers, timeout=timeout) as session:
|
328 |
-
async with session.post(cls.api_endpoint, json=data, proxy=proxy) as response:
|
329 |
-
response.raise_for_status()
|
330 |
-
logger.info(f"Received response with status {response.status}")
|
331 |
-
if model == 'ImageGeneration':
|
332 |
-
response_text = await response.text()
|
333 |
-
url_match = re.search(r'https://storage\.googleapis\.com/[^\s\)]+', response_text)
|
334 |
-
if url_match:
|
335 |
-
image_url = url_match.group(0)
|
336 |
-
logger.info(f"Image URL found.")
|
337 |
-
yield ImageResponse(image_url, alt=messages[-1]['content'])
|
338 |
-
else:
|
339 |
-
logger.error("Image URL not found in the response.")
|
340 |
-
raise Exception("Image URL not found in the response")
|
341 |
-
else:
|
342 |
-
full_response = ""
|
343 |
-
search_results_json = ""
|
344 |
-
try:
|
345 |
-
async for chunk, _ in response.content.iter_chunks():
|
346 |
-
if chunk:
|
347 |
-
decoded_chunk = chunk.decode(errors='ignore')
|
348 |
-
decoded_chunk = re.sub(r'\$@\$v=[^$]+\$@\$', '', decoded_chunk)
|
349 |
-
if decoded_chunk.strip():
|
350 |
-
if '$~~~$' in decoded_chunk:
|
351 |
-
search_results_json += decoded_chunk
|
352 |
-
else:
|
353 |
-
full_response += decoded_chunk
|
354 |
-
yield decoded_chunk
|
355 |
-
logger.info("Finished streaming response chunks.")
|
356 |
-
except Exception as e:
|
357 |
-
logger.exception("Error while iterating over response chunks.")
|
358 |
-
raise e
|
359 |
-
if data["webSearchMode"] and search_results_json:
|
360 |
-
match = re.search(r'\$~~~\$(.*?)\$~~~\$', search_results_json, re.DOTALL)
|
361 |
-
if match:
|
362 |
-
try:
|
363 |
-
search_results = json.loads(match.group(1))
|
364 |
-
formatted_results = "\n\n**Sources:**\n"
|
365 |
-
for i, result in enumerate(search_results[:5], 1):
|
366 |
-
formatted_results += f"{i}. [{result['title']}]({result['link']})\n"
|
367 |
-
logger.info("Formatted search results.")
|
368 |
-
yield formatted_results
|
369 |
-
except json.JSONDecodeError as je:
|
370 |
-
logger.error("Failed to parse search results JSON.")
|
371 |
-
raise je
|
372 |
-
break # Exit the retry loop if successful
|
373 |
-
except ClientError as ce:
|
374 |
-
logger.error(f"Client error occurred: {ce}. Retrying attempt {attempt + 1}/{retry_attempts}")
|
375 |
-
if attempt == retry_attempts - 1:
|
376 |
-
raise HTTPException(status_code=502, detail="Error communicating with the external API.")
|
377 |
-
except asyncio.TimeoutError:
|
378 |
-
logger.error(f"Request timed out. Retrying attempt {attempt + 1}/{retry_attempts}")
|
379 |
-
if attempt == retry_attempts - 1:
|
380 |
-
raise HTTPException(status_code=504, detail="External API request timed out.")
|
381 |
-
except Exception as e:
|
382 |
-
logger.error(f"Unexpected error: {e}. Retrying attempt {attempt + 1}/{retry_attempts}")
|
383 |
-
if attempt == retry_attempts - 1:
|
384 |
-
raise HTTPException(status_code=500, detail=str(e))
|
385 |
-
|
386 |
-
# FastAPI app setup
|
387 |
-
app = FastAPI()
|
388 |
-
|
389 |
-
# Add the cleanup task when the app starts
|
390 |
-
@app.on_event("startup")
|
391 |
-
async def startup_event():
|
392 |
-
asyncio.create_task(cleanup_rate_limit_stores())
|
393 |
-
logger.info("Started rate limit store cleanup task.")
|
394 |
-
|
395 |
-
# Middleware to enhance security and enforce Content-Type for specific endpoints
|
396 |
-
@app.middleware("http")
|
397 |
-
async def security_middleware(request: Request, call_next):
|
398 |
-
client_ip = request.client.host
|
399 |
-
# Enforce that POST requests to /v1/chat/completions must have Content-Type: application/json
|
400 |
-
if request.method == "POST" and request.url.path == "/v1/chat/completions":
|
401 |
-
content_type = request.headers.get("Content-Type")
|
402 |
-
if content_type != "application/json":
|
403 |
-
logger.warning(f"Invalid Content-Type from IP: {client_ip} for path: {request.url.path}")
|
404 |
-
return JSONResponse(
|
405 |
-
status_code=400,
|
406 |
-
content={
|
407 |
-
"error": {
|
408 |
-
"message": "Content-Type must be application/json",
|
409 |
-
"type": "invalid_request_error",
|
410 |
-
"param": None,
|
411 |
-
"code": None
|
412 |
-
}
|
413 |
-
},
|
414 |
-
)
|
415 |
-
response = await call_next(request)
|
416 |
-
return response
|
417 |
-
|
418 |
-
# Request Models
|
419 |
-
class Message(BaseModel):
|
420 |
-
role: str
|
421 |
-
content: str
|
422 |
-
|
423 |
-
class ChatRequest(BaseModel):
|
424 |
-
model: str
|
425 |
-
messages: List[Message]
|
426 |
-
temperature: Optional[float] = 1.0
|
427 |
-
top_p: Optional[float] = 1.0
|
428 |
-
n: Optional[int] = 1
|
429 |
-
stream: Optional[bool] = False
|
430 |
-
stop: Optional[Union[str, List[str]]] = None
|
431 |
-
max_tokens: Optional[int] = None
|
432 |
-
presence_penalty: Optional[float] = 0.0
|
433 |
-
frequency_penalty: Optional[float] = 0.0
|
434 |
-
logit_bias: Optional[Dict[str, float]] = None
|
435 |
-
user: Optional[str] = None
|
436 |
-
webSearchMode: Optional[bool] = False # Custom parameter
|
437 |
-
|
438 |
-
def create_response(content: str, model: str, finish_reason: Optional[str] = None) -> Dict[str, Any]:
|
439 |
-
return {
|
440 |
-
"id": f"chatcmpl-{uuid.uuid4()}",
|
441 |
-
"object": "chat.completion.chunk",
|
442 |
-
"created": int(datetime.now().timestamp()),
|
443 |
-
"model": model,
|
444 |
-
"choices": [
|
445 |
-
{
|
446 |
-
"index": 0,
|
447 |
-
"delta": {"content": content, "role": "assistant"},
|
448 |
-
"finish_reason": finish_reason,
|
449 |
-
}
|
450 |
-
],
|
451 |
-
"usage": None,
|
452 |
-
}
|
453 |
-
|
454 |
-
@app.post("/v1/chat/completions", dependencies=[Depends(rate_limiter_per_ip)])
|
455 |
-
async def chat_completions(request: ChatRequest, req: Request, api_key: str = Depends(get_api_key)):
|
456 |
-
client_ip = req.client.host
|
457 |
-
# Redact user messages only for logging purposes
|
458 |
-
redacted_messages = [{"role": msg.role, "content": "[redacted]"} for msg in request.messages]
|
459 |
-
|
460 |
-
logger.info(f"Received chat completions request from API key: {api_key} | IP: {client_ip} | Model: {request.model} | Messages: {redacted_messages}")
|
461 |
-
|
462 |
-
try:
|
463 |
-
# Validate that the requested model is available
|
464 |
-
if request.model not in Blackbox.models and request.model not in Blackbox.model_aliases:
|
465 |
-
logger.warning(f"Attempt to use unavailable model: {request.model} from IP: {client_ip}")
|
466 |
-
raise HTTPException(status_code=400, detail="Requested model is not available.")
|
467 |
-
|
468 |
-
# Process the request with actual message content, but don't log it
|
469 |
-
async_generator = Blackbox.create_async_generator(
|
470 |
-
model=request.model,
|
471 |
-
messages=[{"role": msg.role, "content": msg.content} for msg in request.messages], # Actual message content used here
|
472 |
-
image=None,
|
473 |
-
image_name=None,
|
474 |
-
webSearchMode=request.webSearchMode
|
475 |
-
)
|
476 |
-
|
477 |
-
if request.stream:
|
478 |
-
async def generate():
|
479 |
-
try:
|
480 |
-
async for chunk in async_generator:
|
481 |
-
if isinstance(chunk, ImageResponse):
|
482 |
-
image_markdown = f""
|
483 |
-
response_chunk = create_response(image_markdown, request.model)
|
484 |
-
else:
|
485 |
-
response_chunk = create_response(chunk, request.model)
|
486 |
-
|
487 |
-
yield f"data: {json.dumps(response_chunk)}\n\n"
|
488 |
-
|
489 |
-
yield "data: [DONE]\n\n"
|
490 |
-
except HTTPException as he:
|
491 |
-
error_response = {"error": he.detail}
|
492 |
-
yield f"data: {json.dumps(error_response)}\n\n"
|
493 |
-
except Exception as e:
|
494 |
-
logger.exception(f"Error during streaming response generation from IP: {client_ip}.")
|
495 |
-
error_response = {"error": str(e)}
|
496 |
-
yield f"data: {json.dumps(error_response)}\n\n"
|
497 |
-
|
498 |
-
return StreamingResponse(generate(), media_type="text/event-stream")
|
499 |
-
else:
|
500 |
-
response_content = ""
|
501 |
-
async for chunk in async_generator:
|
502 |
-
if isinstance(chunk, ImageResponse):
|
503 |
-
response_content += f"\n"
|
504 |
-
else:
|
505 |
-
response_content += chunk
|
506 |
-
|
507 |
-
logger.info(f"Completed non-streaming response generation for API key: {api_key} | IP: {client_ip}")
|
508 |
-
return {
|
509 |
-
"id": f"chatcmpl-{uuid.uuid4()}",
|
510 |
-
"object": "chat.completion",
|
511 |
-
"created": int(datetime.now().timestamp()),
|
512 |
-
"model": request.model,
|
513 |
-
"choices": [
|
514 |
-
{
|
515 |
-
"message": {
|
516 |
-
"role": "assistant",
|
517 |
-
"content": response_content
|
518 |
-
},
|
519 |
-
"finish_reason": "stop",
|
520 |
-
"index": 0
|
521 |
-
}
|
522 |
-
],
|
523 |
-
"usage": {
|
524 |
-
"prompt_tokens": sum(len(msg.content.split()) for msg in request.messages),
|
525 |
-
"completion_tokens": len(response_content.split()),
|
526 |
-
"total_tokens": sum(len(msg.content.split()) for msg in request.messages) + len(response_content.split())
|
527 |
-
},
|
528 |
-
}
|
529 |
-
except ModelNotWorkingException as e:
|
530 |
-
logger.warning(f"Model not working: {e} | IP: {client_ip}")
|
531 |
-
raise HTTPException(status_code=503, detail=str(e))
|
532 |
-
except HTTPException as he:
|
533 |
-
logger.warning(f"HTTPException: {he.detail} | IP: {client_ip}")
|
534 |
-
raise he
|
535 |
-
except Exception as e:
|
536 |
-
logger.exception(f"An unexpected error occurred while processing the chat completions request from IP: {client_ip}.")
|
537 |
-
raise HTTPException(status_code=500, detail=str(e))
|
538 |
-
|
539 |
-
# Re-added endpoints without API key authentication
|
540 |
-
|
541 |
-
# Endpoint: POST /v1/tokenizer
|
542 |
-
class TokenizerRequest(BaseModel):
|
543 |
-
text: str
|
544 |
-
|
545 |
-
@app.post("/v1/tokenizer", dependencies=[Depends(rate_limiter_per_ip)])
|
546 |
-
async def tokenizer(request: TokenizerRequest, req: Request):
|
547 |
-
client_ip = req.client.host
|
548 |
-
text = request.text
|
549 |
-
token_count = len(text.split())
|
550 |
-
logger.info(f"Tokenizer requested from IP: {client_ip} | Text length: {len(text)}")
|
551 |
-
return {"text": text, "tokens": token_count}
|
552 |
-
|
553 |
-
# Endpoint: GET /v1/models
|
554 |
-
@app.get("/v1/models", dependencies=[Depends(rate_limiter_per_ip)])
|
555 |
-
async def get_models(req: Request):
|
556 |
-
client_ip = req.client.host
|
557 |
-
logger.info(f"Fetching available models from IP: {client_ip}")
|
558 |
-
return {"data": [{"id": model, "object": "model"} for model in Blackbox.models]}
|
559 |
-
|
560 |
-
# Endpoint: GET /v1/models/{model}/status
|
561 |
-
@app.get("/v1/models/{model}/status", dependencies=[Depends(rate_limiter_per_ip)])
|
562 |
-
async def model_status(model: str, req: Request):
|
563 |
-
client_ip = req.client.host
|
564 |
-
logger.info(f"Model status requested for '{model}' from IP: {client_ip}")
|
565 |
-
if model in Blackbox.models:
|
566 |
-
return {"model": model, "status": "available"}
|
567 |
-
elif model in Blackbox.model_aliases and Blackbox.model_aliases[model] in Blackbox.models:
|
568 |
-
actual_model = Blackbox.model_aliases[model]
|
569 |
-
return {"model": actual_model, "status": "available via alias"}
|
570 |
-
else:
|
571 |
-
logger.warning(f"Model not found: {model} from IP: {client_ip}")
|
572 |
-
raise HTTPException(status_code=404, detail="Model not found")
|
573 |
-
|
574 |
-
# Endpoint: GET /v1/health
|
575 |
-
@app.get("/v1/health", dependencies=[Depends(rate_limiter_per_ip)])
|
576 |
-
async def health_check(req: Request):
|
577 |
-
client_ip = req.client.host
|
578 |
-
logger.info(f"Health check requested from IP: {client_ip}")
|
579 |
-
return {"status": "ok"}
|
580 |
-
|
581 |
-
# Endpoint: GET /v1/chat/completions (GET method)
|
582 |
-
@app.get("/v1/chat/completions")
|
583 |
-
async def chat_completions_get(req: Request):
|
584 |
-
client_ip = req.client.host
|
585 |
-
logger.info(f"GET request made to /v1/chat/completions from IP: {client_ip}, redirecting to 'about:blank'")
|
586 |
-
return RedirectResponse(url='about:blank')
|
587 |
-
|
588 |
-
# Custom exception handler to match OpenAI's error format
|
589 |
-
@app.exception_handler(HTTPException)
|
590 |
-
async def http_exception_handler(request: Request, exc: HTTPException):
|
591 |
-
client_ip = request.client.host
|
592 |
-
logger.error(f"HTTPException: {exc.detail} | Path: {request.url.path} | IP: {client_ip}")
|
593 |
-
return JSONResponse(
|
594 |
-
status_code=exc.status_code,
|
595 |
-
content={
|
596 |
-
"error": {
|
597 |
-
"message": exc.detail,
|
598 |
-
"type": "invalid_request_error",
|
599 |
-
"param": None,
|
600 |
-
"code": None
|
601 |
-
}
|
602 |
-
},
|
603 |
-
)
|
604 |
|
605 |
if __name__ == "__main__":
|
606 |
-
|
607 |
-
uvicorn.run(app, host="0.0.0.0", port=8000)
|
|
|
1 |
+
import uvicorn
|
2 |
+
from api.app import app
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
3 |
|
4 |
if __name__ == "__main__":
|
5 |
+
uvicorn.run(app, host="0.0.0.0", port=8001)
|
|