Update main.py

main.py

@@ -16,15 +16,19 @@ from fastapi.responses import JSONResponse, StreamingResponse
 from pydantic import BaseModel
 from datetime import datetime
 
-# Configure logging
+# =====================
+# 1. Configure Logging
+# =====================
 logging.basicConfig(
-    level=logging.DEBUG,  # Set to DEBUG to capture detailed logs
+    level=logging.DEBUG,  # Set to DEBUG for detailed logs during development
     format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
     handlers=[logging.StreamHandler()]
 )
 logger = logging.getLogger(__name__)
 
-# Load environment variables
+# ============================
+# 2. 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
 
@@ -32,19 +36,22 @@ 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
+# ====================================
+# 3. Define Rate Limiting Structures
+# ====================================
 rate_limit_store = defaultdict(lambda: {"count": 0, "timestamp": time.time()})
 
 # Define cleanup interval and window
 CLEANUP_INTERVAL = 60  # seconds
 RATE_LIMIT_WINDOW = 60  # seconds
 
-# Define ImageResponse using Pydantic
-class ImageResponse(BaseModel):
+# ========================
+# 4. Define Pydantic Models
+# ========================
+class ImageResponseModel(BaseModel):
     images: str
     alt: str
 
-# Request Models
 class Message(BaseModel):
     role: str
     content: str
@@ -61,14 +68,18 @@ class ChatRequest(BaseModel):
     logit_bias: Optional[Dict[str, float]] = None
     user: Optional[str] = None
 
-# Custom exception for model not working
+# ===============================
+# 5. Define Custom Exceptions
+# ===============================
 class ModelNotWorkingException(Exception):
     def __init__(self, model: str):
         self.model = model
         self.message = f"The model '{model}' is currently not working. Please try another model or wait for it to be fixed."
         super().__init__(self.message)
 
-# Updated Blackbox class with new models and functionalities
+# =======================
+# 6. Define the Blackbox
+# =======================
 class Blackbox:
     label = "Blackbox AI"
     url = "https://www.blackbox.ai"
@@ -344,7 +355,7 @@ class Blackbox:
         proxy: Optional[str] = None,
         websearch: bool = False,
         **kwargs
-    ) -> AsyncGenerator[Union[str, ImageResponse], None]:
+    ) -> AsyncGenerator[Union[str, ImageResponseModel], None]:
         """
         Creates an asynchronous generator for streaming responses from Blackbox AI.
 
@@ -356,7 +367,7 @@ class Blackbox:
             **kwargs: Additional keyword arguments.
 
         Yields:
-            Union[str, ImageResponse]: Segments of the generated response or ImageResponse objects.
+            Union[str, ImageResponseModel]: Segments of the generated response or ImageResponseModel objects.
         """
         model = cls.get_model(model)
 
@@ -432,7 +443,7 @@ class Blackbox:
             "clickedForceWebSearch": False,
             "visitFromDelta": False,
             "mobileClient": False,
-            "webSearchMode": False,
+            "webSearchMode": websearch,
             "userSelectedModel": cls.userSelectedModel.get(model, model)
         }
 
@@ -466,7 +477,7 @@ class Blackbox:
                         match = re.search(r'!\[.*?\]\((https?://[^\)]+)\)', cleaned_response)
                         if match:
                             image_url = match.group(1)
-                            image_response = ImageResponse(images=image_url, alt="Generated Image")
+                            image_response = ImageResponseModel(images=image_url, alt="Generated Image")
                             yield image_response
                         else:
                             yield cleaned_response
@@ -533,201 +544,214 @@ class Blackbox:
             except Exception as e:
                 yield f"Unexpected error during /chat/{chat_id} 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
-
-    async def cleanup_rate_limit_stores():
-        """
-        Periodically cleans up stale entries in the rate_limit_store to prevent memory bloat.
-        """
-        while True:
-            current_time = time.time()
-            ips_to_delete = [ip for ip, value in rate_limit_store.items() if current_time - value["timestamp"] > RATE_LIMIT_WINDOW * 2]
-            for ip in ips_to_delete:
-                del rate_limit_store[ip]
-                logger.debug(f"Cleaned up rate_limit_store for IP: {ip}")
-            await asyncio.sleep(CLEANUP_INTERVAL)
-
-    async def rate_limiter_per_ip(request: Request):
-        """
-        Rate limiter that enforces a limit based on the client's IP address.
-        """
-        client_ip = request.client.host
-        current_time = time.time()
-
-        # Initialize or update the count and timestamp
-        if current_time - rate_limit_store[client_ip]["timestamp"] > RATE_LIMIT_WINDOW:
-            rate_limit_store[client_ip] = {"count": 1, "timestamp": current_time}
-        else:
-            if rate_limit_store[client_ip]["count"] >= RATE_LIMIT:
-                logger.warning(f"Rate limit exceeded for IP address: {client_ip}")
-                raise HTTPException(status_code=429, detail='Rate limit exceeded for IP address | NiansuhAI')
-            rate_limit_store[client_ip]["count"] += 1
-
-    async def get_api_key(request: Request, authorization: str = Header(None)) -> str:
-        """
-        Dependency to extract and validate the API key from the Authorization header.
-        """
-        client_ip = request.client.host
-        if authorization is None or not authorization.startswith('Bearer '):
-            logger.warning(f"Invalid or missing authorization header from IP: {client_ip}")
-            raise HTTPException(status_code=401, detail='Invalid authorization header format')
-        api_key = authorization[7:]
-        if api_key not in API_KEYS:
-            logger.warning(f"Invalid API key attempted: {api_key} from IP: {client_ip}")
-            raise HTTPException(status_code=401, detail='Invalid API key')
-        return api_key
-
-    @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 {
-                "id": f"chatcmpl-{uuid.uuid4()}",
-                "object": "chat.completion",
-                "created": int(datetime.now().timestamp()),
-                "model": request.model,
-                "choices": [
-                    {
-                        "index": 0,
-                        "message": {
-                            "role": "assistant",
-                            "content": response_content
-                        },
-                        "finish_reason": "stop"
+# ============================
+# 7. Initialize FastAPI App
+# ============================
+app = FastAPI()
+
+# ====================================
+# 8. Define Middleware and Dependencies
+# ====================================
+@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
                     }
-                ],
-                "usage": {
-                    "prompt_tokens": sum(len(msg.content.split()) for msg in request.messages),
-                    "completion_tokens": len(response_content.split()),
-                    "total_tokens": sum(len(msg.content.split()) for msg in request.messages) + len(response_content.split())
                 },
-            }
-        except ModelNotWorkingException as e:
-            logger.warning(f"Model not working: {e} | IP: {client_ip}")
-            raise HTTPException(status_code=503, detail=str(e))
-        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))
-
-    @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]}
-
-    @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
+            )
+    response = await call_next(request)
+    return response
+
+async def cleanup_rate_limit_stores():
+    """
+    Periodically cleans up stale entries in the rate_limit_store to prevent memory bloat.
+    """
+    while True:
+        current_time = time.time()
+        ips_to_delete = [ip for ip, value in rate_limit_store.items() if current_time - value["timestamp"] > RATE_LIMIT_WINDOW * 2]
+        for ip in ips_to_delete:
+            del rate_limit_store[ip]
+            logger.debug(f"Cleaned up rate_limit_store for IP: {ip}")
+        await asyncio.sleep(CLEANUP_INTERVAL)
+
+async def rate_limiter_per_ip(request: Request):
+    """
+    Rate limiter that enforces a limit based on the client's IP address.
+    """
+    client_ip = request.client.host
+    current_time = time.time()
+
+    # Initialize or update the count and timestamp
+    if current_time - rate_limit_store[client_ip]["timestamp"] > RATE_LIMIT_WINDOW:
+        rate_limit_store[client_ip] = {"count": 1, "timestamp": current_time}
+    else:
+        if rate_limit_store[client_ip]["count"] >= RATE_LIMIT:
+            logger.warning(f"Rate limit exceeded for IP address: {client_ip}")
+            raise HTTPException(status_code=429, detail='Rate limit exceeded for IP address | NiansuhAI')
+        rate_limit_store[client_ip]["count"] += 1
+
+async def get_api_key(request: Request, authorization: str = Header(None)) -> str:
+    """
+    Dependency to extract and validate the API key from the Authorization header.
+    """
+    client_ip = request.client.host
+    if authorization is None or not authorization.startswith('Bearer '):
+        logger.warning(f"Invalid or missing authorization header from IP: {client_ip}")
+        raise HTTPException(status_code=401, detail='Invalid authorization header format')
+    api_key = authorization[7:]
+    if api_key not in API_KEYS:
+        logger.warning(f"Invalid API key attempted: {api_key} from IP: {client_ip}")
+        raise HTTPException(status_code=401, detail='Invalid API key')
+    return api_key
+
+# =====================================
+# 9. Define FastAPI Event Handlers
+# =====================================
+@app.on_event("startup")
+async def startup_event():
+    asyncio.create_task(cleanup_rate_limit_stores())
+    logger.info("Started rate limit store cleanup task.")
+
+# ==========================================
+# 10. Define FastAPI Endpoints
+# ==========================================
+@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 {
+            "id": f"chatcmpl-{uuid.uuid4()}",
+            "object": "chat.completion",
+            "created": int(datetime.now().timestamp()),
+            "model": request.model,
+            "choices": [
+                {
+                    "index": 0,
+                    "message": {
+                        "role": "assistant",
+                        "content": response_content
+                    },
+                    "finish_reason": "stop"
                 }
+            ],
+            "usage": {
+                "prompt_tokens": sum(len(msg.content.split()) for msg in request.messages),
+                "completion_tokens": len(response_content.split()),
+                "total_tokens": sum(len(msg.content.split()) for msg in request.messages) + len(response_content.split())
             },
+        }
+    except ModelNotWorkingException as e:
+        logger.warning(f"Model not working: {e} | IP: {client_ip}")
+        raise HTTPException(status_code=503, detail=str(e))
+    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))
+
+@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]}
+
+@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"}
+
+# ===============================
+# 11. Define Custom Exception Handler
+# ===============================
+@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
+            }
+        },
+    )
+
+# ============================
+# 12. Optional: Streaming Endpoint
+# ============================
+@app.post("/v1/chat/completions/stream", dependencies=[Depends(rate_limiter_per_ip)])
+async def chat_completions_stream(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 streaming 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.")
+
+        # Create an asynchronous generator for the response
+        async_generator = Blackbox.create_async_generator(
+            model=request.model,
+            messages=[{"role": msg.role, "content": msg.content} for msg in request.messages],
+            temperature=request.temperature,
+            max_tokens=request.max_tokens
         )
 
-    # Optional: Additional Endpoint for Streaming Responses (Using create_async_generator)
-    # This endpoint leverages the new create_async_generator method for streaming responses.
-    # Note: Streaming responses may require clients that support Server-Sent Events (SSE) or WebSockets.
-
-    @app.post("/v1/chat/completions/stream", dependencies=[Depends(rate_limiter_per_ip)])
-    async def chat_completions_stream(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 streaming 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.")
-
-            # Create an asynchronous generator for the response
-            async_generator = Blackbox.create_async_generator(
-                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"Started streaming response for API key: {api_key} | IP: {client_ip}")
-            return StreamingResponse(async_generator, media_type="text/event-stream")
-        except ModelNotWorkingException as e:
-            logger.warning(f"Model not working: {e} | IP: {client_ip}")
-            raise HTTPException(status_code=503, detail=str(e))
-        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 streaming chat completions request from IP: {client_ip}.")
-            raise HTTPException(status_code=500, detail=str(e))
-
-    if __name__ == "__main__":
-        import uvicorn
-        uvicorn.run(app, host="0.0.0.0", port=8000)
+        logger.info(f"Started streaming response for API key: {api_key} | IP: {client_ip}")
+        return StreamingResponse(async_generator, media_type="text/event-stream")
+    except ModelNotWorkingException as e:
+        logger.warning(f"Model not working: {e} | IP: {client_ip}")
+        raise HTTPException(status_code=503, detail=str(e))
+    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 streaming chat completions request from IP: {client_ip}.")
+        raise HTTPException(status_code=500, detail=str(e))
+
+# ========================================
+# 13. Run the Application with Uvicorn
+# ========================================
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)