Update api/utils.py

api/utils.py

@@ -1,10 +1,10 @@
+import re
 from datetime import datetime
 import json
 import uuid
 import asyncio
 import random
 import string
-import re
 from typing import Any, Dict, Optional
 import httpx
 from fastapi import HTTPException
@@ -22,10 +22,8 @@ from api.models import ChatRequest
 from api.logger import setup_logger
 from api.validate import getHid  # Import the asynchronous getHid function
 
-# Setup logger
 logger = setup_logger(__name__)
 
-# Define the blocked message
 BLOCKED_MESSAGE = "Generated by BLACKBOX.AI, try unlimited chat https://www.blackbox.ai"
 
 # Helper function to create chat completion data
@@ -53,7 +51,6 @@ def message_to_dict(message, model_prefix: Optional[str] = None):
     if model_prefix:
         content = f"{model_prefix} {content}"
     if isinstance(message.content, list) and len(message.content) == 2 and "image_url" in message.content[1]:
-        # Ensure base64 images are always included for all models
         image_base64 = message.content[1]["image_url"]["url"]
         return {
             "role": message.role,
@@ -86,16 +83,16 @@ def remove_message_between_special_tags(content: str) -> str:
     content = re.sub(r'\$~~~\$.*?\$~~~\$', '', content, flags=re.DOTALL)
     return content
 
-# Function to remove search results but keep the links intact
+# Improved function to remove search results but keep the links intact
 def remove_search_results(content: str) -> str:
     """Remove search result sections, keeping the URLs intact."""
-    # Regex to match search result JSON (structured with `title`, `snippet`, `link`)
-    search_result_pattern = re.compile(r'(\{.*?"link":\s*"https?://[^\}]+.*?\})', re.DOTALL)
+    # Regex pattern to match search result JSON structures (looking for "link" and capturing it)
+    search_result_pattern = r'\$~~~\$.*?"link":\s*"(https?://[^\"]+).*?\$~~~\$'
 
-    # Remove all search results but keep URLs intact by using a regex
+    # Remove content inside `$~~~$` tags that contains search result data, keeping only the URLs
     cleaned_content = re.sub(search_result_pattern, '', content)
     
-    # Optional: clean up any unnecessary newlines or extra spaces left by the removal
+    # Remove unnecessary newlines or extra spaces left by the removal
     cleaned_content = cleaned_content.strip()
     
     return cleaned_content
@@ -212,16 +209,10 @@ async def process_non_streaming_response(request: ChatRequest):
     trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
     model_prefix = MODEL_PREFIXES.get(request.model, "")
 
-    # Adjust headers_api_chat and headers_chat since referer_url is removed
     headers_api_chat = get_headers_api_chat(BASE_URL)
-    headers_chat = get_headers_chat(
-        BASE_URL,
-        next_action=str(uuid.uuid4()),
-        next_router_state_tree=json.dumps([""]),
-    )
 
     if request.model == 'o1-preview':
-        delay_seconds = random.randint(20, 60)
+        delay_seconds = random.randint(1, 60)
         logger.info(
             f"Introducing a delay of {delay_seconds} seconds for model 'o1-preview' "
             f"(Request ID: {request_id})"