Update main.py

main.py

@@ -20,9 +20,7 @@ from pydantic import BaseModel
 logging.basicConfig(
     level=logging.INFO,
     format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
-    handlers=[
-        logging.StreamHandler()
-    ]
+    handlers=[logging.StreamHandler()]
 )
 logger = logging.getLogger(__name__)
 
@@ -38,10 +36,6 @@ if not API_KEYS or API_KEYS == ['']:
 rate_limit_store = defaultdict(lambda: {"count": 0, "timestamp": time.time()})
 
 async def get_api_key(authorization: str = Header(...)) -> str:
-    """
-    Dependency to extract and validate the API key from the Authorization header.
-    Expects the header in the format: Authorization: Bearer <API_KEY>
-    """
     if not authorization.startswith('Bearer '):
         logger.warning("Invalid authorization header format.")
         raise HTTPException(status_code=401, detail='Invalid authorization header format')
@@ -52,14 +46,9 @@ async def get_api_key(authorization: str = Header(...)) -> str:
     return api_key
 
 async def rate_limiter(api_key: str = Depends(get_api_key)):
-    """
-    Dependency to enforce rate limiting per API key.
-    Raises HTTP 429 if the rate limit is exceeded.
-    """
     current_time = time.time()
     window_start = rate_limit_store[api_key]["timestamp"]
     if current_time - window_start > 60:
-        # Reset the count and timestamp after the time window
         rate_limit_store[api_key] = {"count": 1, "timestamp": current_time}
     else:
         if rate_limit_store[api_key]["count"] >= RATE_LIMIT:
@@ -244,6 +233,10 @@ class Blackbox:
         random_id = ''.join(random.choices(string.ascii_letters + string.digits, k=7))
         messages[-1]['id'] = random_id
         messages[-1]['role'] = 'user'
+
+        # Don't log the full message content for privacy
+        logger.debug(f"Generated message ID: {random_id} for model: {model}")
+
         if image is not None:
             messages[-1]['data'] = {
                 'fileText': '',
@@ -283,7 +276,7 @@ class Blackbox:
             data["trendingAgentMode"] = cls.trendingAgentMode[model]
         elif model in cls.userSelectedModel:
             data["userSelectedModel"] = cls.userSelectedModel[model]
-        logger.info(f"Sending request to {cls.api_endpoint} with data: {data}")
+        logger.info(f"Sending request to {cls.api_endpoint} with data (excluding messages).")
 
         timeout = ClientTimeout(total=60)  # Set an appropriate timeout
         retry_attempts = 10  # Set the number of retry attempts
@@ -299,7 +292,7 @@ class Blackbox:
                             url_match = re.search(r'https://storage\.googleapis\.com/[^\s\)]+', response_text)
                             if url_match:
                                 image_url = url_match.group(0)
-                                logger.info(f"Image URL found: {image_url}")
+                                logger.info(f"Image URL found.")
                                 yield ImageResponse(image_url, alt=messages[-1]['content'])
                             else:
                                 logger.error("Image URL not found in the response.")
@@ -366,7 +359,7 @@ def create_response(content: str, model: str, finish_reason: Optional[str] = Non
     return {
         "id": f"chatcmpl-{uuid.uuid4()}",
         "object": "chat.completion.chunk",
-        "created": int(datetime.now().timestamp()),  # datetime is now correctly imported
+        "created": int(datetime.now().timestamp()),
         "model": model,
         "choices": [
             {
@@ -380,22 +373,18 @@ def create_response(content: str, model: str, finish_reason: Optional[str] = Non
 
 @app.post("/niansuhai/v1/chat/completions", dependencies=[Depends(rate_limiter)])
 async def chat_completions(request: ChatRequest, req: Request, api_key: str = Depends(get_api_key)):
-    """
-    Endpoint to handle chat completions.
-    Protected by API key and rate limiter.
-    """
-    logger.info(f"Received chat completions request from API key: {api_key} | Request: {request}")
+    logger.info(f"Received chat completions request from API key: {api_key} | 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.")
 
-        messages = [{"role": msg.role, "content": msg.content} for msg in request.messages]
-        
+        # Process the request but do not log sensitive content
         async_generator = Blackbox.create_async_generator(
             model=request.model,
-            messages=messages,
+            messages=[{"role": msg.role, "content": "[redacted]"} for msg in request.messages],  # Redact user messages in logs
             image=None,
             image_name=None,
             webSearchMode=request.webSearchMode
@@ -411,10 +400,8 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
                         else:
                             response_chunk = create_response(chunk, request.model)
                         
-                        # Yield each chunk in SSE format
                         yield f"data: {json.dumps(response_chunk)}\n\n"
                     
-                    # Signal the end of the stream
                     yield "data: [DONE]\n\n"
                 except HTTPException as he:
                     error_response = {"error": he.detail}
@@ -437,7 +424,7 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
             return {
                 "id": f"chatcmpl-{uuid.uuid4()}",
                 "object": "chat.completion",
-                "created": int(datetime.now().timestamp()),  # datetime is now correctly imported
+                "created": int(datetime.now().timestamp()),
                 "model": request.model,
                 "choices": [
                     {
@@ -450,9 +437,9 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
                     }
                 ],
                 "usage": {
-                    "prompt_tokens": sum(len(msg['content'].split()) for msg in messages),
+                    "prompt_tokens": sum(len(msg['content'].split()) for msg in request.messages),
                     "completion_tokens": len(response_content.split()),
-                    "total_tokens": sum(len(msg['content'].split()) for msg in messages) + len(response_content.split())
+                    "total_tokens": sum(len(msg['content'].split()) for msg in request.messages) + len(response_content.split())
                 },
             }
     except ModelNotWorkingException as e:
@@ -467,23 +454,17 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
 
 @app.get("/niansuhai/v1/models", dependencies=[Depends(rate_limiter)])
 async def get_models(api_key: str = Depends(get_api_key)):
-    """
-    Endpoint to fetch available models.
-    Protected by API key and rate limiter.
-    """
     logger.info(f"Fetching available models for API key: {api_key}")
     return {"data": [{"id": model} for model in Blackbox.models]}
 
 # Additional endpoints for better functionality
 @app.get("/niansuhai/v1/health", dependencies=[Depends(rate_limiter)])
 async def health_check(api_key: str = Depends(get_api_key)):
-    """Health check endpoint to verify the service is running."""
     logger.info(f"Health check requested by API key: {api_key}")
     return {"status": "ok"}
 
 @app.get("/niansuhai/v1/models/{model}/status", dependencies=[Depends(rate_limiter)])
 async def model_status(model: str, api_key: str = Depends(get_api_key)):
-    """Check if a specific model is available."""
     logger.info(f"Model status requested for '{model}' by API key: {api_key}")
     if model in Blackbox.models:
         return {"model": model, "status": "available"}