Update api/utils.py

api/utils.py

@@ -3,18 +3,15 @@ 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,
-    get_headers_chat,
     BASE_URL,
     AGENT_MODE,
     TRENDING_AGENT_MODE,
-    ALLOWED_MODELS
 )
 from api.models import ChatRequest
 from api.logger import setup_logger
@@ -26,54 +23,16 @@ def generate_chat_id(length: int = 7) -> str:
     characters = string.ascii_letters + string.digits
     return ''.join(random.choices(characters, k=length))
 
-# 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,
-    }
-
-# Function to convert message to dictionary format, ensuring base64 data and optional model prefix
-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]:
-        # Ensure base64 images are always included for all models
-        return {
-            "role": message.role,
-            "content": content,
-            "data": {
-                "imageBase64": message.content[1]["image_url"]["url"],
-                "fileText": "",
-                "title": "snapshot",
-            },
-        }
-    return {"role": message.role, "content": content}
-
-# Function to strip model prefix from content if present
-def strip_model_prefix(content: str) -> str:
-    return content.strip()
+# 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}"
 
 # Process streaming response with headers from config.py
 async def process_streaming_response(request: ChatRequest):
     chat_id = generate_chat_id()
-    referer_url = f"{BASE_URL}/chat/{chat_id}?model={request.model}"
+    referer_url = get_referer_url(chat_id, 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, {})
 
@@ -124,8 +83,7 @@ async def process_streaming_response(request: ChatRequest):
                         content = line
                         if content.startswith("$@$v=undefined-rv1$@$"):
                             content = content[21:]
-                        cleaned_content = strip_model_prefix(content)
-                        yield f"data: {json.dumps(create_chat_completion_data(cleaned_content, request.model, timestamp))}\n\n"
+                        yield f"data: {json.dumps(create_chat_completion_data(content, 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"
@@ -139,17 +97,13 @@ async def process_streaming_response(request: ChatRequest):
 # Process non-streaming response with headers from config.py
 async def process_non_streaming_response(request: ChatRequest):
     chat_id = generate_chat_id()
-    referer_url = f"{BASE_URL}/chat/{chat_id}?model={request.model}"
+    referer_url = get_referer_url(chat_id, 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, {})
 
     headers_api_chat = get_headers_api_chat(referer_url)
-    headers_chat = get_headers_chat(referer_url, next_action=str(uuid.uuid4()), next_router_state_tree=json.dumps([""]))
 
     if request.model == 'o1-preview':
         delay_seconds = random.randint(20, 60)
@@ -198,8 +152,6 @@ async def process_non_streaming_response(request: ChatRequest):
     if full_response.startswith("$@$v=undefined-rv1$@$"):
         full_response = full_response[21:]
 
-    cleaned_full_response = strip_model_prefix(full_response)
-
     return {
         "id": f"chatcmpl-{uuid.uuid4()}",
         "object": "chat.completion",
@@ -208,9 +160,9 @@ async def process_non_streaming_response(request: ChatRequest):
         "choices": [
             {
                 "index": 0,
-                "message": {"role": "assistant", "content": cleaned_full_response},
+                "message": {"role": "assistant", "content": full_response},
                 "finish_reason": "stop",
             }
         ],
         "usage": None,
-    }
+    }