Update api/utils.py

api/utils.py

@@ -25,9 +25,13 @@ 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]:
+    """
+    Create a dictionary representing a chat completion chunk.
+    """
     return {
         "id": f"chatcmpl-{uuid.uuid4()}",
         "object": "chat.completion.chunk",
@@ -43,13 +47,18 @@ def create_chat_completion_data(
         "usage": None,
     }
 
+
 def verify_app_secret(credentials: HTTPAuthorizationCredentials = Depends(bearer_scheme)):
+    """
+    Verify the APP_SECRET from the authorization credentials.
+    """
     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.
@@ -84,9 +93,35 @@ def message_to_dict(message):
 
     return message_dict
 
+
+def get_agent_mode(model: str) -> Dict[str, Any]:
+    """
+    Retrieves the agent mode configuration.
+    Falls back to an empty dictionary if not found.
+    """
+    agent_mode = AGENT_MODE.get(model, {})
+    if not agent_mode:
+        logger.warning(f"No AGENT_MODE configuration found for model: {model}")
+    return agent_mode
+
+
+def get_trending_agent_mode(model: str) -> Dict[str, Any]:
+    """
+    Retrieves the trending agent mode configuration.
+    Falls back to an empty dictionary if not found.
+    """
+    trending_agent_mode = TRENDING_AGENT_MODE.get(model, {})
+    if not trending_agent_mode:
+        logger.warning(f"No TRENDING_AGENT_MODE configuration found for model: {model}")
+    return trending_agent_mode
+
+
 async def process_streaming_response(request: ChatRequest):
-    agent_mode = AGENT_MODE.get(request.model, {})
-    trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
+    """
+    Process a streaming response for a chat completion request.
+    """
+    agent_mode = get_agent_mode(request.model)
+    trending_agent_mode = get_trending_agent_mode(request.model)
 
     # Log reduced information
     logger.info(
@@ -149,10 +184,17 @@ async def process_streaming_response(request: ChatRequest):
         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.")
+
 
 async def process_non_streaming_response(request: ChatRequest):
-    agent_mode = AGENT_MODE.get(request.model, {})
-    trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
+    """
+    Process a non-streaming response for a chat completion request.
+    """
+    agent_mode = get_agent_mode(request.model)
+    trending_agent_mode = get_trending_agent_mode(request.model)
 
     # Log reduced information
     logger.info(