Update api/utils.py

api/utils.py

@@ -3,7 +3,6 @@ import json
 import uuid
 import asyncio
 import random
-import string
 from typing import Any, Dict, Optional
 
 import httpx
@@ -16,9 +15,8 @@ from api.config import (
     AGENT_MODE,
     TRENDING_AGENT_MODE,
     MODEL_PREFIXES,
-    MODEL_REFERERS
 )
-from api.models import ChatRequest
+from api.models import ChatRequest, Message  # Ensure Message is imported
 from api.logger import setup_logger
 from api.validate import getHid  # Import the asynchronous getHid function
 
@@ -27,7 +25,7 @@ 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
+# Function to create chat completion data
 def create_chat_completion_data(
     content: str, model: str, timestamp: int, finish_reason: Optional[str] = None
 ) -> Dict[str, Any]:
@@ -48,12 +46,30 @@ def create_chat_completion_data(
 
 # 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"]
+    content = message.content
     if model_prefix:
         content = f"{model_prefix} {content}"
     message_dict = {"role": message.role, "content": content}
-    if hasattr(message, 'data') and message.data:
-        message_dict['data'] = message.data
+
+    # Check if the message includes an image in base64 format
+    if hasattr(message, 'image_base64') and message.image_base64:
+        # Ensure the base64 image data has the correct data URI prefix
+        image_base64 = message.image_base64
+        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 or name
+        file_name = f"image_{uuid.uuid4()}.jpg"
+        message_dict['data'] = {
+            'imagesData': [
+                {
+                    'filePath': file_name,
+                    'contents': image_base64
+                }
+            ],
+            'fileText': '',
+            'title': ''
+        }
     return message_dict
 
 # Function to strip model prefix from content if present
@@ -64,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 with headers from config.py
+# Process streaming response
 async def process_streaming_response(request: ChatRequest):
     # Generate a unique ID for this request
     request_id = f"chatcmpl-{uuid.uuid4()}"
@@ -148,7 +164,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 with headers from config.py
+# Process non-streaming response
 async def process_non_streaming_response(request: ChatRequest):
     # Generate a unique ID for this request
     request_id = f"chatcmpl-{uuid.uuid4()}"