Update api/utils.py

api/utils.py

@@ -1,5 +1,6 @@
-# main.py or your main application file
+# utils.py or your main application file
 
+import os  # Import os to access environment variables
 from datetime import datetime
 import json
 import uuid
@@ -29,6 +30,9 @@ logger = setup_logger(__name__)
 # Define the blocked message
 BLOCKED_MESSAGE = "Generated by BLACKBOX.AI, try unlimited chat https://www.blackbox.ai"
 
+# Get the custom message from the environment variable
+FREE_USER_MESSAGE = os.getenv('FREE_USER_MESSAGE', '')
+
 # Helper function to create chat completion data
 def create_chat_completion_data(
     content: str, model: str, timestamp: int, finish_reason: Optional[str] = None
@@ -147,7 +151,12 @@ async def process_streaming_response(request: ChatRequest):
                         cleaned_content = strip_model_prefix(content, model_prefix)
                         yield f"data: {json.dumps(create_chat_completion_data(cleaned_content, request.model, timestamp))}\n\n"
 
-                yield f"data: {json.dumps(create_chat_completion_data('', request.model, timestamp, 'stop'))}\n\n"
+                # Append the custom message for free users
+                if request.is_free_user and FREE_USER_MESSAGE:
+                    timestamp = int(datetime.now().timestamp())
+                    yield f"data: {json.dumps(create_chat_completion_data(FREE_USER_MESSAGE, request.model, timestamp, 'stop'))}\n\n"
+                else:
+                    yield f"data: {json.dumps(create_chat_completion_data('', request.model, timestamp, 'stop'))}\n\n"
                 yield "data: [DONE]\n\n"
         except httpx.HTTPStatusError as e:
             logger.error(f"HTTP error occurred for Request ID {request_id}: {e}")
@@ -233,6 +242,10 @@ async def process_non_streaming_response(request: ChatRequest):
 
     cleaned_full_response = strip_model_prefix(full_response, model_prefix)
 
+    # Append the custom message for free users
+    if request.is_free_user and FREE_USER_MESSAGE:
+        cleaned_full_response += f"\n\n{FREE_USER_MESSAGE}"
+
     return {
         "id": f"chatcmpl-{uuid.uuid4()}",
         "object": "chat.completion",