From 19e62f85fcfb4c1d9f961906f01e7e580b245ca5 Mon Sep 17 00:00:00 2001 From: cauvang32 Date: Tue, 26 Aug 2025 23:36:47 +0700 Subject: [PATCH] refactor: add user preference for tool execution display and enhance message handling for search and scraping activities --- src/commands/commands.py | 25 ++ src/database/db_handler.py | 15 ++ src/module/message_handler.py | 488 ++++++++++++++++++---------------- 3 files changed, 295 insertions(+), 233 deletions(-) diff --git a/src/commands/commands.py b/src/commands/commands.py index 1e4dd1b..0bdaa90 100644 --- a/src/commands/commands.py +++ b/src/commands/commands.py @@ -376,12 +376,37 @@ def setup_commands(bot: commands.Bot, db_handler, openai_client, image_generator "/search `` - Search Google and send results to the AI model.\n" "/web `` - Scrape a webpage and send the data to the AI model.\n" "/generate `` - Generate an image from a text prompt.\n" + "/toggle_tools - Toggle display of tool execution details (code, input, output).\n" "/reset - Reset your chat history.\n" "/user_stat - Get information about your input tokens, output tokens, and current model.\n" "/help - Display this help message.\n" ) await interaction.response.send_message(help_message, ephemeral=True) + @tree.command(name="toggle_tools", description="Toggle the display of tool execution details (code, input, output).") + @check_blacklist() + async def toggle_tools(interaction: discord.Interaction): + """Toggle the display of tool execution details for the user.""" + await interaction.response.defer(ephemeral=True) + + user_id = interaction.user.id + current_setting = await db_handler.get_user_tool_display(user_id) + new_setting = not current_setting + + await db_handler.set_user_tool_display(user_id, new_setting) + + status = "enabled" if new_setting else "disabled" + description = ( + "You will now see detailed execution information including code, input, and output when tools are used." + if new_setting else + "Tool execution details are now hidden. You'll only see the final results." + ) + + await interaction.followup.send( + f"🔧 **Tool Display {status.title()}**\n{description}", + ephemeral=True + ) + @tree.command(name="stop", description="Stop any process or queue of the user. Admins can stop other users' tasks by providing their ID.") @app_commands.describe(user_id="The Discord user ID to stop tasks for (admin only)") @check_blacklist() diff --git a/src/database/db_handler.py b/src/database/db_handler.py index 9b288ad..efcfa2f 100644 --- a/src/database/db_handler.py +++ b/src/database/db_handler.py @@ -112,6 +112,20 @@ class DatabaseHandler: upsert=True ) + # Tool display preferences + async def get_user_tool_display(self, user_id: int) -> bool: + """Get user's tool display preference (default: False - disabled)""" + user_data = await self.db.user_preferences.find_one({'user_id': user_id}) + return user_data.get('show_tool_execution', False) if user_data else False + + async def set_user_tool_display(self, user_id: int, show_tools: bool) -> None: + """Set user's tool display preference""" + await self.db.user_preferences.update_one( + {'user_id': user_id}, + {'$set': {'show_tool_execution': show_tools}}, + upsert=True + ) + # Admin and permissions management with caching async def is_admin(self, user_id: int) -> bool: """Check if the user is an admin (no caching for security)""" @@ -163,6 +177,7 @@ class DatabaseHandler: """Create indexes for better query performance""" await self.db.user_histories.create_index("user_id") await self.db.user_models.create_index("user_id") + await self.db.user_preferences.create_index("user_id") await self.db.whitelist.create_index("user_id") await self.db.blacklist.create_index("user_id") diff --git a/src/module/message_handler.py b/src/module/message_handler.py index 37f1165..529d523 100644 --- a/src/module/message_handler.py +++ b/src/module/message_handler.py @@ -241,115 +241,119 @@ class MessageHandler: # Display the executed code information in Discord (but not save to history) if discord_message and code_to_execute: - try: - # Clean up the code display (remove file context comments) - code_lines = code_to_execute.split('\n') - clean_code_lines = [] - for line in code_lines: - if not (line.strip().startswith('# Data file available:') or - line.strip().startswith('# File path:') or - line.strip().startswith('# You can access this file using:')): - clean_code_lines.append(line) - - clean_code = '\n'.join(clean_code_lines).strip() - - # Check if code is too long for Discord message (3000 chars limit) - if len(clean_code) > 3000: - # Send code as file attachment - code_file = discord.File( - io.StringIO(clean_code), - filename="executed_code.py" - ) + # Check user's tool display preference + show_execution_details = await self.db_handler.get_user_tool_display(user_id) if user_id else False + + if show_execution_details: + try: + # Clean up the code display (remove file context comments) + code_lines = code_to_execute.split('\n') + clean_code_lines = [] + for line in code_lines: + if not (line.strip().startswith('# Data file available:') or + line.strip().startswith('# File path:') or + line.strip().startswith('# You can access this file using:')): + clean_code_lines.append(line) - # Create display text without code - execution_display = "**🐍 Python Code Execution**\n\n" + clean_code = '\n'.join(clean_code_lines).strip() - # Show packages to install if any - if packages_to_install: - execution_display += f"**📦 Installing packages:** {', '.join(packages_to_install)}\n\n" - - # Show input data if any - if input_data: - execution_display += "**📥 Input:**\n```\n" - execution_display += input_data[:500] # Limit input length - if len(input_data) > 500: - execution_display += "\n... (input truncated)" - execution_display += "\n```\n\n" - - execution_display += "**💻 Code:** *Attached as file (too long to display)*\n\n" - - # Show the output - if execute_result and execute_result.get("success"): - output = execute_result.get("output", "") - # Remove package installation info from output if it exists - if output and "Installed packages:" in output: - lines = output.split('\n') - output = '\n'.join(lines[2:]) if len(lines) > 2 else "" + # Check if code is too long for Discord message (3000 chars limit) + if len(clean_code) > 3000: + # Send code as file attachment + code_file = discord.File( + io.StringIO(clean_code), + filename="executed_code.py" + ) - if output and output.strip(): - execution_display += "**📤 Output:**\n```\n" - execution_display += output[:2000] # More space for output when code is attached - if len(output) > 2000: - execution_display += "\n... (output truncated)" - execution_display += "\n```" - else: - execution_display += "**📤 Output:** *(No output)*" - else: - error_msg = execute_result.get("error", "Unknown error") if execute_result else "Execution failed" - execution_display += f"**❌ Error:**\n```\n{error_msg[:1000]}\n```" - if len(error_msg) > 1000: - execution_display += "*(Error message truncated)*" - - # Send with file attachment - await discord_message.channel.send(execution_display, file=code_file) - else: - # Use normal display for shorter code - execution_display = "**🐍 Python Code Execution**\n\n" - - # Show packages to install if any - if packages_to_install: - execution_display += f"**📦 Installing packages:** {', '.join(packages_to_install)}\n\n" - - # Show input data if any - if input_data: - execution_display += "**📥 Input:**\n```\n" - execution_display += input_data[:500] # Limit input length - if len(input_data) > 500: - execution_display += "\n... (input truncated)" - execution_display += "\n```\n\n" - - # Show the actual code - execution_display += "**💻 Code:**\n```python\n" - execution_display += clean_code - execution_display += "\n```\n\n" - - # Show the output - if execute_result and execute_result.get("success"): - output = execute_result.get("output", "") - # Remove package installation info from output if it exists - if output and "Installed packages:" in output: - lines = output.split('\n') - output = '\n'.join(lines[2:]) if len(lines) > 2 else "" + # Create display text without code + execution_display = "**🐍 Python Code Execution**\n\n" - if output and output.strip(): - execution_display += "**📤 Output:**\n```\n" - execution_display += output[:1000] # Limit output length for Discord - if len(output) > 1000: - execution_display += "\n... (output truncated)" - execution_display += "\n```" + # Show packages to install if any + if packages_to_install: + execution_display += f"**📦 Installing packages:** {', '.join(packages_to_install)}\n\n" + + # Show input data if any + if input_data: + execution_display += "**📥 Input:**\n```\n" + execution_display += input_data[:500] # Limit input length + if len(input_data) > 500: + execution_display += "\n... (input truncated)" + execution_display += "\n```\n\n" + + execution_display += "**💻 Code:** *Attached as file (too long to display)*\n\n" + + # Show the output + if execute_result and execute_result.get("success"): + output = execute_result.get("output", "") + # Remove package installation info from output if it exists + if output and "Installed packages:" in output: + lines = output.split('\n') + output = '\n'.join(lines[2:]) if len(lines) > 2 else "" + + if output and output.strip(): + execution_display += "**📤 Output:**\n```\n" + execution_display += output[:2000] # More space for output when code is attached + if len(output) > 2000: + execution_display += "\n... (output truncated)" + execution_display += "\n```" + else: + execution_display += "**📤 Output:** *(No output)*" else: - execution_display += "**📤 Output:** *(No output)*" + error_msg = execute_result.get("error", "Unknown error") if execute_result else "Execution failed" + execution_display += f"**❌ Error:**\n```\n{error_msg[:1000]}\n```" + if len(error_msg) > 1000: + execution_display += "*(Error message truncated)*" + + # Send with file attachment + await discord_message.channel.send(execution_display, file=code_file) else: - error_msg = execute_result.get("error", "Unknown error") if execute_result else "Execution failed" - execution_display += f"**❌ Error:**\n```\n{error_msg[:800]}\n```" - if len(error_msg) > 800: - execution_display += "*(Error message truncated)*" + # Use normal display for shorter code + execution_display = "**🐍 Python Code Execution**\n\n" + + # Show packages to install if any + if packages_to_install: + execution_display += f"**📦 Installing packages:** {', '.join(packages_to_install)}\n\n" + + # Show input data if any + if input_data: + execution_display += "**📥 Input:**\n```\n" + execution_display += input_data[:500] # Limit input length + if len(input_data) > 500: + execution_display += "\n... (input truncated)" + execution_display += "\n```\n\n" + + # Show the actual code + execution_display += "**💻 Code:**\n```python\n" + execution_display += clean_code + execution_display += "\n```\n\n" + + # Show the output + if execute_result and execute_result.get("success"): + output = execute_result.get("output", "") + # Remove package installation info from output if it exists + if output and "Installed packages:" in output: + lines = output.split('\n') + output = '\n'.join(lines[2:]) if len(lines) > 2 else "" + + if output and output.strip(): + execution_display += "**📤 Output:**\n```\n" + execution_display += output[:1000] # Limit output length for Discord + if len(output) > 1000: + execution_display += "\n... (output truncated)" + execution_display += "\n```" + else: + execution_display += "**📤 Output:** *(No output)*" + else: + error_msg = execute_result.get("error", "Unknown error") if execute_result else "Execution failed" + execution_display += f"**❌ Error:**\n```\n{error_msg[:800]}\n```" + if len(error_msg) > 800: + execution_display += "*(Error message truncated)*" + + # Send the execution display to Discord as a separate message + await discord_message.channel.send(execution_display) - # Send the execution display to Discord as a separate message - await discord_message.channel.send(execution_display) - - except Exception as e: - logging.error(f"Error displaying code execution: {str(e)}") + except Exception as e: + logging.error(f"Error displaying code execution: {str(e)}") # If there are visualizations, handle them if execute_result and execute_result.get("visualizations"): @@ -1349,6 +1353,11 @@ class MessageHandler: async def _google_search(self, args: Dict[str, Any]): """Perform a Google search with Discord display""" try: + # Find user_id from current task context + user_id = args.get("user_id") + if not user_id: + user_id = self._find_user_id_from_current_task() + # Get the Discord message to display search activity discord_message = self._get_discord_message_from_current_task() @@ -1360,83 +1369,87 @@ class MessageHandler: from src.utils.web_utils import google_search result = await google_search(args) - # Display the search activity in Discord + # Display the search activity in Discord (only if user has enabled tool display) if discord_message and query: - try: - # Parse the result to get structured data - import json - search_data = json.loads(result) if isinstance(result, str) else result - - # Get the combined content - combined_content = search_data.get('combined_content', '') - - # Check if content is too long for Discord message (3000 chars limit) - if len(combined_content) > 3000: - # Send content as file attachment - content_file = discord.File( - io.StringIO(combined_content), - filename="search_results.txt" - ) - - # Create display text without full content - search_display = "**🔍 Google Search**\n\n" - search_display += f"**📝 Query:** `{query}`\n" - search_display += f"**📊 Results:** {num_results} requested\n\n" - - # Show search results with links - if 'results' in search_data and search_data['results']: - search_display += "**🔗 Found Links:**\n" - for i, item in enumerate(search_data['results'][:5], 1): - title = item.get('title', 'No title')[:80] - link = item.get('link', '') - used_mark = "✅" if item.get('used_for_content', False) else "📄" - search_display += f"{i}. {used_mark} [{title}]({link})\n" - search_display += "\n" - - search_display += "**📄 Content:** *Attached as file (too long to display)*" - - if 'error' in search_data: - search_display += f"\n**❌ Error:** {search_data['error'][:300]}" - - # Send with file attachment - await discord_message.channel.send(search_display, file=content_file) - else: - # Use normal display for shorter content - search_display = "**🔍 Google Search**\n\n" - search_display += f"**📝 Query:** `{query}`\n" - search_display += f"**📊 Results:** {num_results} requested\n\n" - - # Show search results with links - if 'results' in search_data and search_data['results']: - search_display += "**🔗 Found Links:**\n" - for i, item in enumerate(search_data['results'][:5], 1): - title = item.get('title', 'No title')[:80] - link = item.get('link', '') - used_mark = "✅" if item.get('used_for_content', False) else "📄" - search_display += f"{i}. {used_mark} [{title}]({link})\n" - search_display += "\n" - - # Show content preview - if combined_content.strip(): - search_display += "**📄 Content:**\n```\n" - search_display += combined_content - search_display += "\n```" - else: - search_display += "**📄 Content:** *(No content retrieved)*" - - if 'error' in search_data: - search_display += f"\n**❌ Error:** {search_data['error']}" - - # Send the search display to Discord - await discord_message.channel.send(search_display) - - except Exception as e: - logging.error(f"Error displaying Google search: {str(e)}") - # Fallback: just send a simple message to prevent bot from getting stuck + # Check user's tool display preference + show_search_details = await self.db_handler.get_user_tool_display(user_id) if user_id else False + + if show_search_details: try: - await discord_message.channel.send(f"🔍 Google search completed for: `{query}`") - except: - pass + # Parse the result to get structured data + import json + search_data = json.loads(result) if isinstance(result, str) else result + + # Get the combined content + combined_content = search_data.get('combined_content', '') + + # Check if content is too long for Discord message (3000 chars limit) + if len(combined_content) > 3000: + # Send content as file attachment + content_file = discord.File( + io.StringIO(combined_content), + filename="search_results.txt" + ) + + # Create display text without full content + search_display = "**🔍 Google Search**\n\n" + search_display += f"**📝 Query:** `{query}`\n" + search_display += f"**📊 Results:** {num_results} requested\n\n" + + # Show search results with links + if 'results' in search_data and search_data['results']: + search_display += "**🔗 Found Links:**\n" + for i, item in enumerate(search_data['results'][:5], 1): + title = item.get('title', 'No title')[:80] + link = item.get('link', '') + used_mark = "✅" if item.get('used_for_content', False) else "📄" + search_display += f"{i}. {used_mark} [{title}]({link})\n" + search_display += "\n" + + search_display += "**📄 Content:** *Attached as file (too long to display)*" + + if 'error' in search_data: + search_display += f"\n**❌ Error:** {search_data['error'][:300]}" + + # Send with file attachment + await discord_message.channel.send(search_display, file=content_file) + else: + # Use normal display for shorter content + search_display = "**🔍 Google Search**\n\n" + search_display += f"**📝 Query:** `{query}`\n" + search_display += f"**📊 Results:** {num_results} requested\n\n" + + # Show search results with links + if 'results' in search_data and search_data['results']: + search_display += "**🔗 Found Links:**\n" + for i, item in enumerate(search_data['results'][:5], 1): + title = item.get('title', 'No title')[:80] + link = item.get('link', '') + used_mark = "✅" if item.get('used_for_content', False) else "📄" + search_display += f"{i}. {used_mark} [{title}]({link})\n" + search_display += "\n" + + # Show content preview + if combined_content.strip(): + search_display += "**📄 Content:**\n```\n" + search_display += combined_content + search_display += "\n```" + else: + search_display += "**📄 Content:** *(No content retrieved)*" + + if 'error' in search_data: + search_display += f"\n**❌ Error:** {search_data['error']}" + + # Send the search display to Discord + await discord_message.channel.send(search_display) + + except Exception as e: + logging.error(f"Error displaying Google search: {str(e)}") + # Fallback: just send a simple message to prevent bot from getting stuck + try: + await discord_message.channel.send(f"🔍 Google search completed for: `{query}`") + except: + pass return result except Exception as e: @@ -1446,6 +1459,11 @@ class MessageHandler: async def _scrape_webpage(self, args: Dict[str, Any]): """Scrape a webpage with Discord display""" try: + # Find user_id from current task context + user_id = args.get("user_id") + if not user_id: + user_id = self._find_user_id_from_current_task() + # Get the Discord message to display scraping activity discord_message = self._get_discord_message_from_current_task() @@ -1457,64 +1475,68 @@ class MessageHandler: from src.utils.web_utils import scrape_webpage result = await scrape_webpage(args) - # Display the scraping activity in Discord + # Display the scraping activity in Discord (only if user has enabled tool display) if discord_message and url: - try: - # Parse the result to get structured data - import json - scrape_data = json.loads(result) if isinstance(result, str) else result - - # Get the scraped content - content = scrape_data.get('content', '') if scrape_data.get('success') else '' - - # Check if content is too long for Discord message (3000 chars limit) - if len(content) > 3000: - # Send content as file attachment - content_file = discord.File( - io.StringIO(content), - filename="scraped_content.txt" - ) - - # Create display text without full content - scrape_display = "**🌐 Webpage Scraping**\n\n" - scrape_display += f"**🔗 URL:** {url}\n" - scrape_display += f"**⚙️ Max Tokens:** {max_tokens}\n\n" - scrape_display += "**📄 Content:** *Attached as file (too long to display)*" - - if 'error' in scrape_data: - scrape_display += f"\n**❌ Error:** {scrape_data['error'][:300]}" - elif content: - scrape_display += f"\n**✅ Success:** Scraped {len(content)} characters" - - # Send with file attachment - await discord_message.channel.send(scrape_display, file=content_file) - else: - # Use normal display for shorter content - scrape_display = "**🌐 Webpage Scraping**\n\n" - scrape_display += f"**🔗 URL:** {url}\n" - scrape_display += f"**⚙️ Max Tokens:** {max_tokens}\n\n" - - # Show content - if content.strip(): - scrape_display += "**📄 Content:**\n```\n" - scrape_display += content - scrape_display += "\n```" - else: - scrape_display += "**📄 Content:** *(No content retrieved)*" - - if 'error' in scrape_data: - scrape_display += f"\n**❌ Error:** {scrape_data['error']}" - - # Send the scraping display to Discord - await discord_message.channel.send(scrape_display) - - except Exception as e: - logging.error(f"Error displaying webpage scraping: {str(e)}") - # Fallback: just send a simple message to prevent bot from getting stuck + # Check user's tool display preference + show_scrape_details = await self.db_handler.get_user_tool_display(user_id) if user_id else False + + if show_scrape_details: try: - await discord_message.channel.send(f"🌐 Webpage scraping completed for: {url}") - except: - pass + # Parse the result to get structured data + import json + scrape_data = json.loads(result) if isinstance(result, str) else result + + # Get the scraped content + content = scrape_data.get('content', '') if scrape_data.get('success') else '' + + # Check if content is too long for Discord message (3000 chars limit) + if len(content) > 3000: + # Send content as file attachment + content_file = discord.File( + io.StringIO(content), + filename="scraped_content.txt" + ) + + # Create display text without full content + scrape_display = "**🌐 Webpage Scraping**\n\n" + scrape_display += f"**🔗 URL:** {url}\n" + scrape_display += f"**⚙️ Max Tokens:** {max_tokens}\n\n" + scrape_display += "**📄 Content:** *Attached as file (too long to display)*" + + if 'error' in scrape_data: + scrape_display += f"\n**❌ Error:** {scrape_data['error'][:300]}" + elif content: + scrape_display += f"\n**✅ Success:** Scraped {len(content)} characters" + + # Send with file attachment + await discord_message.channel.send(scrape_display, file=content_file) + else: + # Use normal display for shorter content + scrape_display = "**🌐 Webpage Scraping**\n\n" + scrape_display += f"**🔗 URL:** {url}\n" + scrape_display += f"**⚙️ Max Tokens:** {max_tokens}\n\n" + + # Show content + if content.strip(): + scrape_display += "**📄 Content:**\n```\n" + scrape_display += content + scrape_display += "\n```" + else: + scrape_display += "**📄 Content:** *(No content retrieved)*" + + if 'error' in scrape_data: + scrape_display += f"\n**❌ Error:** {scrape_data['error']}" + + # Send the scraping display to Discord + await discord_message.channel.send(scrape_display) + + except Exception as e: + logging.error(f"Error displaying webpage scraping: {str(e)}") + # Fallback: just send a simple message to prevent bot from getting stuck + try: + await discord_message.channel.send(f"🌐 Webpage scraping completed for: {url}") + except: + pass return result except Exception as e: