code/05-agents/planning_reflection.py

"""同一个多步任务，比较三种做法：直接做、先列计划再做、做完之后自我检查并修改一次。

在 AI-Course/code/05-agents 目录下运行：python planning_reflection.py
"""
import json
import time

from agent_loop import SYSTEM, TOOLS, RealModel

model = RealModel()
SCHEMAS = [t["schema"] for t in TOOLS.values()]
TASK = ("对比 httpx 的 Client 和 AsyncClient 在三件事上的配置方式是否一样：超时、代理、HTTP/2。"
        "列成一张表，每一格都注明文档出处（文件名和行号）。")


def loop(messages, max_steps=12):
    """执行智能体循环，直到模型不再调用工具。返回最终回答、完整消息列表和统计。"""
    stats = {"steps": 0, "tool_calls": 0, "tokens": 0}
    for _ in range(max_steps):
        message, usage = model(messages, SCHEMAS)
        stats["steps"] += 1
        stats["tokens"] += usage.prompt_tokens + usage.completion_tokens
        if not message.tool_calls:
            return message.content, messages, stats
        messages.append(message.model_dump(exclude_none=True))
        for call in message.tool_calls:
            stats["tool_calls"] += 1
            try:
                result = TOOLS[call.function.name]["fn"](**json.loads(call.function.arguments or "{}"))
            except Exception as e:
                result = f"错误：{type(e).__name__}: {e}"
            messages.append({"role": "tool", "tool_call_id": call.id, "content": result[:3000]})
    return None, messages, stats


def plain():
    return loop([{"role": "system", "content": SYSTEM}, {"role": "user", "content": TASK}])


def with_plan():
    # 第一步：只列计划，不调用工具
    plan, usage = model([{"role": "system", "content": SYSTEM}, {"role": "user", "content":
                         TASK + "\n\n先不要调用工具。列出你打算怎么查，编号列出每一步要找什么，不超过 6 步。"}], None)
    print("  计划：" + plan.content.replace("\n", " | ")[:300])
    answer, messages, stats = loop([
        {"role": "system", "content": SYSTEM},
        {"role": "user", "content": TASK + "\n\n按这个计划执行，执行中发现计划不对可以调整：\n" + plan.content},
    ])
    stats["tokens"] += usage.prompt_tokens + usage.completion_tokens
    return answer, messages, stats


def with_reflection():
    answer, messages, stats = plain()
    evidence = "\n\n".join(m["content"] for m in messages if m["role"] == "tool")
    review, usage = model([{"role": "user", "content": f"""下面是一份回答和查到的全部原文。逐格检查回答里的表格：
每一格的说法，原文里有没有依据？注明的出处（文件和行号）对不对？
只列出有问题的格子和原因。全部没问题就只回复"没有问题"。

回答：
{answer}

原文：
{evidence[:20000]}"""}], None)
    stats["tokens"] += usage.prompt_tokens + usage.completion_tokens
    print("  检查意见：" + review.content.replace("\n", " | ")[:400])
    if "没有问题" in review.content[:20]:
        return answer, messages, stats
    # 有问题就把检查意见交回给智能体，让它继续查、修改回答
    messages += [{"role": "assistant", "content": answer},
                 {"role": "user", "content": "有人检查了你的回答，意见如下。需要的话继续查文档，然后给出修改后的完整回答。\n\n" + review.content}]
    revised, messages, more = loop(messages)
    for k in stats:
        stats[k] += more[k]
    return revised, messages, stats


for name, run in [("直接做", plain), ("先列计划", with_plan), ("做完再检查", with_reflection)]:
    print(f"===== {name}")
    start = time.time()
    answer, _, stats = run()
    print(f"  {stats['steps']} 次模型调用，{stats['tool_calls']} 次工具调用，{stats['tokens']} 词元，{time.time() - start:.0f} 秒")
    print("  回答：\n" + (answer or "（没有完成）"))
    print()