step2_first_tool.py

"""
第 2 步：给模型第一只手 —— 手动完成一次工具调用。

运行：
    python tutorial/step2_first_tool.py

这一步刻意**不写循环**。整个工具调用的过程全部摊开、手动执行一遍，
让你看清中间到底传了什么。看懂这个文件，就看懂了 Agent 的全部机制。

一次工具调用分成四拍：
    1. 我们把「工具说明书」（schema）连同问题一起发给模型
    2. 模型不直接回答，而是回一个 tool_calls：「请帮我调用 calculator，参数是……」
    3. 我们**在本地**真的执行那个 Python 函数，拿到结果
    4. 我们把结果作为一条 role="tool" 的消息塞回对话，再问模型一次

关键认知：**模型自己不会执行任何东西**。
它只会说「我想调用这个函数、参数是这些」，真正动手的永远是你的代码。
所谓 Agent 框架，核心就是这个「传话 + 代跑」的过程。
"""

import ast
import json
import operator
import pathlib
import sys

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

import config  # noqa: E402

# ----------------------------------------------------------------------
# 工具的两半：说明书给模型看，函数给我们自己跑
# ----------------------------------------------------------------------

# 第一半：schema —— 用 OpenAI function calling 格式描述这个工具。
# 模型只能看到这段文字，所以 description 写得好不好，直接决定模型会不会用、用得对不对。
CALCULATOR_SCHEMA = {
    "type": "function",
    "function": {
        "name": "calculator",
        "description": "计算一个数学表达式，返回精确结果。需要算数时必须用它，不要心算。",
        "parameters": {
            "type": "object",
            "properties": {
                "expression": {
                    "type": "string",
                    "description": "合法的数学表达式，例如 '(12+8)*3'",
                }
            },
            "required": ["expression"],
        },
    },
}


# 第二半：handler —— 真正干活的 Python 函数。模型永远看不到这里的代码。
def calculator(expression: str) -> str:
    """教学版计算器：解析 AST 并只对白名单节点求值。

    为什么不直接写 `eval(expression)`？
        因为 expression 是**模型生成的字符串**。用 eval 等于把「在你机器上
        执行任意 Python」这个权限交给了模型。这是 Agent 工具最经典的漏洞：
        工具的输入永远来自模型，而模型的输入可能来自任何人。
        完整实现（还加了长度和幂次限制）见 tools/calculator.py。
    """
    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")))


def main() -> None:
    client = config.get_client()

    messages = [
        {"role": "system", "content": "你是一个助手，需要算数时必须使用 calculator 工具。"},
        {"role": "user", "content": "帮我算一下 (12+8)*3 等于多少？"},
    ]

    # ---- 第 1 拍：带着工具说明书提问 ----
    print("【第 1 拍】把问题 + 工具说明书发给模型")
    response = client.chat.completions.create(
        model=config.MODEL,
        messages=messages,
        tools=[CALCULATOR_SCHEMA],  # 和 step1 唯一的区别，就是多了这一行
    )
    message = response.choices[0].message

    # ---- 第 2 拍：模型没有回答，而是要求调用工具 ----
    if not message.tool_calls:
        print("模型直接回答了，没有用工具：", message.content)
        print("（小模型有时会这样。换个更大的模型，或者把 description 写得更强硬。）")
        return

    tool_call = message.tool_calls[0]
    print("【第 2 拍】模型没有回答，而是要求调用工具：")
    print(f"    工具名：{tool_call.function.name}")
    print(f"    参数  ：{tool_call.function.arguments}   ← 注意这是 JSON 字符串，不是 dict")

    # ---- 第 3 拍：我们在本地真的执行它 ----
    arguments = json.loads(tool_call.function.arguments)
    result = calculator(arguments["expression"])
    print(f"【第 3 拍】我们本地执行 calculator({arguments['expression']!r}) → {result}")

    # ---- 第 4 拍：把结果塞回对话，再问一次 ----
    # 顺序很重要：assistant(tool_calls) 必须在前，紧跟着对应的 tool 消息。
    # tool_call_id 就是把这两条配对起来的钥匙，缺了它 API 会直接报错。
    messages.append({
        "role": "assistant",
        "content": message.content or "",
        "tool_calls": [tool_call.model_dump()],
    })
    messages.append({
        "role": "tool",
        "tool_call_id": tool_call.id,
        "content": result,
    })

    print("【第 4 拍】把工具结果塞回对话，再问模型一次")
    final = client.chat.completions.create(
        model=config.MODEL,
        messages=messages,
        tools=[CALCULATOR_SCHEMA],
    )
    print()
    print("最终回答：", final.choices[0].message.content)
    print()
    print("现在数一下：一次工具调用要发两次请求。")
    print("如果模型想连着用两个工具呢？三个呢？——那就得循环。下一步（step3）。")


if __name__ == "__main__":
    main()