zk-data-agent/src/agent_tool_specs/execution.py

from __future__ import annotations

from collections.abc import Mapping

from ..agent_tool_core import AgentTool, ToolHandler
from ._builder import resolve_handler


def build_execution_tools(handlers: Mapping[str, ToolHandler]) -> list[AgentTool]:
    """构建本地命令、Python 执行与短等待工具声明。"""

    return [
        AgentTool(
            name='python_package',
            description=(
                '在当前用户独立 Python venv 中检查或安装 Python 包。'
                '用于 python_exec 缺少 pandas、pyarrow、openpyxl 等分析依赖时；'
                '不要通过 bash 执行 pip，也不要安装到项目 .venv。'
            ),
            parameters={
                'type': 'object',
                'properties': {
                    'action': {
                        'type': 'string',
                        'enum': ['show', 'install'],
                        'description': 'show 查看当前 Python 环境和 pip；install 安装 packages。',
                    },
                    'packages': {
                        'type': 'array',
                        'items': {'type': 'string'},
                        'description': '要安装的包名或 pip requirement spec，例如 pandas、pyarrow==15.0.0。',
                    },
                    'timeout_seconds': {
                        'type': 'number',
                        'minimum': 1,
                        'maximum': 600,
                        'description': '安装超时时间，默认 120 秒。',
                    },
                    'max_output_chars': {
                        'type': 'integer',
                        'minimum': 1000,
                        'maximum': 50000,
                    },
                },
                'required': ['action'],
            },
            handler=resolve_handler(handlers, 'python_package', 'execution'),
        ),
        AgentTool(
            name='bash',
            description=(
                '运行真实 shell 命令，包括系统命令、进程控制、git 只读检查、用户已确认的依赖安装，'
                '以及所有 Python 脚本执行。'
                '执行 Python 时，把代码先写到 scratchpad 下的 .py 文件，'
                '再用本工具跑 `python3 -u <script> 2>&1 | tee <log>`；'
                '`-u` 关掉 stdout 缓冲，`tee` 让进度行实时落盘，超时被 kill 时仍能 `tail` 日志看到死在哪个 phase。'
                '不要把多行 Python 代码以 `-c` 字符串方式塞进 command。'
                'Python 包安装必须优先使用 python_package；不要通过 bash 执行 pip。'
                '需要等待远端长任务（>30s 的训练、文件落盘、外部 API 回调）时，'
                '设置 run_in_background=true 让命令在后台 detach 跑、立刻返回 task_id；'
                '默认 wait_for_completion=true 时进程结束后会自动起新一轮把结果送回。'
                '后台任务输出落在 <remote_workspace>/.bg_tasks/<task_id>/output。'
            ),
            parameters={
                'type': 'object',
                'properties': {
                    'command': {'type': 'string'},
                    'run_in_background': {
                        'type': 'boolean',
                        'description': (
                            '后台 detach 跑。立刻返回 task_id，不等命令结束。'
                            '适合 >30s 的远端任务、需要等异步事件的场景。'
                        ),
                    },
                    'wait_for_completion': {
                        'type': 'boolean',
                        'description': (
                            '仅在 run_in_background=true 时生效。'
                            'true（默认）：进程结束后由后端自动起新一轮把结果作为 system 消息送回；'
                            'false：不自动续，agent 需自行调 bash_status 取结果。'
                        ),
                    },
                },
                'required': ['command'],
            },
            handler=resolve_handler(handlers, 'bash', 'execution'),
        ),
        AgentTool(
            name='bash_status',
            description=(
                '查询通过 bash(run_in_background=true) 启动的后台任务状态。'
                'block=true 会阻塞到任务终止或 timeout_seconds 到期再返回；'
                'block=false（默认）立刻返回当前快照。'
                '返回 status (running/completed/failed/cancelled)、exit_code 和最近输出预览。'
            ),
            parameters={
                'type': 'object',
                'properties': {
                    'task_id': {
                        'type': 'string',
                        'description': 'bash run_in_background 启动时返回的 task_id。',
                    },
                    'block': {
                        'type': 'boolean',
                        'description': '是否等到任务终止再返回；默认 false。',
                    },
                    'timeout_seconds': {
                        'type': 'number',
                        'minimum': 0,
                        'maximum': 600,
                        'description': 'block=true 时最长等待秒数，默认 30，上限 600。',
                    },
                },
                'required': ['task_id'],
            },
            handler=resolve_handler(handlers, 'bash_status', 'execution'),
        ),
        AgentTool(
            name='bash_kill',
            description=(
                '主动终止通过 bash(run_in_background=true) 启动的后台任务。'
                '注意：用户在前端按 stop 取消当前 run 不会自动杀后台进程（detach 的本意），'
                '需要明确调用本工具才能终止远端进程。'
            ),
            parameters={
                'type': 'object',
                'properties': {
                    'task_id': {
                        'type': 'string',
                        'description': 'bash run_in_background 启动时返回的 task_id。',
                    },
                },
                'required': ['task_id'],
            },
            handler=resolve_handler(handlers, 'bash_kill', 'execution'),
        ),
        AgentTool(
            name='sleep',
            description='Pause execution briefly for bounded local wait flows.',
            parameters={
                'type': 'object',
                'properties': {
                    'seconds': {'type': 'number', 'minimum': 0.0, 'maximum': 5.0},
                },
                'required': ['seconds'],
            },
            handler=resolve_handler(handlers, 'sleep', 'execution'),
        ),
    ]