Update api/utils.py

api/utils.py

@@ -13,7 +13,8 @@ from api.config import (
     get_headers_chat,
     BASE_URL,
     AGENT_MODE,
-    TRENDING_AGENT_MODE
+    TRENDING_AGENT_MODE,
+    ALLOWED_MODELS
 )
 from api.models import ChatRequest
 from api.logger import setup_logger
@@ -70,6 +71,9 @@ async def process_streaming_response(request: ChatRequest):
     referer_url = f"{BASE_URL}/chat/{chat_id}?model={request.model}"
     logger.info(f"Generated Chat ID: {chat_id} - Model: {request.model} - URL: {referer_url}")
 
+    if request.model not in ALLOWED_MODELS:
+        raise HTTPException(status_code=400, detail="Invalid model specified.")
+
     agent_mode = AGENT_MODE.get(request.model, {})
     trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
 
@@ -138,6 +142,9 @@ async def process_non_streaming_response(request: ChatRequest):
     referer_url = f"{BASE_URL}/chat/{chat_id}?model={request.model}"
     logger.info(f"Generated Chat ID: {chat_id} - Model: {request.model} - URL: {referer_url}")
 
+    if request.model not in ALLOWED_MODELS:
+        raise HTTPException(status_code=400, detail="Invalid model specified.")
+
     agent_mode = AGENT_MODE.get(request.model, {})
     trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})