smolagents · 中文源码精读教程

本文集文档发布于灏天文库

第 5 章 · 01 Tool 基类与 @tool 装饰器

第 5 章 · 01 Tool 基类与 @tool 装饰器

本节摘要:本节精读工具体系的根基:Tool 基类——四个类属性(name/description/inputs/output_type)加一个 forward 方法就构成一个工具,__init_subclass__ 挂钩自动校验;PipelineTool 用三个类属性把 transformers pipeline 包成工具;@tool 装饰器把普通函数一键变工具,docstring 当 description、类型提示自动生成参数 schema;幕后功臣是 _function_type_hints_utils.py(get_json_schema 把类型提示与 Google 风格 docstring 解析成 JSON schema)与 tool_validation.py(用 AST 静态校验工具类定义的合法性,为序列化与远程执行把关)。

内容来源:src/smolagents/tools.py、src/smolagents/_function_type_hints_utils.py、src/smolagents/tool_validation.py

⚠️ 注意:工具的 description 不是写给人看的注释——它会被原样拼进系统提示词,LLM 靠它决定「什么时候用这个工具、传什么参数」。写含糊的 description 等于给 LLM 发错误的说明书。

学习目标

  1. 掌握 Tool 基类四件套与 forward 约定,理解自动校验机制。
  2. 会用 @tool 装饰器三分钟写一个工具。
  3. 读懂 get_json_schema:类型提示 + docstring → JSON schema。
  4. 理解 tool_validation 的 AST 静态检查为什么存在。

一、Tool 基类:四件套 + forward

Tool 类的完整约定浓缩在类属性声明里:

class Tool(BaseTool):
    name: str
    description: str
    inputs: dict[str, dict[str, str | type | bool]]
    output_type: str
    output_schema: dict[str, Any] | None = None

    def forward(self, *args, **kwargs):
        raise NotImplementedError("Write this method in your subclass of `Tool`.")
​
  • name:工具名,LLM 在代码里就用这个名字调用它;
  • description:一句话说清「做什么、吃什么参数、吐什么结果」;
  • inputs:参数字典,每个参数必须有 type 和 description 两键;
  • output_type:返回类型(string/integer/image/audio/any 等);
  • output_schema:可选,结构化输出的 JSON schema,提示 LLM 返回值长什么样。

实际逻辑写在 forward 里,但外部统一走 __call__:

def __call__(self, *args, sanitize_inputs_outputs: bool = False, **kwargs):
    if not self.is_initialized:
        self.setup()
    # Handle the arguments might be passed as a single dictionary
    if len(args) == 1 and len(kwargs) == 0 and isinstance(args[0], dict):
        potential_kwargs = args[0]
        if all(key in self.inputs for key in potential_kwargs):
            args = ()
            kwargs = potential_kwargs
    outputs = self.forward(*args, **kwargs)  # sanitize_inputs_outputs 时还做输入/输出类型转换
    return outputs
​

两个设计点:懒初始化——setup()(加载大模型等重活)推迟到第一次调用;字典参数宽容——LLM 传一个 dict 进来时,若键恰好匹配 inputs 就自动摊开成关键字参数。另外 to_code_prompt() 会把工具渲染成「带签名的 Python 函数定义 + docstring」的样子拼进 CodeAgent 系统提示,这就是 LLM 眼中「工具即函数」的由来。

子类化时不必手动检查属性写没写对。__init_subclass__ 钩子配合 validate_after_init 装饰器,把每个子类的 __init__ 包一层——实例化后自动追加执行 self.validate_arguments()。

validate_arguments 检查:四个必填属性是否齐全且类型正确、name 是否为合法 Python 标识符、每个 input 是否都有 type/description 且 type 在 AUTHORIZED_TYPES 白名单里、forward 的参数签名是否与 inputs 的键完全一致。错误配置在实例化当场爆掉,而不是等到 LLM 调用时才 mysteriously 失败。

二、PipelineTool:三个类属性包一个 transformers 模型

PipelineTool 专治「把 transformers 模型变成工具」:

class PipelineTool(Tool):
    pre_processor_class = None
    model_class = None
    post_processor_class = None
    default_checkpoint = None
​

它把调用流程固化为 encode → forward → decode 三段,setup() 负责按 checkpoint 惰性加载:

def setup(self):
    if isinstance(self.pre_processor, str):
        self.pre_processor = self.pre_processor_class.from_pretrained(self.pre_processor, **self.hub_kwargs)
    if isinstance(self.model, str):
        self.model = self.model_class.from_pretrained(self.model, **self.model_kwargs, **self.hub_kwargs)
    ...
    self.model.to(self.device)
    super().setup()
​

三个字段传字符串 checkpoint 也行、传现成实例也行(from_pretrained 只在收到字符串时触发)。子类只需声明三个类属性(pre_processor_class/model_class/default_checkpoint)并按需覆写 encode/decode——比如第 5.3 节的 SpeechToTextTool(Whisper 语音转文字)总共不到 30 行。设备选择交给 accelerate 的 PartialState 自动探测 GPU。

三、@tool 装饰器:普通函数一秒变工具

不想写类?@tool 装饰器直接吃函数:

from smolagents import tool

@tool
def get_travel_duration(start: str, end: str, mode: str) -> str:
    """Gets the travel time in seconds between two places.

    Args:
        start: the start place
        end: the destination place
        mode: the transportation mode, one of 'driving', 'walking', 'bicycling' or 'transit'
    """
    ...  # 实际逻辑
    return f"Travel time by {mode}: {duration}"
​

装饰器源码的主干:

def tool(tool_function: Callable) -> Tool:
    tool_json_schema = get_json_schema(tool_function)["function"]
    if "return" not in tool_json_schema:
        if len(tool_json_schema["parameters"]["properties"]) == 0:
            tool_json_schema["return"] = {"type": "null"}
        else:
            raise TypeHintParsingException(
                "Tool return type not found: make sure your function has a return type hint!"
            )

    class SimpleTool(Tool):
        def __init__(self):
            self.is_initialized = True

    SimpleTool.name = tool_json_schema["name"]
    SimpleTool.description = tool_json_schema["description"]
    SimpleTool.inputs = tool_json_schema["parameters"]["properties"]
    SimpleTool.output_type = tool_json_schema["return"]["type"]
    SimpleTool.forward = staticmethod(tool_function)  # 实际经 wraps 包装
    ...
​

流程:先调 get_json_schema 生成完整 schema,再动态建一个 SimpleTool 类,把 schema 里的 name/description/inputs/output_type 灌进类属性,原函数包成 forward。返回类型提示必须有(无参数工具除外)——因为 output_type 必须告知 LLM。

四、_function_type_hints_utils.py:schema 自动生成的核心

get_json_schema 是 @tool 的引擎,做两件事:解析 docstring、解析类型提示。

def get_json_schema(func: Callable) -> dict:
    doc = inspect.getdoc(func)
    if not doc:
        raise DocstringParsingException(
            f"Cannot generate JSON schema for {func.__name__} because it has no docstring!"
        )
    main_doc, param_descriptions, return_doc = _parse_google_format_docstring(doc)
    json_schema = _convert_type_hints_to_json_schema(func)
    for arg, schema in json_schema["properties"].items():
        if arg not in param_descriptions:
            raise DocstringParsingException(... "the docstring has no description for the argument '{arg}'")
        schema["description"] = param_descriptions[arg]
    output = {"name": func.__name__, "description": main_doc, "parameters": json_schema}
    ...
​

三个正则把 Google 风格 docstring 切成三段(节选自三个预编译正则):

description_re = re.compile(r"^(.*?)(?=\n\s*(Args:|Returns:|Raises:)|\Z)", re.DOTALL)  # 函数描述段
args_re = re.compile(r"\n\s*Args:\n\s*(.*?)[\n\s]*(Returns:|Raises:|\Z)", re.DOTALL)   # Args 段
args_split_re = re.compile(                                                            # 再按参数切分
    r"(?:^|\n)"
    r"\s*(\w+)\s*(?:\([^)]*?\))?:\s*"   # 捕获参数名(忽略括号里的类型)
    r"(.*?)\s*"                          # 捕获参数描述(可跨行)
    r"(?=\n\s*\w+\s*(?:\([^)]*?\))?:|\Z)",
    re.DOTALL | re.VERBOSE,
)
​

没有 docstring、有参数却没写参数描述,都会直接抛异常——tool 的「文档即接口」被强制执行。描述末尾还支持 (choices: ["tea", "coffee"]) 语法,解析成 schema 的 enum 字段。

类型侧,_convert_type_hints_to_json_schema 用 get_type_hints 取注解,inspect.signature 取默认值:无默认值的参数进 required,有默认值的打上 nullable: True。_parse_type_hint 递归处理容器类型:

elif origin is list:
    if not args:
        return {"type": "array"}
    else:
        return {"type": "array", "items": _parse_type_hint(args[0])}
elif origin is dict:
    out = {"type": "object"}
    if len(args) == 2:
        out["additionalProperties"] = _parse_type_hint(args[1])
    return out
elif origin is Literal:
    ...
    final_type.update({"enum": [arg for arg in args if arg is not None]})
​

基础类型映射表很直白:int→integer、float→number、str→string、bool→boolean、list→array、dict→object。Union 类型(X | None)拆成子类型列表并标记 nullable,多类型并集则降级为 any。_get_json_schema_type 还特判了 PIL 的 Image → image 类型、torch 的 Tensor → audio 类型,让多模态工具的参数声明保持自然。

五、tool_validation.py:AST 静态校验

为什么需要静态校验?因为第 4.2 节的 save/push_to_hub 与第 6 章的远程执行器都要「把工具源码抽出来,在别处重建」。抽出来的代码若引用了闭包变量、本地模块或复杂初始化参数,到沙箱里必然崩。validate_tool_attributes 用 AST 在不运行代码的前提下排查隐患,docstring 列了规则清单:__init__ 参数必须有默认值(初始化参数无法追溯重建)、类属性只能是字符串或字典等简单字面量(复杂对象放 __init__)、方法里的 import 必须来自包而非本地文件、所有方法必须自包含(self-contained)。

实现由两个 AST 访问者组成。ClassLevelChecker 查类层面:类属性必须是简单字面量、__init__ 的参数必须带字面量默认值、name 必须是合法 Python 标识符。MethodChecker 查方法体——核心是「引用了未定义的名字」检测:

def visit_Name(self, node):
    if isinstance(node.ctx, ast.Load):
        if not (
            node.id in _BUILTIN_NAMES
            or node.id in BASE_BUILTIN_MODULES
            or node.id in self.arg_names
            or node.id == "self"
            or node.id in self.class_attributes
            or node.id in self.imports
            or node.id in self.from_imports
            or node.id in self.assigned_names
            ...
        ):
            self.errors.append(f"Name '{node.id}' is undefined.")
​

一个名字只有在「内置名/基础模块/函数参数/self/类属性/显式 import/方法内赋值」之列才算合法——visit_Assign、visit_For、visit_With、visit_ExceptHandler、推导式的 generators 等钩子专门负责收集各种赋值形式(with X as y、except E as e、循环变量都算)。这条防线服务两个场景:LLM 生成工具代码时的快速质检,以及远程执行前确认工具自包含。校验失败抛出的 ValueError 会把所有错误一次性列全,方便修复。

💡 阶梯要点:本阶装上「造工具」的装备。三层心智模型:Tool 基类是手工精修路线(全控制),PipelineTool 是 transformers 快车道,@tool 是函数极简路线(类型提示与 docstring 即 schema)。而 tool_validation 保证无论哪条路线造出的工具,都能被抽成源码、异地重建——这正是通向第 6 章远程沙箱的前提。

本节要点回顾

  • Tool 四件套:name/description/inputs/output_type,逻辑写 forward,外部调 __call__(懒 setup + 字典参数摊开)。
  • __init_subclass__ + validate_after_init:实例化即校验四件套、forward 签名与 inputs 键一致。
  • PipelineTool:pre_processor_class/model_class/default_checkpoint 三属性,setup 惰性加载,encode→forward→decode。
  • @tool:get_json_schema 生成 schema,动态建 SimpleTool 类;返回类型提示必填;docstring 必须有且每个参数都要有描述。
  • get_json_schema = Google docstring 解析(三正则,支持 choices→enum)+ 类型提示解析(递归容器类型,int→integer 映射,Union→nullable)。
  • tool_validation:AST 静态检查类属性简单性、__init__ 默认值、方法内名字引用自包含——为源码抽取与远程执行把关。

下一节:02 生态集成四件套与 ToolCollection——from_hub/from_space/from_mcp/from_langchain,让别人的工具直接为你的 agent 所用。

作者与出处
原作者： 灏天文库
整理： 灏天文库整理
本站整理收录，版权归原作者/开源协议所有；欢迎通过原文链接访问源仓库。
发布者: 作者: 灏天文库 转发
评论区 (0)
U
0/500