Update api/utils.py

api/utils.py

@@ -3,6 +3,7 @@ import json
 import uuid
 import asyncio
 import random
+import string
 from typing import Any, Dict, Optional
 
 import httpx
@@ -15,6 +16,7 @@ from api.config import (
     AGENT_MODE,
     TRENDING_AGENT_MODE,
     MODEL_PREFIXES,
+    MODEL_REFERERS
 )
 from api.models import ChatRequest
 from api.logger import setup_logger
@@ -70,18 +72,38 @@ def strip_model_prefix(content: str, model_prefix: Optional[str] = None) -> str:
         return content[len(model_prefix):].strip()
     return content
 
-# Helper function to build JSON data for the request
-def build_json_data(request: ChatRequest, h_value: str, model_prefix: Optional[str]):
+# Process streaming response with headers from config.py
+async def process_streaming_response(request: ChatRequest):
+    # Generate a unique ID for this request
+    request_id = f"chatcmpl-{uuid.uuid4()}"
+    logger.info(f"Processing request with ID: {request_id} - Model: {request.model}")
+
     agent_mode = AGENT_MODE.get(request.model, {})
     trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
-    return {
+    model_prefix = MODEL_PREFIXES.get(request.model, "")
+
+    # Adjust headers_api_chat since referer_url is removed
+    headers_api_chat = get_headers_api_chat(BASE_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' (Request ID: {request_id})")
+        await asyncio.sleep(delay_seconds)
+
+    # Fetch the h-value for the 'validated' field
+    h_value = await getHid()
+    if not h_value:
+        logger.error("Failed to retrieve h-value for validation.")
+        raise HTTPException(status_code=500, detail="Validation failed due to missing h-value.")
+
+    json_data = {
         "agentMode": agent_mode,
         "clickedAnswer2": False,
         "clickedAnswer3": False,
         "clickedForceWebSearch": False,
         "codeModelMode": True,
         "githubToken": None,
-        "id": None,
+        "id": None,  # Using request_id instead of chat_id
         "isChromeExt": False,
         "isMicMode": False,
         "maxTokens": request.max_tokens,
@@ -94,35 +116,10 @@ def build_json_data(request: ChatRequest, h_value: str, model_prefix: Optional[s
         "userId": None,
         "userSelectedModel": MODEL_MAPPING.get(request.model, request.model),
         "userSystemPrompt": None,
-        "validated": "00f37b34-a166-4efb-bce5-1312d87f2f94",
+        "validated": h_value,  # Dynamically set the validated field
         "visitFromDelta": False,
     }
 
-# Process streaming response with headers from config.py
-async def process_streaming_response(request: ChatRequest):
-    logger.info(f"Processing request - Model: {request.model}")
-
-    model_prefix = MODEL_PREFIXES.get(request.model, "")
-
-    # Adjust headers_api_chat since referer_url is removed
-    headers_api_chat = get_headers_api_chat(BASE_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'")
-        await asyncio.sleep(delay_seconds)
-
-    # Fetch the h-value for the 'validated' field
-    h_value = await getHid()
-    if not h_value:
-        logger.error("Failed to retrieve h-value for validation.")
-        raise HTTPException(status_code=500, detail="Validation failed due to missing h-value.")
-
-    json_data = build_json_data(request, h_value, model_prefix)
-
-    # Initialize rolling buffer to handle BLOCKED_MESSAGE split across chunks
-    rolling_buffer = ""
-
     async with httpx.AsyncClient() as client:
         try:
             async with client.stream(
@@ -134,68 +131,46 @@ async def process_streaming_response(request: ChatRequest):
             ) as response:
                 response.raise_for_status()
                 async for chunk in response.aiter_text():
+                    timestamp = int(datetime.now().timestamp())
                     if chunk:
-                        # Combine rolling buffer with current chunk
-                        combined_chunk = rolling_buffer + chunk
-
-                        # Remove any occurrence of BLOCKED_MESSAGE in combined_chunk
-                        if BLOCKED_MESSAGE in combined_chunk:
-                            logger.info("Blocked message detected in response.")
-                            combined_chunk = combined_chunk.replace(BLOCKED_MESSAGE, '')
-
-                        # Remove model prefix if present
-                        cleaned_content = strip_model_prefix(combined_chunk, model_prefix)
-
-                        # Yield the cleaned content
-                        timestamp = int(datetime.now().timestamp())
+                        content = chunk
+                        if content.startswith("$@$v=undefined-rv1$@$"):
+                            content = content[21:]
+                        # Remove the blocked message if present
+                        if BLOCKED_MESSAGE in content:
+                            logger.info(f"Blocked message detected in response for Request ID {request_id}.")
+                            content = content.replace(BLOCKED_MESSAGE, '').strip()
+                            if not content:
+                                continue  # Skip if content is empty after removal
+                        cleaned_content = strip_model_prefix(content, model_prefix)
                         yield f"data: {json.dumps(create_chat_completion_data(cleaned_content, request.model, timestamp))}\n\n"
 
-                        # Update rolling buffer with the end of the combined_chunk
-                        # Keep only the last len(BLOCKED_MESSAGE) - 1 characters
-                        rolling_buffer = combined_chunk[-(len(BLOCKED_MESSAGE) - 1):]
-                    else:
-                        # If chunk is empty, reset rolling buffer
-                        rolling_buffer = ""
-
-                # After streaming is done, check if any remaining content is in the rolling buffer
-                if rolling_buffer:
-                    # Remove any occurrence of BLOCKED_MESSAGE in rolling buffer
-                    if BLOCKED_MESSAGE in rolling_buffer:
-                        logger.info("Blocked message detected in remaining buffer.")
-                        rolling_buffer = rolling_buffer.replace(BLOCKED_MESSAGE, '')
-
-                    cleaned_content = strip_model_prefix(rolling_buffer, model_prefix)
-                    timestamp = int(datetime.now().timestamp())
-                    yield f"data: {json.dumps(create_chat_completion_data(cleaned_content, request.model, timestamp))}\n\n"
-
-                # Signal the end of the streaming
-                timestamp = int(datetime.now().timestamp())
                 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:
-            logger.error(f"HTTP error occurred: {e}")
+            logger.error(f"HTTP error occurred for Request ID {request_id}: {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}")
+            logger.error(f"Error occurred during request for Request ID {request_id}: {e}")
             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):
-    logger.info(f"Processing request - Model: {request.model}")
+    # Generate a unique ID for this request
+    request_id = f"chatcmpl-{uuid.uuid4()}"
+    logger.info(f"Processing request with ID: {request_id} - Model: {request.model}")
 
+    agent_mode = AGENT_MODE.get(request.model, {})
+    trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
     model_prefix = MODEL_PREFIXES.get(request.model, "")
 
     # Adjust headers_api_chat and headers_chat since referer_url is removed
     headers_api_chat = get_headers_api_chat(BASE_URL)
-    headers_chat = get_headers_chat(
-        BASE_URL,
-        next_action=str(uuid.uuid4()),
-        next_router_state_tree=json.dumps([""])
-    )
+    headers_chat = get_headers_chat(BASE_URL, next_action=str(uuid.uuid4()), next_router_state_tree=json.dumps([""]))
 
     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'")
+        logger.info(f"Introducing a delay of {delay_seconds} seconds for model 'o1-preview' (Request ID: {request_id})")
         await asyncio.sleep(delay_seconds)
 
     # Fetch the h-value for the 'validated' field
@@ -204,33 +179,52 @@ async def process_non_streaming_response(request: ChatRequest):
         logger.error("Failed to retrieve h-value for validation.")
         raise HTTPException(status_code=500, detail="Validation failed due to missing h-value.")
 
-    json_data = build_json_data(request, h_value, model_prefix)
+    json_data = {
+        "agentMode": agent_mode,
+        "clickedAnswer2": False,
+        "clickedAnswer3": False,
+        "clickedForceWebSearch": False,
+        "codeModelMode": True,
+        "githubToken": None,
+        "id": None,  # Using request_id instead of chat_id
+        "isChromeExt": False,
+        "isMicMode": False,
+        "maxTokens": request.max_tokens,
+        "messages": [message_to_dict(msg, model_prefix=model_prefix) for msg in request.messages],
+        "mobileClient": False,
+        "playgroundTemperature": request.temperature,
+        "playgroundTopP": request.top_p,
+        "previewToken": None,
+        "trendingAgentMode": trending_agent_mode,
+        "userId": None,
+        "userSelectedModel": MODEL_MAPPING.get(request.model, request.model),
+        "userSystemPrompt": None,
+        "validated": h_value,  # Dynamically set the validated field
+        "visitFromDelta": False,
+    }
 
     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,
+                method="POST", url=f"{BASE_URL}/api/chat", headers=headers_api_chat, json=json_data
             ) as response:
                 response.raise_for_status()
                 async for chunk in response.aiter_text():
                     full_response += chunk
         except httpx.HTTPStatusError as e:
-            logger.error(f"HTTP error occurred: {e}")
+            logger.error(f"HTTP error occurred for Request ID {request_id}: {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}")
+            logger.error(f"Error occurred during request for Request ID {request_id}: {e}")
             raise HTTPException(status_code=500, detail=str(e))
-
+    
     if full_response.startswith("$@$v=undefined-rv1$@$"):
         full_response = full_response[21:]
 
     # Remove the blocked message if present
     if BLOCKED_MESSAGE in full_response:
-        logger.info("Blocked message detected in response.")
+        logger.info(f"Blocked message detected in response for Request ID {request_id}.")
         full_response = full_response.replace(BLOCKED_MESSAGE, '').strip()
         if not full_response:
             raise HTTPException(status_code=500, detail="Blocked message detected in response.")
@@ -250,4 +244,4 @@ async def process_non_streaming_response(request: ChatRequest):
             }
         ],
         "usage": None,
-    }
+    }