Niansuh commited on
Commit
370b45c
·
verified ·
1 Parent(s): ca2768f

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +204 -418
main.py CHANGED
@@ -229,10 +229,10 @@ class Blackbox:
229
  cleaned_text = re.sub(pattern, '', text)
230
  try:
231
  response_json = json.loads(cleaned_text)
232
- # Assuming the response is in {"response": "Your answer here."} format
233
- return response_json.get("response", cleaned_text)
234
  except json.JSONDecodeError:
235
- return cleaned_text
236
 
237
  @classmethod
238
  async def generate_response(
@@ -252,6 +252,7 @@ class Blackbox:
252
 
253
  prefix = cls.model_prefixes.get(model, "")
254
 
 
255
  formatted_prompt = ""
256
  for message in messages:
257
  role = message.get('role', '').capitalize()
@@ -297,30 +298,12 @@ class Blackbox:
297
  "role": "user"
298
  }
299
  ],
300
- "id": chat_id,
301
- "previewToken": None,
302
- "userId": None,
303
- "codeModelMode": True,
304
- "agentMode": agent_mode,
305
- "trendingAgentMode": trending_agent_mode,
306
- "isMicMode": False,
307
- "userSystemPrompt": None,
308
- "maxTokens": 1024,
309
- "playgroundTopP": 0.9,
310
- "playgroundTemperature": 0.5,
311
- "isChromeExt": False,
312
- "githubToken": None,
313
- "clickedAnswer2": False,
314
- "clickedAnswer3": False,
315
- "clickedForceWebSearch": False,
316
- "visitFromDelta": False,
317
- "mobileClient": False,
318
- "webSearchMode": False,
319
- "userSelectedModel": cls.userSelectedModel.get(model, model)
320
  }
321
 
322
  async with ClientSession(headers=common_headers) as session:
323
  try:
 
324
  async with session.post(
325
  cls.api_endpoint,
326
  headers=headers_api_chat_combined,
@@ -347,411 +330,214 @@ class Blackbox:
347
  logger.exception(f"Unexpected error during /api/chat request: {str(e)}")
348
  return f"Unexpected error during /api/chat request: {str(e)}"
349
 
350
- @classmethod
351
- async def create_async_generator(
352
- cls,
353
- model: str,
354
- messages: List[Dict[str, str]],
355
- proxy: Optional[str] = None,
356
- websearch: bool = False,
357
- **kwargs
358
- ) -> AsyncGenerator[Union[str, ImageResponseModel], None]:
359
- """
360
- Creates an asynchronous generator for streaming responses from Blackbox AI.
361
-
362
- Parameters:
363
- model (str): Model to use for generating responses.
364
- messages (List[Dict[str, str]]): Message history.
365
- proxy (Optional[str]): Proxy URL, if needed.
366
- websearch (bool): Enables or disables web search mode.
367
- **kwargs: Additional keyword arguments.
 
 
 
 
 
 
 
 
 
 
 
368
 
369
- Yields:
370
- Union[str, ImageResponseModel]: Segments of the generated response or ImageResponseModel objects.
371
  """
372
- model = cls.get_model(model)
373
-
374
- chat_id = cls.generate_random_string()
375
- next_action = cls.generate_next_action()
376
- next_router_state_tree = cls.generate_next_router_state_tree()
377
-
378
- agent_mode = cls.agentMode.get(model, {})
379
- trending_agent_mode = cls.trendingAgentMode.get(model, {})
380
-
381
- prefix = cls.model_prefixes.get(model, "")
382
-
383
- formatted_prompt = ""
384
- for message in messages:
385
- role = message.get('role', '').capitalize()
386
- content = message.get('content', '')
387
- if role and content:
388
- formatted_prompt += f"{role}: {content}\n"
389
-
390
- if prefix:
391
- formatted_prompt = f"{prefix} {formatted_prompt}".strip()
392
-
393
- referer_path = cls.model_referers.get(model, f"/?model={model}")
394
- referer_url = f"{cls.url}{referer_path}"
395
-
396
- common_headers = {
397
- 'accept': '*/*',
398
- 'accept-language': 'en-US,en;q=0.9',
399
- 'cache-control': 'no-cache',
400
- 'origin': cls.url,
401
- 'pragma': 'no-cache',
402
- 'priority': 'u=1, i',
403
- 'sec-ch-ua': '"Chromium";v="129", "Not=A?Brand";v="8"',
404
- 'sec-ch-ua-mobile': '?0',
405
- 'sec-ch-ua-platform': '"Linux"',
406
- 'sec-fetch-dest': 'empty',
407
- 'sec-fetch-mode': 'cors',
408
- 'sec-fetch-site': 'same-origin',
409
- 'user-agent': 'Mozilla/5.0 (X11; Linux x86_64) '
410
- 'AppleWebKit/537.36 (KHTML, like Gecko) '
411
- 'Chrome/129.0.0.0 Safari/537.36'
412
- }
413
-
414
- headers_api_chat = {
415
- 'Content-Type': 'application/json',
416
- 'Referer': referer_url
417
- }
418
- headers_api_chat_combined = {**common_headers, **headers_api_chat}
419
-
420
- payload_api_chat = {
421
- "messages": [
422
- {
423
- "id": chat_id,
424
- "content": formatted_prompt,
425
- "role": "user"
426
- }
427
- ],
428
- "id": chat_id,
429
- "previewToken": None,
430
- "userId": None,
431
- "codeModelMode": True,
432
- "agentMode": agent_mode,
433
- "trendingAgentMode": trending_agent_mode,
434
- "isMicMode": False,
435
- "userSystemPrompt": None,
436
- "maxTokens": 1024,
437
- "playgroundTopP": 0.9,
438
- "playgroundTemperature": 0.5,
439
- "isChromeExt": False,
440
- "githubToken": None,
441
- "clickedAnswer2": False,
442
- "clickedAnswer3": False,
443
- "clickedForceWebSearch": False,
444
- "visitFromDelta": False,
445
- "mobileClient": False,
446
- "webSearchMode": websearch,
447
- "userSelectedModel": cls.userSelectedModel.get(model, model)
448
- }
449
-
450
- headers_chat = {
451
- 'Accept': 'text/x-component',
452
- 'Content-Type': 'text/plain;charset=UTF-8',
453
- 'Referer': f'{cls.url}/chat/{chat_id}?model={model}',
454
- 'next-action': next_action,
455
- 'next-router-state-tree': next_router_state_tree,
456
- 'next-url': '/'
457
- }
458
- headers_chat_combined = {**common_headers, **headers_chat}
459
-
460
- data_chat = '[]'
461
-
462
- async with ClientSession(headers=common_headers) as session:
463
- try:
464
- async with session.post(
465
- cls.api_endpoint,
466
- headers=headers_api_chat_combined,
467
- json=payload_api_chat,
468
- proxy=proxy
469
- ) as response_api_chat:
470
- response_api_chat.raise_for_status()
471
- text = await response_api_chat.text()
472
- logger.debug(f"Raw response from Blackbox API: {text}") # Log raw response
473
- cleaned_response = cls.clean_response(text)
474
- logger.debug(f"Cleaned response: {cleaned_response}") # Log cleaned response
475
-
476
- if model in cls.image_models:
477
- match = re.search(r'!\[.*?\]\((https?://[^\)]+)\)', cleaned_response)
478
- if match:
479
- image_url = match.group(1)
480
- image_response = ImageResponseModel(images=image_url, alt="Generated Image")
481
- yield image_response
482
- else:
483
- yield cleaned_response
484
- else:
485
- if websearch:
486
- match = re.search(r'\$~~~\$(.*?)\$~~~\$', cleaned_response, re.DOTALL)
487
- if match:
488
- source_part = match.group(1).strip()
489
- answer_part = cleaned_response[match.end():].strip()
490
- try:
491
- sources = json.loads(source_part)
492
- source_formatted = "**Source:**\n"
493
- for item in sources:
494
- title = item.get('title', 'No Title')
495
- link = item.get('link', '#')
496
- position = item.get('position', '')
497
- source_formatted += f"{position}. [{title}]({link})\n"
498
- final_response = f"{answer_part}\n\n{source_formatted}"
499
- except json.JSONDecodeError:
500
- final_response = f"{answer_part}\n\nSource information is unavailable."
501
- else:
502
- final_response = cleaned_response
503
- else:
504
- if '$~~~$' in cleaned_response:
505
- final_response = cleaned_response.split('$~~~$')[0].strip()
506
- else:
507
- final_response = cleaned_response
508
-
509
- yield final_response
510
- except ClientResponseError as e:
511
- error_text = f"Error {e.status}: {e.message}"
512
- try:
513
- error_response = await e.response.text()
514
- cleaned_error = cls.clean_response(error_response)
515
- error_text += f" - {cleaned_error}"
516
- logger.error(f"Blackbox API ClientResponseError: {error_text}")
517
- except Exception:
518
- pass
519
- yield error_text
520
- except Exception as e:
521
- yield f"Unexpected error during /api/chat request: {str(e)}"
522
 
523
- chat_url = f'{cls.url}/chat/{chat_id}?model={model}'
 
 
 
 
 
 
 
524
 
525
- try:
526
- async with session.post(
527
- chat_url,
528
- headers=headers_chat_combined,
529
- data=data_chat,
530
- proxy=proxy
531
- ) as response_chat:
532
- response_chat.raise_for_status()
533
- pass
534
- except ClientResponseError as e:
535
- error_text = f"Error {e.status}: {e.message}"
536
- try:
537
- error_response = await e.response.text()
538
- cleaned_error = cls.clean_response(error_response)
539
- error_text += f" - {cleaned_error}"
540
- logger.error(f"Blackbox API ClientResponseError during chat URL request: {error_text}")
541
- except Exception:
542
- pass
543
- yield error_text
544
- except Exception as e:
545
- yield f"Unexpected error during /chat/{chat_id} request: {str(e)}"
 
 
 
 
 
 
 
 
 
 
 
546
 
547
- # ============================
548
- # 7. Initialize FastAPI App
549
- # ============================
550
- app = FastAPI()
 
 
 
 
 
 
 
 
 
551
 
552
- # ====================================
553
- # 8. Define Middleware and Dependencies
554
- # ====================================
555
- @app.middleware("http")
556
- async def security_middleware(request: Request, call_next):
557
- client_ip = request.client.host
558
- # Enforce that POST requests to /v1/chat/completions must have Content-Type: application/json
559
- if request.method == "POST" and request.url.path == "/v1/chat/completions":
560
- content_type = request.headers.get("Content-Type")
561
- if content_type != "application/json":
562
- logger.warning(f"Invalid Content-Type from IP: {client_ip} for path: {request.url.path}")
563
- return JSONResponse(
564
- status_code=400,
565
- content={
566
- "error": {
567
- "message": "Content-Type must be application/json",
568
- "type": "invalid_request_error",
569
- "param": None,
570
- "code": None
571
  }
 
 
 
 
 
572
  },
573
- )
574
- response = await call_next(request)
575
- return response
576
-
577
- async def cleanup_rate_limit_stores():
578
- """
579
- Periodically cleans up stale entries in the rate_limit_store to prevent memory bloat.
580
- """
581
- while True:
582
- current_time = time.time()
583
- ips_to_delete = [ip for ip, value in rate_limit_store.items() if current_time - value["timestamp"] > RATE_LIMIT_WINDOW * 2]
584
- for ip in ips_to_delete:
585
- del rate_limit_store[ip]
586
- logger.debug(f"Cleaned up rate_limit_store for IP: {ip}")
587
- await asyncio.sleep(CLEANUP_INTERVAL)
588
-
589
- async def rate_limiter_per_ip(request: Request):
590
- """
591
- Rate limiter that enforces a limit based on the client's IP address.
592
- """
593
- client_ip = request.client.host
594
- current_time = time.time()
595
-
596
- # Initialize or update the count and timestamp
597
- if current_time - rate_limit_store[client_ip]["timestamp"] > RATE_LIMIT_WINDOW:
598
- rate_limit_store[client_ip] = {"count": 1, "timestamp": current_time}
599
- else:
600
- if rate_limit_store[client_ip]["count"] >= RATE_LIMIT:
601
- logger.warning(f"Rate limit exceeded for IP address: {client_ip}")
602
- raise HTTPException(status_code=429, detail='Rate limit exceeded for IP address | NiansuhAI')
603
- rate_limit_store[client_ip]["count"] += 1
604
-
605
- async def get_api_key(request: Request, authorization: str = Header(None)) -> str:
606
- """
607
- Dependency to extract and validate the API key from the Authorization header.
608
- """
609
- client_ip = request.client.host
610
- if authorization is None or not authorization.startswith('Bearer '):
611
- logger.warning(f"Invalid or missing authorization header from IP: {client_ip}")
612
- raise HTTPException(status_code=401, detail='Invalid authorization header format')
613
- api_key = authorization[7:]
614
- if api_key not in API_KEYS:
615
- logger.warning(f"Invalid API key attempted: {api_key} from IP: {client_ip}")
616
- raise HTTPException(status_code=401, detail='Invalid API key')
617
- return api_key
618
-
619
- # =====================================
620
- # 9. Define FastAPI Event Handlers
621
- # =====================================
622
- @app.on_event("startup")
623
- async def startup_event():
624
- asyncio.create_task(cleanup_rate_limit_stores())
625
- logger.info("Started rate limit store cleanup task.")
626
-
627
- # ==========================================
628
- # 10. Define FastAPI Endpoints
629
- # ==========================================
630
- @app.post("/v1/chat/completions", dependencies=[Depends(rate_limiter_per_ip)])
631
- async def chat_completions(request: ChatRequest, req: Request, api_key: str = Depends(get_api_key)):
632
- client_ip = req.client.host
633
- # Redact user messages only for logging purposes
634
- redacted_messages = [{"role": msg.role, "content": "[redacted]"} for msg in request.messages]
635
-
636
- logger.info(f"Received chat completions request from API key: {api_key} | IP: {client_ip} | Model: {request.model} | Messages: {redacted_messages}")
637
-
638
- try:
639
- # Validate that the requested model is available
640
- if request.model not in Blackbox.models and request.model not in Blackbox.model_aliases:
641
- logger.warning(f"Attempt to use unavailable model: {request.model} from IP: {client_ip}")
642
- raise HTTPException(status_code=400, detail="Requested model is not available.")
643
-
644
- # Process the request with actual message content, but don't log it
645
- response_content = await Blackbox.generate_response(
646
- model=request.model,
647
- messages=[{"role": msg.role, "content": msg.content} for msg in request.messages],
648
- temperature=request.temperature,
649
- max_tokens=request.max_tokens
650
- )
651
-
652
- logger.info(f"Completed response generation for API key: {api_key} | IP: {client_ip}")
653
- return {
654
- "id": f"chatcmpl-{uuid.uuid4()}",
655
- "object": "chat.completion",
656
- "created": int(datetime.now().timestamp()),
657
- "model": request.model,
658
- "choices": [
659
- {
660
- "index": 0,
661
- "message": {
662
- "role": "assistant",
663
- "content": response_content
664
- },
665
- "finish_reason": "stop"
666
  }
667
- ],
668
- "usage": {
669
- "prompt_tokens": sum(len(msg.content.split()) for msg in request.messages),
670
- "completion_tokens": len(response_content.split()),
671
- "total_tokens": sum(len(msg.content.split()) for msg in request.messages) + len(response_content.split())
672
  },
673
- }
674
- except ModelNotWorkingException as e:
675
- logger.warning(f"Model not working: {e} | IP: {client_ip}")
676
- raise HTTPException(status_code=503, detail=str(e))
677
- except HTTPException as he:
678
- logger.warning(f"HTTPException: {he.detail} | IP: {client_ip}")
679
- raise he
680
- except Exception as e:
681
- logger.exception(f"An unexpected error occurred while processing the chat completions request from IP: {client_ip}.")
682
- raise HTTPException(status_code=500, detail=str(e))
683
-
684
- @app.get("/v1/models", dependencies=[Depends(rate_limiter_per_ip)])
685
- async def get_models(req: Request):
686
- client_ip = req.client.host
687
- logger.info(f"Fetching available models from IP: {client_ip}")
688
- return {"data": [{"id": model, "object": "model"} for model in Blackbox.models]}
689
-
690
- @app.get("/v1/health", dependencies=[Depends(rate_limiter_per_ip)])
691
- async def health_check(req: Request):
692
- client_ip = req.client.host
693
- logger.info(f"Health check requested from IP: {client_ip}")
694
- return {"status": "ok"}
695
 
696
- # ===============================
697
- # 11. Define Custom Exception Handler
698
- # ===============================
699
- @app.exception_handler(HTTPException)
700
- async def http_exception_handler(request: Request, exc: HTTPException):
701
- client_ip = request.client.host
702
- logger.error(f"HTTPException: {exc.detail} | Path: {request.url.path} | IP: {client_ip}")
703
- return JSONResponse(
704
- status_code=exc.status_code,
705
- content={
706
- "error": {
707
- "message": exc.detail,
708
- "type": "invalid_request_error",
709
- "param": None,
710
- "code": None
711
- }
712
- },
713
- )
714
 
715
- # ============================
716
- # 12. Optional: Streaming Endpoint
717
- # ============================
718
- @app.post("/v1/chat/completions/stream", dependencies=[Depends(rate_limiter_per_ip)])
719
- async def chat_completions_stream(request: ChatRequest, req: Request, api_key: str = Depends(get_api_key)):
720
- client_ip = req.client.host
721
- # Redact user messages only for logging purposes
722
- redacted_messages = [{"role": msg.role, "content": "[redacted]"} for msg in request.messages]
723
-
724
- logger.info(f"Received streaming chat completions request from API key: {api_key} | IP: {client_ip} | Model: {request.model} | Messages: {redacted_messages}")
725
-
726
- try:
727
- # Validate that the requested model is available
728
- if request.model not in Blackbox.models and request.model not in Blackbox.model_aliases:
729
- logger.warning(f"Attempt to use unavailable model: {request.model} from IP: {client_ip}")
730
- raise HTTPException(status_code=400, detail="Requested model is not available.")
731
-
732
- # Create an asynchronous generator for the response
733
- async_generator = Blackbox.create_async_generator(
734
- model=request.model,
735
- messages=[{"role": msg.role, "content": msg.content} for msg in request.messages],
736
- temperature=request.temperature,
737
- max_tokens=request.max_tokens
738
- )
739
 
740
- logger.info(f"Started streaming response for API key: {api_key} | IP: {client_ip}")
741
- return StreamingResponse(async_generator, media_type="text/event-stream")
742
- except ModelNotWorkingException as e:
743
- logger.warning(f"Model not working: {e} | IP: {client_ip}")
744
- raise HTTPException(status_code=503, detail=str(e))
745
- except HTTPException as he:
746
- logger.warning(f"HTTPException: {he.detail} | IP: {client_ip}")
747
- raise he
748
- except Exception as e:
749
- logger.exception(f"An unexpected error occurred while processing the streaming chat completions request from IP: {client_ip}.")
750
- raise HTTPException(status_code=500, detail=str(e))
751
-
752
- # ========================================
753
- # 13. Run the Application with Uvicorn
754
- # ========================================
755
- if __name__ == "__main__":
756
- import uvicorn
757
- uvicorn.run(app, host="0.0.0.0", port=8000)
 
229
  cleaned_text = re.sub(pattern, '', text)
230
  try:
231
  response_json = json.loads(cleaned_text)
232
+ # Adjust based on actual response structure
233
+ return response_json.get("response", response_json.get("data", cleaned_text))
234
  except json.JSONDecodeError:
235
+ return cleaned_text.strip()
236
 
237
  @classmethod
238
  async def generate_response(
 
252
 
253
  prefix = cls.model_prefixes.get(model, "")
254
 
255
+ # Construct the prompt
256
  formatted_prompt = ""
257
  for message in messages:
258
  role = message.get('role', '').capitalize()
 
298
  "role": "user"
299
  }
300
  ],
301
+ "model": model # Simplified payload for testing
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
302
  }
303
 
304
  async with ClientSession(headers=common_headers) as session:
305
  try:
306
+ logger.debug(f"Payload sent to Blackbox API: {json.dumps(payload_api_chat)}")
307
  async with session.post(
308
  cls.api_endpoint,
309
  headers=headers_api_chat_combined,
 
330
  logger.exception(f"Unexpected error during /api/chat request: {str(e)}")
331
  return f"Unexpected error during /api/chat request: {str(e)}"
332
 
333
+ # ===============================
334
+ # 7. Initialize FastAPI App
335
+ # ===============================
336
+ app = FastAPI()
337
+
338
+ # ====================================
339
+ # 8. Define Middleware and Dependencies
340
+ # ====================================
341
+ @app.middleware("http")
342
+ async def security_middleware(request: Request, call_next):
343
+ client_ip = request.client.host
344
+ # Enforce that POST requests to /v1/chat/completions must have Content-Type: application/json
345
+ if request.method == "POST" and request.url.path == "/v1/chat/completions":
346
+ content_type = request.headers.get("Content-Type")
347
+ if content_type != "application/json":
348
+ logger.warning(f"Invalid Content-Type from IP: {client_ip} for path: {request.url.path}")
349
+ return JSONResponse(
350
+ status_code=400,
351
+ content={
352
+ "error": {
353
+ "message": "Content-Type must be application/json",
354
+ "type": "invalid_request_error",
355
+ "param": None,
356
+ "code": None
357
+ }
358
+ },
359
+ )
360
+ response = await call_next(request)
361
+ return response
362
 
363
+ async def cleanup_rate_limit_stores():
 
364
  """
365
+ Periodically cleans up stale entries in the rate_limit_store to prevent memory bloat.
366
+ """
367
+ while True:
368
+ current_time = time.time()
369
+ ips_to_delete = [ip for ip, value in rate_limit_store.items() if current_time - value["timestamp"] > RATE_LIMIT_WINDOW * 2]
370
+ for ip in ips_to_delete:
371
+ del rate_limit_store[ip]
372
+ logger.debug(f"Cleaned up rate_limit_store for IP: {ip}")
373
+ await asyncio.sleep(CLEANUP_INTERVAL)
374
+
375
+ async def rate_limiter_per_ip(request: Request):
376
+ """
377
+ Rate limiter that enforces a limit based on the client's IP address.
378
+ """
379
+ client_ip = request.client.host
380
+ current_time = time.time()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
381
 
382
+ # Initialize or update the count and timestamp
383
+ if current_time - rate_limit_store[client_ip]["timestamp"] > RATE_LIMIT_WINDOW:
384
+ rate_limit_store[client_ip] = {"count": 1, "timestamp": current_time}
385
+ else:
386
+ if rate_limit_store[client_ip]["count"] >= RATE_LIMIT:
387
+ logger.warning(f"Rate limit exceeded for IP address: {client_ip}")
388
+ raise HTTPException(status_code=429, detail='Rate limit exceeded for IP address | NiansuhAI')
389
+ rate_limit_store[client_ip]["count"] += 1
390
 
391
+ async def get_api_key(request: Request, authorization: str = Header(None)) -> str:
392
+ """
393
+ Dependency to extract and validate the API key from the Authorization header.
394
+ """
395
+ client_ip = request.client.host
396
+ if authorization is None or not authorization.startswith('Bearer '):
397
+ logger.warning(f"Invalid or missing authorization header from IP: {client_ip}")
398
+ raise HTTPException(status_code=401, detail='Invalid authorization header format')
399
+ api_key = authorization[7:]
400
+ if api_key not in API_KEYS:
401
+ logger.warning(f"Invalid API key attempted: {api_key} from IP: {client_ip}")
402
+ raise HTTPException(status_code=401, detail='Invalid API key')
403
+ return api_key
404
+
405
+ # =====================================
406
+ # 9. Define FastAPI Event Handlers
407
+ # =====================================
408
+ @app.on_event("startup")
409
+ async def startup_event():
410
+ asyncio.create_task(cleanup_rate_limit_stores())
411
+ logger.info("Started rate limit store cleanup task.")
412
+
413
+ # ==========================================
414
+ # 10. Define FastAPI Endpoints
415
+ # ==========================================
416
+ @app.post("/v1/chat/completions", dependencies=[Depends(rate_limiter_per_ip)])
417
+ async def chat_completions(request: ChatRequest, req: Request, api_key: str = Depends(get_api_key)):
418
+ client_ip = req.client.host
419
+ # Redact user messages only for logging purposes
420
+ redacted_messages = [{"role": msg.role, "content": "[redacted]"} for msg in request.messages]
421
+
422
+ logger.info(f"Received chat completions request from API key: {api_key} | IP: {client_ip} | Model: {request.model} | Messages: {redacted_messages}")
423
 
424
+ try:
425
+ # Validate that the requested model is available
426
+ if request.model not in Blackbox.models and request.model not in Blackbox.model_aliases:
427
+ logger.warning(f"Attempt to use unavailable model: {request.model} from IP: {client_ip}")
428
+ raise HTTPException(status_code=400, detail="Requested model is not available.")
429
+
430
+ # Process the request with actual message content, but don't log it
431
+ response_content = await Blackbox.generate_response(
432
+ model=request.model,
433
+ messages=[{"role": msg.role, "content": msg.content} for msg in request.messages],
434
+ temperature=request.temperature,
435
+ max_tokens=request.max_tokens
436
+ )
437
 
438
+ logger.info(f"Completed response generation for API key: {api_key} | IP: {client_ip}")
439
+ return {
440
+ "id": f"chatcmpl-{uuid.uuid4()}",
441
+ "object": "chat.completion",
442
+ "created": int(datetime.now().timestamp()),
443
+ "model": request.model,
444
+ "choices": [
445
+ {
446
+ "index": 0,
447
+ "message": {
448
+ "role": "assistant",
449
+ "content": response_content
450
+ },
451
+ "finish_reason": "stop"
 
 
 
 
 
452
  }
453
+ ],
454
+ "usage": {
455
+ "prompt_tokens": sum(len(msg.content.split()) for msg in request.messages),
456
+ "completion_tokens": len(response_content.split()),
457
+ "total_tokens": sum(len(msg.content.split()) for msg in request.messages) + len(response_content.split())
458
  },
459
+ }
460
+ except ModelNotWorkingException as e:
461
+ logger.warning(f"Model not working: {e} | IP: {client_ip}")
462
+ raise HTTPException(status_code=503, detail=str(e))
463
+ except HTTPException as he:
464
+ logger.warning(f"HTTPException: {he.detail} | IP: {client_ip}")
465
+ raise he
466
+ except Exception as e:
467
+ logger.exception(f"An unexpected error occurred while processing the chat completions request from IP: {client_ip}.")
468
+ raise HTTPException(status_code=500, detail=str(e))
469
+
470
+ @app.get("/v1/models", dependencies=[Depends(rate_limiter_per_ip)])
471
+ async def get_models(req: Request):
472
+ client_ip = req.client.host
473
+ logger.info(f"Fetching available models from IP: {client_ip}")
474
+ return {"data": [{"id": model, "object": "model"} for model in Blackbox.models]}
475
+
476
+ @app.get("/v1/health", dependencies=[Depends(rate_limiter_per_ip)])
477
+ async def health_check(req: Request):
478
+ client_ip = req.client.host
479
+ logger.info(f"Health check requested from IP: {client_ip}")
480
+ return {"status": "ok"}
481
+
482
+ # ========================================
483
+ # 11. Define Custom Exception Handler
484
+ # ========================================
485
+ @app.exception_handler(HTTPException)
486
+ async def http_exception_handler(request: Request, exc: HTTPException):
487
+ client_ip = request.client.host
488
+ logger.error(f"HTTPException: {exc.detail} | Path: {request.url.path} | IP: {client_ip}")
489
+ return JSONResponse(
490
+ status_code=exc.status_code,
491
+ content={
492
+ "error": {
493
+ "message": exc.detail,
494
+ "type": "invalid_request_error",
495
+ "param": None,
496
+ "code": None
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
497
  }
 
 
 
 
 
498
  },
499
+ )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
500
 
501
+ # ============================
502
+ # 12. Optional: Streaming Endpoint
503
+ # ============================
504
+ @app.post("/v1/chat/completions/stream", dependencies=[Depends(rate_limiter_per_ip)])
505
+ async def chat_completions_stream(request: ChatRequest, req: Request, api_key: str = Depends(get_api_key)):
506
+ client_ip = req.client.host
507
+ # Redact user messages only for logging purposes
508
+ redacted_messages = [{"role": msg.role, "content": "[redacted]"} for msg in request.messages]
 
 
 
 
 
 
 
 
 
 
509
 
510
+ logger.info(f"Received streaming chat completions request from API key: {api_key} | IP: {client_ip} | Model: {request.model} | Messages: {redacted_messages}")
511
+
512
+ try:
513
+ # Validate that the requested model is available
514
+ if request.model not in Blackbox.models and request.model not in Blackbox.model_aliases:
515
+ logger.warning(f"Attempt to use unavailable model: {request.model} from IP: {client_ip}")
516
+ raise HTTPException(status_code=400, detail="Requested model is not available.")
517
+
518
+ # Create an asynchronous generator for the response
519
+ async_generator = Blackbox.create_async_generator(
520
+ model=request.model,
521
+ messages=[{"role": msg.role, "content": msg.content} for msg in request.messages],
522
+ temperature=request.temperature,
523
+ max_tokens=request.max_tokens
524
+ )
 
 
 
 
 
 
 
 
 
525
 
526
+ logger.info(f"Started streaming response for API key: {api_key} | IP: {client_ip}")
527
+ return StreamingResponse(async_generator, media_type="text/event-stream")
528
+ except ModelNotWorkingException as e:
529
+ logger.warning(f"Model not working: {e} | IP: {client_ip}")
530
+ raise HTTPException(status_code=503, detail=str(e))
531
+ except HTTPException as he:
532
+ logger.warning(f"HTTPException: {he.detail} | IP: {client_ip}")
533
+ raise he
534
+ except Exception as e:
535
+ logger.exception(f"An unexpected error occurred while processing the streaming chat completions request from IP: {client_ip}.")
536
+ raise HTTPException(status_code=500, detail=str(e))
537
+
538
+ # ========================================
539
+ # 13. Run the Application with Uvicorn
540
+ # ========================================
541
+ if __name__ == "__main__":
542
+ import uvicorn
543
+ uvicorn.run(app, host="0.0.0.0", port=8000)