Update api/utils.py

api/utils.py

@@ -6,29 +6,40 @@ import uuid
 
 import httpx
 from api import validate
-from api.config import MODEL_MAPPING, headers, AGENT_MODE, TRENDING_AGENT_MODE
-from fastapi import Depends, security
-from fastapi.security import HTTPAuthorizationCredentials
+from api.config import (
+    MODEL_MAPPING,
+    headers,
+    AGENT_MODE,
+    TRENDING_AGENT_MODE,
+    MODEL_PREFIXES,
+    APP_SECRET,
+    BASE_URL,
+)
+from fastapi import Depends, HTTPException
+from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
 
-from api.config import APP_SECRET, BASE_URL
 from api.models import ChatRequest
-
 from api.logger import setup_logger
 
 logger = setup_logger(__name__)
 
+# Initialize HTTPBearer for security
+security = HTTPBearer()
+
 def create_chat_completion_data(
     content: str, model: str, timestamp: int, finish_reason: Optional[str] = None
 ) -> Dict[str, Any]:
+    prefix = MODEL_PREFIXES.get(model, "")
+    full_content = f"{prefix} {content}" if prefix else content
     return {
         "id": f"chatcmpl-{uuid.uuid4()}",
         "object": "chat.completion.chunk",
         "created": timestamp,
-        "model": model,
+        "model": f"{prefix}{model}" if prefix else model,
         "choices": [
             {
                 "index": 0,
-                "delta": {"content": content, "role": "assistant"},
+                "delta": {"content": full_content, "role": "assistant"},
                 "finish_reason": finish_reason,
             }
         ],
@@ -41,28 +52,34 @@ def verify_app_secret(credentials: HTTPAuthorizationCredentials = Depends(securi
     return credentials.credentials
 
 def message_to_dict(message):
+    prefix = MODEL_PREFIXES.get(message.model, "") if hasattr(message, 'model') else ""
+    
     if isinstance(message.content, str):
-        return {"role": message.role, "content": message.content}
+        content = f"{prefix} {message.content}" if prefix else message.content
     elif isinstance(message.content, list) and len(message.content) == 2:
-        return {
-            "role": message.role,
-            "content": message.content[0]["text"],
-            "data": {
-                "imageBase64": message.content[1]["image_url"]["url"],
-                "fileText": "",
-                "title": "snapshot",
-            },
-        }
+        content = f"{prefix} {message.content[0]['text']}" if prefix else message.content[0]['text']
     else:
-        return {"role": message.role, "content": message.content}
+        content = message.content
+
+    message_dict = {"role": message.role, "content": content}
+
+    if isinstance(message.content, list) and len(message.content) == 2:
+        message_dict["data"] = {
+            "imageBase64": message.content[1]["image_url"]["url"],
+            "fileText": "",
+            "title": "snapshot",
+        }
+
+    return message_dict
 
 async def process_streaming_response(request: ChatRequest):
     agent_mode = AGENT_MODE.get(request.model, {})
     trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
+    model_prefix = MODEL_PREFIXES.get(request.model, "")
 
-    # Log reduced information
+    # Log reduced information with prefix
     logger.info(
-        f"Streaming request for model: '{request.model}', "
+        f"Streaming request for model: '{model_prefix}{request.model}', "
         f"agent mode: {agent_mode}, trending agent mode: {trending_agent_mode}"
     )
 
@@ -86,7 +103,7 @@ async def process_streaming_response(request: ChatRequest):
         "visitFromDelta": False,
         "mobileClient": False,
         "userSelectedModel": MODEL_MAPPING.get(request.model),
-        "validated": validate.getHid()
+        "validated": validate.getHid(),
     }
 
     async with httpx.AsyncClient() as client:
@@ -125,10 +142,11 @@ 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_prefix = MODEL_PREFIXES.get(request.model, "")
 
-    # Log reduced information
+    # Log reduced information with prefix
     logger.info(
-        f"Non-streaming request for model: '{request.model}', "
+        f"Non-streaming request for model: '{model_prefix}{request.model}', "
         f"agent mode: {agent_mode}, trending agent mode: {trending_agent_mode}"
     )
 
@@ -152,7 +170,7 @@ async def process_non_streaming_response(request: ChatRequest):
         "visitFromDelta": False,
         "mobileClient": False,
         "userSelectedModel": MODEL_MAPPING.get(request.model),
-        "validated": validate.getHid()
+        "validated": validate.getHid(),
     }
 
     full_response = ""
@@ -162,16 +180,23 @@ async def process_non_streaming_response(request: ChatRequest):
         ) as response:
             async for chunk in response.aiter_text():
                 full_response += chunk
+
     if "https://www.blackbox.ai" in full_response:
         validate.getHid(True)
         full_response = "hid已刷新，重新对话即可"
+
     if full_response.startswith("$@$v=undefined-rv1$@$"):
         full_response = full_response[21:]
+
+    # Prepend model prefix to the final response content
+    if model_prefix:
+        full_response = f"{model_prefix} {full_response}"
+
     return {
         "id": f"chatcmpl-{uuid.uuid4()}",
         "object": "chat.completion",
         "created": int(datetime.now().timestamp()),
-        "model": request.model,
+        "model": f"{model_prefix}{request.model}" if model_prefix else request.model,
         "choices": [
             {
                 "index": 0,