Update api/utils.py

api/utils.py

@@ -3,10 +3,10 @@ import json
 from typing import Any, Dict, Optional
 
 import httpx
-import aiohttp  # Add this import
+import aiohttp
 from api.config import (
-    DDG_API_ENDPOINT,  # Updated
-    DDG_STATUS_URL,    # Updated
+    DDG_API_ENDPOINT,
+    DDG_STATUS_URL,
     MODEL_MAPPING,
     ALLOWED_MODELS,
     AGENT_MODE,
@@ -23,9 +23,6 @@ import uuid
 
 logger = setup_logger(__name__)
 
-# Remove Blackbox-specific functions and add DDG-specific functions
-
-# New DDG Integration
 async def get_ddg_vqd():
     status_url = DDG_STATUS_URL
 
@@ -39,7 +36,10 @@ async def get_ddg_vqd():
         try:
             async with session.get(status_url, headers=headers) as response:
                 if response.status == 200:
-                    return response.headers.get("x-vqd-4")
+                    vqd = response.headers.get("x-vqd-4")
+                    if not vqd:
+                        logger.error("VQD token not found in response headers.")
+                    return vqd
                 else:
                     logger.error(f"Error: Status code {response.status} when fetching VQD")
                     return None
@@ -47,31 +47,12 @@ async def get_ddg_vqd():
             logger.error(f"Error getting VQD: {e}")
             return None
 
-def format_prompt(messages):
-    """Helper function to format messages into a prompt string."""
-    return "\n".join([f"{msg.role}: {msg.content}" for msg in messages])
-
-def message_to_dict_ddg(message, model_prefix: Optional[str] = None):
+def message_to_dict_ddg(message):
     if isinstance(message.content, str):
-        content = message.content
-        if model_prefix:
-            content = f"{model_prefix} {content}"
-        return {"role": message.role, "content": content}
-    elif isinstance(message.content, list) and len(message.content) == 2:
-        content = message.content[0]["text"]
-        if model_prefix:
-            content = f"{model_prefix} {content}"
-        return {
-            "role": message.role,
-            "content": content,
-            "data": {
-                "imageBase64": message.content[1]["image_url"]["url"],
-                "fileText": "",
-                "title": "snapshot",
-            },
-        }
-    else:
         return {"role": message.role, "content": message.content}
+    else:
+        # If DDG API does not support non-string content, raise an error
+        raise ValueError("Message content must be a string.")
 
 def strip_model_prefix(content: str, model_prefix: Optional[str] = None) -> str:
     """Remove the model prefix from the response content if present."""
@@ -81,28 +62,53 @@ def strip_model_prefix(content: str, model_prefix: Optional[str] = None) -> str:
     logger.debug("No prefix to strip from content.")
     return content
 
+def create_chat_completion_data(
+    content: str, model: str, timestamp: int, finish_reason: Optional[str] = None
+) -> Dict[str, Any]:
+    return {
+        "id": f"chatcmpl-{uuid.uuid4()}",
+        "object": "chat.completion.chunk",
+        "created": timestamp,
+        "model": model,
+        "choices": [
+            {
+                "index": 0,
+                "delta": {"content": content, "role": "assistant"},
+                "finish_reason": finish_reason,
+            }
+        ],
+        "usage": None,
+    }
+
 async def process_ddg_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, "")
-    referer_path = MODEL_REFERERS.get(request.model, f"/?model={request.model}")
-    referer_url = f"https://duckduckgo.com{referer_path}"  # Updated BASE_URL
+    # referer_path and referer_url are not required unless specified by DDG API
+
+    # Obtain VQD token
+    vqd = await get_ddg_vqd()
+    if not vqd:
+        raise HTTPException(status_code=500, detail="Failed to obtain VQD token")
 
-    # Update headers with dynamic Referer
+    # Build headers
     dynamic_headers = {
         'accept': 'text/event-stream',
         'content-type': 'application/json',
         'user-agent': 'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/128.0.0.0 Safari/537.36',
-        'x-vqd-4': await get_ddg_vqd()
+        'x-vqd-4': vqd
     }
 
-    if not dynamic_headers['x-vqd-4']:
-        raise HTTPException(status_code=500, detail="Failed to obtain VQD token")
+    # Construct message history
+    try:
+        message_history = [message_to_dict_ddg(msg) for msg in request.messages]
+    except ValueError as ve:
+        logger.error(f"Invalid message format: {ve}")
+        raise HTTPException(status_code=400, detail=str(ve))
 
     json_data = {
         "model": MODEL_MAPPING.get(request.model, request.model),
-        "messages": [message_to_dict_ddg(msg, model_prefix=model_prefix) for msg in request.messages],
-        # Add additional fields if required by DDG API
+        "messages": message_history,
     }
 
     async with httpx.AsyncClient() as client:
@@ -121,7 +127,7 @@ async def process_ddg_streaming_response(request: ChatRequest):
                         content = line
                         if content.startswith("$@$v=undefined-rv1$@$"):
                             content = content[21:]
-                        # Strip the model prefix from the response content
+                        # No model_prefix handling as messages are correctly formatted
                         cleaned_content = strip_model_prefix(content, model_prefix)
                         yield f"data: {json.dumps(create_chat_completion_data(cleaned_content, request.model, timestamp))}\n\n"
 
@@ -138,27 +144,33 @@ async def process_ddg_non_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, "")
-    referer_path = MODEL_REFERERS.get(request.model, f"/?model={request.model}")
-    referer_url = f"https://duckduckgo.com{referer_path}"  # Updated BASE_URL
+    # referer_path and referer_url are not required unless specified by DDG API
+
+    # Obtain VQD token
+    vqd = await get_ddg_vqd()
+    if not vqd:
+        raise HTTPException(status_code=500, detail="Failed to obtain VQD token")
 
-    # Update headers with dynamic Referer
+    # Build headers
     dynamic_headers = {
         'accept': 'application/json',
         'content-type': 'application/json',
         'user-agent': 'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/128.0.0.0 Safari/537.36',
-        'x-vqd-4': await get_ddg_vqd()
+        'x-vqd-4': vqd
     }
 
-    if not dynamic_headers['x-vqd-4']:
-        raise HTTPException(status_code=500, detail="Failed to obtain VQD token")
+    # Construct message history
+    try:
+        message_history = [message_to_dict_ddg(msg) for msg in request.messages]
+    except ValueError as ve:
+        logger.error(f"Invalid message format: {ve}")
+        raise HTTPException(status_code=400, detail=str(ve))
 
     json_data = {
         "model": MODEL_MAPPING.get(request.model, request.model),
-        "messages": [message_to_dict_ddg(msg, model_prefix=model_prefix) for msg in request.messages],
-        # Add additional fields if required by DDG API
+        "messages": message_history,
     }
 
-    full_response = ""
     async with httpx.AsyncClient() as client:
         try:
             response = await client.post(
@@ -179,7 +191,6 @@ async def process_ddg_non_streaming_response(request: ChatRequest):
     if full_response.startswith("$@$v=undefined-rv1$@$"):
         full_response = full_response[21:]
 
-    # Strip the model prefix from the full response
     cleaned_full_response = strip_model_prefix(full_response, model_prefix)
 
     return {
@@ -196,21 +207,3 @@ async def process_ddg_non_streaming_response(request: ChatRequest):
         ],
         "usage": None,
     }
-
-def create_chat_completion_data(
-    content: str, model: str, timestamp: int, finish_reason: Optional[str] = None
-) -> Dict[str, Any]:
-    return {
-        "id": f"chatcmpl-{uuid.uuid4()}",
-        "object": "chat.completion.chunk",
-        "created": timestamp,
-        "model": model,
-        "choices": [
-            {
-                "index": 0,
-                "delta": {"content": content, "role": "assistant"},
-                "finish_reason": finish_reason,
-            }
-        ],
-        "usage": None,
-    }