Update main.py

main.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import os
 import re
 import random
@@ -53,7 +55,8 @@ async def cleanup_rate_limit_stores():
     """
     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]
+        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}")
@@ -118,6 +121,7 @@ class Blackbox:
     models = [
         default_model,
         'blackboxai-pro',
+        *image_models,
         "llama-3.1-8b",
         'llama-3.1-70b',
         'llama-3.1-405b',
@@ -138,8 +142,6 @@ class Blackbox:
         'ReactAgent',
         'XcodeAgent',
         'AngularJSAgent',
-        *image_models,
-        'Niansuh',
     ]
 
     # Filter models based on AVAILABLE_MODELS
@@ -148,7 +150,6 @@ class Blackbox:
 
     agentMode = {
         'ImageGeneration': {'mode': True, 'id': "ImageGenerationLV45LJp", 'name': "Image Generation"},
-        'Niansuh': {'mode': True, 'id': "NiansuhAIk1HgESy", 'name': "Niansuh"},
     }
     trendingAgentMode = {
         "blackboxai": {},
@@ -197,33 +198,65 @@ class Blackbox:
         'AngularJSAgent': '@AngularJS Agent',
         'blackboxai-pro': '@BLACKBOXAI-PRO',
         'ImageGeneration': '@Image Generation',
-        'Niansuh': '@Niansuh',
     }
 
     model_referers = {
-        "blackboxai": f"{url}/?model=blackboxai",
-        "gpt-4o": f"{url}/?model=gpt-4o",
-        "gemini-pro": f"{url}/?model=gemini-pro",
-        "claude-sonnet-3.5": f"{url}/?model=claude-sonnet-3.5"
+        "blackboxai": "/?model=blackboxai",
+        "gpt-4o": "/?model=gpt-4o",
+        "gemini-pro": "/?model=gemini-pro",
+        "claude-sonnet-3.5": "/?model=claude-sonnet-3.5"
     }
 
     model_aliases = {
         "gemini-flash": "gemini-1.5-flash",
         "claude-3.5-sonnet": "claude-sonnet-3.5",
         "flux": "ImageGeneration",
-        "niansuh": "Niansuh",
     }
 
     @classmethod
-    def get_model(cls, model: str) -> Optional[str]:
+    def get_model(cls, model: str) -> str:
         if model in cls.models:
             return model
-        elif model in cls.userSelectedModel and cls.userSelectedModel[model] in cls.models:
-            return cls.userSelectedModel[model]
-        elif model in cls.model_aliases and cls.model_aliases[model] in cls.models:
+        elif model in cls.model_aliases:
             return cls.model_aliases[model]
         else:
-            return cls.default_model if cls.default_model in cls.models else None
+            return cls.default_model
+
+    @staticmethod
+    def generate_random_string(length: int = 7) -> str:
+        characters = string.ascii_letters + string.digits
+        return ''.join(random.choices(characters, k=length))
+
+    @staticmethod
+    def generate_next_action() -> str:
+        return uuid.uuid4().hex
+
+    @staticmethod
+    def generate_next_router_state_tree() -> str:
+        router_state = [
+            "",
+            {
+                "children": [
+                    "(chat)",
+                    {
+                        "children": [
+                            "__PAGE__",
+                            {}
+                        ]
+                    }
+                ]
+            },
+            None,
+            None,
+            True
+        ]
+        return json.dumps(router_state)
+
+    @staticmethod
+    def clean_response(text: str) -> str:
+        pattern = r'^\$\@\$v=undefined-rv1\$\@\$'
+        cleaned_text = re.sub(pattern, '', text)
+        return cleaned_text
 
     @classmethod
     async def create_async_generator(
@@ -231,73 +264,87 @@ class Blackbox:
         model: str,
         messages: List[Dict[str, str]],
         proxy: Optional[str] = None,
-        image: Any = None,
-        image_name: Optional[str] = None,
-        webSearchMode: bool = False,
+        websearch: bool = False,
         **kwargs
-    ) -> AsyncGenerator[Any, None]:
+    ) -> AsyncGenerator[Union[str, ImageResponse], None]:
+        """
+        Creates an asynchronous generator for streaming responses from Blackbox AI.
+
+        Parameters:
+            model (str): Model to use for generating responses.
+            messages (List[Dict[str, str]]): Message history.
+            proxy (Optional[str]): Proxy URL, if needed.
+            websearch (bool): Enables or disables web search mode.
+            **kwargs: Additional keyword arguments.
+
+        Yields:
+            Union[str, ImageResponse]: Segments of the generated response or ImageResponse objects.
+        """
         model = cls.get_model(model)
-        if model is None:
-            logger.error(f"Model {model} is not available.")
-            raise ModelNotWorkingException(model)
 
-        logger.info(f"Selected model: {model}")
+        chat_id = cls.generate_random_string()
+        next_action = cls.generate_next_action()
+        next_router_state_tree = cls.generate_next_router_state_tree()
 
-        if not cls.working or model not in cls.models:
-            logger.error(f"Model {model} is not working or not supported.")
-            raise ModelNotWorkingException(model)
-        
-        headers = {
-            "accept": "*/*",
-            "accept-language": "en-US,en;q=0.9",
-            "cache-control": "no-cache",
-            "content-type": "application/json",
-            "origin": cls.url,
-            "pragma": "no-cache",
-            "priority": "u=1, i",
-            "referer": cls.model_referers.get(model, cls.url),
-            "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",
-        }
+        agent_mode = cls.agentMode.get(model, {})
+        trending_agent_mode = cls.trendingAgentMode.get(model, {})
 
-        if model in cls.model_prefixes:
-            prefix = cls.model_prefixes[model]
-            if not messages[0]['content'].startswith(prefix):
-                logger.debug(f"Adding prefix '{prefix}' to the first message.")
-                messages[0]['content'] = f"{prefix} {messages[0]['content']}"
+        prefix = cls.model_prefixes.get(model, "")
         
-        random_id = ''.join(random.choices(string.ascii_letters + string.digits, k=7))
-        messages[-1]['id'] = random_id
-        messages[-1]['role'] = 'user'
-
-        # Don't log the full message content for privacy
-        logger.debug(f"Generated message ID: {random_id} for model: {model}")
-
-        if image is not None:
-            messages[-1]['data'] = {
-                'fileText': '',
-                'imageBase64': to_data_uri(image),
-                'title': image_name
-            }
-            messages[-1]['content'] = 'FILE:BB\n$#$\n\n$#$\n' + messages[-1]['content']
-            logger.debug("Image data added to the message.")
+        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"
         
-        data = {
-            "messages": messages,
-            "id": random_id,
+        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'
+        }
+
+        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": {},
-            "trendingAgentMode": {},
+            "agentMode": agent_mode,
+            "trendingAgentMode": trending_agent_mode,
             "isMicMode": False,
             "userSystemPrompt": None,
-            "maxTokens": 99999999,
+            "maxTokens": 1024,
             "playgroundTopP": 0.9,
             "playgroundTemperature": 0.5,
             "isChromeExt": False,
@@ -307,81 +354,102 @@ class Blackbox:
             "clickedForceWebSearch": False,
             "visitFromDelta": False,
             "mobileClient": False,
-            "userSelectedModel": None,
-            "webSearchMode": webSearchMode,
+            "webSearchMode": websearch,
+            "userSelectedModel": cls.userSelectedModel.get(model, model)
         }
 
-        if model in cls.agentMode:
-            data["agentMode"] = cls.agentMode[model]
-        elif model in cls.trendingAgentMode:
-            data["trendingAgentMode"] = cls.trendingAgentMode[model]
-        elif model in cls.userSelectedModel:
-            data["userSelectedModel"] = cls.userSelectedModel[model]
-        logger.info(f"Sending request to {cls.api_endpoint} with data (excluding messages).")
+        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}
 
-        timeout = ClientTimeout(total=60)  # Set an appropriate timeout
-        retry_attempts = 10  # Set the number of retry attempts
+        data_chat = '[]'
 
-        for attempt in range(retry_attempts):
+        async with ClientSession(headers=common_headers) as session:
             try:
-                async with ClientSession(headers=headers, timeout=timeout) as session:
-                    async with session.post(cls.api_endpoint, json=data, proxy=proxy) as response:
-                        response.raise_for_status()
-                        logger.info(f"Received response with status {response.status}")
-                        if model == 'ImageGeneration':
-                            response_text = await response.text()
-                            url_match = re.search(r'https://storage\.googleapis\.com/[^\s\)]+', response_text)
-                            if url_match:
-                                image_url = url_match.group(0)
-                                logger.info(f"Image URL found.")
-                                yield ImageResponse(image_url, alt=messages[-1]['content'])
+                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()
+                    cleaned_response = cls.clean_response(text)
+
+                    if model in cls.image_models:
+                        match = re.search(r'!\[.*?\]\((https?://[^\)]+)\)', cleaned_response)
+                        if match:
+                            image_url = match.group(1)
+                            image_response = ImageResponse(url=image_url, alt="Generated Image")
+                            yield image_response
+                        else:
+                            yield cleaned_response
+                    else:
+                        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 = "**Sources:**\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:
-                                logger.error("Image URL not found in the response.")
-                                raise Exception("Image URL not found in the response")
+                                final_response = cleaned_response
                         else:
-                            full_response = ""
-                            search_results_json = ""
-                            try:
-                                async for chunk, _ in response.content.iter_chunks():
-                                    if chunk:
-                                        decoded_chunk = chunk.decode(errors='ignore')
-                                        decoded_chunk = re.sub(r'\$@\$v=[^$]+\$@\$', '', decoded_chunk)
-                                        if decoded_chunk.strip():
-                                            if '$~~~$' in decoded_chunk:
-                                                search_results_json += decoded_chunk
-                                            else:
-                                                full_response += decoded_chunk
-                                                yield decoded_chunk
-                                logger.info("Finished streaming response chunks.")
-                            except Exception as e:
-                                logger.exception("Error while iterating over response chunks.")
-                                raise e
-                            if data["webSearchMode"] and search_results_json:
-                                match = re.search(r'\$~~~\$(.*?)\$~~~\$', search_results_json, re.DOTALL)
-                                if match:
-                                    try:
-                                        search_results = json.loads(match.group(1))
-                                        formatted_results = "\n\n**Sources:**\n"
-                                        for i, result in enumerate(search_results[:5], 1):
-                                            formatted_results += f"{i}. [{result['title']}]({result['link']})\n"
-                                        logger.info("Formatted search results.")
-                                        yield formatted_results
-                                    except json.JSONDecodeError as je:
-                                        logger.error("Failed to parse search results JSON.")
-                                        raise je
-                break  # Exit the retry loop if successful
-            except ClientError as ce:
-                logger.error(f"Client error occurred: {ce}. Retrying attempt {attempt + 1}/{retry_attempts}")
-                if attempt == retry_attempts - 1:
-                    raise HTTPException(status_code=502, detail="Error communicating with the external API.")
-            except asyncio.TimeoutError:
-                logger.error(f"Request timed out. Retrying attempt {attempt + 1}/{retry_attempts}")
-                if attempt == retry_attempts - 1:
-                    raise HTTPException(status_code=504, detail="External API request timed out.")
+                            if '$~~~$' in cleaned_response:
+                                final_response = cleaned_response.split('$~~~$')[0].strip()
+                            else:
+                                final_response = cleaned_response
+
+                        yield 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}"
+                except Exception:
+                    pass
+                yield error_text
             except Exception as e:
-                logger.error(f"Unexpected error: {e}. Retrying attempt {attempt + 1}/{retry_attempts}")
-                if attempt == retry_attempts - 1:
-                    raise HTTPException(status_code=500, detail=str(e))
+                yield 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}"
+                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()
@@ -466,6 +534,7 @@ def create_response(content: str, model: str, finish_reason: Optional[str] = Non
         "usage": None,  # To be filled in non-streaming responses
     }
 
+# Existing /v1/chat/completions Endpoint
 @app.post("/v1/chat/completions", dependencies=[Depends(rate_limiter_per_ip)])
 async def chat_completions(request: ChatRequest, req: Request, api_key: str = Depends(get_api_key)):
     client_ip = req.client.host
@@ -486,7 +555,7 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
             messages=[{"role": msg.role, "content": msg.content} for msg in request.messages],  # Actual message content used here
             image=None,
             image_name=None,
-            webSearchMode=request.webSearchMode
+            websearch=request.webSearchMode
         )
 
         if request.stream:
@@ -606,6 +675,45 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
         logger.exception(f"An unexpected error occurred while processing the chat completions request from IP: {client_ip}.")
         raise HTTPException(status_code=500, detail=str(e))
 
+# New Web Search Endpoint
+class WebSearchRequest(BaseModel):
+    query: str
+    num_results: Optional[int] = 5
+
+@app.post("/v1/websearch", dependencies=[Depends(rate_limiter_per_ip)])
+async def web_search(request: WebSearchRequest, req: Request, api_key: str = Depends(get_api_key)):
+    """
+    Handles web search requests.
+
+    Parameters:
+        query (str): The search query.
+        num_results (int): Number of search results to return.
+
+    Returns:
+        JSONResponse: Contains the search results.
+    """
+    client_ip = req.client.host
+    logger.info(f"Received web search request from API key: {api_key} | IP: {client_ip} | Query: {request.query} | Num Results: {request.num_results}")
+
+    # Implement your web search logic here.
+    # This is a mock implementation. Replace it with actual web search integration.
+
+    try:
+        # Mock search results
+        search_results = []
+        for i in range(1, request.num_results + 1):
+            search_results.append({
+                "position": i,
+                "title": f"Sample Search Result {i} for '{request.query}'",
+                "link": f"https://www.example.com/search-result-{i}"
+            })
+
+        logger.info(f"Web search completed for query: {request.query}")
+        return {"results": search_results}
+    except Exception as e:
+        logger.exception(f"An error occurred during web search from IP: {client_ip}.")
+        raise HTTPException(status_code=500, detail="An error occurred during web search.")
+
 # Endpoint: POST /v1/tokenizer
 @app.post("/v1/tokenizer", dependencies=[Depends(rate_limiter_per_ip)])
 async def tokenizer(request: TokenizerRequest, req: Request):