step4_memory.py

"""
第 4 步：记忆 —— 让 Agent 记住上一轮发生了什么。

运行：
    python tutorial/step4_memory.py

step3 的循环每次只处理一个目标，跑完就散了。
但真实对话是连续的：「算一下 12+8」→「把结果写到文件里」——
第二句话里的「结果」，只有在模型能看见第一轮的对话时才有意义。

这一步做的事，说穿了非常朴素：
    **把 messages 列表在多轮之间留着，不要每次都新建。**

这就是所谓「记忆」的全部。没有向量数据库，没有魔法，就是一个 list。

这个文件会跑两遍同一个第二轮请求做对比：
    有记忆 → 模型看得到 "20"，写进文件的是 20
    无记忆 → 模型不知道「结果」指什么，只能把问题本身写进去

顺带一个真实世界的坑：对话越长，token 花得越多，还会撞上上下文长度上限。
所以 Memory 需要裁剪策略 —— 见 agent/memory.py 的 _trim()，
那里还处理了一个隐蔽约束：tool 消息必须紧跟在它配对的 assistant 之后，
裁剪时不能把这一对拆散，否则 API 会直接 400。
"""

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

# 教程产生的文件都丢进这个目录，不弄脏项目根
OUTPUT_DIR = pathlib.Path(__file__).resolve().parent / "tutorial_output"

TOOL_SCHEMAS = [
    {
        "type": "function",
        "function": {
            "name": "calculator",
            "description": "计算数学表达式。需要算数时必须用它。",
            "parameters": {
                "type": "object",
                "properties": {"expression": {"type": "string"}},
                "required": ["expression"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "file_io",
            "description": "把内容写入文件。",
            "parameters": {
                "type": "object",
                "properties": {
                    "action": {"type": "string", "enum": ["write"]},
                    "filename": {"type": "string"},
                    "content": {"type": "string"},
                },
                "required": ["action", "filename", "content"],
            },
        },
    },
]


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")))


def file_io(action: str, filename: str, content: str = "") -> str:
    """极简写文件工具。所有写入都关在 tutorial_output/ 里。

    真实版本见 tools/file_io.py —— 那里还要防路径穿越、防软链接逃逸、防误删。
    """
    if action != "write":
        return f"[错误] 本教程只实现了 write，收到: {action!r}"
    OUTPUT_DIR.mkdir(exist_ok=True)
    # 只取文件名部分，丢掉任何目录成分，防止 ../../ 跑出去
    target = OUTPUT_DIR / pathlib.Path(filename).name
    target.write_text(content, encoding="utf-8")
    return f"已写入 {target.name}（{len(content)} 字符）"


HANDLERS = {"calculator": calculator, "file_io": file_io}


class Agent:
    """和 step3 的 react() 几乎一样，唯一区别：messages 是实例属性。

    messages 活在 self 上而不是函数局部变量里 —— 就这一个改动，
    Agent 就有了跨轮次的记忆。
    """

    def __init__(self, system_prompt: str):
        self.messages = [{"role": "system", "content": system_prompt}]

    def run(self, goal: str, max_steps: int = 5) -> str:
        client = config.get_client()
        self.messages.append({"role": "user", "content": goal})

        for _ in range(max_steps):
            response = client.chat.completions.create(
                model=config.MODEL,
                messages=self.messages,
                tools=TOOL_SCHEMAS,
            )
            message = response.choices[0].message
            self.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:
                return message.content

            for tool_call in message.tool_calls:
                arguments = json.loads(tool_call.function.arguments)
                try:
                    result = HANDLERS[tool_call.function.name](**arguments)
                except Exception as exc:  # noqa: BLE001 — 教学演示
                    result = f"[工具出错] {type(exc).__name__}: {exc}"
                print(f"    [Act] {tool_call.function.name}({tool_call.function.arguments})")
                print(f"    [Observe] {result}")
                self.messages.append({
                    "role": "tool",
                    "tool_call_id": tool_call.id,
                    "content": result,
                })

        return f"[已达最大步数 {max_steps}]"


SYSTEM = "你是一个助手，必须用工具完成任务，不要心算也不要假装写了文件。"
FOLLOW_UP = "把上一步的结果写到 note.txt"


def show_note() -> str:
    """读回刚写的文件内容 —— 两个场景写的是同一个文件名，必须马上看，不然会被覆盖。"""
    note = OUTPUT_DIR / "note.txt"
    return note.read_text(encoding="utf-8") if note.exists() else "(文件不存在)"


def main() -> None:
    print("=" * 64)
    print("场景 A：有记忆 —— 同一个 Agent 连续处理两轮")
    print("=" * 64)
    agent = Agent(SYSTEM)

    print("\n第 1 轮：计算 12+8")
    print("  →", agent.run("计算 12+8"))

    print(f"\n第 2 轮：{FOLLOW_UP}")
    agent.run(FOLLOW_UP)
    with_memory = show_note()
    print(f"\n此时 messages 里已经攒了 {len(agent.messages)} 条消息，")
    print("第一轮的计算结果就在里面 —— 所以模型知道「上一步的结果」是什么。")

    print()
    print("=" * 64)
    print("场景 B：无记忆 —— 换一个全新 Agent，只给它第二轮的请求")
    print("=" * 64)
    amnesiac = Agent(SYSTEM)
    print(f"\n只说：{FOLLOW_UP}")
    amnesiac.run(FOLLOW_UP)
    without_memory = show_note()

    print()
    print("-" * 64)
    print("note.txt 两次被写入的内容：")
    print(f"  有记忆 → {with_memory!r}")
    print(f"  无记忆 → {without_memory!r}")
    print()
    print("无记忆时模型根本不知道「结果」指的是什么，只好把问题原样抄进去。")
    print("这就是 Memory 的作用 —— 它不神秘，就是那个没被丢掉的 list。")
    print("下一步（step5）：多个工具怎么管理。")


if __name__ == "__main__":
    main()