Soulter
79 lines
2.9 KiB
Python
79 lines
2.9 KiB
Python
import json
|
|
from dataclasses import dataclass, field
|
|
|
|
from astrbot.api import FunctionTool
|
|
from astrbot.core.agent.run_context import ContextWrapper
|
|
from astrbot.core.agent.tool import ToolExecResult
|
|
from astrbot.core.astr_agent_context import AstrAgentContext
|
|
|
|
from ..computer_client import get_booter, get_local_booter
|
|
|
|
|
|
def _check_admin_permission(context: ContextWrapper[AstrAgentContext]) -> str | None:
|
|
cfg = context.context.context.get_config(
|
|
umo=context.context.event.unified_msg_origin
|
|
)
|
|
provider_settings = cfg.get("provider_settings", {})
|
|
require_admin = provider_settings.get("computer_use_require_admin", True)
|
|
if require_admin and context.context.event.role != "admin":
|
|
return (
|
|
"error: Permission denied. Shell execution is only allowed for admin users. "
|
|
"Tell user to set admins in `AstrBot WebUI -> Config -> General Config` by adding their user ID to the admins list if they need this feature."
|
|
f"User's ID is: {context.context.event.get_sender_id()}. User's ID can be found by using /sid command."
|
|
)
|
|
return None
|
|
|
|
|
|
@dataclass
|
|
class ExecuteShellTool(FunctionTool):
|
|
name: str = "astrbot_execute_shell"
|
|
description: str = "Execute a command in the shell."
|
|
parameters: dict = field(
|
|
default_factory=lambda: {
|
|
"type": "object",
|
|
"properties": {
|
|
"command": {
|
|
"type": "string",
|
|
"description": "The bash command to execute. Equal to 'cd {working_dir} && {your_command}'.",
|
|
},
|
|
"background": {
|
|
"type": "boolean",
|
|
"description": "Whether to run the command in the background.",
|
|
"default": False,
|
|
},
|
|
"env": {
|
|
"type": "object",
|
|
"description": "Optional environment variables to set for the file creation process.",
|
|
"additionalProperties": {"type": "string"},
|
|
"default": {},
|
|
},
|
|
},
|
|
"required": ["command"],
|
|
}
|
|
)
|
|
|
|
is_local: bool = False
|
|
|
|
async def call(
|
|
self,
|
|
context: ContextWrapper[AstrAgentContext],
|
|
command: str,
|
|
background: bool = False,
|
|
env: dict = {},
|
|
) -> ToolExecResult:
|
|
if permission_error := _check_admin_permission(context):
|
|
return permission_error
|
|
|
|
if self.is_local:
|
|
sb = get_local_booter()
|
|
else:
|
|
sb = await get_booter(
|
|
context.context.context,
|
|
context.context.event.unified_msg_origin,
|
|
)
|
|
try:
|
|
result = await sb.shell.exec(command, background=background, env=env)
|
|
return json.dumps(result)
|
|
except Exception as e:
|
|
return f"Error executing command: {str(e)}"
|