Update main.py

main.py

@@ -10,19 +10,32 @@ import time
 from collections import defaultdict
 from typing import List, Dict, Any, Optional, AsyncGenerator, Union
 from datetime import datetime
+import contextvars  # Added for context variables
 
 from aiohttp import ClientSession, ClientTimeout, ClientError
 from fastapi import FastAPI, HTTPException, Request, Depends, Header
 from fastapi.responses import StreamingResponse, JSONResponse, RedirectResponse
 from pydantic import BaseModel
 
+# Define a context variable for client_ip
+client_ip_var = contextvars.ContextVar("client_ip", default="N/A")
+
+# Custom logging filter to inject client_ip from context variable
+class ContextFilter(logging.Filter):
+    def filter(self, record):
+        record.client_ip = client_ip_var.get()
+        return True
+
 # Configure logging
 logging.basicConfig(
     level=logging.INFO,
     format="%(asctime)s [%(levelname)s] %(name)s [IP: %(client_ip)s]: %(message)s",
     handlers=[logging.StreamHandler()]
 )
+
+# Get the root logger and add the ContextFilter
 logger = logging.getLogger(__name__)
+logger.addFilter(ContextFilter())
 
 # Load environment variables
 API_KEYS = os.getenv('API_KEYS', '').split(',')  # Comma-separated API keys
@@ -377,9 +390,8 @@ app = FastAPI()
 @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)
+    # Set the client_ip in the context variable
+    client_ip_var.set(client_ip)
     
     # Enforce that POST requests to sensitive endpoints must have a valid Content-Type
     if request.method == "POST" and request.url.path == "/v1/chat/completions":
@@ -397,11 +409,15 @@ 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):
@@ -441,16 +457,15 @@ 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} | Client IP: {client_ip} | Model: {request.model} | Messages: {redacted_messages}")
+    logger.info(f"Received chat completions request from API key: {api_key} | 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} | Client IP: {client_ip}")
+            logger.warning(f"Attempt to use unavailable model: {request.model}")
             raise HTTPException(status_code=400, detail="Requested model is not available. | NiansuhAI")
 
         # Process the request with actual message content, but don't log it
@@ -492,7 +507,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} | Client IP: {client_ip}")
+            logger.info(f"Completed non-streaming response generation for API key: {api_key} | Model: {request.model}")
             return {
                 "id": f"chatcmpl-{uuid.uuid4()}",
                 "object": "chat.completion",
@@ -515,10 +530,10 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
                 },
             }
     except ModelNotWorkingException as e:
-        logger.warning(f"Model not working: {e} | Client IP: {client_ip}")
+        logger.warning(f"Model not working: {e}")
         raise HTTPException(status_code=503, detail=str(e))
     except HTTPException as he:
-        logger.warning(f"HTTPException: {he.detail} | Client IP: {client_ip}")
+        logger.warning(f"HTTPException: {he.detail}")
         raise he
     except Exception as e:
         logger.exception("An unexpected error occurred while processing the chat completions request.")
@@ -535,27 +550,27 @@ async def chat_completions_get():
 @app.get("/v1/models")
 async def get_models(req: Request):
     client_ip = req.client.host
-    logger.info(f"Fetching available models | Client IP: {client_ip}")
+    logger.info(f"Fetching available models")
     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):
     client_ip = req.client.host
-    logger.info(f"Health check requested | Client IP: {client_ip}")
+    logger.info(f"Health check requested")
     return {"status": "ok"}
 
 @app.get("/v1/models/{model}/status")
 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}")
+    logger.info(f"Model status requested for '{model}'")
     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} | Client IP: {client_ip}")
+        logger.warning(f"Model not found: {model}")
         raise HTTPException(status_code=404, detail="Model not found")
 
 # Custom exception handler to match OpenAI's error format
@@ -584,7 +599,7 @@ 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}")
+    logger.info(f"Tokenizer called | Tokens: {token_count}")
     return {"text": text, "tokens": token_count}
 
 if __name__ == "__main__":