Update main.py

main.py

@@ -64,21 +64,11 @@ async def cleanup_rate_limit_stores():
 
         await asyncio.sleep(CLEANUP_INTERVAL)
 
-async def get_api_key(request: Request, authorization: str = Header(None)) -> str:
-    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
-
-async def rate_limiter(request: Request, api_key: str = Depends(get_api_key)):
+# Define rate limiter for endpoints requiring API key
+async def rate_limiter_with_api_key(request: Request, api_key: str = Depends(lambda request: get_api_key(request, authorization=request.headers.get('Authorization')))):
     client_ip = request.client.host
     current_time = time.time()
-    
+
     # Rate limiting per API key
     window_start = rate_limit_store[api_key]["timestamp"]
     if current_time - window_start > RATE_LIMIT_WINDOW:
@@ -99,6 +89,32 @@ async def rate_limiter(request: Request, api_key: str = Depends(get_api_key)):
             raise HTTPException(status_code=429, detail='Rate limit exceeded for IP address')
         ip_rate_limit_store[client_ip]["count"] += 1
 
+# Define rate limiter for endpoints NOT requiring API key
+async def rate_limiter_per_ip(request: Request):
+    client_ip = request.client.host
+    current_time = time.time()
+
+    # Rate limiting per IP address
+    window_start_ip = ip_rate_limit_store[client_ip]["timestamp"]
+    if current_time - window_start_ip > RATE_LIMIT_WINDOW:
+        ip_rate_limit_store[client_ip] = {"count": 1, "timestamp": current_time}
+    else:
+        if ip_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')
+        ip_rate_limit_store[client_ip]["count"] += 1
+
+async def get_api_key(request: Request, authorization: str = Header(None)) -> str:
+    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
+
 # Custom exception for model not working
 class ModelNotWorkingException(Exception):
     def __init__(self, model: str):
@@ -399,11 +415,11 @@ app = FastAPI()
 # Middleware to enhance security
 @app.middleware("http")
 async def security_middleware(request: Request, call_next):
+    client_ip = request.client.host
     # Enforce that POST requests to sensitive endpoints must have a valid Content-Type
-    if request.method == "POST" and request.url.path in ["/v1/chat/completions", "/v1/completions"]:
+    if request.method == "POST" and request.url.path == "/v1/chat/completions":
         content_type = request.headers.get("Content-Type")
         if content_type != "application/json":
-            client_ip = request.client.host
             logger.warning(f"Invalid Content-Type from IP: {client_ip} for path: {request.url.path}")
             return JSONResponse(
                 status_code=400,
@@ -454,7 +470,7 @@ def create_response(content: str, model: str, finish_reason: Optional[str] = Non
         "usage": None,
     }
 
-@app.post("/v1/chat/completions", dependencies=[Depends(rate_limiter)])
+@app.post("/v1/chat/completions", dependencies=[Depends(rate_limiter_with_api_key)])
 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
@@ -539,27 +555,29 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
         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))
 
-# Return 'about:blank' when accessing the endpoint via GET
-@app.get("/v1/chat/completions")
-async def chat_completions_get(req: Request):
+# Re-added endpoints without API key authentication
+
+# Endpoint: POST /v1/tokenizer
+class TokenizerRequest(BaseModel):
+    text: str
+
+@app.post("/v1/tokenizer", dependencies=[Depends(rate_limiter_per_ip)])
+async def tokenizer(request: TokenizerRequest, req: Request):
     client_ip = req.client.host
-    logger.info(f"GET request made to /v1/chat/completions from IP: {client_ip}, redirecting to 'about:blank'")
-    return RedirectResponse(url='about:blank')
+    text = request.text
+    token_count = len(text.split())
+    logger.info(f"Tokenizer requested from IP: {client_ip} | Text length: {len(text)}")
+    return {"text": text, "tokens": token_count}
 
-@app.get("/v1/models")
+# 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]}
 
-# Additional endpoints for better functionality
-@app.get("/v1/health", dependencies=[Depends(rate_limiter)])
-async def health_check(req: Request, api_key: str = Depends(get_api_key)):
-    client_ip = req.client.host
-    logger.info(f"Health check requested by API key: {api_key} | IP: {client_ip}")
-    return {"status": "ok"}
-
-@app.get("/v1/models/{model}/status")
+# Endpoint: GET /v1/models/{model}/status
+@app.get("/v1/models/{model}/status", dependencies=[Depends(rate_limiter_per_ip)])
 async def model_status(model: str, req: Request):
     client_ip = req.client.host
     logger.info(f"Model status requested for '{model}' from IP: {client_ip}")
@@ -572,6 +590,20 @@ async def model_status(model: str, req: Request):
         logger.warning(f"Model not found: {model} from IP: {client_ip}")
         raise HTTPException(status_code=404, detail="Model not found")
 
+# 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"}
+
+# Endpoint: GET /v1/chat/completions (GET method)
+@app.get("/v1/chat/completions")
+async def chat_completions_get(req: Request):
+    client_ip = req.client.host
+    logger.info(f"GET request made to /v1/chat/completions from IP: {client_ip}, redirecting to 'about:blank'")
+    return RedirectResponse(url='about:blank')
+
 # Custom exception handler to match OpenAI's error format
 @app.exception_handler(HTTPException)
 async def http_exception_handler(request: Request, exc: HTTPException):
@@ -589,76 +621,6 @@ async def http_exception_handler(request: Request, exc: HTTPException):
         },
     )
 
-# New endpoint: /v1/tokenizer to calculate token counts
-class TokenizerRequest(BaseModel):
-    text: str
-
-@app.post("/v1/tokenizer")
-async def tokenizer(request: TokenizerRequest, req: Request, api_key: str = Depends(get_api_key)):
-    client_ip = req.client.host
-    text = request.text
-    token_count = len(text.split())
-    logger.info(f"Tokenizer requested by API key: {api_key} | IP: {client_ip} | Text length: {len(text)}")
-    return {"text": text, "tokens": token_count}
-
-# New endpoint: /v1/completions to support text completions
-class CompletionRequest(BaseModel):
-    model: str
-    prompt: str
-    max_tokens: Optional[int] = 16
-    temperature: Optional[float] = 1.0
-    top_p: Optional[float] = 1.0
-    n: Optional[int] = 1
-    stream: Optional[bool] = False
-    stop: Optional[Union[str, List[str]]] = None
-    logprobs: Optional[int] = None
-    echo: Optional[bool] = False
-    presence_penalty: Optional[float] = 0.0
-    frequency_penalty: Optional[float] = 0.0
-    best_of: Optional[int] = 1
-    logit_bias: Optional[Dict[str, float]] = None
-    user: Optional[str] = None
-
-@app.post("/v1/completions", dependencies=[Depends(rate_limiter)])
-async def completions(request: CompletionRequest, req: Request, api_key: str = Depends(get_api_key)):
-    client_ip = req.client.host
-    logger.info(f"Received completion request from API key: {api_key} | IP: {client_ip} | Model: {request.model}")
-
-    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.")
-
-        # Simulate a simple completion by echoing the prompt
-        completion_text = f"{request.prompt} [Completed by {request.model}]"
-
-        return {
-            "id": f"cmpl-{uuid.uuid4()}",
-            "object": "text_completion",
-            "created": int(datetime.now().timestamp()),
-            "model": request.model,
-            "choices": [
-                {
-                    "text": completion_text,
-                    "index": 0,
-                    "logprobs": None,
-                    "finish_reason": "length"
-                }
-            ],
-            "usage": {
-                "prompt_tokens": len(request.prompt.split()),
-                "completion_tokens": len(completion_text.split()),
-                "total_tokens": len(request.prompt.split()) + len(completion_text.split())
-            }
-        }
-    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 completions request from IP: {client_ip}.")
-        raise HTTPException(status_code=500, detail=str(e))
-
 # Add the cleanup task when the app starts
 @app.on_event("startup")
 async def startup_event():