Update api/utils.py

api/utils.py

@@ -3,12 +3,14 @@ import json
 from typing import Any, Dict, Optional
 
 import httpx
+import aiohttp  # Add this import
 from api.config import (
+    DDG_API_ENDPOINT,  # Updated
+    DDG_STATUS_URL,    # Updated
     MODEL_MAPPING,
-    headers,
+    ALLOWED_MODELS,
     AGENT_MODE,
     TRENDING_AGENT_MODE,
-    BASE_URL,
     MODEL_PREFIXES,
     MODEL_REFERERS
 )
@@ -18,30 +20,38 @@ from api.models import ChatRequest
 from api.logger import setup_logger
 
 import uuid
-import asyncio
-import random  # Newly added imports
 
 logger = setup_logger(__name__)
 
-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,
+# Remove Blackbox-specific functions and add DDG-specific functions
+
+# New DDG Integration
+async def get_ddg_vqd():
+    status_url = DDG_STATUS_URL
+
+    headers = {
+        'User-Agent': 'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/128.0.0.0 Safari/537.36',
+        'Accept': 'text/event-stream',
+        'x-vqd-accept': '1'
     }
 
-def message_to_dict(message, model_prefix: Optional[str] = None):
+    async with aiohttp.ClientSession() as session:
+        try:
+            async with session.get(status_url, headers=headers) as response:
+                if response.status == 200:
+                    return response.headers.get("x-vqd-4")
+                else:
+                    logger.error(f"Error: Status code {response.status} when fetching VQD")
+                    return None
+        except Exception as e:
+            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):
     if isinstance(message.content, str):
         content = message.content
         if model_prefix:
@@ -71,50 +81,35 @@ def strip_model_prefix(content: str, model_prefix: Optional[str] = None) -> str:
     logger.debug("No prefix to strip from content.")
     return content
 
-async def process_streaming_response(request: ChatRequest):
+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"{BASE_URL}{referer_path}"
+    referer_url = f"https://duckduckgo.com{referer_path}"  # Updated BASE_URL
 
     # Update headers with dynamic Referer
-    dynamic_headers = headers.copy()
-    dynamic_headers['Referer'] = referer_url
+    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()
+    }
 
-    # Introduce delay for 'o1-preview' model
-    if request.model == 'o1-preview':
-        delay_seconds = random.randint(20, 60)
-        logger.info(f"Introducing a delay of {delay_seconds} seconds for model 'o1-preview'")
-        await asyncio.sleep(delay_seconds)
+    if not dynamic_headers['x-vqd-4']:
+        raise HTTPException(status_code=500, detail="Failed to obtain VQD token")
 
     json_data = {
-        "messages": [message_to_dict(msg, model_prefix=model_prefix) for msg in request.messages],
-        "previewToken": None,
-        "userId": None,
-        "codeModelMode": True,
-        "agentMode": agent_mode,
-        "trendingAgentMode": trending_agent_mode,
-        "isMicMode": False,
-        "userSystemPrompt": None,
-        "maxTokens": request.max_tokens,
-        "playgroundTopP": request.top_p,
-        "playgroundTemperature": request.temperature,
-        "isChromeExt": False,
-        "githubToken": None,
-        "clickedAnswer2": False,
-        "clickedAnswer3": False,
-        "clickedForceWebSearch": False,
-        "visitFromDelta": False,
-        "mobileClient": False,
-        "userSelectedModel": MODEL_MAPPING.get(request.model, request.model),
+        "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
     }
 
     async with httpx.AsyncClient() as client:
         try:
             async with client.stream(
                 "POST",
-                f"{BASE_URL}/api/chat",
+                DDG_API_ENDPOINT,
                 headers=dynamic_headers,
                 json=json_data,
                 timeout=100,
@@ -139,59 +134,48 @@ async def process_streaming_response(request: ChatRequest):
             logger.error(f"Error occurred during request: {e}")
             raise HTTPException(status_code=500, detail=str(e))
 
-async def process_non_streaming_response(request: ChatRequest):
+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"{BASE_URL}{referer_path}"
+    referer_url = f"https://duckduckgo.com{referer_path}"  # Updated BASE_URL
 
     # Update headers with dynamic Referer
-    dynamic_headers = headers.copy()
-    dynamic_headers['Referer'] = referer_url
+    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()
+    }
 
-    # Introduce delay for 'o1-preview' model
-    if request.model == 'o1-preview':
-        delay_seconds = random.randint(20, 60)
-        logger.info(f"Introducing a delay of {delay_seconds} seconds for model 'o1-preview'")
-        await asyncio.sleep(delay_seconds)
+    if not dynamic_headers['x-vqd-4']:
+        raise HTTPException(status_code=500, detail="Failed to obtain VQD token")
 
     json_data = {
-        "messages": [message_to_dict(msg, model_prefix=model_prefix) for msg in request.messages],
-        "previewToken": None,
-        "userId": None,
-        "codeModelMode": True,
-        "agentMode": agent_mode,
-        "trendingAgentMode": trending_agent_mode,
-        "isMicMode": False,
-        "userSystemPrompt": None,
-        "maxTokens": request.max_tokens,
-        "playgroundTopP": request.top_p,
-        "playgroundTemperature": request.temperature,
-        "isChromeExt": False,
-        "githubToken": None,
-        "clickedAnswer2": False,
-        "clickedAnswer3": False,
-        "clickedForceWebSearch": False,
-        "visitFromDelta": False,
-        "mobileClient": False,
-        "userSelectedModel": MODEL_MAPPING.get(request.model, request.model),
+        "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
     }
+
     full_response = ""
     async with httpx.AsyncClient() as client:
         try:
-            async with client.stream(
-                method="POST", url=f"{BASE_URL}/api/chat", headers=dynamic_headers, json=json_data
-            ) as response:
-                response.raise_for_status()
-                async for chunk in response.aiter_text():
-                    full_response += chunk
+            response = await client.post(
+                DDG_API_ENDPOINT,
+                headers=dynamic_headers,
+                json=json_data,
+                timeout=100
+            )
+            response.raise_for_status()
+            full_response = response.text
         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 full_response.startswith("$@$v=undefined-rv1$@$"):
         full_response = full_response[21:]
 
@@ -212,3 +196,21 @@ async def process_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,
+    }