Update api/utils.py

api/utils.py

@@ -18,7 +18,6 @@ from api.config import (
     MODEL_PREFIXES,
     MODEL_REFERERS
 )
-from api.models import ChatRequest
 from api.logger import (
     log_generated_chat_id,
     log_model_delay,
@@ -84,23 +83,27 @@ def strip_model_prefix(content: str, model_prefix: Optional[str] = None) -> str:
 
 # Process streaming response with headers from config.py
 async def process_streaming_response(request: ChatRequest):
-    chat_id = generate_chat_id()
-    chat_url = f"/chat/{chat_id}?model={request.model}"
-    log_generated_chat_id(chat_id, chat_url)
-    
+    # Determine if a Chat ID and referer URL should be generated
+    if request.model in MODEL_REFERERS:
+        chat_id = generate_chat_id()
+        chat_url = f"/chat/{chat_id}?model={request.model}"
+        log_generated_chat_id(chat_id, chat_url)
+        
+        referer_path = MODEL_REFERERS[request.model]
+        referer_url = f"{BASE_URL}{referer_path}"
+    else:
+        chat_id = None
+        referer_url = BASE_URL  # Use base URL for models not in MODEL_REFERERS
+
     agent_mode = AGENT_MODE.get(request.model, {})
     trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
     model_prefix = MODEL_PREFIXES.get(request.model, "")
-    
-    # Retrieve referer based on the model specified in the request
-    referer_path = MODEL_REFERERS.get(request.model, "/?model=default")
-    referer_url = f"{BASE_URL}{referer_path}"
 
     headers_api_chat = get_headers_api_chat(referer_url)
 
     if request.model == 'o1-preview':
         delay_seconds = random.randint(20, 60)
-        log_model_delay(delay_seconds, chat_id, request.model)
+        log_model_delay(delay_seconds, chat_id, request.model if chat_id else "unknown")
         await asyncio.sleep(delay_seconds)
 
     json_data = {
@@ -110,7 +113,7 @@ async def process_streaming_response(request: ChatRequest):
         "clickedForceWebSearch": False,
         "codeModelMode": True,
         "githubToken": None,
-        "id": chat_id,
+        "id": chat_id if chat_id else "unknown",
         "isChromeExt": False,
         "isMicMode": False,
         "maxTokens": request.max_tokens,
@@ -149,32 +152,35 @@ async def process_streaming_response(request: ChatRequest):
                 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:
-            log_http_error(chat_id, e)
+            log_http_error(chat_id if chat_id else "unknown", e)
             raise HTTPException(status_code=e.response.status_code, detail=str(e))
         except httpx.RequestError as e:
-            log_request_error(chat_id, e)
+            log_request_error(chat_id if chat_id else "unknown", 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):
-    chat_id = generate_chat_id()
-    chat_url = f"/chat/{chat_id}?model={request.model}"
-    log_generated_chat_id(chat_id, chat_url)
-    
+    if request.model in MODEL_REFERERS:
+        chat_id = generate_chat_id()
+        chat_url = f"/chat/{chat_id}?model={request.model}"
+        log_generated_chat_id(chat_id, chat_url)
+        
+        referer_path = MODEL_REFERERS[request.model]
+        referer_url = f"{BASE_URL}{referer_path}"
+    else:
+        chat_id = None
+        referer_url = BASE_URL
+
     agent_mode = AGENT_MODE.get(request.model, {})
     trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
     model_prefix = MODEL_PREFIXES.get(request.model, "")
-    
-    # Retrieve referer based on the model specified in the request
-    referer_path = MODEL_REFERERS.get(request.model, "/?model=default")
-    referer_url = f"{BASE_URL}{referer_path}"
 
     headers_api_chat = get_headers_api_chat(referer_url)
-    headers_chat = get_headers_chat(chat_url, next_action=str(uuid.uuid4()), next_router_state_tree=json.dumps([""]))
+    headers_chat = get_headers_chat(chat_url, next_action=str(uuid.uuid4()), next_router_state_tree=json.dumps([""])) if chat_id else None
 
     if request.model == 'o1-preview':
         delay_seconds = random.randint(20, 60)
-        log_model_delay(delay_seconds, chat_id, request.model)
+        log_model_delay(delay_seconds, chat_id if chat_id else "unknown", request.model)
         await asyncio.sleep(delay_seconds)
 
     json_data = {
@@ -184,7 +190,7 @@ async def process_non_streaming_response(request: ChatRequest):
         "clickedForceWebSearch": False,
         "codeModelMode": True,
         "githubToken": None,
-        "id": chat_id,
+        "id": chat_id if chat_id else "unknown",
         "isChromeExt": False,
         "isMicMode": False,
         "maxTokens": request.max_tokens,
@@ -211,10 +217,10 @@ async def process_non_streaming_response(request: ChatRequest):
                 async for chunk in response.aiter_text():
                     full_response += chunk
         except httpx.HTTPStatusError as e:
-            log_http_error(chat_id, e)
+            log_http_error(chat_id if chat_id else "unknown", e)
             raise HTTPException(status_code=e.response.status_code, detail=str(e))
         except httpx.RequestError as e:
-            log_request_error(chat_id, e)
+            log_request_error(chat_id if chat_id else "unknown", e)
             raise HTTPException(status_code=500, detail=str(e))
     if full_response.startswith("$@$v=undefined-rv1$@$"):
         full_response = full_response[21:]