Update api/utils.py

api/utils.py

@@ -5,34 +5,21 @@ import uuid
 import re
 
 import httpx
-from api.config import MODEL_MAPPING, headers, AGENT_MODE, TRENDING_AGENT_MODE, BASE_URL
+from api.config import (
+    MODEL_MAPPING,
+    MODEL_ALIASES,
+    headers,
+    AGENT_MODE,
+    TRENDING_AGENT_MODE,
+    BASE_URL
+)
 from fastapi import HTTPException
-from api.models import ChatRequest, ImageResponse
+from api.models import ChatRequest
 
 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
@@ -59,9 +46,9 @@ def message_to_dict(message):
     elif isinstance(message.content, list) and len(message.content) == 2:
         return {
             "role": message.role,
-            "content": message.content[0].get("text", ""),
+            "content": message.content[0]["text"],
             "data": {
-                "imageBase64": message.content[1].get("imageBase64", ""),
+                "imageBase64": message.content[1]["image_url"]["url"],
                 "fileText": "",
                 "title": "snapshot",
             },
@@ -71,29 +58,16 @@ def message_to_dict(message):
 
 
 async def process_streaming_response(request: ChatRequest):
-    agent_mode = AGENT_MODE.get(request.model, {})
-    trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
-    model = MODEL_MAPPING.get(request.model, request.model)
-    
-    # Add prefix if applicable
-    prefix = model_prefixes.get(request.model, "")
-    
-    # Format messages with prefix
-    formatted_messages = []
-    for msg in request.messages:
-        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": data
-        })
+    # Map the requested model to the actual model used by the API
+    model = MODEL_MAPPING.get(request.model, MODEL_ALIASES.get(request.model, "blackboxai"))
+
+    logger.info(f"Using model: {model}")
+
+    agent_mode = AGENT_MODE.get(model, {})
+    trending_agent_mode = TRENDING_AGENT_MODE.get(model, {})
 
     json_data = {
-        "messages": formatted_messages,
+        "messages": [message_to_dict(msg) for msg in request.messages],
         "previewToken": None,
         "userId": None,
         "codeModelMode": True,
@@ -111,10 +85,11 @@ async def process_streaming_response(request: ChatRequest):
         "clickedForceWebSearch": False,
         "visitFromDelta": False,
         "mobileClient": False,
-        "webSearchMode": False,  # Set to True if web search is needed
         "userSelectedModel": model,
     }
 
+    logger.debug(f"Payload for streaming request: {json.dumps(json_data)}")
+
     async with httpx.AsyncClient() as client:
         try:
             async with client.stream(
@@ -131,17 +106,11 @@ async def process_streaming_response(request: ChatRequest):
                         content = line
                         if content.startswith("$@$v=undefined-rv1$@$"):
                             content = content[21:]
-                        
-                        # 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 f"data: {json.dumps(create_chat_completion_data(content, model, timestamp))}\n\n"
+
+                # Indicate the end of the stream
+                timestamp = int(datetime.now().timestamp())
+                yield f"data: {json.dumps(create_chat_completion_data('', model, timestamp, 'stop'))}\n\n"
                 yield "data: [DONE]\n\n"
         except httpx.HTTPStatusError as e:
             logger.error(f"HTTP error occurred: {e}")
@@ -152,29 +121,16 @@ async def process_streaming_response(request: ChatRequest):
 
 
 async def process_non_streaming_response(request: ChatRequest):
-    agent_mode = AGENT_MODE.get(request.model, {})
-    trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
-    model = MODEL_MAPPING.get(request.model, request.model)
-    
-    # Add prefix if applicable
-    prefix = model_prefixes.get(request.model, "")
-    
-    # Format messages with prefix
-    formatted_messages = []
-    for msg in request.messages:
-        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": data
-        })
+    # Map the requested model to the actual model used by the API
+    model = MODEL_MAPPING.get(request.model, MODEL_ALIASES.get(request.model, "blackboxai"))
+
+    logger.info(f"Using model: {model}")
+
+    agent_mode = AGENT_MODE.get(model, {})
+    trending_agent_mode = TRENDING_AGENT_MODE.get(model, {})
 
     json_data = {
-        "messages": formatted_messages,
+        "messages": [message_to_dict(msg) for msg in request.messages],
         "previewToken": None,
         "userId": None,
         "codeModelMode": True,
@@ -192,20 +148,23 @@ async def process_non_streaming_response(request: ChatRequest):
         "clickedForceWebSearch": False,
         "visitFromDelta": False,
         "mobileClient": False,
-        "webSearchMode": False,  # Set to True if web search is needed
         "userSelectedModel": model,
     }
+
+    logger.debug(f"Payload for non-streaming request: {json.dumps(json_data)}")
+
     full_response = ""
     async with httpx.AsyncClient() as client:
         try:
-            async with client.post(
-                f"{BASE_URL}/api/chat",
+            async with client.stream(
+                method="POST",
+                url=f"{BASE_URL}/api/chat",
                 headers=headers,
                 json=json_data,
-                timeout=100,
             ) as response:
                 response.raise_for_status()
-                full_response = await response.text()
+                async for chunk in response.aiter_text():
+                    full_response += chunk
         except httpx.HTTPStatusError as e:
             logger.error(f"HTTP error occurred: {e}")
             raise HTTPException(status_code=e.response.status_code, detail=str(e))
@@ -216,32 +175,13 @@ 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
-        }
+    logger.debug(f"Full non-streaming response: {full_response}")
 
     return {
         "id": f"chatcmpl-{uuid.uuid4()}",
         "object": "chat.completion",
         "created": int(datetime.now().timestamp()),
-        "model": request.model,
+        "model": model,
         "choices": [
             {
                 "index": 0,