Update api/utils.py

api/utils.py

@@ -3,6 +3,7 @@ import json
 import uuid
 import asyncio
 import random
+import string
 from typing import Any, Dict, Optional
 
 import httpx
@@ -15,8 +16,9 @@ from api.config import (
     AGENT_MODE,
     TRENDING_AGENT_MODE,
     MODEL_PREFIXES,
+    MODEL_REFERERS
 )
-from api.models import ChatRequest, Message  # Ensure Message is imported
+from api.models import ChatRequest
 from api.logger import setup_logger
 from api.validate import getHid  # Import the asynchronous getHid function
 
@@ -49,28 +51,26 @@ 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}"
-    message_dict = {"role": message.role, "content": content}
-    
-    # Check if the message includes image data
-    if hasattr(message, 'image_base64') and message.image_base64:
-        image_base64 = message.image_base64
-        # Ensure the base64 image data has the correct data URI prefix
-        if not image_base64.startswith('data:image/'):
-            # Assuming JPEG format; adjust if necessary
-            image_base64 = f"data:image/jpeg;base64,{image_base64}"
-        # Generate a unique file path
-        file_name = f"MultipleFiles/{uuid.uuid4().hex}.jpg"
-        message_dict['data'] = {
-            'imagesData': [
-                {
-                    'filePath': file_name,
-                    'contents': image_base64
-                }
-            ],
-            'fileText': '',
-            'title': ''
+    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,
+            "content": content,
+            "data": {
+                "imageBase64": image_base64,
+                "fileText": "",
+                "title": "snapshot",
+                # Added imagesData field here
+                "imagesData": [
+                    {
+                        "filePath": f"MultipleFiles/{uuid.uuid4().hex}.jpg",
+                        "contents": image_base64
+                    }
+                ],
+            },
         }
-    return message_dict
+    return {"role": message.role, "content": content}
 
 # Function to strip model prefix from content if present
 def strip_model_prefix(content: str, model_prefix: Optional[str] = None) -> str:
@@ -80,7 +80,7 @@ def strip_model_prefix(content: str, model_prefix: Optional[str] = None) -> str:
         return content[len(model_prefix):].strip()
     return content
 
-# Process streaming response
+# 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()}"
@@ -95,14 +95,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,
@@ -115,7 +120,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,
@@ -124,7 +131,7 @@ async def process_streaming_response(request: ChatRequest):
         "userId": None,
         "userSelectedModel": MODEL_MAPPING.get(request.model, request.model),
         "userSystemPrompt": None,
-        "validated": h_value,
+        "validated": h_value,  # Dynamically set the validated field
         "visitFromDelta": False,
         "webSearchModePrompt": False,
         "imageGenerationMode": False,  # Added this line
@@ -148,7 +155,9 @@ 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
@@ -161,10 +170,12 @@ async def process_streaming_response(request: ChatRequest):
             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
+# 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()}"
@@ -176,18 +187,27 @@ async def process_non_streaming_response(request: ChatRequest):
 
     # 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([""]))
+    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' "
+            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,
@@ -200,7 +220,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,
@@ -209,7 +231,7 @@ async def process_non_streaming_response(request: ChatRequest):
         "userId": None,
         "userSelectedModel": MODEL_MAPPING.get(request.model, request.model),
         "userSystemPrompt": None,
-        "validated": h_value,
+        "validated": h_value,  # Dynamically set the validated field
         "visitFromDelta": False,
         "webSearchModePrompt": False,
         "imageGenerationMode": False,  # Added this line
@@ -219,7 +241,10 @@ 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", headers=headers_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():
@@ -228,18 +253,24 @@ async def process_non_streaming_response(request: ChatRequest):
             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:]
 
     # 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}.")
+        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.")
+            raise HTTPException(
+                status_code=500, detail="Blocked message detected in response."
+            )
 
     cleaned_full_response = strip_model_prefix(full_response, model_prefix)