Update api/utils.py

api/utils.py

@@ -1,5 +1,3 @@
-# utils.py
-
 from datetime import datetime
 import json
 import uuid
@@ -12,18 +10,23 @@ import httpx
 from fastapi import HTTPException
 from api.config import (
     MODEL_MAPPING,
+    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
 
 logger = setup_logger(__name__)
 
-# Helper function to create a UUID-based chat ID
-def generate_chat_id() -> str:
-    return str(uuid.uuid4())
+# Helper function to create a random alphanumeric chat ID
+def generate_chat_id(length: int = 7) -> str:
+    characters = string.ascii_letters + string.digits
+    return ''.join(random.choices(characters, k=length))
 
 # Helper function to create chat completion data
 def create_chat_completion_data(
@@ -44,15 +47,11 @@ def create_chat_completion_data(
         "usage": None,
     }
 
-# Function to convert message to dictionary format, ensuring base64 data
-def message_to_dict(message):
-    if isinstance(message.content, str):
-        content = message.content
-    elif isinstance(message.content, list) and len(message.content) > 0:
-        content = message.content[0].get("text", "")
-    else:
-        content = ""
-    
+# Function to convert message to dictionary format, ensuring base64 data and optional model prefix
+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]:
         # Ensure base64 images are always included for all models
         return {
@@ -67,18 +66,21 @@ def message_to_dict(message):
     return {"role": message.role, "content": content}
 
 # Function to strip model prefix from content if present
-def strip_model_prefix(content: str) -> str:
-    """Function retained but no model_prefix is used."""
-    # Since model_prefix is removed, this function can simply return content
+def strip_model_prefix(content: str, model_prefix: Optional[str] = None) -> str:
+    """Remove the model prefix from the response content if present."""
+    if model_prefix and content.startswith(model_prefix):
+        logger.debug(f"Stripping prefix '{model_prefix}' from content.")
+        return content[len(model_prefix):].strip()
     return content
 
 # Function to get the correct referer URL for logging
 def get_referer_url(chat_id: str, model: str) -> str:
     """Generate the referer URL based on specific models listed in MODEL_REFERERS."""
-    # MODEL_REFERERS has been removed; referer_url defaults to BASE_URL
+    if model in MODEL_REFERERS:
+        return f"{BASE_URL}/chat/{chat_id}?model={model}"
     return BASE_URL
 
-# Process streaming response without model prefixes and referers
+# Process streaming response with headers from config.py
 async def process_streaming_response(request: ChatRequest):
     chat_id = generate_chat_id()
     referer_url = get_referer_url(chat_id, request.model)
@@ -86,22 +88,9 @@ async def process_streaming_response(request: ChatRequest):
 
     agent_mode = AGENT_MODE.get(request.model, {})
     trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
-    
-    headers_api_chat = {
-        'Accept': '*/*',
-        'Accept-Language': 'en-US,en;q=0.9',
-        'Cache-Control': 'no-cache',
-        'Origin': BASE_URL,
-        'Pragma': 'no-cache',
-        'Priority': 'u=1, i',
-        'User-Agent': (
-            'Mozilla/5.0 (Windows NT 10.0; Win64; x64) '
-            'AppleWebKit/537.36 (KHTML, like Gecko) '
-            'Chrome/130.0.0.0 Safari/537.36'
-        ),
-        'Content-Type': 'application/json',
-        'Referer': referer_url,  # Retain referer if necessary
-    }
+    model_prefix = MODEL_PREFIXES.get(request.model, "")
+
+    headers_api_chat = get_headers_api_chat(referer_url)
 
     if request.model == 'o1-preview':
         delay_seconds = random.randint(1, 60)
@@ -119,7 +108,7 @@ async def process_streaming_response(request: ChatRequest):
         "isChromeExt": False,
         "isMicMode": False,
         "maxTokens": request.max_tokens,
-        "messages": [message_to_dict(msg) 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,
@@ -148,7 +137,7 @@ async def process_streaming_response(request: ChatRequest):
                         content = line
                         if content.startswith("$@$v=undefined-rv1$@$"):
                             content = content[21:]
-                        cleaned_content = strip_model_prefix(content)
+                        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"
@@ -160,7 +149,7 @@ async def process_streaming_response(request: ChatRequest):
             logger.error(f"Error occurred during request for Chat ID {chat_id}: {e}")
             raise HTTPException(status_code=500, detail=str(e))
 
-# Process non-streaming response without model prefixes and referers
+# Process non-streaming response with headers from config.py
 async def process_non_streaming_response(request: ChatRequest):
     chat_id = generate_chat_id()
     referer_url = get_referer_url(chat_id, request.model)
@@ -168,35 +157,10 @@ async def process_non_streaming_response(request: ChatRequest):
 
     agent_mode = AGENT_MODE.get(request.model, {})
     trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
-    
-    headers_api_chat = {
-        'Accept': '*/*',
-        'Accept-Language': 'en-US,en;q=0.9',
-        'Cache-Control': 'no-cache',
-        'Origin': BASE_URL,
-        'Pragma': 'no-cache',
-        'Priority': 'u=1, i',
-        'User-Agent': (
-            'Mozilla/5.0 (Windows NT 10.0; Win64; x64) '
-            'AppleWebKit/537.36 (KHTML, like Gecko) '
-            'Chrome/130.0.0.0 Safari/537.36'
-        ),
-        'Content-Type': 'application/json',
-        'Referer': referer_url,  # Retain referer if necessary
-    }
-    headers_chat = {
-        'Accept': '*/*',
-        'Accept-Language': 'en-US,en;q=0.9',
-        'Cache-Control': 'no-cache',
-        'Content-Type': 'text/plain;charset=UTF-8',
-        'Origin': BASE_URL,
-        'Pragma': 'no-cache',
-        'Priority': 'u=1, i',
-        'Referer': referer_url,
-        'Next-Action': str(uuid.uuid4()),
-        'Next-Router-State-Tree': json.dumps([""]),
-        'Next-URL': '/',
-    }
+    model_prefix = MODEL_PREFIXES.get(request.model, "")
+
+    headers_api_chat = get_headers_api_chat(referer_url)
+    headers_chat = get_headers_chat(referer_url, next_action=str(uuid.uuid4()), next_router_state_tree=json.dumps([""]))
 
     if request.model == 'o1-preview':
         delay_seconds = random.randint(20, 60)
@@ -214,7 +178,7 @@ async def process_non_streaming_response(request: ChatRequest):
         "isChromeExt": False,
         "isMicMode": False,
         "maxTokens": request.max_tokens,
-        "messages": [message_to_dict(msg) 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,
@@ -242,11 +206,10 @@ async def process_non_streaming_response(request: ChatRequest):
         except httpx.RequestError as e:
             logger.error(f"Error occurred during request for Chat ID {chat_id}: {e}")
             raise HTTPException(status_code=500, detail=str(e))
-    
     if full_response.startswith("$@$v=undefined-rv1$@$"):
         full_response = full_response[21:]
 
-    cleaned_full_response = strip_model_prefix(full_response)
+    cleaned_full_response = strip_model_prefix(full_response, model_prefix)
 
     return {
         "id": f"chatcmpl-{uuid.uuid4()}",
@@ -261,4 +224,4 @@ async def process_non_streaming_response(request: ChatRequest):
             }
         ],
         "usage": None,
-    }
+    }