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

This commit is contained in:
2025-08-26 23:36:47 +07:00
parent eaaef0676a
commit 19e62f85fc
3 changed files with 295 additions and 233 deletions

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"
"/web `<url>` - Scrape a webpage and send the data to the AI model.\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"
"/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()

View File

@@ -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")

View File

@@ -241,6 +241,10 @@ class MessageHandler:
# Display the executed code information in Discord (but not save to history)
if discord_message and code_to_execute:
# 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')
@@ -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,8 +1369,12 @@ 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:
# 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:
# Parse the result to get structured data
import json
@@ -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,8 +1475,12 @@ 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:
# 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:
# Parse the result to get structured data
import json