Niansuh commited on
Commit
8b177d4
·
verified ·
1 Parent(s): 0253f2a

Update main.py

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