Update api/utils.py

api/utils.py

@@ -2,16 +2,37 @@ from datetime import datetime
 import json
 from typing import Any, Dict, Optional
 import uuid
+import re
 
 import httpx
 from api.config import MODEL_MAPPING, headers, AGENT_MODE, TRENDING_AGENT_MODE, BASE_URL
 from fastapi import HTTPException
-from api.models import ChatRequest
+from api.models import ChatRequest, ImageResponse
 
 from api.logger import setup_logger
 
 logger = setup_logger(__name__)
 
+model_prefixes = {
+    'gpt-4o': '@GPT-4o',
+    'gemini-pro': '@Gemini-PRO',
+    'claude-sonnet-3.5': '@Claude-Sonnet-3.5',
+    'PythonAgent': '@Python Agent',
+    'JavaAgent': '@Java Agent',
+    'JavaScriptAgent': '@JavaScript Agent',
+    'HTMLAgent': '@HTML Agent',
+    'GoogleCloudAgent': '@Google Cloud Agent',
+    'AndroidDeveloper': '@Android Developer',
+    'SwiftDeveloper': '@Swift Developer',
+    'Next.jsAgent': '@Next.js Agent',
+    'MongoDBAgent': '@MongoDB Agent',
+    'PyTorchAgent': '@PyTorch Agent',
+    'ReactAgent': '@React Agent',
+    'XcodeAgent': '@Xcode Agent',
+    'AngularJSAgent': '@AngularJS Agent',
+    'blackboxai-pro': '@BLACKBOXAI-PRO',
+    'ImageGeneration': '@Image Generation',
+}
 
 def create_chat_completion_data(
     content: str, model: str, timestamp: int, finish_reason: Optional[str] = None
@@ -38,9 +59,9 @@ def message_to_dict(message):
     elif isinstance(message.content, list) and len(message.content) == 2:
         return {
             "role": message.role,
-            "content": message.content[0]["text"],
+            "content": message.content[0].get("text", ""),
             "data": {
-                "imageBase64": message.content[1]["image_url"]["url"],
+                "imageBase64": message.content[1].get("imageBase64", ""),
                 "fileText": "",
                 "title": "snapshot",
             },
@@ -55,27 +76,6 @@ async def process_streaming_response(request: ChatRequest):
     model = MODEL_MAPPING.get(request.model, request.model)
     
     # Add prefix if applicable
-    prefix = ""
-    model_prefixes = {
-        'gpt-4o': '@GPT-4o',
-        'gemini-pro': '@Gemini-PRO',
-        'claude-sonnet-3.5': '@Claude-Sonnet-3.5',
-        'PythonAgent': '@Python Agent',
-        'JavaAgent': '@Java Agent',
-        'JavaScriptAgent': '@JavaScript Agent',
-        'HTMLAgent': '@HTML Agent',
-        'GoogleCloudAgent': '@Google Cloud Agent',
-        'AndroidDeveloper': '@Android Developer',
-        'SwiftDeveloper': '@Swift Developer',
-        'Next.jsAgent': '@Next.js Agent',
-        'MongoDBAgent': '@MongoDB Agent',
-        'PyTorchAgent': '@PyTorch Agent',
-        'ReactAgent': '@React Agent',
-        'XcodeAgent': '@Xcode Agent',
-        'AngularJSAgent': '@AngularJS Agent',
-        'blackboxai-pro': '@BLACKBOXAI-PRO',
-        'ImageGeneration': '@Image Generation',
-    }
     prefix = model_prefixes.get(request.model, "")
     
     # Format messages with prefix
@@ -84,10 +84,12 @@ async def process_streaming_response(request: ChatRequest):
         formatted_content = msg.content
         if prefix and msg.role.lower() == "user":
             formatted_content = f"{prefix} {formatted_content}"
+        # Ensure 'data' is included if present
+        data = msg.content.get('data') if isinstance(msg.content, dict) else None
         formatted_messages.append({
             "role": msg.role,
             "content": formatted_content,
-            "data": msg.content.get('data') if isinstance(msg.content, dict) else None
+            "data": data
         })
 
     json_data = {
@@ -129,7 +131,15 @@ async def process_streaming_response(request: ChatRequest):
                         content = line
                         if content.startswith("$@$v=undefined-rv1$@$"):
                             content = content[21:]
-                        yield f"data: {json.dumps(create_chat_completion_data(content, request.model, timestamp))}\n\n"
+                        
+                        # Check for Base64 image data
+                        match_data_uri = re.search(r'!\[.*?\]\((data:image/[^;]+;base64,[^\)]+)\)', content)
+                        if match_data_uri:
+                            image_data = match_data_uri.group(1)
+                            image_response = ImageResponse(images=image_data, alt="Generated Image")
+                            yield image_response
+                        else:
+                            yield f"data: {json.dumps(create_chat_completion_data(content, request.model, timestamp))}\n\n"
 
                 yield f"data: {json.dumps(create_chat_completion_data('', request.model, timestamp, 'stop'))}\n\n"
                 yield "data: [DONE]\n\n"
@@ -147,27 +157,6 @@ async def process_non_streaming_response(request: ChatRequest):
     model = MODEL_MAPPING.get(request.model, request.model)
     
     # Add prefix if applicable
-    prefix = ""
-    model_prefixes = {
-        'gpt-4o': '@GPT-4o',
-        'gemini-pro': '@Gemini-PRO',
-        'claude-sonnet-3.5': '@Claude-Sonnet-3.5',
-        'PythonAgent': '@Python Agent',
-        'JavaAgent': '@Java Agent',
-        'JavaScriptAgent': '@JavaScript Agent',
-        'HTMLAgent': '@HTML Agent',
-        'GoogleCloudAgent': '@Google Cloud Agent',
-        'AndroidDeveloper': '@Android Developer',
-        'SwiftDeveloper': '@Swift Developer',
-        'Next.jsAgent': '@Next.js Agent',
-        'MongoDBAgent': '@MongoDB Agent',
-        'PyTorchAgent': '@PyTorch Agent',
-        'ReactAgent': '@React Agent',
-        'XcodeAgent': '@Xcode Agent',
-        'AngularJSAgent': '@AngularJS Agent',
-        'blackboxai-pro': '@BLACKBOXAI-PRO',
-        'ImageGeneration': '@Image Generation',
-    }
     prefix = model_prefixes.get(request.model, "")
     
     # Format messages with prefix
@@ -176,10 +165,12 @@ async def process_non_streaming_response(request: ChatRequest):
         formatted_content = msg.content
         if prefix and msg.role.lower() == "user":
             formatted_content = f"{prefix} {formatted_content}"
+        # Ensure 'data' is included if present
+        data = msg.content.get('data') if isinstance(msg.content, dict) else None
         formatted_messages.append({
             "role": msg.role,
             "content": formatted_content,
-            "data": msg.content.get('data') if isinstance(msg.content, dict) else None
+            "data": data
         })
 
     json_data = {
@@ -204,7 +195,6 @@ async def process_non_streaming_response(request: ChatRequest):
         "webSearchMode": False,  # Set to True if web search is needed
         "userSelectedModel": model,
     }
-
     full_response = ""
     async with httpx.AsyncClient() as client:
         try:
@@ -226,6 +216,27 @@ async def process_non_streaming_response(request: ChatRequest):
     if full_response.startswith("$@$v=undefined-rv1$@$"):
         full_response = full_response[21:]
 
+    # Handle Base64 image data in the response
+    match_data_uri = re.search(r'!\[.*?\]\((data:image/[^;]+;base64,[^\)]+)\)', full_response)
+    if match_data_uri:
+        image_data = match_data_uri.group(1)
+        # Optionally, you can convert data URI to a more usable format or pass it directly
+        return {
+            "id": f"chatcmpl-{uuid.uuid4()}",
+            "object": "chat.completion",
+            "created": int(datetime.now().timestamp()),
+            "model": request.model,
+            "choices": [
+                {
+                    "index": 0,
+                    "message": {"role": "assistant", "content": full_response},
+                    "finish_reason": "stop",
+                }
+            ],
+            "usage": None,
+            "imageData": image_data  # Add image data to the response if needed
+        }
+
     return {
         "id": f"chatcmpl-{uuid.uuid4()}",
         "object": "chat.completion",