feat: Add shell command extraction and execution handling in Teto AI responses
This commit is contained in:
parent
8c451f1a48
commit
b643e14a9e
GPG Key ID:
13E498CE010AC6FD
103
cogs/teto_cog.py
103
cogs/teto_cog.py
@ -4,7 +4,6 @@ from discord import app_commands
|
||||
import re
|
||||
import base64
|
||||
import io
|
||||
import json
|
||||
import asyncio
|
||||
import subprocess
|
||||
|
||||
@ -15,6 +14,27 @@ def strip_think_blocks(text):
|
||||
def encode_image_to_base64(image_data):
|
||||
return base64.b64encode(image_data).decode('utf-8')
|
||||
|
||||
def extract_shell_command(text):
|
||||
"""
|
||||
Extracts shell commands from text using the custom format:
|
||||
```shell-command
|
||||
command
|
||||
```
|
||||
|
||||
Returns a tuple of (command, text_without_command) if a command is found,
|
||||
or (None, original_text) if no command is found.
|
||||
"""
|
||||
pattern = r"```shell-command\n(.*?)\n```"
|
||||
match = re.search(pattern, text, re.DOTALL)
|
||||
|
||||
if match:
|
||||
command = match.group(1).strip()
|
||||
# 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
|
||||
|
||||
# In-memory conversation history for Kasane Teto AI (keyed by channel id)
|
||||
_teto_conversations = {}
|
||||
|
||||
@ -28,24 +48,6 @@ class TetoCog(commands.Cog):
|
||||
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:
|
||||
@ -72,10 +74,9 @@ class TetoCog(commands.Cog):
|
||||
return f"Error executing command: {e}"
|
||||
|
||||
|
||||
async def _teto_reply_ai_with_messages(self, messages, system_mode="reply"):
|
||||
async def _teto_reply_ai_with_messages(self, messages):
|
||||
"""
|
||||
Use OpenRouter AI to generate a Kasane Teto-style response.
|
||||
system_mode: "reply" for replying as Kasane Teto.
|
||||
"""
|
||||
api_key = os.getenv("AI_API_KEY")
|
||||
if not api_key:
|
||||
@ -101,18 +102,17 @@ class TetoCog(commands.Cog):
|
||||
"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.\n\n"
|
||||
"When using the shell command tool, do not be in character and write the command properly. Avoid running obviously dangerous commands."
|
||||
"IMPORTANT: When you need to execute a shell command, include it in your response using this exact format:\n"
|
||||
"```shell-command\n"
|
||||
"your_command_here\n"
|
||||
"```\n"
|
||||
"Do not be in character when writing the command itself - write it properly. Avoid running obviously dangerous commands, regardless of your character's personality. \n"
|
||||
)
|
||||
|
||||
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,
|
||||
"tools": tools_to_send # Include tools in the payload
|
||||
"max_tokens": 2000
|
||||
}
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.post(url, headers=headers, json=payload) as resp:
|
||||
@ -127,38 +127,23 @@ class TetoCog(commands.Cog):
|
||||
|
||||
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
|
||||
# Get the AI's text response
|
||||
ai_content = response_message.get("content", "")
|
||||
|
||||
# Check for custom tool call format in the response
|
||||
if self._allow_shell_commands:
|
||||
command, _ = extract_shell_command(ai_content)
|
||||
if command:
|
||||
# Execute the shell command
|
||||
tool_result = await self._execute_shell_command(command)
|
||||
|
||||
# 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)
|
||||
|
||||
return ai_content
|
||||
|
||||
else:
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user