Update api/utils.py

api/utils.py

@@ -1,19 +1,24 @@
+# main.py or your main application file
+
 from datetime import datetime
 import json
 import uuid
 import asyncio
 import random
+import string
 from typing import Any, Dict, Optional
 
 import httpx
 from fastapi import HTTPException
 from api.config import (
     MODEL_MAPPING,
-    common_headers,  # Import common_headers directly
+    get_headers_api_chat,
+    get_headers_chat,
     BASE_URL,
     AGENT_MODE,
     TRENDING_AGENT_MODE,
     MODEL_PREFIXES,
+    MODEL_REFERERS
 )
 from api.models import ChatRequest
 from api.logger import setup_logger
@@ -21,6 +26,9 @@ from api.validate import getHid  # Import the asynchronous getHid function
 
 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
 def create_chat_completion_data(
     content: str, model: str, timestamp: int, finish_reason: Optional[str] = None
@@ -45,11 +53,7 @@ def message_to_dict(message, model_prefix: Optional[str] = None):
     content = message.content if isinstance(message.content, str) else message.content[0]["text"]
     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]
-    ):
+    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
         return {
             "role": message.role,
@@ -70,7 +74,7 @@ def strip_model_prefix(content: str, model_prefix: Optional[str] = None) -> str:
         return content[len(model_prefix):].strip()
     return content
 
-# Process streaming response without headers
+# Process streaming response with headers from config.py
 async def process_streaming_response(request: ChatRequest):
     # Generate a unique ID for this request
     request_id = f"chatcmpl-{uuid.uuid4()}"
@@ -80,11 +84,12 @@ async def process_streaming_response(request: ChatRequest):
     trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
     model_prefix = MODEL_PREFIXES.get(request.model, "")
 
+    # Adjust headers_api_chat since referer_url is removed
+    headers_api_chat = get_headers_api_chat(BASE_URL)
+
     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' (Request ID: {request_id})")
         await asyncio.sleep(delay_seconds)
 
     # Fetch the h-value for the 'validated' field
@@ -104,9 +109,7 @@ 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,
@@ -124,6 +127,7 @@ async def process_streaming_response(request: ChatRequest):
             async with client.stream(
                 "POST",
                 f"{BASE_URL}/api/chat",
+                headers=headers_api_chat,
                 json=json_data,
                 timeout=100,
             ) as response:
@@ -134,6 +138,12 @@ async def process_streaming_response(request: ChatRequest):
                         content = line
                         if content.startswith("$@$v=undefined-rv1$@$"):
                             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}.")
+                            content = content.replace(BLOCKED_MESSAGE, '').strip()
+                            if not content:
+                                continue  # Skip if content is empty after removal
                         cleaned_content = strip_model_prefix(content, model_prefix)
                         yield f"data: {json.dumps(create_chat_completion_data(cleaned_content, request.model, timestamp))}\n\n"
 
@@ -146,7 +156,7 @@ async def process_streaming_response(request: ChatRequest):
             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 without headers
+# Process non-streaming response with headers from config.py
 async def process_non_streaming_response(request: ChatRequest):
     # Generate a unique ID for this request
     request_id = f"chatcmpl-{uuid.uuid4()}"
@@ -156,11 +166,13 @@ 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)
-        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' (Request ID: {request_id})")
         await asyncio.sleep(delay_seconds)
 
     # Fetch the h-value for the 'validated' field
@@ -180,9 +192,7 @@ 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,
@@ -199,9 +209,7 @@ async def process_non_streaming_response(request: ChatRequest):
     async with httpx.AsyncClient() as client:
         try:
             async with client.stream(
-                method="POST",
-                url=f"{BASE_URL}/api/chat",
-                json=json_data,
+                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():
@@ -212,10 +220,17 @@ async def process_non_streaming_response(request: ChatRequest):
         except httpx.RequestError as 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:]
 
+    # Remove the blocked message if present
+    if BLOCKED_MESSAGE in full_response:
+        logger.info(f"Blocked message detected in response for Request ID {request_id}.")
+        full_response = full_response.replace(BLOCKED_MESSAGE, '').strip()
+        if not full_response:
+            raise HTTPException(status_code=500, detail="Blocked message detected in response.")
+
     cleaned_full_response = strip_model_prefix(full_response, model_prefix)
 
     return {