Update main.py

main.py

@@ -19,7 +19,7 @@ from pydantic import BaseModel
 # Configure logging
 logging.basicConfig(
     level=logging.INFO,
-    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
+    format="%(asctime)s [%(levelname)s] %(name)s [IP: %(client_ip)s]: %(message)s",
     handlers=[logging.StreamHandler()]
 )
 logger = logging.getLogger(__name__)
@@ -373,13 +373,19 @@ class Blackbox:
 # FastAPI app setup
 app = FastAPI()
 
-# Middleware to enhance security
+# Middleware to enhance security and log client IP
 @app.middleware("http")
 async def security_middleware(request: Request, call_next):
+    client_ip = request.client.host
+    # Enrich the logger with client_ip
+    extra = {"client_ip": client_ip}
+    logger = logging.LoggerAdapter(logging.getLogger(__name__), extra)
+    
     # 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":
+            logger.warning("Invalid Content-Type for /v1/chat/completions")
             return JSONResponse(
                 status_code=400,
                 content={
@@ -391,7 +397,11 @@ async def security_middleware(request: Request, call_next):
                     }
                 },
             )
+    # Log the incoming request
+    logger.info(f"Incoming request: {request.method} {request.url.path}")
     response = await call_next(request)
+    # Log the response status
+    logger.info(f"Response status: {response.status_code}")
     return response
 
 class Message(BaseModel):
@@ -431,15 +441,16 @@ def create_response(content: str, model: str, finish_reason: Optional[str] = Non
 
 @app.post("/v1/chat/completions", dependencies=[Depends(rate_limiter)])
 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} | Model: {request.model} | Messages: {redacted_messages}")
+    logger.info(f"Received chat completions request from API key: {api_key} | Client 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}")
+            logger.warning(f"Attempt to use unavailable model: {request.model} | Client IP: {client_ip}")
             raise HTTPException(status_code=400, detail="Requested model is not available. | NiansuhAI")
 
         # Process the request with actual message content, but don't log it
@@ -481,7 +492,7 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
                 else:
                     response_content += chunk
 
-            logger.info(f"Completed non-streaming response generation for API key: {api_key}")
+            logger.info(f"Completed non-streaming response generation for API key: {api_key} | Client IP: {client_ip}")
             return {
                 "id": f"chatcmpl-{uuid.uuid4()}",
                 "object": "chat.completion",
@@ -504,15 +515,17 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
                 },
             }
     except ModelNotWorkingException as e:
-        logger.warning(f"Model not working: {e}")
+        logger.warning(f"Model not working: {e} | Client IP: {client_ip}")
         raise HTTPException(status_code=503, detail=str(e))
     except HTTPException as he:
-        logger.warning(f"HTTPException: {he.detail}")
+        logger.warning(f"HTTPException: {he.detail} | Client IP: {client_ip}")
         raise he
     except Exception as e:
         logger.exception("An unexpected error occurred while processing the chat completions request.")
         raise HTTPException(status_code=500, detail=str(e))
 
+# Removed the /v1/completions endpoint as per user request
+
 # Return 'about:blank' when accessing the endpoint via GET
 @app.get("/v1/chat/completions")
 async def chat_completions_get():
@@ -520,31 +533,36 @@ async def chat_completions_get():
     return RedirectResponse(url='about:blank')
 
 @app.get("/v1/models")
-async def get_models():
-    logger.info("Fetching available models")
+async def get_models(req: Request):
+    client_ip = req.client.host
+    logger.info(f"Fetching available models | Client IP: {client_ip}")
     return {"data": [{"id": model, "object": "model"} for model in Blackbox.models]}
 
 # Additional endpoints for better functionality
 @app.get("/v1/health")
 async def health_check(req: Request):
-    logger.info("Health check requested")
+    client_ip = req.client.host
+    logger.info(f"Health check requested | Client IP: {client_ip}")
     return {"status": "ok"}
 
 @app.get("/v1/models/{model}/status")
-async def model_status(model: str):
-    logger.info(f"Model status requested for '{model}'")
+async def model_status(model: str, req: Request):
+    client_ip = req.client.host
+    logger.info(f"Model status requested for '{model}' | Client IP: {client_ip}")
     if model in Blackbox.models:
         return {"model": model, "status": "available"}
     elif model in Blackbox.model_aliases and Blackbox.model_aliases[model] in Blackbox.models:
         actual_model = Blackbox.model_aliases[model]
         return {"model": actual_model, "status": "available via alias"}
     else:
-        logger.warning(f"Model not found: {model}")
+        logger.warning(f"Model not found: {model} | Client IP: {client_ip}")
         raise HTTPException(status_code=404, detail="Model not found")
 
 # 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} | Client IP: {client_ip}")
     return JSONResponse(
         status_code=exc.status_code,
         content={
@@ -562,68 +580,13 @@ class TokenizerRequest(BaseModel):
     text: str
 
 @app.post("/v1/tokenizer")
-async def tokenizer(request: TokenizerRequest):
+async def tokenizer(request: TokenizerRequest, req: Request):
+    client_ip = req.client.host
     text = request.text
     token_count = len(text.split())
+    logger.info(f"Tokenizer called | Client IP: {client_ip} | Tokens: {token_count}")
     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")
-async def completions(request: CompletionRequest, req: Request):
-    logger.info(f"Received completion request | 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}")
-            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}")
-        raise he
-    except Exception as e:
-        logger.exception("An unexpected error occurred while processing the completions request.")
-        raise HTTPException(status_code=500, detail=str(e))
-
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=8000)