Update api/utils.py

api/utils.py

@@ -4,16 +4,16 @@ import uuid
 import asyncio
 import random
 import string
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, List
 
 import httpx
 from fastapi import HTTPException
 from api.config import (
     MODEL_MAPPING,
+    AGENT_MODE,
+    TRENDING_AGENT_MODE,
     get_headers,
     BASE_URL,
-    AGENT_MODE,
-    TRENDING_AGENT_MODE
 )
 from api.models import ChatRequest
 from api.logger import setup_logger
@@ -44,9 +44,11 @@ def create_chat_completion_data(
         "usage": None,
     }
 
-# Function to convert message to dictionary format
-def message_to_dict(message):
+# Function to convert message to dictionary format for API
+def message_to_dict(message, model_prefix: Optional[str] = None) -> Dict[str, Any]:
     content = message.content if isinstance(message.content, str) else message.content[0]["text"]
+    if model_prefix:
+        content = f"{model_prefix} {content}"
     if isinstance(message.content, list) and len(message.content) == 2 and "image_url" in message.content[1]:
         return {
             "role": message.role,
@@ -59,16 +61,25 @@ def message_to_dict(message):
         }
     return {"role": message.role, "content": content}
 
-# Process streaming response
+# Function to get agent modes for specific models
+def get_agent_modes(model: str) -> Dict[str, Any]:
+    return {
+        "agentMode": AGENT_MODE.get(model, {}),
+        "trendingAgentMode": TRENDING_AGENT_MODE.get(model, {}),
+    }
+
+# Process streaming response with headers from config.py
 async def process_streaming_response(request: ChatRequest):
     chat_id = generate_chat_id()
-    agent_mode = AGENT_MODE.get(request.model, {})
-    trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
+    logger.info(f"Generated Chat ID: {chat_id} - Model: {request.model}")
 
+    # Retrieve agent modes based on the model
+    agent_modes = get_agent_modes(request.model)
     headers = get_headers()
 
     json_data = {
-        "agentMode": agent_mode,
+        "agentMode": agent_modes["agentMode"],
+        "trendingAgentMode": agent_modes["trendingAgentMode"],
         "clickedAnswer2": False,
         "clickedAnswer3": False,
         "clickedForceWebSearch": False,
@@ -83,7 +94,6 @@ async def process_streaming_response(request: ChatRequest):
         "playgroundTemperature": request.temperature,
         "playgroundTopP": request.top_p,
         "previewToken": None,
-        "trendingAgentMode": trending_agent_mode,
         "userId": None,
         "userSelectedModel": MODEL_MAPPING.get(request.model, request.model),
         "userSystemPrompt": None,
@@ -91,9 +101,16 @@ async def process_streaming_response(request: ChatRequest):
         "visitFromDelta": False,
     }
 
+    logger.debug(f"Data Payload: {json_data}")  # Debugging line to inspect payload
     async with httpx.AsyncClient() as client:
         try:
-            async with client.stream("POST", f"{BASE_URL}/api/chat", headers=headers, json=json_data, timeout=100) as response:
+            async with client.stream(
+                "POST",
+                f"{BASE_URL}/api/chat",
+                headers=headers,
+                json=json_data,
+                timeout=100,
+            ) as response:
                 response.raise_for_status()
                 async for line in response.aiter_lines():
                     timestamp = int(datetime.now().timestamp())
@@ -108,16 +125,18 @@ async def process_streaming_response(request: ChatRequest):
             logger.error(f"Error occurred during request for Chat ID {chat_id}: {e}")
             raise HTTPException(status_code=500, detail=str(e))
 
-# Process non-streaming response
+# Process non-streaming response with headers from config.py
 async def process_non_streaming_response(request: ChatRequest):
     chat_id = generate_chat_id()
-    agent_mode = AGENT_MODE.get(request.model, {})
-    trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
+    logger.info(f"Generated Chat ID: {chat_id} - Model: {request.model}")
 
+    # Retrieve agent modes based on the model
+    agent_modes = get_agent_modes(request.model)
     headers = get_headers()
 
     json_data = {
-        "agentMode": agent_mode,
+        "agentMode": agent_modes["agentMode"],
+        "trendingAgentMode": agent_modes["trendingAgentMode"],
         "clickedAnswer2": False,
         "clickedAnswer3": False,
         "clickedForceWebSearch": False,
@@ -132,7 +151,6 @@ async def process_non_streaming_response(request: ChatRequest):
         "playgroundTemperature": request.temperature,
         "playgroundTopP": request.top_p,
         "previewToken": None,
-        "trendingAgentMode": trending_agent_mode,
         "userId": None,
         "userSelectedModel": MODEL_MAPPING.get(request.model, request.model),
         "userSystemPrompt": None,
@@ -140,10 +158,13 @@ async def process_non_streaming_response(request: ChatRequest):
         "visitFromDelta": False,
     }
 
+    logger.debug(f"Data Payload: {json_data}")  # Debugging line to inspect payload
     full_response = ""
     async with httpx.AsyncClient() as client:
         try:
-            async with client.stream("POST", f"{BASE_URL}/api/chat", headers=headers, json=json_data) as response:
+            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