动手实现 Agent（二）：工具系统增强

笔者承诺本文除直接的 LLM 消息原文外，其余内容的 AI 含量为 0。

本系列的源代码可以在 https://github.com/Eslzzyl/tinyagent 找到。

在上一篇 https://www.cnblogs.com/eslzzyl/p/20211623 中，我们实现了一个最简的 Agent。在本文中，我们将主要对它的工具系统进行增强。

优化输入体验

但在这之前，我们先用一个经典的库 prompt toolkit 替换原本用于读取用户输入的 input() 内置方法。input() 在处理中文删除时会有些问题，prompt toolkit 提供了更可靠的处理。

from prompt_toolkit import prompt

request = prompt("Your Message:")

工具系统增强

Schema 的自动解析

此前，我们需要手动为所有的工具定义 Schema：

tools = [
    {
        "type": "function",
        "function": {
            "name": "read",
            "description": "read a file from filesystem",
            "parameters": {
                "type": "object",
                "properties": {
                    "path": {
                        "type": "string",
                        "description": "file path",
                    },
                },
                "required": ["path"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "write",
            "description": "write a file to filesystem",
            "parameters": {
                "type": "object",
                "properties": {
                    "path": {
                        "type": "string",
                        "description": "file path",
                    },
                    "content": {
                        "type": "string",
                        "description": "content to write",
                    },
                },
                "required": ["content"],
            },
        },
    },
]

每当我们更新工具时，都需要对应更新这里的 Schema，这相当麻烦。不过，我们可以通过解析函数签名和 docstring 来自动生成 Schema。

docstring_parser 库（https://github.com/rr-/docstring_parser ）提供了解析常见格式 docstring 的功能，可供我们使用。这里，我们使用一种 docstring 和函数签名混合的解析方法来获得所有需要的数据。

import inspect

from docstring_parser import parse

# Python 类型到 JSON Schema 类型
_PY_TO_JSON = {
    "str": "string",
    "int": "integer",
    "float": "number",
    "bool": "boolean",
    "list": "array",
    "dict": "object",
}


def generate_tool_schema(func) -> dict:
    # 获取函数签名
    signature = inspect.signature(func)
    # 函数名
    name = func.__name__
    # raw docstring
    doc = inspect.getdoc(func) or ""
    parsed_doc = parse(doc)
    # 短描述
    description = parsed_doc.short_description
    # 参数
    params = parsed_doc.params
    properties = {}
    # 将参数转换成 openai tool schema 格式
    for param in params:
        properties[param.arg_name] = {
            "type": _PY_TO_JSON[param.type_name] if param.type_name else "null",
            "description": param.description,
        }
    # 根据函数签名中的参数默认值确定某个参数是否可选
    required = []
    for name, param in signature.parameters.items():
        # 如果默认参数是空，即没有默认参数，就说明是必须的
        if param.default is inspect.Parameter.empty:
            required.append(name)

    return {
        "type": "function",
        "function": {
            "name": name,
            "description": description,
            "parameters": {"type": "object", "properties": properties},
            "required": required,
        },
    }

接下来我们需要为工具函数补上 docstring：

def read(path: str) -> str:
    """
    read a file from filesystem.

    Args:
        path (str) : the file path
    """
    with open(path, mode="r", encoding="utf-8") as f:
        lines = f.readlines()
        return "".join(lines)


def write(path: str, content: str) -> str:
    """
    write a file to filesystem.

    Args:
        path (str) : the file path
        content (str): the content to write
    """
    with open(path, mode="w", encoding="utf-8") as f:
        f.write(content)
    return f"Wrote file {path}"

有了这些之后，我们就可以删除 Agent 类构造函数中的 tool_spec_list 参数，并在需要时自动生成 Schema 了。

response = self.client.call_with_tools(
    messages=messages,
    tool_spec_list=[generate_tool_schema(tool) for tool in self.tools],
)

增强 read 工具

读取目录

首先，我们需要支持读取目录。主流的 Agent 产品中，通常不对列出目录内容设置一个独立的工具，而是让 read 工具兼顾这个功能。

import os
from pathlib import Path

def read(path: str) -> str:
    """
    Read a file or folder from filesystem. If `path` is a folder, list all contents in it.

    Args:
        path (str) : the file or folder path
    """
    _path = Path(path)
    if _path.is_dir():
        return "\n".join(os.listdir(_path))
    else:
        with open(path, mode="r", encoding="utf-8") as f:
            lines = f.readlines()
            return "".join(lines)

试一下效果：

Your Message:当前目录下都有哪些文件？
Assistant:
Tool: read({'path': '.'})
Assistant: 当前目录下包含以下文件和文件夹：

**文件：**
- `main.py` - 主程序文件
- `pyproject.toml` - Python 项目配置
- `README.md` - 项目说明文档
- `uv.lock` - uv 包管理器锁定文件
- `.env` - 环境变量配置
- `.env.example` - 环境变量示例
- `.python-version` - Python 版本配置
- `.gitignore` - Git 忽略文件配置

**目录：**
- `src` - 源代码目录
- `.venv` - 虚拟环境目录
- `.git` - Git 版本控制目录
- `.ruff_cache` - Ruff 代码检查缓存目录

这是一个典型的 Python 项目结构，使用了 uv 包管理器和 Ruff 代码检查工具。

附加行号

这相当简单：

with open(path, mode="r", encoding="utf-8") as f:
    lines = f.readlines()
    lines = [f"{index}:{lines[index]}" for index in range(1, len(lines))]

    return "".join(lines)

允许指定 offset 和 limit

offset 用于指定读取的起始位置，而 limit 用于指定读取的行数。

为了实现简单，我们还是一次性读取整个文件，然后在内存中对行进行截取，应用这两个参数。

def read(path: str, offset: int | None = None, limit: int | None = None) -> str:
    """
    Read a file or folder from filesystem. If `path` is a folder, list all contents in it.

    Args:
        path (str): The file or folder path
        offset (int, optional): The starting line number to read (starting from 1).
        limit (int, optional): Number of lines to read
    """
    _path = Path(path)
    if _path.is_dir():
        return "\n".join(os.listdir(_path))
    else:
        if offset and offset < 1:
            return f"Error: invalid offset: {offset}"
        if limit and limit < 1:
            return f"Error: invalid limit: {limit}"

        with open(path, mode="r", encoding="utf-8") as f:
            lines = f.readlines()
            lines = [f"{index}:{lines[index]}" for index in range(1, len(lines))]
            if not offset:
                offset = 1
            if offset - 1 > len(lines):
                return f"Error: offset {offset} is larger than file line counts {len(lines)}"
            # 截取从 offset 开始的所有行
            lines = lines[offset - 1 :]
            if limit:
                if limit > len(lines):
                    return f"Error: limit {offset} is larger than file line counts {len(lines)} from offset {offset}"
                lines = lines[:limit]

            return "".join(lines)

实际上，这里还应该做行数和字节数的超长截断。但是为了实现简单没有做。

增强 write 工具

write 工具好像没有太大的问题，我们可以跳过它。

edit 工具

Agent 通常采用文本替换的方法来编辑文件。

一共 4 个参数，一个 path，一个旧文本，一个新文本，一个可选的 replace_all 参数。

def edit(path: str, old_text: str, new_text: str, replace_all: bool = False) -> str:
    """
    Edit a file. replace old_text with new_text.

    Args:
        path (str) : the file path
        old_text (str): old content to replace
        new_text (str): new content
        replace_all (bool, optional): whether to replace all matches, default to false.
    """
    with open(path, mode="w", encoding="utf-8") as f:
        content = f.read()
        if content.find(old_text) == -1:
            return "Error: could not find old_text in file"
        if not replace_all and content.count(old_text) > 1:
            return "Error: found multiple matches in file. provider more context or set replace_all to true"
        content.replace(old_text, new_text)
        return f"Edited file {path}"

bash 工具

获取命令和 timeout 参数，然后通过 sh -c 调用。原则上应该做超长内容的截断，但是为了实现简单没有做。

def bash(command: str, timeout: int) -> str:
    """
    Run shell command.

    Args:
        command (str): the command to run
        timeout (int): command timeout
    """
    result = subprocess.run(
        ["sh", "-c", command], capture_output=True, text=True, timeout=timeout
    )
    return result.stdout

错误处理

我们的代码可能无法覆盖所有的错误。但是我们至少可以将错误反馈给模型，让它调整。为了做到这一点，我们只需要给整个工具调用加一个 try-except，然后把捕获到的错误返回即可。

try:
    result = getattr(tool, name)(**arguments)
except Exception as e:
    result = f"Error: {e}"