Update main.py

main.py

@@ -29,7 +29,7 @@ API_KEYS = os.getenv('API_KEYS', '').split(',')  # Comma-separated API keys
 RATE_LIMIT = int(os.getenv('RATE_LIMIT', '60'))  # Requests per minute
 
 if not API_KEYS or API_KEYS == ['']:
-    logger.error("No API keys found. Please set the API_KEYS environment variable.")
+    logger.error("No API keys found. Please set the API_KEYS environment variable. | NiansuhAI")
     raise Exception("API_KEYS environment variable not set.")
 
 # Simple in-memory rate limiter
@@ -295,8 +295,8 @@ class Blackbox:
                                 logger.info(f"Image URL found.")
                                 yield ImageResponse(image_url, alt=messages[-1]['content'])
                             else:
-                                logger.error("Image URL not found in the response.")
-                                raise Exception("Image URL not found in the response")
+                                logger.error("Image URL not found in the response. | NiansuhAI")
+                                raise Exception("Image URL not found in the response | NiansuhAI")
                         else:
                             full_response = ""
                             search_results_json = ""
@@ -313,7 +313,7 @@ class Blackbox:
                                                 yield decoded_chunk
                                 logger.info("Finished streaming response chunks.")
                             except Exception as e:
-                                logger.exception("Error while iterating over response chunks.")
+                                logger.exception("Error while iterating over response chunks. | NiansuhAI")
                                 raise e
                             if data["webSearchMode"] and search_results_json:
                                 match = re.search(r'\$~~~\$(.*?)\$~~~\$', search_results_json, re.DOTALL)
@@ -326,7 +326,7 @@ class Blackbox:
                                         logger.info("Formatted search results.")
                                         yield formatted_results
                                     except json.JSONDecodeError as je:
-                                        logger.error("Failed to parse search results JSON.")
+                                        logger.error("Failed to parse search results JSON. | NiansuhAI")
                                         raise je
                 break  # Exit the retry loop if successful
             except ClientError as ce:
@@ -382,7 +382,7 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
         # 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.")
+            raise HTTPException(status_code=400, detail="Requested model is not available. | NiansuhAI")
 
         # Process the request with actual message content, but don't log it
         async_generator = Blackbox.create_async_generator(
@@ -410,7 +410,7 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
                     error_response = {"error": he.detail}
                     yield f"data: {json.dumps(error_response)}\n\n"
                 except Exception as e:
-                    logger.exception("Error during streaming response generation.")
+                    logger.exception("Error during streaming response generation. | NiansuhAI")
                     error_response = {"error": str(e)}
                     yield f"data: {json.dumps(error_response)}\n\n"
 
@@ -452,7 +452,7 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
         logger.warning(f"HTTPException: {he.detail}")
         raise he
     except Exception as e:
-        logger.exception("An unexpected error occurred while processing the chat completions request.")
+        logger.exception("An unexpected error occurred while processing the chat completions request. | NiansuhAI")
         raise HTTPException(status_code=500, detail=str(e))
 
 @app.get("/niansuhai/v1/models", dependencies=[Depends(rate_limiter)])
@@ -470,13 +470,13 @@ async def health_check(api_key: str = Depends(get_api_key)):
 async def model_status(model: str, api_key: str = Depends(get_api_key)):
     logger.info(f"Model status requested for '{model}' by API key: {api_key}")
     if model in Blackbox.models:
-        return {"model": model, "status": "available"}
+        return {"model": model, "status": "available | NiansuhAI"}
     elif model in Blackbox.model_aliases:
         actual_model = Blackbox.model_aliases[model]
-        return {"model": actual_model, "status": "available via alias"}
+        return {"model": actual_model, "status": "available via alias | NiansuhAI"}
     else:
         logger.warning(f"Model not found: {model}")
-        raise HTTPException(status_code=404, detail="Model not found")
+        raise HTTPException(status_code=404, detail="Model not found | NiansuhAI")
 
 if __name__ == "__main__":
     import uvicorn