Update main.py

main.py

@@ -21,7 +21,7 @@ from pydantic import BaseModel
 
 # Configure logging
 logging.basicConfig(
-    level=logging.INFO,  # Set to DEBUG for more detailed logs
+    level=logging.DEBUG,  # Set to DEBUG for detailed logs; change to INFO in production
     format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
     handlers=[logging.StreamHandler()]
 )
@@ -237,6 +237,12 @@ class Blackbox:
         websearch: bool = False,
         **kwargs
     ) -> Union[str, ImageResponseModel]:
+        """
+        Generates a response by interacting with the Blackbox API.
+        Performs two POST requests:
+        1. Initiates the chat and obtains a chat_id.
+        2. Retrieves the chat response using the chat_id.
+        """
         model = cls.get_model(model)
         chat_id = cls.generate_random_string()
         next_action = cls.generate_next_action()
@@ -328,6 +334,7 @@ class Blackbox:
 
         async with ClientSession(headers=common_headers) as session:
             try:
+                # First POST request to initiate the chat and get chat_id
                 async with session.post(
                     cls.api_endpoint,
                     headers=headers_api_chat_combined,
@@ -336,12 +343,26 @@ class Blackbox:
                 ) 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
+                    logger.debug(f"Raw response from Blackbox API (initiate chat): {text}")  # Added logging
                     cleaned_response = cls.clean_response(text)
-                    logger.debug(f"Cleaned response: {cleaned_response}")  # Added logging
+                    logger.debug(f"Cleaned response (initiate chat): {cleaned_response}")  # Added logging
+
+                # Second POST request to retrieve the chat response using chat_id
+                chat_url = f'{cls.url}/chat/{chat_id}?model={model}'
+                async with session.post(
+                    chat_url,
+                    headers=headers_chat_combined,
+                    data=data_chat,
+                    proxy=proxy
+                ) as response_chat:
+                    response_chat.raise_for_status()
+                    chat_response_text = await response_chat.text()
+                    logger.debug(f"Raw response from Blackbox API (retrieve chat): {chat_response_text}")  # Added logging
+                    cleaned_chat_response = cls.clean_response(chat_response_text)
+                    logger.debug(f"Cleaned response (retrieve chat): {cleaned_chat_response}")  # Added logging
 
                     if model in cls.image_models:
-                        match = re.search(r'!\[.*?\]\((https?://[^\)]+)\)', cleaned_response)
+                        match = re.search(r'!\[.*?\]\((https?://[^\)]+)\)', cleaned_chat_response)
                         if match:
                             image_url = match.group(1)
                             image_response = ImageResponseModel(images=image_url, alt="Generated Image")
@@ -349,13 +370,13 @@ class Blackbox:
                             return image_response
                         else:
                             logger.debug("No image URL found in the response.")  # Added logging
-                            return cleaned_response
+                            return cleaned_chat_response
                     else:
                         if websearch:
-                            match = re.search(r'\$~~~\$(.*?)\$~~~\$', cleaned_response, re.DOTALL)
+                            match = re.search(r'\$~~~\$(.*?)\$~~~\$', cleaned_chat_response, re.DOTALL)
                             if match:
                                 source_part = match.group(1).strip()
-                                answer_part = cleaned_response[match.end():].strip()
+                                answer_part = cleaned_chat_response[match.end():].strip()
                                 try:
                                     sources = json.loads(source_part)
                                     source_formatted = "**Source:**\n"
@@ -368,12 +389,12 @@ class Blackbox:
                                 except json.JSONDecodeError:
                                     final_response = f"{answer_part}\n\nSource information is unavailable."
                             else:
-                                final_response = cleaned_response
+                                final_response = cleaned_chat_response
                         else:
-                            if '$~~~$' in cleaned_response:
-                                final_response = cleaned_response.split('$~~~$')[0].strip()
+                            if '$~~~$' in cleaned_chat_response:
+                                final_response = cleaned_chat_response.split('$~~~$')[0].strip()
                             else:
-                                final_response = cleaned_response
+                                final_response = cleaned_chat_response
 
                         logger.debug(f"Final response to return: {final_response}")  # Added logging
                         return final_response
@@ -391,31 +412,6 @@ class Blackbox:
                 logger.exception(f"Unexpected error during /api/chat request: {str(e)}")  # Added logging
                 return f"Unexpected error during /api/chat request: {str(e)}"
 
-        # The following block is unreachable due to the return statements above
-        # It can be removed unless it's intended for future use
-        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
-            return error_text
-        except Exception as e:
-            return f"Unexpected error during /chat/{chat_id} request: {str(e)}"
-
     @classmethod
     async def create_async_generator(
         cls,
@@ -427,194 +423,22 @@ class Blackbox:
     ) -> AsyncGenerator[Union[str, ImageResponseModel], 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, ImageResponseModel]: Segments of the generated response or ImageResponse objects.
         """
-        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'
-        }
-
-        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:
-                            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:
-                            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:
-                            if '$~~~$' in cleaned_response:
-                                final_response = cleaned_response.split('$~~~$')[0].strip()
-                            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}"
-                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
-                yield error_text
-            except Exception as e:
-                yield f"Unexpected error during /api/chat request: {str(e)}"
-
-        # The following block is unreachable due to the yield statements above
-        # It can be removed unless it's intended for future use
-        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
+            response = await cls.generate_response(
+                model=model,
+                messages=messages,
+                proxy=proxy,
+                websearch=websearch,
+                **kwargs
+            )
+            if isinstance(response, ImageResponseModel):
+                yield response
+            else:
+                yield response
         except Exception as e:
-            yield f"Unexpected error during /chat/{chat_id} request: {str(e)}"
+            logger.exception(f"Error in create_async_generator: {str(e)}")
+            yield f"Unexpected error: {str(e)}"
 
 # ----------------------------- FastAPI App Setup -----------------------------