slipstream/discordbot
slipstream/discordbot
1
1
Fork 0
Code Issues Pull Requests 3 Actions Packages Projects Releases Wiki Activity

feat: Enhance shell command extraction and response formatting in Teto AI

Browse Source
This commit is contained in:
Slipstream 2025-05-19 22:34:50 -06:00
parent c747bd9f2f
commit 909616fb60
Signed by: slipstream
GPG Key ID: 13E498CE010AC6FD
1 changed files with 70 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

84
cogs/teto_cog.py
View File

@ -21,8 +21,8 @@ def extract_shell_command(text):
command
```
Returns a tuple of (command, text_without_command) if a command is found,
or (None, original_text) if no command is found.
Returns a tuple of (command, text_without_command, text_before_command) if a command is found,
or (None, original_text, None) if no command is found.
"""
pattern = r"```shell-command\n(.*?)\n```"
match = re.search(pattern, text, re.DOTALL)
@ -30,11 +30,17 @@ def extract_shell_command(text):
if match:
print(f"[TETO DEBUG] Found shell command: {match.group(1)}")
command = match.group(1).strip()
# Get the text before the command block
start_idx = match.start()
text_before_command = text[:start_idx].strip() if start_idx > 0 else None
# Remove the command block from the text
text_without_command = re.sub(pattern, "", text, flags=re.DOTALL).strip()
return command, text_without_command
return None, text
return command, text_without_command, text_before_command
return None, text, None
# In-memory conversation history for Kasane Teto AI (keyed by channel id)
_teto_conversations = {}
@ -50,7 +56,7 @@ class TetoCog(commands.Cog):
self._allow_shell_commands = False # Flag to control shell command tool usage
async def _execute_shell_command(self, command: str) -> str:
"""Executes a shell command and returns its output."""
"""Executes a shell command and returns its output, limited to first 5 lines."""
try:
# Use subprocess.run for simple command execution
# Consider security implications of running arbitrary commands
@ -63,9 +69,20 @@ class TetoCog(commands.Cog):
output = ""
if stdout:
output += f"Stdout:\n{stdout.decode()}\n"
# Limit stdout to first 5 lines
stdout_lines = stdout.decode().splitlines()
limited_stdout = "\n".join(stdout_lines[:5])
if len(stdout_lines) > 5:
limited_stdout += "\n... (output truncated, showing first 5 lines)"
output += f"Stdout:\n{limited_stdout}\n"
if stderr:
output += f"Stderr:\n{stderr.decode()}\n"
# Limit stderr to first 5 lines
stderr_lines = stderr.decode().splitlines()
limited_stderr = "\n".join(stderr_lines[:5])
if len(stderr_lines) > 5:
limited_stderr += "\n... (output truncated, showing first 5 lines)"
output += f"Stderr:\n{limited_stderr}\n"
if not output:
output = "Command executed successfully with no output."
@ -136,7 +153,9 @@ class TetoCog(commands.Cog):
"Reply to the user in a short, conversational manner, staying in character.\n\n"
"You have the ability to run any arbitrary shell commands. \n"
"If a user asks you to show the entire output of a command, you can and should do so. \n"
"IMPORTANT: When you need to execute a shell command, include it in your response using this exact format:\n"
"IMPORTANT: When you need to execute a shell command, you can include a message before the command block. \n"
"For example, you can say 'Let me check that for you!' and then include the command block. \n"
"Include the command in your response using this exact format:\n"
"```shell-command\n"
"your_command_here\n"
"```\n"
@ -166,20 +185,38 @@ class TetoCog(commands.Cog):
# Check for custom tool call format in the response
if self._allow_shell_commands:
command, _ = extract_shell_command(ai_content)
command, _, text_before_command = extract_shell_command(ai_content)
if command:
if self._is_dangerous_command(command):
tool_result = "Error: Execution was blocked due to a potentially dangerous command."
tool_result = "❌ Error: Execution was blocked due to a potentially dangerous command."
else:
# Execute the shell command
tool_result = await self._execute_shell_command(command)
# Format the response with the AI's message before the command (if any)
formatted_response = ai_content
if text_before_command:
# Replace the original AI content with just the text before the command
# plus a formatted command execution message
if self._is_dangerous_command(command):
formatted_response = f"{text_before_command}\n\n*❌ Command \"{command}\" blocked (potentially dangerous)*\n\n{tool_result}"
else:
formatted_response = f"{text_before_command}\n\n*✅ Command \"{command}\" executed successfully*\n\n{tool_result}"
else:
# If there was no text before the command, just show the command execution message
if self._is_dangerous_command(command):
formatted_response = f"*❌ Command \"{command}\" blocked (potentially dangerous)*\n\n{tool_result}"
else:
formatted_response = f"*✅ Command \"{command}\" executed successfully*\n\n{tool_result}"
# Append the original message and tool result to the conversation
messages.append({"role": "assistant", "content": ai_content})
messages.append({"role": "user", "content": f"Command output:\n{tool_result}"})
# Make another API call with the tool result
return await self._teto_reply_ai_with_messages(messages)
# Make another API call with the tool result, but return the formatted response
# to be displayed in Discord
ai_follow_up = await self._teto_reply_ai_with_messages(messages)
return formatted_response + "\n\n" + ai_follow_up
return ai_content
@ -305,7 +342,17 @@ class TetoCog(commands.Cog):
ai_reply = await self._teto_reply_ai_with_messages(messages=convo)
ai_reply = strip_think_blocks(ai_reply)
await message.reply(ai_reply)
convo.append({"role": "assistant", "content": ai_reply})
# Extract the original AI content (without command execution formatting)
# for storing in conversation history
command, content_without_command, _ = extract_shell_command(ai_reply)
if command:
# If there was a command, store the original AI content without the formatted execution message
convo.append({"role": "assistant", "content": content_without_command if content_without_command else ai_reply})
else:
# If there was no command, store the full reply
convo.append({"role": "assistant", "content": ai_reply})
_teto_conversations[convo_key] = convo[-10:] # Keep last 10 interactions
log.info("[TETO DEBUG] AI reply sent successfully.")
except Exception as e:
@ -365,7 +412,16 @@ async def teto_context_menu_ai_reply(interaction: discord.Interaction, message:
ai_reply = strip_think_blocks(ai_reply)
await message.reply(ai_reply)
await interaction.followup.send("Teto AI replied desu~", ephemeral=True)
convo.append({"role": "assistant", "content": ai_reply})
# Extract the original AI content (without command execution formatting)
# for storing in conversation history
command, content_without_command, _ = extract_shell_command(ai_reply)
if command:
# If there was a command, store the original AI content without the formatted execution message
convo.append({"role": "assistant", "content": content_without_command if content_without_command else ai_reply})
else:
# If there was no command, store the full reply
convo.append({"role": "assistant", "content": ai_reply})
_teto_conversations[convo_key] = convo[-10:]
except Exception as e:
await interaction.followup.send(f"Teto AI reply failed: {e} desu~", ephemeral=True)