Update api/utils.py

api/utils.py

@@ -68,16 +68,13 @@ def message_to_dict(message, model_prefix: Optional[str] = None):
         return {"role": message.role, "content": content}
 
 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, "")
+    model = request.model  # Use the requested model directly
+    agent_mode = AGENT_MODE.get(model, {})
+    trending_agent_mode = TRENDING_AGENT_MODE.get(model, {})
+    model_prefix = MODEL_PREFIXES.get(model, "")
 
     hid = validate.getHid()
-    logger.info(f"Using hid: {hid} for model: {request.model}")
-
-    user_selected_model = MODEL_MAPPING.get(request.model, request.model)
-    logger.info(f"Processing request for model: {request.model}")
-    logger.info(f"Using userSelectedModel: {user_selected_model}")
+    logger.info(f"Using hid: {hid} for model: {model}")
 
     json_data = {
         "messages": [message_to_dict(msg, model_prefix=model_prefix) for msg in request.messages],
@@ -98,7 +95,7 @@ async def process_streaming_response(request: ChatRequest):
         "clickedForceWebSearch": False,
         "visitFromDelta": False,
         "mobileClient": False,
-        "userSelectedModel": user_selected_model,
+        "userSelectedModel": model,  # Use the requested model
         "validated": hid
     }
 
@@ -120,13 +117,14 @@ async def process_streaming_response(request: ChatRequest):
                             validate.getHid(True)
                             content = "hid已刷新,重新对话即可"
                             logger.info(f"hid refreshed due to content: {content}")
-                            yield f"data: {json.dumps(create_chat_completion_data(content, request.model, timestamp))}\n\n"
+                            yield f"data: {json.dumps(create_chat_completion_data(content, model, timestamp))}\n\n"
                             break
                         if content.startswith("$@$v=undefined-rv1$@$"):
                             content = content[21:]
-                        yield f"data: {json.dumps(create_chat_completion_data(content, request.model, timestamp))}\n\n"
+                        # Do not strip model prefix or modify content
+                        yield f"data: {json.dumps(create_chat_completion_data(content, 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('', model, timestamp, 'stop'))}\n\n"
                 yield "data: [DONE]\n\n"
         except httpx.HTTPStatusError as e:
             logger.error(f"HTTP error occurred: {e}")
@@ -135,4 +133,72 @@ async def process_streaming_response(request: ChatRequest):
             logger.error(f"Error occurred during request: {e}")
             raise HTTPException(status_code=500, detail=str(e))
 
-# Similar updates for process_non_streaming_response
+async def process_non_streaming_response(request: ChatRequest):
+    model = request.model  # Use the requested model directly
+    agent_mode = AGENT_MODE.get(model, {})
+    trending_agent_mode = TRENDING_AGENT_MODE.get(model, {})
+    model_prefix = MODEL_PREFIXES.get(model, "")
+
+    hid = validate.getHid()
+    logger.info(f"Using hid: {hid} for model: {model}")
+
+    json_data = {
+        "messages": [message_to_dict(msg, model_prefix=model_prefix) for msg in request.messages],
+        "previewToken": None,
+        "userId": None,
+        "codeModelMode": True,
+        "agentMode": agent_mode,
+        "trendingAgentMode": trending_agent_mode,
+        "isMicMode": False,
+        "userSystemPrompt": None,
+        "maxTokens": request.max_tokens,
+        "playgroundTopP": request.top_p,
+        "playgroundTemperature": request.temperature,
+        "isChromeExt": False,
+        "githubToken": None,
+        "clickedAnswer2": False,
+        "clickedAnswer3": False,
+        "clickedForceWebSearch": False,
+        "visitFromDelta": False,
+        "mobileClient": False,
+        "userSelectedModel": model,  # Use the requested model
+        "validated": hid
+    }
+    full_response = ""
+    async with httpx.AsyncClient() as client:
+        try:
+            async with client.stream(
+                method="POST", url=f"{BASE_URL}/api/chat", headers=headers, json=json_data
+            ) as response:
+                response.raise_for_status()
+                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))
+        except httpx.RequestError as e:
+            logger.error(f"Error occurred during request: {e}")
+            raise HTTPException(status_code=500, detail=str(e))
+
+    if "https://www.blackbox.ai" in full_response:
+        validate.getHid(True)
+        full_response = "hid已刷新，重新对话即可"
+        logger.info("hid refreshed due to response content")
+
+    if full_response.startswith("$@$v=undefined-rv1$@$"):
+        full_response = full_response[21:]
+    # Do not strip model prefix or modify content
+    return {
+        "id": f"chatcmpl-{uuid.uuid4()}",
+        "object": "chat.completion",
+        "created": int(datetime.now().timestamp()),
+        "model": model,
+        "choices": [
+            {
+                "index": 0,
+                "message": {"role": "assistant", "content": full_response},
+                "finish_reason": "stop",
+            }
+        ],
+        "usage": None,
+    }