step5_multi_tool.py

"""
第 5 步：多个工具 —— 注册表，以及「加一个工具要改几处代码」。

运行：
    python tutorial/step5_multi_tool.py

前四步一直只有一两个工具，schema 和 handler 是分开手写的。
工具一多，这种写法立刻出问题：

    - schema 列表和 handler 字典是两份，很容易改了一边忘了另一边
    - 每加一个工具，都要回去改主流程的代码
    - 工具的「说明书」和「实现」离得很远，读代码要来回跳

解决办法朴素得有点无聊：**让工具自己注册自己**。
一个 register_tool(schema, handler)，同时往两个容器里各塞一份，
主流程从此再也不用认识任何具体工具。

这就是 tools/registry.py 做的全部事情 —— 20 行，没有别的。
看完这一步，你已经把 MiniAgent 的五个概念全部走完了：
Tool、Memory、Executor、Loop、Registry。
"""

import ast
import json
import operator
import pathlib
import sys
from datetime import datetime

sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent))

import config  # noqa: E402

# ----------------------------------------------------------------------
# 注册表：两个容器 + 一个注册函数
# ----------------------------------------------------------------------

SCHEMAS: list[dict] = []      # 发给模型的「说明书」列表
HANDLERS: dict[str, object] = {}  # 我们本地执行用的「工具名 → 函数」


def register_tool(schema: dict, handler) -> None:
    """把一个工具的两半同时登记进来，保证它们永远不会对不上。"""
    SCHEMAS.append(schema)
    HANDLERS[schema["function"]["name"]] = handler


# ----------------------------------------------------------------------
# 工具一：计算器
# ----------------------------------------------------------------------


def calculator(expression: str) -> str:
    operators = {
        ast.Add: operator.add,
        ast.Sub: operator.sub,
        ast.Mult: operator.mul,
        ast.Div: operator.truediv,
        ast.USub: operator.neg,
    }

    def walk(node):
        if isinstance(node, ast.Expression):
            return walk(node.body)
        if isinstance(node, ast.Constant) and isinstance(node.value, (int, float)):
            return node.value
        if isinstance(node, ast.BinOp) and type(node.op) in operators:
            return operators[type(node.op)](walk(node.left), walk(node.right))
        if isinstance(node, ast.UnaryOp) and type(node.op) in operators:
            return operators[type(node.op)](walk(node.operand))
        raise ValueError(f"不支持的表达式：{expression!r}")

    return str(walk(ast.parse(expression, mode="eval")))


register_tool(
    {
        "type": "function",
        "function": {
            "name": "calculator",
            "description": "计算数学表达式。需要算数时必须用它，不要心算。",
            "parameters": {
                "type": "object",
                "properties": {"expression": {"type": "string", "description": "如 '15*4'"}},
                "required": ["expression"],
            },
        },
    },
    calculator,
)


# ----------------------------------------------------------------------
# 工具二：当前时间
#
# 注意这个工具「没有参数」—— properties 是空的，也没有 required。
# 模型照样能调用它。很多人以为工具必须有参数，其实不必。
# ----------------------------------------------------------------------


def get_time() -> str:
    return datetime.now().strftime("%Y-%m-%d %H:%M:%S")


register_tool(
    {
        "type": "function",
        "function": {
            "name": "get_time",
            "description": "返回当前的本地日期和时间。问到时间必须用它，不要猜。",
            "parameters": {"type": "object", "properties": {}},
        },
    },
    get_time,
)


# ----------------------------------------------------------------------
# 主循环：和 step3 一模一样，但它已经不认识任何具体工具了
# ----------------------------------------------------------------------


def react(goal: str, max_steps: int = 6) -> str:
    client = config.get_client()
    messages = [
        {"role": "system", "content": "你是一个助手，必须用工具完成任务，不要自己猜答案。"},
        {"role": "user", "content": goal},
    ]

    for step in range(1, max_steps + 1):
        print(f"\n--- 第 {step} 圈 ---")
        response = client.chat.completions.create(
            model=config.MODEL,
            messages=messages,
            tools=SCHEMAS,  # ← 从注册表来，主流程不需要知道里面有什么
        )
        message = response.choices[0].message
        messages.append({
            "role": "assistant",
            "content": message.content or "",
            "tool_calls": [tc.model_dump() for tc in message.tool_calls or []],
        })

        if not message.tool_calls:
            print(f"[Answer] {message.content}")
            return message.content

        for tool_call in message.tool_calls:
            name = tool_call.function.name
            print(f"[Act]     {name}({tool_call.function.arguments})")

            # 模型有时会编一个不存在的工具名。别崩，把可用列表告诉它，它通常会改。
            if name not in HANDLERS:
                result = f"[错误] 没有叫 {name!r} 的工具。可用的有：{list(HANDLERS)}"
            else:
                try:
                    result = HANDLERS[name](**json.loads(tool_call.function.arguments))
                except Exception as exc:  # noqa: BLE001 — 教学演示
                    result = f"[工具出错] {type(exc).__name__}: {exc}"

            print(f"[Observe] {result}")
            messages.append({
                "role": "tool",
                "tool_call_id": tool_call.id,
                "content": result,
            })

    return f"[已达最大步数 {max_steps}，任务未完成]"


if __name__ == "__main__":
    print("已注册的工具：", list(HANDLERS))
    print("=" * 64)
    react("现在几点？顺便帮我算一下 15*4")
    print("=" * 64)
    print()
    print("教程到这里就结束了。回头看你会发现：")
    print("  step3 的循环  ≈  agent/loop.py")
    print("  step4 的 messages  ≈  agent/memory.py")
    print("  step5 的 register_tool  ≈  tools/registry.py")
    print()
    print("现在去读项目根目录的那几个文件，应该没有任何一行是陌生的了。")
    print("然后试着自己加一个工具 —— README 的「扩展」一节有模板。")