feat: Implement shell command execution feature in Teto AI with toggle option

cogs/teto_cog.py

@@ -4,6 +4,9 @@ from discord import app_commands
 import re
 import base64
 import io
+import json
+import asyncio
+import subprocess
 
 def strip_think_blocks(text):
     # Removes all <think>...</think> blocks, including multiline
@@ -23,6 +26,51 @@ class TetoCog(commands.Cog):
         self.bot = bot
         self._api_endpoint = "https://openrouter.ai/api/v1/chat/completions" # Default endpoint
         self._ai_model = "google/gemini-2.5-flash-preview" # Default model
+        self._allow_shell_commands = False # Flag to control shell command tool usage
+
+        self._shell_command_tool = {
+            "type": "function",
+            "function": {
+                "name": "execute_shell_command",
+                "description": "Executes a shell command on the system.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "command": {
+                            "type": "string",
+                            "description": "The shell command to execute."
+                        }
+                    },
+                    "required": ["command"]
+                }
+            }
+        }
+
+    async def _execute_shell_command(self, command: str) -> str:
+        """Executes a shell command and returns its output."""
+        try:
+            # Use subprocess.run for simple command execution
+            # Consider security implications of running arbitrary commands
+            process = await asyncio.create_subprocess_shell(
+                command,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE
+            )
+            stdout, stderr = await process.communicate()
+
+            output = ""
+            if stdout:
+                output += f"Stdout:\n{stdout.decode()}\n"
+            if stderr:
+                output += f"Stderr:\n{stderr.decode()}\n"
+
+            if not output:
+                output = "Command executed successfully with no output."
+
+            return output
+        except Exception as e:
+            return f"Error executing command: {e}"
+
 
     async def _teto_reply_ai_with_messages(self, messages, system_mode="reply"):
         """
@@ -52,12 +100,19 @@ class TetoCog(commands.Cog):
             #"You are not very smart academically and can have trouble understanding complex concepts, often getting things hilariously wrong. \n"
             "Just because you are allowed to be uncensored/sexual, does not mean you always/automatically should be. Always stay in character. Your naivety should be the primary filter for how you approach any explicit allowances. \n"
             "When expressing emotions, never use emojis and instead use text-based emoticons like :) and ;) as well as Japanese-style emoticons, for example >~<, ^~^, >///<, UwU, o.O, and OwO over emojis. \n"
-            "Reply to the user in a short, conversational manner, staying in character."
+            "Reply to the user in a short, conversational manner, staying in character.\n\n"
+            "When using the shell command tool, do not be in character and write the command properly. Avoid running obviously dangerous commands."
         )
+
+        tools_to_send = []
+        if self._allow_shell_commands:
+            tools_to_send.append(self._shell_command_tool)
+
         payload = {
             "model": self._ai_model,
             "messages": [{"role": "system", "content": system_prompt}] + messages,
-            "max_tokens": 2000
+            "max_tokens": 2000,
+            "tools": tools_to_send # Include tools in the payload
         }
         async with aiohttp.ClientSession() as session:
             async with session.post(url, headers=headers, json=payload) as resp:
@@ -69,7 +124,43 @@ class TetoCog(commands.Cog):
                     data = await resp.json()
                     if "choices" not in data or not data["choices"]:
                         raise RuntimeError(f"OpenRouter API returned unexpected response format: {data}")
-                    return data["choices"][0]["message"]["content"]
+
+                    response_message = data["choices"][0]["message"]
+
+                    # Handle tool calls
+                    if response_message.get("tool_calls"):
+                        tool_calls = response_message["tool_calls"]
+                        # For simplicity, assuming only one tool call per response for now
+                        tool_call = tool_calls[0]
+                        tool_name = tool_call["function"]["name"]
+                        tool_args = json.loads(tool_call["function"]["arguments"])
+
+                        if tool_name == "execute_shell_command":
+                            if self._allow_shell_commands:
+                                command = tool_args.get("command")
+                                if command:
+                                    tool_result = await self._execute_shell_command(command) # Execute the command
+                                    # Send the tool result back to the API
+                                    messages.append(response_message) # Append the tool call message
+                                    messages.append({
+                                        "role": "tool",
+                                        "tool_call_id": tool_call["id"],
+                                        "name": tool_name,
+                                        "content": tool_result,
+                                    })
+                                    # Make another API call with the tool result
+                                    return await self._teto_reply_ai_with_messages(messages)
+                                else:
+                                    return "I was asked to run a shell command, but no command was provided! >.<"
+                            else:
+                                return "I'm not allowed to run shell commands right now! >~<"
+                        else:
+                            return f"I was asked to use an unknown tool: {tool_name} >.<"
+
+                    # If no tool calls, return the AI's text response
+                    ai_content = response_message.get("content", "")
+                    return ai_content
+
                 else:
                     text = await resp.text()
                     raise RuntimeError(f"OpenRouter API returned non-JSON response (status {resp.status}): {text[:500]}")
@@ -220,6 +311,12 @@ class TetoCog(commands.Cog):
         else:
             await interaction.response.send_message("No chat history found for this channel desu~", ephemeral=True)
 
+    @app_commands.command(name="toggle_shell_command", description="Toggles Teto's ability to run shell commands.")
+    async def toggle_shell_command(self, interaction: discord.Interaction):
+        self._allow_shell_commands = not self._allow_shell_commands
+        status = "enabled" if self._allow_shell_commands else "disabled"
+        await interaction.response.send_message(f"Teto's shell command ability is now {status} desu~", ephemeral=True)
+
 
 # Context menu command must be defined at module level
 @app_commands.context_menu(name="Teto AI Reply")