Update api/utils.py

api/utils.py

@@ -22,6 +22,7 @@ from api.config import (
 )
 from api.models import ChatRequest
 from api.logger import setup_logger
+from api.image import ImageResponse, to_data_uri  # Assuming image utilities are here
 
 logger = setup_logger(__name__)
 
@@ -76,18 +77,21 @@ def get_referer_url(chat_id: str, model: str) -> str:
     return f"{BASE_URL}/chat/{chat_id}?model={model}"
 
 # Process streaming response with headers from config.py
-async def process_streaming_response(request: ChatRequest):
+async def process_streaming_response(request: ChatRequest, web_search: bool = False):
     chat_id = generate_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, {})
-    model_prefix = MODEL_PREFIXES.get(request.model, "")
+    # Resolve model aliases
+    resolved_model = MODEL_ALIASES.get(request.model, request.model)
+
+    agent_mode = AGENT_MODE.get(resolved_model, {})
+    trending_agent_mode = TRENDING_AGENT_MODE.get(resolved_model, {})
+    model_prefix = MODEL_PREFIXES.get(resolved_model, "")
 
     headers_api_chat = get_headers_api_chat()
 
-    if request.model == 'o1-preview':
+    if resolved_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)
@@ -110,10 +114,11 @@ async def process_streaming_response(request: ChatRequest):
         "previewToken": None,
         "trendingAgentMode": trending_agent_mode,
         "userId": None,
-        "userSelectedModel": request.model if request.model in MODELS else None,
+        "userSelectedModel": resolved_model if resolved_model in MODELS else None,
         "userSystemPrompt": None,
         "validated": "69783381-2ce4-4dbd-ac78-35e9063feabc",
         "visitFromDelta": False,
+        "webSearchMode": web_search,  # Include web search mode
     }
 
     async with httpx.AsyncClient() as client:
@@ -133,9 +138,9 @@ async def process_streaming_response(request: ChatRequest):
                         if content.startswith("$@$v=undefined-rv1$@$"):
                             content = content[21:]
                         cleaned_content = strip_model_prefix(content, model_prefix)
-                        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(cleaned_content, resolved_model, timestamp))}\n\n"
 
-                yield f"data: {json.dumps(create_chat_completion_data('', request.model, timestamp, 'stop'))}\n\n"
+                yield f"data: {json.dumps(create_chat_completion_data('', resolved_model, timestamp, 'stop'))}\n\n"
                 yield "data: [DONE]\n\n"
         except httpx.HTTPStatusError as e:
             logger.error(f"HTTP error occurred for Chat ID {chat_id}: {e}")
@@ -145,18 +150,21 @@ async def process_streaming_response(request: ChatRequest):
             raise HTTPException(status_code=500, detail=str(e))
 
 # Process non-streaming response with headers from config.py
-async def process_non_streaming_response(request: ChatRequest):
+async def process_non_streaming_response(request: ChatRequest, web_search: bool = False):
     chat_id = generate_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, {})
-    model_prefix = MODEL_PREFIXES.get(request.model, "")
+    # Resolve model aliases
+    resolved_model = MODEL_ALIASES.get(request.model, request.model)
+
+    agent_mode = AGENT_MODE.get(resolved_model, {})
+    trending_agent_mode = TRENDING_AGENT_MODE.get(resolved_model, {})
+    model_prefix = MODEL_PREFIXES.get(resolved_model, "")
 
     headers_api_chat = get_headers_api_chat()
 
-    if request.model == 'o1-preview':
+    if resolved_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)
@@ -179,10 +187,11 @@ async def process_non_streaming_response(request: ChatRequest):
         "previewToken": None,
         "trendingAgentMode": trending_agent_mode,
         "userId": None,
-        "userSelectedModel": request.model if request.model in MODELS else None,
+        "userSelectedModel": resolved_model if resolved_model in MODELS else None,
         "userSystemPrompt": None,
         "validated": "69783381-2ce4-4dbd-ac78-35e9063feabc",
         "visitFromDelta": False,
+        "webSearchMode": web_search,  # Include web search mode
     }
 
     full_response = ""
@@ -214,7 +223,7 @@ async def process_non_streaming_response(request: ChatRequest):
         "id": f"chatcmpl-{uuid.uuid4()}",
         "object": "chat.completion",
         "created": int(datetime.now().timestamp()),
-        "model": request.model,
+        "model": resolved_model,
         "choices": [
             {
                 "index": 0,