Coder-Vippro/ChatGPT-Discord-Bot
1
0
Fork 0
Code Issues Pull Requests 3 Actions Packages Projects Releases 13 Wiki Activity

refactor: add user preference for tool execution display and enhance message handling for search and scraping activities

Browse Source
This commit is contained in:
cauvang32
2025-08-26 23:36:47 +07:00
parent eaaef0676a
commit 19e62f85fc
3 changed files with 295 additions and 233 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
src/commands/commands.py
View File

@@ -376,12 +376,37 @@ def setup_commands(bot: commands.Bot, db_handler, openai_client, image_generator
"/search `<query>` - Search Google and send results to the AI model.\n" "/search `<query>` - Search Google and send results to the AI model.\n"
"/web `<url>` - Scrape a webpage and send the data to the AI model.\n" "/web `<url>` - Scrape a webpage and send the data to the AI model.\n"
"/generate `<prompt>` - Generate an image from a text prompt.\n" "/generate `<prompt>` - 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" "/reset - Reset your chat history.\n"
"/user_stat - Get information about your input tokens, output tokens, and current model.\n" "/user_stat - Get information about your input tokens, output tokens, and current model.\n"
"/help - Display this help message.\n" "/help - Display this help message.\n"
) )
await interaction.response.send_message(help_message, ephemeral=True) 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.") @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)") @app_commands.describe(user_id="The Discord user ID to stop tasks for (admin only)")
@check_blacklist() @check_blacklist()

15
src/database/db_handler.py
View File

@@ -112,6 +112,20 @@ class DatabaseHandler:
upsert=True 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 # Admin and permissions management with caching
async def is_admin(self, user_id: int) -> bool: async def is_admin(self, user_id: int) -> bool:
"""Check if the user is an admin (no caching for security)""" """Check if the user is an admin (no caching for security)"""
@@ -163,6 +177,7 @@ class DatabaseHandler:
"""Create indexes for better query performance""" """Create indexes for better query performance"""
await self.db.user_histories.create_index("user_id") await self.db.user_histories.create_index("user_id")
await self.db.user_models.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.whitelist.create_index("user_id")
await self.db.blacklist.create_index("user_id") await self.db.blacklist.create_index("user_id")

488
src/module/message_handler.py
View File

@@ -241,115 +241,119 @@ class MessageHandler:
# Display the executed code information in Discord (but not save to history) # Display the executed code information in Discord (but not save to history)
if discord_message and code_to_execute: if discord_message and code_to_execute:
try: # Check user's tool display preference
# Clean up the code display (remove file context comments) show_execution_details = await self.db_handler.get_user_tool_display(user_id) if user_id else False
code_lines = code_to_execute.split('\n')
clean_code_lines = [] if show_execution_details:
for line in code_lines: try:
if not (line.strip().startswith('# Data file available:') or # Clean up the code display (remove file context comments)
line.strip().startswith('# File path:') or code_lines = code_to_execute.split('\n')
line.strip().startswith('# You can access this file using:')): clean_code_lines = []
clean_code_lines.append(line) for line in code_lines:
if not (line.strip().startswith('# Data file available:') or
clean_code = '\n'.join(clean_code_lines).strip() line.strip().startswith('# File path:') or
line.strip().startswith('# You can access this file using:')):
# Check if code is too long for Discord message (3000 chars limit) clean_code_lines.append(line)
if len(clean_code) > 3000:
# Send code as file attachment
code_file = discord.File(
io.StringIO(clean_code),
filename="executed_code.py"
)
# Create display text without code clean_code = '\n'.join(clean_code_lines).strip()
execution_display = "**🐍 Python Code Execution**\n\n"
# Show packages to install if any # Check if code is too long for Discord message (3000 chars limit)
if packages_to_install: if len(clean_code) > 3000:
execution_display += f"**📦 Installing packages:** {', '.join(packages_to_install)}\n\n" # Send code as file attachment
code_file = discord.File(
# Show input data if any io.StringIO(clean_code),
if input_data: filename="executed_code.py"
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(): # Create display text without code
execution_display += "**📤 Output:**\n```\n" execution_display = "**🐍 Python Code Execution**\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 ""
if output and output.strip(): # Show packages to install if any
execution_display += "**📤 Output:**\n```\n" if packages_to_install:
execution_display += output[:1000] # Limit output length for Discord execution_display += f"**📦 Installing packages:** {', '.join(packages_to_install)}\n\n"
if len(output) > 1000:
execution_display += "\n... (output truncated)" # Show input data if any
execution_display += "\n```" 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: 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: else:
error_msg = execute_result.get("error", "Unknown error") if execute_result else "Execution failed" # Use normal display for shorter code
execution_display += f"**❌ Error:**\n```\n{error_msg[:800]}\n```" execution_display = "**🐍 Python Code Execution**\n\n"
if len(error_msg) > 800:
execution_display += "*(Error message truncated)*" # 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 except Exception as e:
await discord_message.channel.send(execution_display) 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 there are visualizations, handle them
if execute_result and execute_result.get("visualizations"): if execute_result and execute_result.get("visualizations"):
@@ -1349,6 +1353,11 @@ class MessageHandler:
async def _google_search(self, args: Dict[str, Any]): async def _google_search(self, args: Dict[str, Any]):
"""Perform a Google search with Discord display""" """Perform a Google search with Discord display"""
try: 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 # Get the Discord message to display search activity
discord_message = self._get_discord_message_from_current_task() discord_message = self._get_discord_message_from_current_task()
@@ -1360,83 +1369,87 @@ class MessageHandler:
from src.utils.web_utils import google_search from src.utils.web_utils import google_search
result = await google_search(args) 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: if discord_message and query:
try: # Check user's tool display preference
# Parse the result to get structured data show_search_details = await self.db_handler.get_user_tool_display(user_id) if user_id else False
import json
search_data = json.loads(result) if isinstance(result, str) else result if show_search_details:
# 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: try:
await discord_message.channel.send(f"🔍 Google search completed for: `{query}`") # Parse the result to get structured data
except: import json
pass 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 return result
except Exception as e: except Exception as e:
@@ -1446,6 +1459,11 @@ class MessageHandler:
async def _scrape_webpage(self, args: Dict[str, Any]): async def _scrape_webpage(self, args: Dict[str, Any]):
"""Scrape a webpage with Discord display""" """Scrape a webpage with Discord display"""
try: 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 # Get the Discord message to display scraping activity
discord_message = self._get_discord_message_from_current_task() discord_message = self._get_discord_message_from_current_task()
@@ -1457,64 +1475,68 @@ class MessageHandler:
from src.utils.web_utils import scrape_webpage from src.utils.web_utils import scrape_webpage
result = await scrape_webpage(args) 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: if discord_message and url:
try: # Check user's tool display preference
# Parse the result to get structured data show_scrape_details = await self.db_handler.get_user_tool_display(user_id) if user_id else False
import json
scrape_data = json.loads(result) if isinstance(result, str) else result if show_scrape_details:
# 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: try:
await discord_message.channel.send(f"🌐 Webpage scraping completed for: {url}") # Parse the result to get structured data
except: import json
pass 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 return result
except Exception as e: except Exception as e: