Update api/utils.py

api/utils.py

@@ -4,7 +4,7 @@ import uuid
 import asyncio
 import random
 import string
-import re
+import os
 from typing import Any, Dict, Optional
 
 import httpx
@@ -28,6 +28,9 @@ logger = setup_logger(__name__)
 # Define the blocked message
 BLOCKED_MESSAGE = "Generated by BLACKBOX.AI, try unlimited chat https://www.blackbox.ai"
 
+# Fetch the advertisement text from environment variable
+ADVERTISEMENT_TEXT = os.getenv("ADVERTISEMENT_TEXT", "")
+
 # Helper function to create chat completion data
 def create_chat_completion_data(
     content: str, model: str, timestamp: int, finish_reason: Optional[str] = None
@@ -81,18 +84,6 @@ def strip_model_prefix(content: str, model_prefix: Optional[str] = None) -> str:
         return content[len(model_prefix):].strip()
     return content
 
-# Helper function to remove content starting with N\n\n$~~~$ and ending with $~~~$
-def remove_enclosed_content(message: str) -> str:
-    """
-    Remove content that starts with N\n\n$~~~$ and ends with $~~~$ from the message.
-    If no such content is found, return the original message.
-    """
-    pattern = r'N\n\n\$\~\~\~\$(.*?)\$\~\~\~\$'
-    match = re.search(pattern, message, re.DOTALL)
-    if match:
-        return message[:match.start()] + message[match.end():]
-    return message
-
 # Process streaming response with headers from config.py
 async def process_streaming_response(request: ChatRequest):
     # Generate a unique ID for this request
@@ -108,14 +99,19 @@ async def process_streaming_response(request: ChatRequest):
 
     if request.model == 'o1-preview':
         delay_seconds = random.randint(1, 60)
-        logger.info(f"Introducing a delay of {delay_seconds} seconds for model 'o1-preview' (Request ID: {request_id})")
+        logger.info(
+            f"Introducing a delay of {delay_seconds} seconds for model 'o1-preview' "
+            f"(Request ID: {request_id})"
+        )
         await asyncio.sleep(delay_seconds)
 
     # Fetch the h-value for the 'validated' field
     h_value = await getHid()
     if not h_value:
         logger.error("Failed to retrieve h-value for validation.")
-        raise HTTPException(status_code=500, detail="Validation failed due to missing h-value.")
+        raise HTTPException(
+            status_code=500, detail="Validation failed due to missing h-value."
+        )
 
     json_data = {
         "agentMode": agent_mode,
@@ -128,7 +124,9 @@ async def process_streaming_response(request: ChatRequest):
         "isChromeExt": False,
         "isMicMode": False,
         "maxTokens": request.max_tokens,
-        "messages": [message_to_dict(msg, model_prefix=model_prefix) for msg in request.messages],
+        "messages": [
+            message_to_dict(msg, model_prefix=model_prefix) for msg in request.messages
+        ],
         "mobileClient": False,
         "playgroundTemperature": request.temperature,
         "playgroundTopP": request.top_p,
@@ -146,7 +144,11 @@ async def process_streaming_response(request: ChatRequest):
     async with httpx.AsyncClient() as client:
         try:
             async with client.stream(
-                "POST", f"{BASE_URL}/api/chat", headers=headers_api_chat, json=json_data, timeout=100
+                "POST",
+                f"{BASE_URL}/api/chat",
+                headers=headers_api_chat,
+                json=json_data,
+                timeout=100,
             ) as response:
                 response.raise_for_status()
                 async for chunk in response.aiter_text():
@@ -157,22 +159,29 @@ async def process_streaming_response(request: ChatRequest):
                             content = content[21:]
                         # Remove the blocked message if present
                         if BLOCKED_MESSAGE in content:
-                            logger.info(f"Blocked message detected in response for Request ID {request_id}.")
+                            logger.info(
+                                f"Blocked message detected in response for Request ID {request_id}."
+                            )
                             content = content.replace(BLOCKED_MESSAGE, '').strip()
                             if not content:
                                 continue  # Skip if content is empty after removal
-                        # Remove content starting with N\n\n$~~~$ and ending with $~~~$
-                        cleaned_content = remove_enclosed_content(content)
-                        cleaned_content = strip_model_prefix(cleaned_content, model_prefix)
+                        cleaned_content = strip_model_prefix(content, model_prefix)
+                        
+                        # Append advertisement text
+                        if ADVERTISEMENT_TEXT:
+                            cleaned_content += "\n\n" + ADVERTISEMENT_TEXT
+
                         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('ok', request.model, timestamp))}\n\n"
+                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}")
             raise HTTPException(status_code=e.response.status_code, detail=str(e))
         except httpx.RequestError as e:
-            logger.error(f"Error occurred during request for Request ID {request_id}: {e}")
+            logger.error(
+                f"Error occurred during request for Request ID {request_id}: {e}"
+            )
             raise HTTPException(status_code=500, detail=str(e))
 
 # Process non-streaming response with headers from config.py
@@ -195,14 +204,19 @@ async def process_non_streaming_response(request: ChatRequest):
 
     if request.model == 'o1-preview':
         delay_seconds = random.randint(20, 60)
-        logger.info(f"Introducing a delay of {delay_seconds} seconds for model 'o1-preview' (Request ID: {request_id})")
+        logger.info(
+            f"Introducing a delay of {delay_seconds} seconds for model 'o1-preview' "
+            f"(Request ID: {request_id})"
+        )
         await asyncio.sleep(delay_seconds)
 
     # Fetch the h-value for the 'validated' field
     h_value = await getHid()
     if not h_value:
         logger.error("Failed to retrieve h-value for validation.")
-        raise HTTPException(status_code=500, detail="Validation failed due to missing h-value.")
+        raise HTTPException(
+            status_code=500, detail="Validation failed due to missing h-value."
+        )
 
     json_data = {
         "agentMode": agent_mode,
@@ -215,7 +229,9 @@ async def process_non_streaming_response(request: ChatRequest):
         "isChromeExt": False,
         "isMicMode": False,
         "maxTokens": request.max_tokens,
-        "messages": [message_to_dict(msg, model_prefix=model_prefix) for msg in request.messages],
+        "messages": [
+            message_to_dict(msg, model_prefix=model_prefix) for msg in request.messages
+        ],
         "mobileClient": False,
         "playgroundTemperature": request.temperature,
         "playgroundTopP": request.top_p,
@@ -230,26 +246,27 @@ async def process_non_streaming_response(request: ChatRequest):
         "imageGenerationMode": False,  # Added this line
     }
 
+    full_response = ""
     async with httpx.AsyncClient() as client:
         try:
-            response = await client.post(
-                f"{BASE_URL}/api/chat", headers=headers_api_chat, json=json_data, timeout=100
-            )
-            response.raise_for_status()
-
-            result = response.json()
-            content = result.get("choices", [{}])[0].get("text", "")
-            cleaned_content = remove_enclosed_content(content)
-            cleaned_content = strip_model_prefix(cleaned_content, model_prefix)
-            return {"choices": [{"text": cleaned_content}], "model": request.model}
+            async with client.stream(
+                method="POST",
+                url=f"{BASE_URL}/api/chat",
+                headers=headers_api_chat,
+                json=json_data,
+            ) as response:
+                response.raise_for_status()
+                async for chunk in response.aiter_text():
+                    full_response += chunk
         except httpx.HTTPStatusError as e:
             logger.error(f"HTTP error occurred for Request ID {request_id}: {e}")
             raise HTTPException(status_code=e.response.status_code, detail=str(e))
         except httpx.RequestError as e:
-            logger.error(f"Error occurred during request for Request ID {request_id}: {e}")
+            logger.error(
+                f"Error occurred during request for Request ID {request_id}: {e}"
+            )
             raise HTTPException(status_code=500, detail=str(e))
 
-
     if full_response.startswith("$@$v=undefined-rv1$@$"):
         full_response = full_response[21:]
 
@@ -264,14 +281,12 @@ async def process_non_streaming_response(request: ChatRequest):
                 status_code=500, detail="Blocked message detected in response."
             )
 
-    # Remove content between special tags
-    full_response = remove_message_between_special_tags(full_response)
-
-    # Remove search results but keep the links intact
-    full_response = remove_search_results(full_response)
-
     cleaned_full_response = strip_model_prefix(full_response, model_prefix)
 
+    # Append the advertisement text if available
+    if ADVERTISEMENT_TEXT:
+        cleaned_full_response += "\n\n" + ADVERTISEMENT_TEXT
+
     return {
         "id": f"chatcmpl-{uuid.uuid4()}",
         "object": "chat.completion",