refactor: add user preference for tool execution display and enhance message handling for search and scraping activities
This commit is contained in:
@@ -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()
|
||||
|
||||
@@ -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")
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user