Update main.py

main.py

@@ -473,38 +473,6 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
             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
-
 # Request Models
 class Message(BaseModel):
     role: str
@@ -556,6 +524,38 @@ def create_response(content: str, model: str, finish_reason: Optional[str] = Non
         "usage": None,  # To be filled in non-streaming responses
     }
 
+# Initialize FastAPI app
+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
+
 # FastAPI Endpoints
 
 @app.post("/v1/chat/completions", dependencies=[Depends(rate_limiter_per_ip)])
@@ -610,7 +610,7 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
                         yield f"data: {json.dumps(response_chunk)}\n\n"
                     
                     # After all chunks are sent, send the final message with finish_reason
-                    prompt_tokens = sum(len(msg['content'].split()) for msg in request.messages)
+                    prompt_tokens = sum(len(msg.content.split()) for msg in request.messages)
                     completion_tokens = len(assistant_content.split())
                     total_tokens = prompt_tokens + completion_tokens
                     estimated_cost = calculate_estimated_cost(prompt_tokens, completion_tokens)
@@ -695,7 +695,7 @@ 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))
 
-# Tokenizer Endpoint
+# Endpoint: POST /v1/tokenizer
 @app.post("/v1/tokenizer", dependencies=[Depends(rate_limiter_per_ip)])
 async def tokenizer(request: TokenizerRequest, req: Request):
     client_ip = req.client.host
@@ -704,14 +704,14 @@ async def tokenizer(request: TokenizerRequest, req: Request):
     logger.info(f"Tokenizer requested from IP: {client_ip} | Text length: {len(text)}")
     return {"text": text, "tokens": token_count}
 
-# Get Models Endpoint
+# 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]}
 
-# Model Status Endpoint
+# 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
@@ -725,21 +725,21 @@ 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")
 
-# Health Check Endpoint
+# 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"}
 
-# Redirect GET requests to /v1/chat/completions to 'about:blank'
+# 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
+# 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