Update main.py

main.py

@@ -39,6 +39,22 @@ rate_limit_store = defaultdict(lambda: {"count": 0, "timestamp": time.time()})
 CLEANUP_INTERVAL = 60  # seconds
 RATE_LIMIT_WINDOW = 60  # seconds
 
+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')
+        rate_limit_store[client_ip]["count"] += 1
+
 class Blackbox:
     label = "Blackbox AI"
     url = "https://www.blackbox.ai"
@@ -327,121 +343,3 @@ class Blackbox:
                 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
-                raise HTTPException(status_code=e.status, detail=error_text)
-            except Exception as 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
-            }
-        },
-    )