Update main.py

main.py

@@ -9,16 +9,16 @@ import asyncio
 import time
 from collections import defaultdict
 from typing import List, Dict, Any, Optional, Union, AsyncGenerator
+from datetime import datetime  # <-- Added import
 
 from aiohttp import ClientSession, ClientResponseError
 from fastapi import FastAPI, HTTPException, Request, Depends, Header
 from fastapi.responses import JSONResponse
 from pydantic import BaseModel
-from datetime import datetime
 
 # Configure logging
 logging.basicConfig(
-    level=logging.INFO,
+    level=logging.DEBUG,  # Changed to DEBUG for detailed logs
     format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
     handlers=[logging.StreamHandler()]
 )
@@ -202,167 +202,190 @@ class Blackbox:
         return cleaned_text
 
     @classmethod
-async def generate_response(
-    cls,
-    model: str,
-    messages: List[Dict[str, str]],
-    proxy: Optional[str] = None,
-    websearch: bool = False,
-    **kwargs
-) -> Union[str, ImageResponseModel]:
-    model = cls.get_model(model)
-    chat_id = cls.generate_random_string()
-    next_action = cls.generate_next_action()
-    next_router_state_tree = cls.generate_next_router_state_tree()
-
-    agent_mode = cls.agentMode.get(model, {})
-    trending_agent_mode = cls.trendingAgentMode.get(model, {})
-
-    prefix = cls.model_prefixes.get(model, "")
-    
-    formatted_prompt = ""
-    for message in messages:
-        role = message.get('role', '').capitalize()
-        content = message.get('content', '')
-        if role and content:
-            formatted_prompt += f"{role}: {content}\n"
-    
-    if prefix:
-        formatted_prompt = f"{prefix} {formatted_prompt}".strip()
-
-    referer_path = cls.model_referers.get(model, f"/?model={model}")
-    referer_url = f"{cls.url}{referer_path}"
-
-    common_headers = {
-        'accept': '*/*',
-        'accept-language': 'en-US,en;q=0.9',
-        'cache-control': 'no-cache',
-        'origin': cls.url,
-        'pragma': 'no-cache',
-        'priority': 'u=1, i',
-        'sec-ch-ua': '"Chromium";v="129", "Not=A?Brand";v="8"',
-        'sec-ch-ua-mobile': '?0',
-        'sec-ch-ua-platform': '"Linux"',
-        'sec-fetch-dest': 'empty',
-        'sec-fetch-mode': 'cors',
-        'sec-fetch-site': 'same-origin',
-        'user-agent': 'Mozilla/5.0 (X11; Linux x86_64) '
-                      'AppleWebKit/537.36 (KHTML, like Gecko) '
-                      'Chrome/129.0.0.0 Safari/537.36'
-    }
+    async def generate_response(
+        cls,
+        model: str,
+        messages: List[Dict[str, str]],
+        proxy: Optional[str] = None,
+        websearch: bool = False,
+        **kwargs
+    ) -> Union[str, ImageResponseModel]:
+        model = cls.get_model(model)
+        chat_id = cls.generate_random_string()
+        next_action = cls.generate_next_action()
+        next_router_state_tree = cls.generate_next_router_state_tree()
 
-    headers_api_chat = {
-        'Content-Type': 'application/json',
-        'Referer': referer_url
-    }
-    headers_api_chat_combined = {**common_headers, **headers_api_chat}
+        agent_mode = cls.agentMode.get(model, {})
+        trending_agent_mode = cls.trendingAgentMode.get(model, {})
 
-    payload_api_chat = {
-        "messages": [
-            {
-                "id": chat_id,
-                "content": formatted_prompt,
-                "role": "user"
-            }
-        ],
-        "id": chat_id,
-        "previewToken": None,
-        "userId": None,
-        "codeModelMode": True,
-        "agentMode": agent_mode,
-        "trendingAgentMode": trending_agent_mode,
-        "isMicMode": False,
-        "userSystemPrompt": None,
-        "maxTokens": 1024,
-        "playgroundTopP": 0.9,
-        "playgroundTemperature": 0.5,
-        "isChromeExt": False,
-        "githubToken": None,
-        "clickedAnswer2": False,
-        "clickedAnswer3": False,
-        "clickedForceWebSearch": False,
-        "visitFromDelta": False,
-        "mobileClient": False,
-        "webSearchMode": websearch,
-        "userSelectedModel": cls.userSelectedModel.get(model, model)
-    }
+        prefix = cls.model_prefixes.get(model, "")
+        
+        formatted_prompt = ""
+        for message in messages:
+            role = message.get('role', '').capitalize()
+            content = message.get('content', '')
+            if role and content:
+                formatted_prompt += f"{role}: {content}\n"
+        
+        if prefix:
+            formatted_prompt = f"{prefix} {formatted_prompt}".strip()
 
-    headers_chat = {
-        'Accept': 'text/x-component',
-        'Content-Type': 'text/plain;charset=UTF-8',
-        'Referer': f'{cls.url}/chat/{chat_id}?model={model}',
-        'next-action': next_action,
-        'next-router-state-tree': next_router_state_tree,
-        'next-url': '/'
-    }
-    headers_chat_combined = {**common_headers, **headers_chat}
+        referer_path = cls.model_referers.get(model, f"/?model={model}")
+        referer_url = f"{cls.url}{referer_path}"
 
-    data_chat = '[]'
+        common_headers = {
+            'accept': '*/*',
+            'accept-language': 'en-US,en;q=0.9',
+            'cache-control': 'no-cache',
+            'origin': cls.url,
+            'pragma': 'no-cache',
+            'priority': 'u=1, i',
+            'sec-ch-ua': '"Chromium";v="129", "Not=A?Brand";v="8"',
+            'sec-ch-ua-mobile': '?0',
+            'sec-ch-ua-platform': '"Linux"',
+            'sec-fetch-dest': 'empty',
+            'sec-fetch-mode': 'cors',
+            'sec-fetch-site': 'same-origin',
+            'user-agent': 'Mozilla/5.0 (X11; Linux x86_64) '
+                          'AppleWebKit/537.36 (KHTML, like Gecko) '
+                          'Chrome/129.0.0.0 Safari/537.36'
+        }
 
-    async with ClientSession(headers=common_headers) as session:
-        try:
-            async with session.post(
-                cls.api_endpoint,
-                headers=headers_api_chat_combined,
-                json=payload_api_chat,
-                proxy=proxy
-            ) as response_api_chat:
-                response_api_chat.raise_for_status()
-                text = await response_api_chat.text()
-                logger.debug(f"Raw response from Blackbox API: {text}")  # Added logging
-                cleaned_response = cls.clean_response(text)
-                logger.debug(f"Cleaned response: {cleaned_response}")  # Added logging
-
-                if model in cls.image_models:
-                    match = re.search(r'!\[.*?\]\((https?://[^\)]+)\)', cleaned_response)
-                    if match:
-                        image_url = match.group(1)
-                        image_response = ImageResponseModel(images=image_url, alt="Generated Image")
-                        logger.debug(f"Image URL extracted: {image_url}")  # Added logging
-                        return image_response
-                    else:
-                        logger.debug("No image URL found in the response.")  # Added logging
-                        return cleaned_response
-                else:
-                    if websearch:
-                        match = re.search(r'\$~~~\$(.*?)\$~~~\$', cleaned_response, re.DOTALL)
+        headers_api_chat = {
+            'Content-Type': 'application/json',
+            'Referer': referer_url
+        }
+        headers_api_chat_combined = {**common_headers, **headers_api_chat}
+
+        payload_api_chat = {
+            "messages": [
+                {
+                    "id": chat_id,
+                    "content": formatted_prompt,
+                    "role": "user"
+                }
+            ],
+            "id": chat_id,
+            "previewToken": None,
+            "userId": None,
+            "codeModelMode": True,
+            "agentMode": agent_mode,
+            "trendingAgentMode": trending_agent_mode,
+            "isMicMode": False,
+            "userSystemPrompt": None,
+            "maxTokens": 1024,
+            "playgroundTopP": 0.9,
+            "playgroundTemperature": 0.5,
+            "isChromeExt": False,
+            "githubToken": None,
+            "clickedAnswer2": False,
+            "clickedAnswer3": False,
+            "clickedForceWebSearch": False,
+            "visitFromDelta": False,
+            "mobileClient": False,
+            "webSearchMode": websearch,
+            "userSelectedModel": cls.userSelectedModel.get(model, model)
+        }
+
+        headers_chat = {
+            'Accept': 'text/x-component',
+            'Content-Type': 'text/plain;charset=UTF-8',
+            'Referer': f'{cls.url}/chat/{chat_id}?model={model}',
+            'next-action': next_action,
+            'next-router-state-tree': next_router_state_tree,
+            'next-url': '/'
+        }
+        headers_chat_combined = {**common_headers, **headers_chat}
+
+        data_chat = '[]'
+
+        async with ClientSession(headers=common_headers) as session:
+            try:
+                async with session.post(
+                    cls.api_endpoint,
+                    headers=headers_api_chat_combined,
+                    json=payload_api_chat,
+                    proxy=proxy
+                ) as response_api_chat:
+                    response_api_chat.raise_for_status()
+                    text = await response_api_chat.text()
+                    logger.debug(f"Raw response from Blackbox API: {text}")  # Added logging
+                    cleaned_response = cls.clean_response(text)
+                    logger.debug(f"Cleaned response: {cleaned_response}")  # Added logging
+
+                    if model in cls.image_models:
+                        match = re.search(r'!\[.*?\]\((https?://[^\)]+)\)', cleaned_response)
                         if match:
-                            source_part = match.group(1).strip()
-                            answer_part = cleaned_response[match.end():].strip()
-                            try:
-                                sources = json.loads(source_part)
-                                source_formatted = "**Source:**\n"
-                                for item in sources:
-                                    title = item.get('title', 'No Title')
-                                    link = item.get('link', '#')
-                                    position = item.get('position', '')
-                                    source_formatted += f"{position}. [{title}]({link})\n"
-                                final_response = f"{answer_part}\n\n{source_formatted}"
-                            except json.JSONDecodeError:
-                                final_response = f"{answer_part}\n\nSource information is unavailable."
+                            image_url = match.group(1)
+                            image_response = ImageResponseModel(images=image_url, alt="Generated Image")
+                            logger.debug(f"Image URL extracted: {image_url}")  # Added logging
+                            return image_response
                         else:
-                            final_response = cleaned_response
+                            logger.debug("No image URL found in the response.")  # Added logging
+                            return cleaned_response
                     else:
-                        if '$~~~$' in cleaned_response:
-                            final_response = cleaned_response.split('$~~~$')[0].strip()
+                        if websearch:
+                            match = re.search(r'\$~~~\$(.*?)\$~~~\$', cleaned_response, re.DOTALL)
+                            if match:
+                                source_part = match.group(1).strip()
+                                answer_part = cleaned_response[match.end():].strip()
+                                try:
+                                    sources = json.loads(source_part)
+                                    source_formatted = "**Source:**\n"
+                                    for item in sources:
+                                        title = item.get('title', 'No Title')
+                                        link = item.get('link', '#')
+                                        position = item.get('position', '')
+                                        source_formatted += f"{position}. [{title}]({link})\n"
+                                    final_response = f"{answer_part}\n\n{source_formatted}"
+                                except json.JSONDecodeError:
+                                    final_response = f"{answer_part}\n\nSource information is unavailable."
+                            else:
+                                final_response = cleaned_response
                         else:
-                            final_response = cleaned_response
+                            if '$~~~$' in cleaned_response:
+                                final_response = cleaned_response.split('$~~~$')[0].strip()
+                            else:
+                                final_response = cleaned_response
 
-                    logger.debug(f"Final response to return: {final_response}")  # Added logging
-                    return final_response
+                        logger.debug(f"Final response to return: {final_response}")  # Added logging
+                        return final_response
+            except ClientResponseError as e:
+                error_text = f"Error {e.status}: {e.message}"
+                try:
+                    error_response = await e.response.text()
+                    cleaned_error = cls.clean_response(error_response)
+                    error_text += f" - {cleaned_error}"
+                    logger.error(f"ClientResponseError: {error_text}")  # Added logging
+                except Exception:
+                    pass
+                return error_text
+            except Exception as e:
+                logger.exception(f"Unexpected error during /api/chat request: {str(e)}")  # Added logging
+                return f"Unexpected error during /api/chat request: {str(e)}"
+
+        chat_url = f'{cls.url}/chat/{chat_id}?model={model}'
+
+        try:
+            async with session.post(
+                chat_url,
+                headers=headers_chat_combined,
+                data=data_chat,
+                proxy=proxy
+            ) as response_chat:
+                response_chat.raise_for_status()
+                pass
         except ClientResponseError as e:
             error_text = f"Error {e.status}: {e.message}"
             try:
                 error_response = await e.response.text()
                 cleaned_error = cls.clean_response(error_response)
                 error_text += f" - {cleaned_error}"
-                logger.error(f"ClientResponseError: {error_text}")  # Added logging
             except Exception:
                 pass
             return error_text
         except Exception as e:
-            logger.exception(f"Unexpected error during /api/chat request: {str(e)}")  # Added logging
-            return f"Unexpected error during /api/chat request: {str(e)}"
+            return f"Unexpected error during /chat/{chat_id} request: {str(e)}"
 
     @classmethod
     async def create_async_generator(
@@ -485,15 +508,19 @@ async def generate_response(
                 ) as response_api_chat:
                     response_api_chat.raise_for_status()
                     text = await response_api_chat.text()
+                    logger.debug(f"Raw response from Blackbox API: {text}")  # Added logging
                     cleaned_response = cls.clean_response(text)
+                    logger.debug(f"Cleaned response: {cleaned_response}")  # Added logging
 
                     if model in cls.image_models:
                         match = re.search(r'!\[.*?\]\((https?://[^\)]+)\)', cleaned_response)
                         if match:
                             image_url = match.group(1)
                             image_response = ImageResponseModel(images=image_url, alt="Generated Image")
+                            logger.debug(f"Image URL extracted: {image_url}")  # Added logging
                             yield image_response
                         else:
+                            logger.debug("No image URL found in the response.")  # Added logging
                             yield cleaned_response
                     else:
                         if websearch:
@@ -520,6 +547,7 @@ async def generate_response(
                             else:
                                 final_response = cleaned_response
 
+                        logger.debug(f"Final response to yield: {final_response}")  # Added logging
                         yield final_response
             except ClientResponseError as e:
                 error_text = f"Error {e.status}: {e.message}"
@@ -527,38 +555,52 @@ async def generate_response(
                     error_response = await e.response.text()
                     cleaned_error = cls.clean_response(error_response)
                     error_text += f" - {cleaned_error}"
+                    logger.error(f"ClientResponseError: {error_text}")  # Added logging
                 except Exception:
                     pass
                 yield error_text
             except Exception as e:
                 yield f"Unexpected error during /api/chat request: {str(e)}"
 
-            chat_url = f'{cls.url}/chat/{chat_id}?model={model}'
+        chat_url = f'{cls.url}/chat/{chat_id}?model={model}'
 
+        try:
+            async with session.post(
+                chat_url,
+                headers=headers_chat_combined,
+                data=data_chat,
+                proxy=proxy
+            ) as response_chat:
+                response_chat.raise_for_status()
+                pass
+        except ClientResponseError as e:
+            error_text = f"Error {e.status}: {e.message}"
             try:
-                async with session.post(
-                    chat_url,
-                    headers=headers_chat_combined,
-                    data=data_chat,
-                    proxy=proxy
-                ) as response_chat:
-                    response_chat.raise_for_status()
-                    pass
-            except ClientResponseError as e:
-                error_text = f"Error {e.status}: {e.message}"
-                try:
-                    error_response = await e.response.text()
-                    cleaned_error = cls.clean_response(error_response)
-                    error_text += f" - {cleaned_error}"
-                except Exception:
-                    pass
-                yield error_text
-            except Exception as e:
-                yield f"Unexpected error during /chat/{chat_id} request: {str(e)}"
+                error_response = await e.response.text()
+                cleaned_error = cls.clean_response(error_response)
+                error_text += f" - {cleaned_error}"
+            except Exception:
+                pass
+            yield error_text
+        except Exception as e:
+            yield f"Unexpected error during /chat/{chat_id} request: {str(e)}"
 
 # FastAPI app setup
 app = FastAPI()
 
+# Rate Limiter Cleanup Task
+async def cleanup_rate_limit_stores():
+    """
+    Periodically cleans up stale entries in the rate_limit_store to prevent memory bloat.
+    """
+    while True:
+        current_time = time.time()
+        ips_to_delete = [ip for ip, value in rate_limit_store.items() if current_time - value["timestamp"] > RATE_LIMIT_WINDOW * 2]
+        for ip in ips_to_delete:
+            del rate_limit_store[ip]
+            logger.debug(f"Cleaned up rate_limit_store for IP: {ip}")
+        await asyncio.sleep(CLEANUP_INTERVAL)
+
 # Add the cleanup task when the app starts
 @app.on_event("startup")
 async def startup_event():
@@ -605,19 +647,6 @@ class ChatRequest(BaseModel):
     logit_bias: Optional[Dict[str, float]] = None
     user: Optional[str] = None
 
-# Rate Limiter Cleanup Task
-async def cleanup_rate_limit_stores():
-    """
-    Periodically cleans up stale entries in the rate_limit_store to prevent memory bloat.
-    """
-    while True:
-        current_time = time.time()
-        ips_to_delete = [ip for ip, value in rate_limit_store.items() if current_time - value["timestamp"] > RATE_LIMIT_WINDOW * 2]
-        for ip in ips_to_delete:
-            del rate_limit_store[ip]
-            logger.debug(f"Cleaned up rate_limit_store for IP: {ip}")
-        await asyncio.sleep(CLEANUP_INTERVAL)
-
 # Rate Limiter Dependency
 async def rate_limiter_per_ip(request: Request):
     """
@@ -682,7 +711,7 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
             return {
                 "id": f"chatcmpl-{uuid.uuid4()}",
                 "object": "chat.completion",
-                "created": int(datetime.now().timestamp()),
+                "created": int(datetime.now().timestamp()),  # Fixed: datetime is now imported
                 "model": request.model,
                 "choices": [
                     {
@@ -705,7 +734,7 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
         return {
             "id": f"chatcmpl-{uuid.uuid4()}",
             "object": "chat.completion",
-            "created": int(datetime.now().timestamp()),
+            "created": int(datetime.now().timestamp()),  # Fixed: datetime is now imported
             "model": request.model,
             "choices": [
                 {