Update main.py

main.py

@@ -1,22 +1,45 @@
-from __future__ import annotations
-
-import asyncio
-import aiohttp
+import os
+import re
 import random
 import string
-import json
 import uuid
-import re
-from typing import Optional, AsyncGenerator, Union
+import json
+import logging
+import asyncio
+import time
+from collections import defaultdict
+from typing import List, Dict, Any, Optional, Union, AsyncGenerator
+from datetime import datetime
 
 from aiohttp import ClientSession, ClientResponseError
-
-from ..typing import AsyncResult, Messages
-from .base_provider import AsyncGeneratorProvider, ProviderModelMixin
-from ..image import ImageResponse
-
-
-class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
+from fastapi import FastAPI, HTTPException, Request, Depends, Header
+from fastapi.responses import JSONResponse
+from pydantic import BaseModel
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
+    handlers=[logging.StreamHandler()]
+)
+logger = logging.getLogger(__name__)
+
+# Load environment variables
+API_KEYS = os.getenv('API_KEYS', '').split(',')  # Comma-separated API keys
+RATE_LIMIT = int(os.getenv('RATE_LIMIT', '60'))  # Requests per minute
+
+if not API_KEYS or API_KEYS == ['']:
+    logger.error("No API keys found. Please set the API_KEYS environment variable.")
+    raise Exception("API_KEYS environment variable not set.")
+
+# Simple in-memory rate limiter based solely on IP addresses
+rate_limit_store = defaultdict(lambda: {"count": 0, "timestamp": time.time()})
+
+# Define cleanup interval and window
+CLEANUP_INTERVAL = 60  # seconds
+RATE_LIMIT_WINDOW = 60  # seconds
+
+class Blackbox:
     label = "Blackbox AI"
     url = "https://www.blackbox.ai"
     api_endpoint = "https://www.blackbox.ai/api/chat"
@@ -166,26 +189,26 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
         return cleaned_text
 
     @classmethod
-    async def create_async_generator(
+    async def generate_response(
         cls,
         model: str,
-        messages: Messages,
+        messages: List[Dict[str, str]],
         proxy: Optional[str] = None,
         websearch: bool = False,
         **kwargs
-    ) -> AsyncGenerator[Union[str, ImageResponse], None]:
+    ) -> Dict[str, Any]:
         """
-        Creates an asynchronous generator for streaming responses from Blackbox AI.
+        Generates a response from Blackbox AI for the /v1/chat/completions endpoint.
 
         Parameters:
             model (str): Model to use for generating responses.
-            messages (Messages): Message history.
+            messages (List[Dict[str, str]]): Message history.
             proxy (Optional[str]): Proxy URL, if needed.
             websearch (bool): Enables or disables web search mode.
             **kwargs: Additional keyword arguments.
 
-        Yields:
-            Union[str, ImageResponse]: Segments of the generated response or ImageResponse objects.
+        Returns:
+            Dict[str, Any]: The response dictionary in the format required by /v1/chat/completions.
         """
         model = cls.get_model(model)
 
@@ -265,18 +288,6 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
             "userSelectedModel": cls.userSelectedModel.get(model, model)
         }
 
-        headers_chat = {
-            'Accept': 'text/x-component',
-            'Content-Type': 'text/plain;charset=UTF-8',
-            'Referer': f'{cls.url}/chat/{chat_id}?model={model}',
-            'next-action': next_action,
-            'next-router-state-tree': next_router_state_tree,
-            'next-url': '/'
-        }
-        headers_chat_combined = {**common_headers, **headers_chat}
-
-        data_chat = '[]'
-
         async with ClientSession(headers=common_headers) as session:
             try:
                 async with session.post(
@@ -289,63 +300,29 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
                     text = await response_api_chat.text()
                     cleaned_response = cls.clean_response(text)
 
-                    if model in cls.image_models:
-                        match = re.search(r'!\[.*?\]\((https?://[^\)]+)\)', cleaned_response)
-                        if match:
-                            image_url = match.group(1)
-                            image_response = ImageResponse(images=image_url, alt="Generated Image")
-                            yield image_response
-                        else:
-                            yield cleaned_response
-                    else:
-                        if websearch:
-                            match = re.search(r'\$~~~\$(.*?)\$~~~\$', cleaned_response, re.DOTALL)
-                            if match:
-                                source_part = match.group(1).strip()
-                                answer_part = cleaned_response[match.end():].strip()
-                                try:
-                                    sources = json.loads(source_part)
-                                    source_formatted = "**Source:**\n"
-                                    for item in sources:
-                                        title = item.get('title', 'No Title')
-                                        link = item.get('link', '#')
-                                        position = item.get('position', '')
-                                        source_formatted += f"{position}. [{title}]({link})\n"
-                                    final_response = f"{answer_part}\n\n{source_formatted}"
-                                except json.JSONDecodeError:
-                                    final_response = f"{answer_part}\n\nSource information is unavailable."
-                            else:
-                                final_response = cleaned_response
-                        else:
-                            if '$~~~$' in cleaned_response:
-                                final_response = cleaned_response.split('$~~~$')[0].strip()
-                            else:
-                                final_response = cleaned_response
-
-                        yield final_response
-            except ClientResponseError as e:
-                error_text = f"Error {e.status}: {e.message}"
-                try:
-                    error_response = await e.response.text()
-                    cleaned_error = cls.clean_response(error_response)
-                    error_text += f" - {cleaned_error}"
-                except Exception:
-                    pass
-                yield error_text
-            except Exception as e:
-                yield f"Unexpected error during /api/chat request: {str(e)}"
-
-            chat_url = f'{cls.url}/chat/{chat_id}?model={model}'
+                    response_data = {
+                        "id": f"chatcmpl-{uuid.uuid4()}",
+                        "object": "chat.completion",
+                        "created": int(datetime.now().timestamp()),
+                        "model": model,
+                        "choices": [
+                            {
+                                "index": 0,
+                                "message": {
+                                    "role": "assistant",
+                                    "content": cleaned_response
+                                },
+                                "finish_reason": "stop"
+                            }
+                        ],
+                        "usage": {
+                            "prompt_tokens": sum(len(msg['content'].split()) for msg in messages),
+                            "completion_tokens": len(cleaned_response.split()),
+                            "total_tokens": sum(len(msg['content'].split()) for msg in messages) + len(cleaned_response.split())
+                        }
+                    }
 
-            try:
-                async with session.post(
-                    chat_url,
-                    headers=headers_chat_combined,
-                    data=data_chat,
-                    proxy=proxy
-                ) as response_chat:
-                    response_chat.raise_for_status()
-                    pass
+                    return response_data
             except ClientResponseError as e:
                 error_text = f"Error {e.status}: {e.message}"
                 try:
@@ -354,6 +331,117 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
                     error_text += f" - {cleaned_error}"
                 except Exception:
                     pass
-                yield error_text
+                raise HTTPException(status_code=e.status, detail=error_text)
             except Exception as e:
-                yield f"Unexpected error during /chat/{chat_id} request: {str(e)}"
+                raise HTTPException(status_code=500, detail=f"Unexpected error during /api/chat request: {str(e)}")
+
+# FastAPI app setup
+app = FastAPI()
+
+# Add the cleanup task when the app starts
+@app.on_event("startup")
+async def startup_event():
+    asyncio.create_task(cleanup_rate_limit_stores())
+    logger.info("Started rate limit store cleanup task.")
+
+# Middleware to enhance security and enforce Content-Type for specific endpoints
+@app.middleware("http")
+async def security_middleware(request: Request, call_next):
+    client_ip = request.client.host
+    # Enforce that POST requests to /v1/chat/completions must have Content-Type: application/json
+    if request.method == "POST" and request.url.path == "/v1/chat/completions":
+        content_type = request.headers.get("Content-Type")
+        if content_type != "application/json":
+            logger.warning(f"Invalid Content-Type from IP: {client_ip} for path: {request.url.path}")
+            return JSONResponse(
+                status_code=400,
+                content={
+                    "error": {
+                        "message": "Content-Type must be application/json",
+                        "type": "invalid_request_error",
+                        "param": None,
+                        "code": None
+                    }
+                },
+            )
+    response = await call_next(request)
+    return response
+
+# Request Models
+class Message(BaseModel):
+    role: str
+    content: str
+
+class ChatRequest(BaseModel):
+    model: str
+    messages: List[Message]
+    temperature: Optional[float] = 1.0
+    top_p: Optional[float] = 1.0
+    n: Optional[int] = 1
+    max_tokens: Optional[int] = None
+    presence_penalty: Optional[float] = 0.0
+    frequency_penalty: Optional[float] = 0.0
+    logit_bias: Optional[Dict[str, float]] = None
+    user: Optional[str] = None
+
+@app.post("/v1/chat/completions", dependencies=[Depends(rate_limiter_per_ip)])
+async def chat_completions(request: ChatRequest, req: Request, api_key: str = Depends(get_api_key)):
+    client_ip = req.client.host
+    # Redact user messages only for logging purposes
+    redacted_messages = [{"role": msg.role, "content": "[redacted]"} for msg in request.messages]
+
+    logger.info(f"Received chat completions request from API key: {api_key} | IP: {client_ip} | Model: {request.model} | Messages: {redacted_messages}")
+
+    try:
+        # Validate that the requested model is available
+        if request.model not in Blackbox.models and request.model not in Blackbox.model_aliases:
+            logger.warning(f"Attempt to use unavailable model: {request.model} from IP: {client_ip}")
+            raise HTTPException(status_code=400, detail="Requested model is not available.")
+
+        # Process the request with actual message content, but don't log it
+        response_content = await Blackbox.generate_response(
+            model=request.model,
+            messages=[{"role": msg.role, "content": msg.content} for msg in request.messages],
+            temperature=request.temperature,
+            max_tokens=request.max_tokens
+        )
+
+        logger.info(f"Completed response generation for API key: {api_key} | IP: {client_ip}")
+        return response_content
+    except HTTPException as he:
+        logger.warning(f"HTTPException: {he.detail} | IP: {client_ip}")
+        raise he
+    except Exception as e:
+        logger.exception(f"An unexpected error occurred while processing the chat completions request from IP: {client_ip}.")
+        raise HTTPException(status_code=500, detail=str(e))
+
+# Endpoint: GET /v1/models
+@app.get("/v1/models", dependencies=[Depends(rate_limiter_per_ip)])
+async def get_models(req: Request):
+    client_ip = req.client.host
+    logger.info(f"Fetching available models from IP: {client_ip}")
+    return {"data": [{"id": model, "object": "model"} for model in Blackbox.models]}
+
+# Endpoint: GET /v1/health
+@app.get("/v1/health", dependencies=[Depends(rate_limiter_per_ip)])
+async def health_check(req: Request):
+    client_ip = req.client.host
+    logger.info(f"Health check requested from IP: {client_ip}")
+    return {"status": "ok"}
+
+# Custom exception handler to match OpenAI's error format
+@app.exception_handler(HTTPException)
+async def http_exception_handler(request: Request, exc: HTTPException):
+    client_ip = request.client.host
+    logger.error(f"HTTPException: {exc.detail} | Path: {request.url.path} | IP: {client_ip}")
+    return JSONResponse(
+        status_code=exc.status_code,
+        content={
+            "error": {
+                "message": exc.detail,
+                "type": "invalid_request_error",
+                "param": None,
+                "code": None
+            }
+        },
+    )