Update api/utils.py

api/utils.py

@@ -1,22 +1,30 @@
 from datetime import datetime
-from http.client import HTTPException
 import json
 from typing import Any, Dict, Optional
 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 fastapi import Depends, HTTPException
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 
-from api.config import APP_SECRET, BASE_URL
+from api import validate
+from api.config import (
+    MODEL_MAPPING,
+    headers,
+    AGENT_MODE,
+    TRENDING_AGENT_MODE,
+    APP_SECRET,
+    BASE_URL,
+)
 from api.models import ChatRequest
 
 from api.logger import setup_logger
 
 logger = setup_logger(__name__)
 
+# Initialize the HTTPBearer security scheme
+bearer_scheme = HTTPBearer()
+
 def create_chat_completion_data(
     content: str, model: str, timestamp: int, finish_reason: Optional[str] = None
 ) -> Dict[str, Any]:
@@ -35,26 +43,46 @@ def create_chat_completion_data(
         "usage": None,
     }
 
-def verify_app_secret(credentials: HTTPAuthorizationCredentials = Depends(security)):
+def verify_app_secret(credentials: HTTPAuthorizationCredentials = Depends(bearer_scheme)):
     if credentials.credentials != APP_SECRET:
+        logger.warning("Invalid APP_SECRET provided.")
         raise HTTPException(status_code=403, detail="Invalid APP_SECRET")
+    logger.debug("APP_SECRET verified successfully.")
     return credentials.credentials
 
 def message_to_dict(message):
+    """
+    Convert a message object to a dictionary suitable for the API request.
+    Handles different content types gracefully.
+    """
+    message_dict = {"role": message.role}
+
     if isinstance(message.content, str):
-        return {"role": message.role, "content": 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",
-            },
-        }
+        message_dict["content"] = message.content
+    elif isinstance(message.content, list):
+        # Handle list content more robustly
+        try:
+            if len(message.content) >= 2:
+                # Assuming the first element has 'text' and the second has 'image_url'
+                text_content = message.content[0].get("text", "")
+                image_url = message.content[1].get("image_url", {}).get("url", "")
+                message_dict["content"] = text_content
+                message_dict["data"] = {
+                    "imageBase64": image_url,
+                    "fileText": "",
+                    "title": "snapshot",
+                }
+            else:
+                # Fallback if the list doesn't have expected structure
+                message_dict["content"] = json.dumps(message.content)
+        except (AttributeError, KeyError, TypeError) as e:
+            logger.error(f"Error parsing message content: {e}")
+            message_dict["content"] = "Invalid message format."
     else:
-        return {"role": message.role, "content": message.content}
+        # Fallback for unexpected content types
+        message_dict["content"] = str(message.content)
+
+    return message_dict
 
 async def process_streaming_response(request: ChatRequest):
     agent_mode = AGENT_MODE.get(request.model, {})
@@ -85,8 +113,8 @@ async def process_streaming_response(request: ChatRequest):
         "clickedForceWebSearch": False,
         "visitFromDelta": False,
         "mobileClient": False,
-        "userSelectedModel": MODEL_MAPPING.get(request.model, request.model),  # Added default
-        "validated": validate.getHid()
+        "userSelectedModel": MODEL_MAPPING.get(request.model),
+        "validated": validate.getHid(),
     }
 
     async with httpx.AsyncClient() as client:
@@ -96,31 +124,31 @@ async def process_streaming_response(request: ChatRequest):
                 f"{BASE_URL}/api/chat",
                 headers=headers,
                 json=json_data,
-                timeout=100,
+                timeout=httpx.Timeout(100.0),
             ) as response:
                 response.raise_for_status()
+                timestamp = int(datetime.now().timestamp())
                 async for line in response.aiter_lines():
-                    timestamp = int(datetime.now().timestamp())
                     if line:
-                        content = line + "\n"
+                        content = line.strip() + "\n"
                         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"
                             break
                         if content.startswith("$@$v=undefined-rv1$@$"):
-                            yield f"data: {json.dumps(create_chat_completion_data(content[21:], request.model, timestamp))}\n\n"
-                        else:
-                            yield f"data: {json.dumps(create_chat_completion_data(content, request.model, timestamp))}\n\n"
+                            content = content[21:]
+                        yield f"data: {json.dumps(create_chat_completion_data(content, request.model, timestamp))}\n\n"
 
+                # Indicate the end of the stream
                 yield f"data: {json.dumps(create_chat_completion_data('', request.model, timestamp, 'stop'))}\n\n"
                 yield "data: [DONE]\n\n"
         except httpx.HTTPStatusError as e:
-            logger.error(f"HTTP error occurred: {e}")
-            raise HTTPException(status_code=e.response.status_code, detail=str(e))
+            logger.error(f"HTTP error occurred: {e.response.status_code} - {e.response.text}")
+            raise HTTPException(status_code=e.response.status_code, detail="Error from upstream service.")
         except httpx.RequestError as e:
-            logger.error(f"Error occurred during request: {e}")
-            raise HTTPException(status_code=500, detail=str(e))
+            logger.error(f"Request error occurred: {e}")
+            raise HTTPException(status_code=500, detail="Internal server error.")
 
 async def process_non_streaming_response(request: ChatRequest):
     agent_mode = AGENT_MODE.get(request.model, {})
@@ -151,41 +179,48 @@ async def process_non_streaming_response(request: ChatRequest):
         "clickedForceWebSearch": False,
         "visitFromDelta": False,
         "mobileClient": False,
-        "userSelectedModel": MODEL_MAPPING.get(request.model, request.model),  # Added default
-        "validated": validate.getHid()
+        "userSelectedModel": MODEL_MAPPING.get(request.model),
+        "validated": validate.getHid(),
     }
 
-    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已刷新，重新对话即可"
-    if full_response.startswith("$@$v=undefined-rv1$@$"):
-        full_response = full_response[21:]
-    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,
-    }
+    try:
+        async with httpx.AsyncClient() as client:
+            response = await client.post(
+                f"{BASE_URL}/api/chat",
+                headers=headers,
+                json=json_data,
+                timeout=httpx.Timeout(100.0),
+            )
+            response.raise_for_status()
+            full_response = response.text
+
+        # Process the full response
+        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:]
+
+        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,
+        }
+    except httpx.HTTPStatusError as e:
+        logger.error(f"HTTP error occurred: {e.response.status_code} - {e.response.text}")
+        raise HTTPException(status_code=e.response.status_code, detail="Error from upstream service.")
+    except httpx.RequestError as e:
+        logger.error(f"Request error occurred: {e}")
+        raise HTTPException(status_code=500, detail="Internal server error.")
+    except Exception as e:
+        logger.error(f"Unexpected error: {e}")
+        raise HTTPException(status_code=500, detail="Internal server error.")