Update main.py

main.py

@@ -6,7 +6,6 @@ import uuid
 import json
 import logging
 import asyncio
-import time
 from aiohttp import ClientSession, ClientTimeout, ClientError
 from fastapi import FastAPI, HTTPException, Request, Depends, Header, status
 from fastapi.responses import StreamingResponse, JSONResponse
@@ -56,7 +55,7 @@ app.state.limiter = limiter
 app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
 
 # API Key Authentication
-API_KEYS = set(os.getenv("API_KEYS", "").split(","))  # Load API keys from environment variable
+API_KEYS = set(api_key.strip() for api_key in os.getenv("API_KEYS", "").split(",") if api_key.strip())
 
 async def get_api_key(authorization: Optional[str] = Header(None)):
     """
@@ -426,7 +425,11 @@ def create_response(content: str, model: str, finish_reason: Optional[str] = Non
 # Endpoint: Chat Completions
 @app.post("/v1/chat/completions", response_model=ChatCompletionResponse)
 @limiter.limit("60/minute")  # Example: 60 requests per minute per IP
-async def chat_completions(request: ChatRequest, req: Request, api_key: str = Depends(get_api_key)):
+async def chat_completions(
+    request: ChatRequest,
+    req: Request,
+    api_key: str = Depends(get_api_key)
+):
     logger.info(f"Received chat completions request: {request}")
     try:
         messages = [{"role": msg.role, "content": msg.content} for msg in request.messages]
@@ -509,16 +512,23 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
         raise HTTPException(status_code=500, detail=str(e))
 
 # Endpoint: List Models
-@app.get("/v1/models")
+@app.get("/v1/models", response_model=Dict[str, List[Dict[str, str]]])
 @limiter.limit("60/minute")
-async def get_models(api_key: str = Depends(get_api_key)):
+async def get_models(
+    request: Request,
+    api_key: str = Depends(get_api_key)
+):
     logger.info("Fetching available models.")
     return {"data": [{"id": model} for model in Blackbox.models]}
 
 # Endpoint: Model Status
-@app.get("/v1/models/{model}/status")
+@app.get("/v1/models/{model}/status", response_model=Dict[str, str])
 @limiter.limit("60/minute")
-async def model_status(model: str, api_key: str = Depends(get_api_key)):
+async def model_status(
+    model: str,
+    request: Request,
+    api_key: str = Depends(get_api_key)
+):
     """Check if a specific model is available."""
     if model in Blackbox.models:
         return {"model": model, "status": "available"}
@@ -529,9 +539,9 @@ async def model_status(model: str, api_key: str = Depends(get_api_key)):
         raise HTTPException(status_code=404, detail="Model not found")
 
 # Endpoint: Health Check
-@app.get("/v1/health")
+@app.get("/v1/health", response_model=Dict[str, str])
 @limiter.limit("60/minute")
-async def health_check():
+async def health_check(request: Request):
     """Health check endpoint to verify the service is running."""
     return {"status": "ok"}