Update api/utils.py

api/utils.py

@@ -3,15 +3,17 @@ import json
 import uuid
 import asyncio
 import random
+import string
 from typing import Any, Dict, Optional
 
 import httpx
 from fastapi import HTTPException
 from api.config import (
-    get_headers_api_chat,
+    MODEL_MAPPING,
+    get_headers,
     BASE_URL,
     AGENT_MODE,
-    TRENDING_AGENT_MODE,
+    TRENDING_AGENT_MODE
 )
 from api.models import ChatRequest
 from api.logger import setup_logger
@@ -23,25 +25,47 @@ def generate_chat_id(length: int = 7) -> str:
     characters = string.ascii_letters + string.digits
     return ''.join(random.choices(characters, k=length))
 
-# Function to get the correct referer URL for logging
-def get_referer_url(chat_id: str, model: str) -> str:
-    return f"{ BASE_URL}/chat/{chat_id}"
+# Helper function to create chat completion data
+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,
+    }
 
-# Process streaming response with headers from config.py
+# Function to convert message to dictionary format
+def message_to_dict(message):
+    content = message.content if isinstance(message.content, str) else message.content[0]["text"]
+    if isinstance(message.content, list) and len(message.content) == 2 and "image_url" in message.content[1]:
+        return {
+            "role": message.role,
+            "content": content,
+            "data": {
+                "imageBase64": message.content[1]["image_url"]["url"],
+                "fileText": "",
+                "title": "snapshot",
+            },
+        }
+    return {"role": message.role, "content": content}
+
+# Process streaming response
 async def process_streaming_response(request: ChatRequest):
     chat_id = generate_chat_id()
-    referer_url = get_referer_url(chat_id, request.model)
-    logger.info(f"Generated Chat ID: {chat_id} - Model: {request.model} - URL: {referer_url}")
-
     agent_mode = AGENT_MODE.get(request.model, {})
     trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
 
-    headers_api_chat = get_headers_api_chat(referer_url)
-
-    if request.model == 'o1-preview':
-        delay_seconds = random.randint(1, 60)
-        logger.info(f"Introducing a delay of {delay_seconds} seconds for model 'o1-preview' (Chat ID: {chat_id})")
-        await asyncio.sleep(delay_seconds)
+    headers = get_headers()
 
     json_data = {
         "agentMode": agent_mode,
@@ -61,7 +85,7 @@ async def process_streaming_response(request: ChatRequest):
         "previewToken": None,
         "trendingAgentMode": trending_agent_mode,
         "userId": None,
-        "userSelectedModel": request.model,
+        "userSelectedModel": MODEL_MAPPING.get(request.model, request.model),
         "userSystemPrompt": None,
         "validated": "69783381-2ce4-4dbd-ac78-35e9063feabc",
         "visitFromDelta": False,
@@ -69,22 +93,12 @@ async def process_streaming_response(request: ChatRequest):
 
     async with httpx.AsyncClient() as client:
         try:
-            async with client.stream(
-                "POST",
-                f"{BASE_URL}/api/chat",
-                headers=headers_api_chat,
-                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())
                     if line:
-                        content = line
-                        if content.startswith("$@$v=undefined-rv1$@$"):
-                            content = content[21:]
-                        yield f"data: {json.dumps(create_chat_completion_data(content, request.model, timestamp))}\n\n"
-
+                        yield f"data: {json.dumps(create_chat_completion_data(line, request.model, timestamp))}\n\n"
                 yield f"data: {json.dumps(create_chat_completion_data('', request.model, timestamp, 'stop'))}\n\n"
                 yield "data: [DONE]\n\n"
         except httpx.HTTPStatusError as e:
@@ -94,21 +108,13 @@ 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 with headers from config.py
+# Process non-streaming response
 async def process_non_streaming_response(request: ChatRequest):
     chat_id = generate_chat_id()
-    referer_url = get_referer_url(chat_id, request.model)
-    logger.info(f"Generated Chat ID: {chat_id} - Model: {request.model} - URL: {referer_url}")
-
     agent_mode = AGENT_MODE.get(request.model, {})
     trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
 
-    headers_api_chat = get_headers_api_chat(referer_url)
-
-    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' (Chat ID: {chat_id})")
-        await asyncio.sleep(delay_seconds)
+    headers = get_headers()
 
     json_data = {
         "agentMode": agent_mode,
@@ -128,7 +134,7 @@ async def process_non_streaming_response(request: ChatRequest):
         "previewToken": None,
         "trendingAgentMode": trending_agent_mode,
         "userId": None,
-        "userSelectedModel": request.model,
+        "userSelectedModel": MODEL_MAPPING.get(request.model, request.model),
         "userSystemPrompt": None,
         "validated": "69783381-2ce4-4dbd-ac78-35e9063feabc",
         "visitFromDelta": False,
@@ -137,9 +143,7 @@ async def process_non_streaming_response(request: ChatRequest):
     full_response = ""
     async with httpx.AsyncClient() as client:
         try:
-            async with client.stream(
-                method="POST", url=f"{BASE_URL}/api/chat", headers=headers_api_chat, json=json_data
-            ) as response:
+            async with client.stream("POST", 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
@@ -149,8 +153,6 @@ async def process_non_streaming_response(request: ChatRequest):
         except httpx.RequestError as e:
             logger.error(f"Error occurred during request for Chat ID {chat_id}: {e}")
             raise HTTPException(status_code=500, detail=str(e))
-    if full_response.startswith("$@$v=undefined-rv1$@$"):
-        full_response = full_response[21:]
 
     return {
         "id": f"chatcmpl-{uuid.uuid4()}",
@@ -165,4 +167,4 @@ async def process_non_streaming_response(request: ChatRequest):
             }
         ],
         "usage": None,
-    }
+    }