Update main.py
Browse files
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 |
-
#
|
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 |
-
"
|
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 |
-
|
351 |
-
|
352 |
-
|
353 |
-
|
354 |
-
|
355 |
-
|
356 |
-
|
357 |
-
|
358 |
-
)
|
359 |
-
|
360 |
-
|
361 |
-
|
362 |
-
|
363 |
-
|
364 |
-
|
365 |
-
|
366 |
-
|
367 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
368 |
|
369 |
-
|
370 |
-
Union[str, ImageResponseModel]: Segments of the generated response or ImageResponseModel objects.
|
371 |
"""
|
372 |
-
|
373 |
-
|
374 |
-
|
375 |
-
|
376 |
-
|
377 |
-
|
378 |
-
|
379 |
-
|
380 |
-
|
381 |
-
|
382 |
-
|
383 |
-
|
384 |
-
|
385 |
-
|
386 |
-
|
387 |
-
|
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 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
524 |
|
525 |
-
|
526 |
-
|
527 |
-
|
528 |
-
|
529 |
-
|
530 |
-
|
531 |
-
|
532 |
-
|
533 |
-
|
534 |
-
|
535 |
-
|
536 |
-
|
537 |
-
|
538 |
-
|
539 |
-
|
540 |
-
|
541 |
-
|
542 |
-
|
543 |
-
|
544 |
-
|
545 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
546 |
|
547 |
-
|
548 |
-
#
|
549 |
-
|
550 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
551 |
|
552 |
-
|
553 |
-
|
554 |
-
|
555 |
-
|
556 |
-
|
557 |
-
|
558 |
-
|
559 |
-
|
560 |
-
|
561 |
-
|
562 |
-
|
563 |
-
|
564 |
-
|
565 |
-
|
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 |
-
|
575 |
-
|
576 |
-
|
577 |
-
|
578 |
-
|
579 |
-
|
580 |
-
|
581 |
-
|
582 |
-
|
583 |
-
|
584 |
-
|
585 |
-
|
586 |
-
|
587 |
-
|
588 |
-
|
589 |
-
|
590 |
-
""
|
591 |
-
|
592 |
-
|
593 |
-
|
594 |
-
|
595 |
-
|
596 |
-
#
|
597 |
-
|
598 |
-
|
599 |
-
|
600 |
-
|
601 |
-
|
602 |
-
|
603 |
-
|
604 |
-
|
605 |
-
|
606 |
-
|
607 |
-
|
608 |
-
|
609 |
-
|
610 |
-
|
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 |
-
#
|
698 |
-
#
|
699 |
-
@app.
|
700 |
-
async def
|
701 |
-
|
702 |
-
|
703 |
-
|
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 |
-
|
717 |
-
|
718 |
-
|
719 |
-
|
720 |
-
|
721 |
-
|
722 |
-
|
723 |
-
|
724 |
-
|
725 |
-
|
726 |
-
|
727 |
-
|
728 |
-
|
729 |
-
|
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 |
-
|
741 |
-
|
742 |
-
|
743 |
-
|
744 |
-
|
745 |
-
|
746 |
-
|
747 |
-
|
748 |
-
|
749 |
-
|
750 |
-
|
751 |
-
|
752 |
-
# ========================================
|
753 |
-
# 13. Run the Application with Uvicorn
|
754 |
-
# ========================================
|
755 |
-
if __name__ == "__main__":
|
756 |
-
|
757 |
-
|
|
|
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)
|