Update api/utils.py

api/utils.py

@@ -1,7 +1,8 @@
+# utils.py
+
 from datetime import datetime
 import json
 import uuid
-import asyncio
 import random
 from typing import Any, Dict, Optional
 
@@ -14,6 +15,7 @@ from api.config import (
     TRENDING_AGENT_MODE,
     BASE_URL,
     generate_id,
+    USER_SELECTED_MODELS,
 )
 from api.models import ChatRequest
 from api.logger import setup_logger
@@ -50,10 +52,17 @@ def strip_model_prefix(content: str, model_prefix: Optional[str] = None) -> str:
 
 # 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 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]:
+    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 {
             "role": message.role,
@@ -105,7 +114,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,
@@ -134,7 +145,9 @@ async def process_streaming_response(request: ChatRequest):
                         content = line
                         if "https://www.blackbox.ai" in content:
                             getHid(True)
-                            content = "HID has been refreshed, please start a new conversation.\n"
+                            content = (
+                                "HID has been refreshed, please start a new conversation.\n"
+                            )
                             yield f"data: {json.dumps(create_chat_completion_data(content, request.model, timestamp))}\n\n"
                             break
                         if content.startswith("$@$v=undefined-rv1$@$"):
@@ -190,7 +203,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,
@@ -207,7 +222,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():