Update api/utils.py

api/utils.py

@@ -6,7 +6,13 @@ import uuid
 
 import httpx
 from api import validate
-from api.config import MODEL_MAPPING, headers, AGENT_MODE, TRENDING_AGENT_MODE
+from api.config import (
+    MODEL_MAPPING,
+    MODEL_PREFIXES,  # Import MODEL_PREFIXES
+    headers,
+    AGENT_MODE,
+    TRENDING_AGENT_MODE
+)
 from fastapi import Depends, security
 from fastapi.security import HTTPAuthorizationCredentials
 
@@ -56,6 +62,10 @@ def message_to_dict(message):
     else:
         return {"role": message.role, "content": message.content}
 
+def get_model_prefix(model: str) -> str:
+    """Retrieve the prefix for a given model."""
+    return MODEL_PREFIXES.get(model, "")
+
 async def process_streaming_response(request: ChatRequest):
     agent_mode = AGENT_MODE.get(request.model, {})
     trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
@@ -89,6 +99,11 @@ async def process_streaming_response(request: ChatRequest):
         "validated": validate.getHid()
     }
 
+    # Get the model prefix
+    prefix = get_model_prefix(request.model)
+    if prefix:
+        prefix += " "  # Add space after prefix
+
     async with httpx.AsyncClient() as client:
         try:
             async with client.stream(
@@ -106,14 +121,25 @@ async def process_streaming_response(request: ChatRequest):
                         if "https://www.blackbox.ai" in content:
                             validate.getHid(True)
                             content = "hid已刷新,重新对话即可\n"
-                            yield f"data: {json.dumps(create_chat_completion_data(content, request.model, timestamp))}\n\n"
+                            # Prepend prefix to the content
+                            prefixed_content = prefix + content if prefix else content
+                            yield f"data: {json.dumps(create_chat_completion_data(prefixed_content, request.model, timestamp))}\n\n"
                             break
                         if content.startswith("$@$v=undefined-rv1$@$"):
-                            yield f"data: {json.dumps(create_chat_completion_data(content[21:], request.model, timestamp))}\n\n"
+                            # Remove the special prefix and prepend model prefix
+                            actual_content = content[21:]
+                            prefixed_content = prefix + actual_content if prefix else actual_content
+                            yield f"data: {json.dumps(create_chat_completion_data(prefixed_content, request.model, timestamp))}\n\n"
                         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"
+                            # Prepend model prefix to the content
+                            prefixed_content = prefix + content if prefix else content
+                            yield f"data: {json.dumps(create_chat_completion_data(prefixed_content, request.model, timestamp))}\n\n"
+
+                # Indicate the end of the streaming response with the prefix
+                finish_content = ''
+                if prefix:
+                    finish_content = prefix
+                yield f"data: {json.dumps(create_chat_completion_data(finish_content, request.model, timestamp, 'stop'))}\n\n"
                 yield "data: [DONE]\n\n"
         except httpx.HTTPStatusError as e:
             logger.error(f"HTTP error occurred: {e}")
@@ -155,6 +181,11 @@ async def process_non_streaming_response(request: ChatRequest):
         "validated": validate.getHid()
     }
 
+    # Get the model prefix
+    prefix = get_model_prefix(request.model)
+    if prefix:
+        prefix += " "  # Add space after prefix
+
     full_response = ""
     async with httpx.AsyncClient() as client:
         async with client.stream(
@@ -162,11 +193,18 @@ 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 full response
+    if prefix:
+        full_response = prefix + full_response
+
     return {
         "id": f"chatcmpl-{uuid.uuid4()}",
         "object": "chat.completion",
@@ -180,4 +218,4 @@ async def process_non_streaming_response(request: ChatRequest):
             }
         ],
         "usage": None,
-    }
+    }