Tools 体系
1. 概述
Tools 模块提供了 AKG Agents 的工具执行框架。Agent 通过工具与外部环境交互 —— 读写文件、运行脚本、验证算子等。
核心组件:
- ToolExecutor — 执行工具、解析参数、捕获错误、持久化结果
- Basic Tools — 内置通用工具(文件读写、Shell、代码检查)
- Domain Tools — 场景专用工具(算子验证、性能分析)
- 工具配置 — 基于 YAML 的工具定义(
tools.yaml、domain_tools.yaml)
2. ToolExecutor
ToolExecutor 是 ReActAgent 使用的核心工具执行引擎,它:
- 接收 LLM 输出的工具名称和参数
- 解析参数表达式(如
read_json_file('path')['key']) - 分发到对应的工具函数
- 捕获执行期间的错误(WARNING+ 级别日志、stderr)
- 将结果持久化到 trace 节点路径
核心特性
- 参数解析:通过
arg_resolver支持工具参数中的表达式,如read_json_file('result.json')['output'] - 错误捕获:捕获工具执行期间的 WARNING+ 级别日志和 stderr,将精简错误摘要注入结果
- Agent/Workflow 分发:除常规工具外,还可分发到已注册的 Agent(
call_agent类型)或 Workflow(call_workflow类型)
3. 内置工具(basic_tools)
| 工具 | 说明 | 必需参数 |
|---|---|---|
read_file |
读取文件内容。支持相对路径和绝对路径。 | file_path |
write_file |
将内容写入文件。支持按 op_name 自动命名。 | content |
ask_user |
向用户提问并等待回复。 | message |
check_python_code |
检查 Python 语法并使用 autopep8 自动格式化。 | file_path |
check_markdown |
使用 markdownlint-cli 检查 Markdown 格式。 | file_path |
execute_script |
执行 Shell 或 Python 脚本。 | script_path |
write_file 参数
| 参数 | 类型 | 必需 | 说明 |
|---|---|---|---|
content |
string | 是 | 要写入的文件内容 |
file_path |
string | 否 | 显式文件路径。未提供时根据 op_name 和 file_type 自动生成 |
op_name |
string | 否 | 算子名称,用于自动命名 |
file_type |
string | 否 | 文件类型/扩展名,用于自动命名 |
encoding |
string | 否 | 文件编码(默认 "utf-8") |
overwrite |
bool | 否 | 是否覆盖已有文件(默认 True) |
execute_script 参数
| 参数 | 类型 | 必需 | 说明 |
|---|---|---|---|
script_path |
string | 是 | 要执行的脚本路径 |
args |
list | 否 | 传递给脚本的命令行参数 |
stdin_input |
string | 否 | 传递给脚本的标准输入 |
timeout |
int | 否 | 执行超时时间,秒(默认 300) |
working_dir |
string | 否 | 脚本执行的工作目录 |
返回格式
所有内置工具返回标准化字典:
{
"status": "success" | "error",
"output": "...", # 工具输出
"error_information": "..." # 错误详情(成功时为空)
}
4. 领域工具(domain_tools)
领域工具是面向算子生成与优化场景的专用工具。
verify_kernel
通过对比框架实现(如 PyTorch)和生成实现(如 Triton)的输出结果,验证生成算子代码的正确性。
| 参数 | 类型 | 必需 | 说明 |
|---|---|---|---|
task_code |
string | 是 | 框架实现代码(需包含 class Model) |
generated_code |
string | 是 | 待验证的生成代码(需包含 class ModelNew) |
op_name |
string | 是 | 算子名称(如 "relu"、"matmul") |
task_id |
string | 否 | 任务 ID(默认 "default_task") |
device_id |
int | 否 | 设备 ID(默认 0) |
timeout |
int | 否 | 超时时间,秒(默认 300) |
cur_path |
string | 否 | 当前 trace 节点路径(默认 "") |
framework |
string | 否 | 计算框架(默认 "torch") |
backend |
string | 否 | 硬件后端(默认 "cuda") |
arch |
string | 否 | 硬件架构(默认 "a100") |
dsl |
string | 否 | 目标 DSL(默认 "triton") |
profile_kernel
算子代码性能分析,返回执行时间和加速比。
| 参数 | 类型 | 必需 | 说明 |
|---|---|---|---|
task_code |
string | 是 | 框架实现代码 |
generated_code |
string | 是 | 待分析的生成代码 |
op_name |
string | 是 | 算子名称 |
task_id |
string | 否 | 任务 ID(默认 "default_task") |
device_id |
int | 否 | 设备 ID(默认 0) |
run_times |
int | 否 | 性能测试运行次数(默认 50) |
warmup_times |
int | 否 | 预热次数(默认 5) |
cur_path |
string | 否 | 当前 trace 节点路径(默认 "") |
framework |
string | 否 | 计算框架(默认 "torch") |
backend |
string | 否 | 硬件后端(默认 "cuda") |
arch |
string | 否 | 硬件架构(默认 "a100") |
dsl |
string | 否 | 目标 DSL(默认 "triton") |
注意:
profile_kernel仅做性能分析,不验证正确性。
5. 工具配置
工具通过 YAML 配置文件定义:
tools.yaml— 基础工具定义domain_tools.yaml— 领域专用工具定义
配置格式
tools:
read_file:
type: "basic_tool"
function:
name: "read_file"
description: "读取文件内容。"
parameters:
type: "object"
properties:
file_path:
type: "string"
description: "文件路径"
required: ["file_path"]
verify_kernel:
type: "domain_tool"
function:
name: "verify_kernel"
description: "验证生成的算子代码正确性。"
parameters:
type: "object"
properties:
task_code:
type: "string"
description: "框架实现代码"
generated_code:
type: "string"
description: "待验证的生成代码"
op_name:
type: "string"
description: "算子名称"
required: ["task_code", "generated_code", "op_name"]
工具类型
| 类型 | 说明 |
|---|---|
basic_tool |
通用工具,分发到 basic_tools 模块 |
domain_tool |
领域专用工具,分发到 domain_tools 模块 |
call_agent |
分发到已注册的 Agent(通过 AgentRegistry) |
call_workflow |
分发到已注册的 Workflow(通过 WorkflowRegistry) |
6. 参数解析器
arg_resolver 模块支持工具调用中的动态参数解析。当 LLM 生成的工具参数包含表达式时,会在执行前自动解析。
支持的表达式
# 读取 JSON 文件并提取字段
read_json_file('path/to/result.json')['output']
# 嵌套访问
read_json_file('path/to/result.json')['data']['code']
这使得 Agent 可以链式传递工具结果 —— 例如读取上一个工具的 result.json 输出,作为下一个工具的输入。