大模型 Agent 的原理

"""极简 ReAct 循环：不依赖任何框架，60 行看懂 Agent 的自主性到底在哪。

运行前：
    pip install openai
    export OPENAI_API_KEY=...        # 兼容 OpenAI 协议的任意服务都行
    export OPENAI_BASE_URL=...       # 换供应商只改这一行
"""
import json
import os

from openai import OpenAI

client = OpenAI(
    api_key=os.environ.get("OPENAI_API_KEY"),
    base_url=os.environ.get("OPENAI_BASE_URL"),
)
MODEL = os.environ.get("MODEL_NAME", "gpt-4o-mini")


# ① 两个最普通的 Python 函数，Agent 能做的事全靠它们
def search(keyword: str) -> str:
    fake = {"杭州人口": "杭州市常住人口约 1252 万人",
            "南京人口": "南京市常住人口约 954 万人"}
    return fake.get(keyword, "没有查到相关资料")


def calculate(expression: str) -> str:
    # 真实项目请换成受限求值器，这里为了聚焦主线用最短写法
    return str(eval(expression, {"__builtins__": {}}, {}))


TOOLS = [
    {"type": "function", "function": {
        "name": "search", "description": "按关键词检索资料，返回一句话事实",
        "parameters": {"type": "object", "properties": {
            "keyword": {"type": "string", "description": "检索关键词"}},
            "required": ["keyword"]}}},
    {"type": "function", "function": {
        "name": "calculate", "description": "计算一个纯数字算术表达式",
        "parameters": {"type": "object", "properties": {
            "expression": {"type": "string", "description": "例如 (12-9)/9*100"}},
            "required": ["expression"]}}},
]
REGISTRY = {"search": search, "calculate": calculate}

MAX_STEPS = 6          # ② 循环的刹车：没有它，Agent 可能一直转下去


def run(question: str) -> str:
    messages = [
        {"role": "system", "content": "你是一个严谨的助手。需要事实就去检索，"
                                      "需要算数就去计算，不要凭印象回答。"},
        {"role": "user", "content": question},
    ]
    for step in range(1, MAX_STEPS + 1):
        # ③ 思考：把「目前已知的一切」整个交给模型，让它决定下一步
        reply = client.chat.completions.create(
            model=MODEL, messages=messages, tools=TOOLS, tool_choice="auto"
        ).choices[0].message
        messages.append(reply.model_dump(exclude_none=True))

        # ④ 终止条件：模型这一轮不再要工具，说明它认为可以收尾了
        if not reply.tool_calls:
            print("第 %d 轮收尾" % step)
            return reply.content

        # ⑤ 行动 + 观察：执行工具，把结果作为新事实追加回去，进入下一轮
        for call in reply.tool_calls:
            args = json.loads(call.function.arguments)
            result = REGISTRY[call.function.name](**args)
            print("第 %d 轮 → %s(%s) = %s" % (step, call.function.name, args, result))
            messages.append({"role": "tool", "tool_call_id": call.id,
                             "content": str(result)})

    # ⑥ 步数用尽仍未收敛：必须给出明确结论，不能静默返回空
    return "已达到最大步数 %d，未能得出结论" % MAX_STEPS


if __name__ == "__main__":
    print(run("杭州比南京多多少人口？折算成百分比是多少？"))

"""纯文本版 ReAct：模型不支持结构化工具调用时，靠提示词 + 正则把循环跑起来。

ReAct 论文里的原始形态就是文本：模型输出
    Thought: ...
    Action: 工具名
    Action Input: 参数
程序解析出动作并执行，把结果作为
    Observation: ...
拼回提示词，再让模型续写下一轮。结构化工具调用是后来的工程化改良，
但两者的骨架一模一样：想 → 做 → 看 → 再想。
"""
import os
import re

from openai import OpenAI

client = OpenAI(
    api_key=os.environ.get("OPENAI_API_KEY"),
    base_url=os.environ.get("OPENAI_BASE_URL"),
)
MODEL = os.environ.get("MODEL_NAME", "gpt-4o-mini")


def search(q: str) -> str:
    table = {"珠穆朗玛峰高度": "珠穆朗玛峰海拔 8848.86 米",
             "华山高度": "华山海拔 2154.9 米"}
    return table.get(q.strip(), "未检索到相关资料")


def calculate(expr: str) -> str:
    return str(round(eval(expr, {"__builtins__": {}}, {}), 4))


TOOLBOX = {"检索": search, "计算": calculate}

TEMPLATE = """你可以使用以下工具：
检索：输入一个关键词，返回一句事实
计算：输入一个算术表达式，返回数值

严格按以下格式逐轮输出，一次只输出一轮：
Thought: 你的推理
Action: 检索 或 计算
Action Input: 工具的输入

当信息已经足够时，改为输出：
Thought: 你的推理
Final Answer: 最终答案

问题：{question}
{scratchpad}"""

ACTION_RE = re.compile(r"Action:\s*(\S+)\s*\nAction Input:\s*(.+)")


def run(question: str, max_steps: int = 6) -> str:
    scratchpad = ""                       # 这就是 Agent 的短期记忆
    for step in range(1, max_steps + 1):
        prompt = TEMPLATE.format(question=question, scratchpad=scratchpad)
        text = client.chat.completions.create(
            model=MODEL, temperature=0, stop=["Observation:"],
            messages=[{"role": "user", "content": prompt}],
        ).choices[0].message.content.strip()
        print("── 第 %d 轮 ──\n%s" % (step, text))

        if "Final Answer:" in text:
            return text.split("Final Answer:", 1)[1].strip()

        m = ACTION_RE.search(text)
        if not m:
            # 模型没按格式写，把要求再说一遍，给它一次自我纠正的机会
            scratchpad += text + "\nObservation: 输出格式不正确，请严格按格式重写本轮。\n"
            continue

        name, arg = m.group(1).strip(), m.group(2).strip()
        observation = TOOLBOX.get(name, lambda _: "没有这个工具")(arg)
        # 关键一步：把「这一轮做了什么、看到了什么」原样接在提示词后面
        scratchpad += text + "\nObservation: " + observation + "\n"
    return "超过 %d 轮未得出结论" % max_steps


if __name__ == "__main__":
    print(run("珠穆朗玛峰比华山高多少米？高出的部分相当于华山的百分之多少？"))

"""链的写法：步骤与顺序都由开发者写死，模型只负责每一格里的那点活。

同一个需求「杭州比南京多多少人口、折算成百分比」，先用链实现一遍，
再看 agent_loop.py 用循环实现的版本，两份代码放在一起，差别一眼就出来。
"""
import os
import re

from openai import OpenAI

client = OpenAI(
    api_key=os.environ.get("OPENAI_API_KEY"),
    base_url=os.environ.get("OPENAI_BASE_URL"),
)
MODEL = os.environ.get("MODEL_NAME", "gpt-4o-mini")

CITY_DB = {"杭州": 1252.0, "南京": 954.0}      # 单位：万人


def ask(prompt: str) -> str:
    """一次纯文本问答，没有工具，没有循环。"""
    resp = client.chat.completions.create(
        model=MODEL, messages=[{"role": "user", "content": prompt}], temperature=0
    )
    return resp.choices[0].message.content.strip()


def step1_extract(question: str) -> tuple:
    """第一步：让模型从原话里抽出两个城市名。位置固定，先后固定。"""
    raw = ask("从这句话里抽出两个城市名，只输出两个词、用逗号隔开，"
              "不要任何解释：\n" + question)
    names = [x.strip() for x in re.split(r"[,，]", raw) if x.strip()]
    return names[0], names[1]


def step2_lookup(a: str, b: str) -> tuple:
    """第二步：查数。这一步压根没模型什么事，是写死的字典查询。"""
    if a not in CITY_DB or b not in CITY_DB:
        raise KeyError("资料库里没有 %s 或 %s" % (a, b))
    return CITY_DB[a], CITY_DB[b]


def step3_compute(pa: float, pb: float) -> tuple:
    """第三步：算数。同样写死，模型碰都不碰。"""
    return pa - pb, (pa - pb) / pb * 100


def step4_phrase(a: str, b: str, diff: float, pct: float) -> str:
    """第四步：让模型把数字说成人话。它只做这一件事。"""
    return ask("用一句中文自然地陈述这个结论，不要加任何额外信息：%s 比 %s 多 "
               "%.0f 万人，相当于多出 %.2f%%。" % (a, b, diff, pct))


def run(question: str) -> str:
    """四步固定串行。谁先谁后、跑几步，全写在这里，和模型无关。"""
    a, b = step1_extract(question)
    pa, pb = step2_lookup(a, b)
    diff, pct = step3_compute(pa, pb)
    return step4_phrase(a, b, diff, pct)


if __name__ == "__main__":
    print(run("杭州比南京多多少人口？折算成百分比是多少？"))
    # 换个问法试试：「南京和苏州哪个人口多」——第二步直接抛 KeyError，
    # 因为这条流水线只会按「抽两个城市→查表→相减→润色」这一条路走。
    # 流水线不会改道，能力边界就是它被写死的那几格。

"""同一个需求的循环写法：步骤数和顺序都由模型在运行时决定。

和 chain_hardcoded.py 放在一起读：那边是四个写死的函数按序执行，
这边只有一个 while，工具怎么用、用几次、先用哪个，全在循环里现场决定。
"""
import json
import os

from openai import OpenAI

client = OpenAI(
    api_key=os.environ.get("OPENAI_API_KEY"),
    base_url=os.environ.get("OPENAI_BASE_URL"),
)
MODEL = os.environ.get("MODEL_NAME", "gpt-4o-mini")

CITY_DB = {"杭州": 1252.0, "南京": 954.0, "苏州": 1296.0, "合肥": 985.0}


def lookup_population(city: str) -> str:
    """查一个城市的常住人口，单位万人。"""
    if city not in CITY_DB:
        return "资料库里没有 %s 的人口数据" % city
    return "%s 常住人口 %.0f 万人" % (city, CITY_DB[city])


def calculate(expression: str) -> str:
    """算一个纯数字表达式。"""
    try:
        return str(round(eval(expression, {"__builtins__": {}}, {}), 4))
    except Exception as exc:
        return "表达式无法计算：%s" % exc


TOOLS = [
    {"type": "function", "function": {
        "name": "lookup_population",
        "description": "查询指定城市的常住人口数量，单位为万人",
        "parameters": {"type": "object", "properties": {
            "city": {"type": "string", "description": "城市中文名，例如 杭州"}},
            "required": ["city"]}}},
    {"type": "function", "function": {
        "name": "calculate",
        "description": "计算一个纯数字的算术表达式，返回数值结果",
        "parameters": {"type": "object", "properties": {
            "expression": {"type": "string", "description": "例如 (1252-954)/954*100"}},
            "required": ["expression"]}}},
]
REGISTRY = {"lookup_population": lookup_population, "calculate": calculate}

SYSTEM = ("你是一个严谨的数据助手。凡是涉及具体数字的事实，一律通过工具获取，"
          "不许凭记忆作答；凡是涉及算术，一律交给计算工具。"
          "信息足够时直接给出最终结论，不要再调用工具。")


def run(question: str, max_steps: int = 8) -> str:
    messages = [{"role": "system", "content": SYSTEM},
                {"role": "user", "content": question}]
    for step in range(1, max_steps + 1):
        reply = client.chat.completions.create(
            model=MODEL, messages=messages, tools=TOOLS,
            tool_choice="auto", temperature=0
        ).choices[0].message
        messages.append(reply.model_dump(exclude_none=True))

        if not reply.tool_calls:
            return reply.content

        for call in reply.tool_calls:
            args = json.loads(call.function.arguments)
            out = REGISTRY[call.function.name](**args)
            print("  第 %d 轮 · %s%s → %s" % (step, call.function.name, args, out))
            messages.append({"role": "tool", "tool_call_id": call.id,
                             "content": str(out)})
    return "超过最大步数 %d 仍未收敛，请缩小问题范围" % max_steps


if __name__ == "__main__":
    # 同一份代码，三个问法所需的轮数完全不同：
    #   两城比较 → 查两次 + 算一次 + 收尾，约 3~4 轮
    #   单城查询 → 查一次 + 收尾，2 轮
    #   资料库里没有的城市 → 工具如实返回「没有数据」，模型据此改口而不是硬编
    for q in ["杭州比南京多多少人口？折算成百分比是多少？",
              "苏州有多少人？",
              "拉萨和合肥哪个人口多？"]:
        print("问题：", q)
        print("回答：", run(q), "\n")

"""完整案例一：调研型 Agent —— 轮数由问题难度决定，而不是由代码写死。

同一份代码，问「北京人口多少」两轮就收尾，问「长三角三城谁增速最快」
可能要转六七轮。这正是循环相对流水线的价值所在。
"""
import json
import os
import time

from openai import OpenAI

client = OpenAI(
    api_key=os.environ.get("OPENAI_API_KEY"),
    base_url=os.environ.get("OPENAI_BASE_URL"),
)
MODEL = os.environ.get("MODEL_NAME", "gpt-4o-mini")

# ---------------- 工具区：三个互相独立的能力 ----------------
CITY = {
    "杭州": {"pop": 1252.0, "gdp": 20059.0},
    "南京": {"pop": 954.0, "gdp": 17421.0},
    "苏州": {"pop": 1296.0, "gdp": 24653.0},
}


def city_population(city: str) -> str:
    d = CITY.get(city)
    return "%s 常住人口 %.0f 万人" % (city, d["pop"]) if d else "没有 %s 的数据" % city


def city_gdp(city: str) -> str:
    d = CITY.get(city)
    return "%s 地区生产总值 %.0f 亿元" % (city, d["gdp"]) if d else "没有 %s 的数据" % city


def calculate(expression: str) -> str:
    try:
        return str(round(eval(expression, {"__builtins__": {}}, {}), 2))
    except Exception as exc:
        return "计算失败：%s" % exc


REGISTRY = {"city_population": city_population, "city_gdp": city_gdp,
            "calculate": calculate}

TOOLS = [
    {"type": "function", "function": {
        "name": "city_population", "description": "查询城市常住人口，单位万人",
        "parameters": {"type": "object",
                       "properties": {"city": {"type": "string", "description": "城市中文名"}},
                       "required": ["city"]}}},
    {"type": "function", "function": {
        "name": "city_gdp", "description": "查询城市地区生产总值，单位亿元",
        "parameters": {"type": "object",
                       "properties": {"city": {"type": "string", "description": "城市中文名"}},
                       "required": ["city"]}}},
    {"type": "function", "function": {
        "name": "calculate", "description": "计算纯数字算术表达式",
        "parameters": {"type": "object",
                       "properties": {"expression": {"type": "string",
                                                     "description": "例如 20059/1252"}},
                       "required": ["expression"]}}},
]

SYSTEM = ("你是一名数据调研助手。所有数字必须来自工具，禁止凭记忆作答；"
          "所有算术必须交给计算工具。每次只做一小步，拿到结果后再决定下一步。"
          "信息足够时直接给出结论并说明依据来自哪几次查询。")


def research(question: str, max_steps: int = 10, verbose: bool = True) -> dict:
    messages = [{"role": "system", "content": SYSTEM},
                {"role": "user", "content": question}]
    tool_calls_made, t0 = 0, time.time()

    for step in range(1, max_steps + 1):
        resp = client.chat.completions.create(
            model=MODEL, messages=messages, tools=TOOLS,
            tool_choice="auto", temperature=0)
        msg = resp.choices[0].message
        messages.append(msg.model_dump(exclude_none=True))

        if not msg.tool_calls:
            return {"answer": msg.content, "steps": step,
                    "tool_calls": tool_calls_made,
                    "seconds": round(time.time() - t0, 2)}

        for call in msg.tool_calls:
            name = call.function.name
            args = json.loads(call.function.arguments)
            result = REGISTRY.get(name, lambda **_: "没有这个工具")(**args)
            tool_calls_made += 1
            if verbose:
                print("  第%d轮 %s%s → %s" % (step, name, args, result))
            messages.append({"role": "tool", "tool_call_id": call.id,
                             "content": str(result)})

    return {"answer": "超过 %d 轮仍未收敛" % max_steps, "steps": max_steps,
            "tool_calls": tool_calls_made, "seconds": round(time.time() - t0, 2)}


if __name__ == "__main__":
    for q in ["南京有多少人口？",
              "杭州和南京，哪个人均生产总值更高？高出百分之多少？"]:
        print("【问题】", q)
        out = research(q)
        print("【回答】", out["answer"])
        print("【统计】共 %d 轮，%d 次工具调用，用时 %.2f 秒\n"
              % (out["steps"], out["tool_calls"], out["seconds"]))

"""规划：先拆再做，和边想边做，是两种不同的规划姿势。

ReAct 是「边想边做」——每一轮只决定下一步。
Plan-and-Execute 是「先拆再做」——一次性列出全部子任务，再逐条执行。
两者不是谁取代谁，长任务里常常混用：先拆出大纲，每条大纲内部再走 ReAct 循环。
"""
import json
import os

from openai import OpenAI

client = OpenAI(
    api_key=os.environ.get("OPENAI_API_KEY"),
    base_url=os.environ.get("OPENAI_BASE_URL"),
)
MODEL = os.environ.get("MODEL_NAME", "gpt-4o-mini")

PLAN_PROMPT = """把下面这个目标拆成 2 到 5 个可以独立执行的子任务。
每个子任务必须是一句可以直接动手的指令，不要写「分析一下」这种空话。
只输出 JSON 数组，形如 ["子任务一", "子任务二"]，不要任何解释。

目标：{goal}"""

REVIEW_PROMPT = """目标：{goal}

已经完成的子任务及其结果：
{done}

还没做的子任务：
{todo}

请判断：现有结果是否已经足够达成目标？
如果足够，输出 JSON：{{"finish": true, "answer": "最终答案"}}
如果不够，输出 JSON：{{"finish": false, "next": "下一步应该做什么，可以修改原计划"}}
只输出 JSON。"""


def ask_json(prompt: str):
    """要模型只吐 JSON，还得自己兜底：它很爱用 ```json 包一层。"""
    text = client.chat.completions.create(
        model=MODEL, temperature=0,
        messages=[{"role": "user", "content": prompt}],
    ).choices[0].message.content.strip()
    if text.startswith("```"):
        text = text.strip("`").lstrip("json").strip()
    return json.loads(text)


def make_plan(goal: str):
    """第一步：一次性拆出计划。计划是文本，不是承诺——后面随时可以改。"""
    return ask_json(PLAN_PROMPT.format(goal=goal))


def run(goal: str, executor, max_rounds: int = 6):
    """executor 是你自己的执行器：给它一句子任务，它跑一轮 ReAct 返回结果。

    注意 review 这一步：每执行完一条就回头看一次目标。
    少了它，Agent 会一根筋地把计划执行到底——哪怕第二条就已经能收尾，
    哪怕第一条的结果已经证明整个计划走错了方向。
    """
    todo = make_plan(goal)
    print("初始计划：", todo)
    done = []
    for rnd in range(1, max_rounds + 1):
        if not todo:
            break
        task = todo.pop(0)
        result = executor(task)
        done.append((task, result))
        print("第 %d 轮完成：%s → %s" % (rnd, task, result))

        verdict = ask_json(REVIEW_PROMPT.format(
            goal=goal,
            done="\n".join("- %s：%s" % (t, r) for t, r in done),
            todo="\n".join("- " + t for t in todo) or "（无）"))
        if verdict.get("finish"):
            return verdict["answer"]
        if verdict.get("next"):
            todo.insert(0, verdict["next"])      # 允许中途改计划
    return "计划执行完毕，但未能确认目标达成：%s" % done


if __name__ == "__main__":
    def fake_executor(task):
        return "（示例结果）已处理：" + task

    print(run("统计长三角三个城市的人口并找出最多的那个", fake_executor))

"""完整案例二：带人工确认的 Agent —— 循环可以自主，动作不必全自主。

读操作让它随便跑，写操作必须停下来等人点头。
实现方式很朴素：在执行工具之前插一个确认闸门，不通过就把「被拒绝」
作为一条观察结果交回循环，让模型换个做法，而不是让程序崩掉。
"""
import json
import os

from openai import OpenAI

client = OpenAI(
    api_key=os.environ.get("OPENAI_API_KEY"),
    base_url=os.environ.get("OPENAI_BASE_URL"),
)
MODEL = os.environ.get("MODEL_NAME", "gpt-4o-mini")

INBOX = [{"id": 1, "from": "客户A", "subject": "报价单有误", "urgent": True},
         {"id": 2, "from": "同事B", "subject": "周会挪到周四", "urgent": False}]
SENT = []

READ_ONLY = {"list_mails", "read_mail"}
WRITE_OPS = {"send_reply"}


def list_mails() -> str:
    return json.dumps([{k: m[k] for k in ("id", "from", "subject", "urgent")}
                       for m in INBOX], ensure_ascii=False)


def read_mail(mail_id: int) -> str:
    for m in INBOX:
        if m["id"] == int(mail_id):
            return "来自 %s 的邮件：%s" % (m["from"], m["subject"])
    return "没有编号为 %s 的邮件" % mail_id


def send_reply(mail_id: int, content: str) -> str:
    SENT.append({"mail_id": int(mail_id), "content": content})
    return "已回复邮件 %s" % mail_id


REGISTRY = {"list_mails": list_mails, "read_mail": read_mail, "send_reply": send_reply}

TOOLS = [
    {"type": "function", "function": {
        "name": "list_mails", "description": "列出收件箱里所有邮件的摘要",
        "parameters": {"type": "object", "properties": {}, "required": []}}},
    {"type": "function", "function": {
        "name": "read_mail", "description": "按编号读取一封邮件的正文",
        "parameters": {"type": "object",
                       "properties": {"mail_id": {"type": "integer", "description": "邮件编号"}},
                       "required": ["mail_id"]}}},
    {"type": "function", "function": {
        "name": "send_reply", "description": "对指定邮件发送一条回复，这是对外可见的写操作",
        "parameters": {"type": "object", "properties": {
            "mail_id": {"type": "integer", "description": "邮件编号"},
            "content": {"type": "string", "description": "回复正文"}},
            "required": ["mail_id", "content"]}}},
]


def ask_human(name: str, args: dict) -> bool:
    """确认闸门。演示里走命令行，真实系统里换成工单、审批流或聊天机器人按钮。"""
    print("\n⚠ 需要确认：即将执行 %s" % name)
    print("  参数：%s" % json.dumps(args, ensure_ascii=False))
    return input("  批准执行？(y/N) ").strip().lower() == "y"


def run(goal: str, max_steps: int = 8, auto_approve: bool = False) -> str:
    messages = [
        {"role": "system", "content": "你是邮件助理。先查看收件箱，判断哪封最紧急，"
                                      "读完内容后再起草回复。回复属于写操作，"
                                      "若被拒绝，就把草稿念给用户听，不要反复重试。"},
        {"role": "user", "content": goal}]

    for step in range(1, max_steps + 1):
        msg = client.chat.completions.create(
            model=MODEL, messages=messages, tools=TOOLS,
            tool_choice="auto", temperature=0).choices[0].message
        messages.append(msg.model_dump(exclude_none=True))
        if not msg.tool_calls:
            return msg.content

        for call in msg.tool_calls:
            name = call.function.name
            args = json.loads(call.function.arguments)
            if name in WRITE_OPS and not auto_approve and not ask_human(name, args):
                result = "用户拒绝了这次操作，请不要重复尝试，改为把内容告诉用户"
            else:
                result = REGISTRY[name](**args)
            print("  第%d轮 %s → %s" % (step, name, result))
            messages.append({"role": "tool", "tool_call_id": call.id,
                             "content": str(result)})
    return "超过最大步数 %d" % max_steps


if __name__ == "__main__":
    print(run("看一下收件箱，挑最紧急的那封回一句「已收到，今天内处理」"))
    print("实际发出的回复：", SENT)

"""可复制的 Agent 循环骨架：只改 TODO 处，其余代码不用动。

它把前面拆开讲的四件事合成一份：工具注册、循环刹车、逐轮轨迹、执行前护栏。
从零写一个 Agent，从这份文件开始改是最快的。
"""
import json
import os
import time

from openai import OpenAI

# ---- TODO 1：换成你的模型服务（任何兼容 OpenAI 协议的服务都行） ----
client = OpenAI(
    api_key=os.environ.get("OPENAI_API_KEY"),
    base_url=os.environ.get("OPENAI_BASE_URL"),
)
MODEL = os.environ.get("MODEL_NAME", "gpt-4o-mini")


# ---- TODO 2：换成你自己的工具函数，返回值一律是字符串 ----
def example_tool(keyword: str) -> str:
    """一个示例工具。真实实现可以是查库、调接口、读文件。"""
    return "关于「%s」的示例结果" % keyword


# ---- TODO 3：把函数登记进来，键名必须与下面 TOOLS 里的 name 完全一致 ----
REGISTRY = {"example_tool": example_tool}

# ---- TODO 4：把每个工具描述给模型。description 决定它选不选你这个工具 ----
TOOLS = [
    {"type": "function", "function": {
        "name": "example_tool",
        "description": "按关键词查询资料，返回一句话结果",
        "parameters": {"type": "object", "properties": {
            "keyword": {"type": "string", "description": "要查询的关键词"}},
            "required": ["keyword"]}}},
]

# ---- TODO 5：换成你的人设与任务约束 ----
SYSTEM = ("你是一个严谨的助手。需要外部事实时调用工具，不要凭记忆作答；"
          "信息足够时直接给出结论，不要再调用工具。")

# ---- 循环上限：三条刹车，按你的场景调 ----
MAX_STEPS = 8
MAX_SECONDS = 90
MAX_TOKENS = 30000
REPEAT_LIMIT = 3          # 同一动作连续几次就判定为原地打转


def run(question: str, verbose: bool = True) -> dict:
    messages = [{"role": "system", "content": SYSTEM},
                {"role": "user", "content": question}]
    trace, fingerprints = [], []
    tokens, t0 = 0, time.monotonic()

    for step in range(1, MAX_STEPS + 1):
        if time.monotonic() - t0 > MAX_SECONDS:
            return _stop("超过时间上限", trace, step, tokens, t0)
        if tokens > MAX_TOKENS:
            return _stop("超过 token 预算", trace, step, tokens, t0)
        if len(fingerprints) >= REPEAT_LIMIT and len(set(fingerprints[-REPEAT_LIMIT:])) == 1:
            return _stop("检测到重复动作", trace, step, tokens, t0)

        resp = client.chat.completions.create(
            model=MODEL, messages=messages, tools=TOOLS,
            tool_choice="auto", temperature=0)
        tokens += getattr(resp.usage, "total_tokens", 0) or 0
        msg = resp.choices[0].message
        messages.append(msg.model_dump(exclude_none=True))

        if not msg.tool_calls:
            trace.append({"step": step, "phase": "finish"})
            return {"answer": msg.content, "stop": "模型给出最终答案",
                    "steps": step, "tokens": tokens,
                    "seconds": round(time.monotonic() - t0, 2), "trace": trace}

        for call in msg.tool_calls:
            name, raw = call.function.name, call.function.arguments
            fingerprints.append("%s|%s" % (name, raw))
            result = _safe_invoke(name, raw)
            trace.append({"step": step, "tool": name, "args": raw,
                          "result": str(result)[:200]})
            if verbose:
                print("  第%d轮 %s(%s) → %s" % (step, name, raw, result))
            messages.append({"role": "tool", "tool_call_id": call.id,
                             "content": str(result)})

    return _stop("达到最大步数", trace, MAX_STEPS, tokens, t0)


def _safe_invoke(name: str, raw_args: str) -> str:
    """执行前的最后一道关：工具存在吗、参数是合法 JSON 吗、执行会抛异常吗。

    三种情况都不抛出，而是把原因作为观察结果交回去，让模型有机会自己改正。
    """
    fn = REGISTRY.get(name)
    if fn is None:
        return "没有名为 %s 的工具，可选：%s" % (name, "、".join(REGISTRY))
    try:
        args = json.loads(raw_args or "{}")
    except json.JSONDecodeError as exc:
        return "参数不是合法 JSON：%s" % exc
    try:
        return str(fn(**args))
    except Exception as exc:
        return "工具执行失败：%s: %s" % (type(exc).__name__, exc)


def _stop(reason: str, trace, step, tokens, t0) -> dict:
    return {"answer": "未得出结论（%s）" % reason, "stop": reason, "steps": step,
            "tokens": tokens, "seconds": round(time.monotonic() - t0, 2),
            "trace": trace}


if __name__ == "__main__":
    out = run("帮我查一下关键词「智能体」的资料")
    print("\n答案：", out["answer"])
    print("停止原因：%s；%d 轮，%d token，%.2f 秒"
          % (out["stop"], out["steps"], out["tokens"], out["seconds"]))

"""工具注册表：把「函数本体」和「写给模型看的说明」绑在一起，只登记一次。

循环里最容易长歪的地方，是每加一个工具就要同时改三处：函数、schema、
分发的 if 分支。用一个装饰器把三处收成一处，Agent 的工具集就再也不会对不上。
"""
import inspect
import json


class ToolRegistry:
    def __init__(self):
        self._funcs = {}
        self._schemas = []

    def tool(self, description: str, **param_desc):
        """装饰器：登记一个函数，顺便根据签名生成 JSON Schema。

        param_desc 里写每个参数的中文说明——这句说明决定模型填不填、填什么，
        比函数名重要得多。
        """
        def deco(fn):
            sig = inspect.signature(fn)
            props, required = {}, []
            for name, p in sig.parameters.items():
                kind = {int: "integer", float: "number", bool: "boolean"}.get(
                    p.annotation, "string")
                props[name] = {"type": kind,
                               "description": param_desc.get(name, name)}
                if p.default is inspect.Parameter.empty:
                    required.append(name)
            self._funcs[fn.__name__] = fn
            self._schemas.append({"type": "function", "function": {
                "name": fn.__name__, "description": description,
                "parameters": {"type": "object", "properties": props,
                               "required": required}}})
            return fn
        return deco

    @property
    def schemas(self):
        """直接丢给模型的 tools 参数。"""
        return self._schemas

    def invoke(self, name: str, arguments: str) -> str:
        """按名字执行，并把任何异常翻译成模型看得懂的文字。

        这里是循环稳定性的关键：工具抛异常时如果让它冒泡，整个循环就断了；
        把错误变成一条 observation 交回去，模型才有机会换个参数重试。
        """
        fn = self._funcs.get(name)
        if fn is None:
            return "没有名为 %s 的工具，请从清单里重新选择" % name
        try:
            args = json.loads(arguments) if isinstance(arguments, str) else arguments
        except json.JSONDecodeError:
            return "参数不是合法 JSON，请重新生成：%s" % arguments
        try:
            return str(fn(**args))
        except TypeError as exc:
            return "参数不匹配：%s" % exc
        except Exception as exc:
            return "工具执行失败：%s: %s" % (type(exc).__name__, exc)


registry = ToolRegistry()


@registry.tool("查询指定城市今天的天气状况", city="城市中文名，例如 北京")
def get_weather(city: str) -> str:
    table = {"北京": "晴，18 到 28 度", "上海": "多云，21 到 27 度"}
    return table.get(city, "暂无 %s 的天气数据" % city)


@registry.tool("按关键词检索内部资料库", keyword="检索关键词", top_k="返回条数")
def search_docs(keyword: str, top_k: int = 3) -> str:
    return "关于「%s」检索到 %d 条资料（示例数据）" % (keyword, top_k)


@registry.tool("计算一个纯数字算术表达式", expression="例如 (28-18)/18*100")
def calculate(expression: str) -> str:
    return str(round(eval(expression, {"__builtins__": {}}, {}), 4))


if __name__ == "__main__":
    print(json.dumps(registry.schemas, ensure_ascii=False, indent=2))
    print(registry.invoke("get_weather", '{"city": "北京"}'))
    print(registry.invoke("get_weather", '{"城市": "北京"}'))   # 参数名写错
    print(registry.invoke("不存在的工具", "{}"))                 # 幻觉调用

"""循环控制器：Agent 真正的工程难点不在「怎么转起来」，而在「什么时候停」。

一个 while 很好写，难的是它会不会转到天荒地老、会不会把钱烧光。
这里把四种刹车做成一个可复用的对象，任何循环都能挂上去。
"""
import time


class StopReason:
    DONE = "模型给出最终答案"
    MAX_STEPS = "达到最大步数"
    TIMEOUT = "超过挂钟时间上限"
    BUDGET = "超过 token 预算"
    REPEAT = "检测到重复动作，判定为原地打转"


class LoopController:
    def __init__(self, max_steps=8, max_seconds=60.0, max_tokens=20000,
                 repeat_threshold=3):
        self.max_steps = max_steps
        self.max_seconds = max_seconds
        self.max_tokens = max_tokens
        self.repeat_threshold = repeat_threshold
        self.step = 0
        self.tokens = 0
        self.started = time.monotonic()
        self.history = []          # 每一轮实际执行的 (工具名, 参数) 指纹

    def elapsed(self) -> float:
        return time.monotonic() - self.started

    def check(self):
        """每轮开头调一次；返回停止原因，返回 None 表示可以继续。"""
        if self.step >= self.max_steps:
            return StopReason.MAX_STEPS
        if self.elapsed() > self.max_seconds:
            return StopReason.TIMEOUT
        if self.tokens > self.max_tokens:
            return StopReason.BUDGET
        if self._is_spinning():
            return StopReason.REPEAT
        return None

    def _is_spinning(self) -> bool:
        """同一个工具 + 同一组参数连续出现 N 次，就是死循环的典型指纹。

        模型卡住时的表现通常不是报错，而是一遍遍调同一个工具期待不同结果。
        步数上限能兜住，但那要烧满所有步数；这个检测能提前刹车。
        """
        n = self.repeat_threshold
        return len(self.history) >= n and len(set(self.history[-n:])) == 1

    def record(self, tool_name: str, arguments: str, usage=None):
        """每执行完一轮记一笔：动作指纹 + token 消耗。"""
        self.step += 1
        self.history.append("%s|%s" % (tool_name, arguments))
        if usage is not None:
            self.tokens += getattr(usage, "total_tokens", 0) or 0

    def summary(self, reason: str) -> str:
        return ("停止原因：%s；共 %d 轮，耗时 %.1f 秒，累计 %d token"
                % (reason, self.step, self.elapsed(), self.tokens))


def demo():
    """不联网也能看懂控制器的行为：模拟一个一直重复同一动作的模型。"""
    ctrl = LoopController(max_steps=10, repeat_threshold=3)
    while True:
        reason = ctrl.check()
        if reason:
            print(ctrl.summary(reason))
            break
        ctrl.record("search", '{"keyword": "同一个词"}')
        print("第 %d 轮：调用 search" % ctrl.step)


if __name__ == "__main__":
    demo()

"""Agent 的记忆：循环之内的草稿纸，和循环之外的长期档案，是两回事。

短期记忆决定「这一轮能不能接上上一轮」，长期记忆决定「换个会话还记不记得你」。
把它们混在一个列表里，是新手最常见的设计事故。
"""
import json
import os
import time


class Scratchpad:
    """短期记忆：一次任务之内，把每轮的思考、动作、观察按序摊平。

    它就是循环的状态本体——没有它，每一轮模型都是从零开始，
    所谓「自主决策」也就无从谈起。
    """

    def __init__(self, system: str, question: str, max_chars=6000):
        self.messages = [{"role": "system", "content": system},
                         {"role": "user", "content": question}]
        self.max_chars = max_chars

    def add_assistant(self, message):
        self.messages.append(message)

    def add_observation(self, tool_call_id: str, text: str, limit=1200):
        """工具结果要截断再入库。

        一次网页检索可能返回几万字符，原样塞进去，两三轮就把上下文窗口撑满，
        后面的轮次会因为超长而报错或者被迫丢掉最早的指令。
        """
        if len(text) > limit:
            text = text[:limit] + "……（结果过长，已截断）"
        self.messages.append({"role": "tool", "tool_call_id": tool_call_id,
                              "content": text})

    def compact(self):
        """超长时的退让策略：保留 system 与最初的问题，压缩中间的轮次。

        中间轮次一旦丢弃就不可恢复，所以要留一条摘要占位，
        让模型知道「前面做过事，只是细节省略了」。
        """
        total = sum(len(str(m.get("content") or "")) for m in self.messages)
        if total <= self.max_chars:
            return False
        head, tail = self.messages[:2], self.messages[-4:]
        dropped = len(self.messages) - len(head) - len(tail)
        if dropped <= 0:
            return False
        note = {"role": "user",
                "content": "（此前已执行 %d 条中间步骤，结论已并入后续观察结果）" % dropped}
        self.messages = head + [note] + tail
        return True


class LongTermMemory:
    """长期记忆：跨任务、跨会话的事实档案，用最朴素的 JSON 文件演示。

    生产环境会换成向量库或数据库，但接口形态就是这两个方法：
    写入一条事实、按关键词召回若干条。
    """

    def __init__(self, path="agent_memory.json"):
        self.path = path
        self.items = []
        if os.path.exists(path):
            self.items = json.load(open(path, encoding="utf-8"))

    def remember(self, text: str, tags=()):
        self.items.append({"text": text, "tags": list(tags), "ts": time.time()})
        with open(self.path, "w", encoding="utf-8") as f:
            json.dump(self.items, f, ensure_ascii=False, indent=2)

    def recall(self, keyword: str, top_k=3):
        hit = [x for x in self.items if keyword in x["text"] or keyword in x["tags"]]
        return [x["text"] for x in hit[-top_k:]]


if __name__ == "__main__":
    pad = Scratchpad("你是一个助手", "帮我对比两个城市的人口")
    pad.add_assistant({"role": "assistant", "content": "我先查第一个城市"})
    pad.add_observation("call_1", "杭州常住人口 1252 万人")
    print("短期记忆条数：", len(pad.messages))

    mem = LongTermMemory("/tmp/agent_memory_demo.json")
    mem.remember("这位用户长期关注长三角城市数据", tags=["偏好"])
    print("长期记忆召回：", mem.recall("偏好"))

"""轨迹记录：Agent 出错时，你要能回答「它第几轮开始走偏的」。

链出问题，看一眼输入输出就知道；Agent 出问题，中间可能已经转了七八轮，
没有轨迹就只能靠猜。所以循环里每一轮都要落一条结构化记录。
"""
import json
import time


class Trace:
    def __init__(self, goal: str):
        self.goal = goal
        self.started = time.time()
        self.steps = []

    def log_think(self, step: int, content, tool_calls):
        """记一轮「想」：模型说了什么、打算调哪些工具。"""
        self.steps.append({
            "step": step, "phase": "think", "ts": round(time.time() - self.started, 3),
            "content": (content or "")[:200],
            "tools": [{"name": c.function.name, "args": c.function.arguments}
                      for c in (tool_calls or [])],
        })

    def log_act(self, step: int, name: str, args: str, result: str, cost_ms: int):
        """记一轮「做 + 看」：工具、入参、返回、耗时。"""
        self.steps.append({
            "step": step, "phase": "act", "ts": round(time.time() - self.started, 3),
            "tool": name, "args": args, "result": str(result)[:300],
            "cost_ms": cost_ms,
        })

    def log_stop(self, reason: str, answer: str = ""):
        self.steps.append({"phase": "stop", "reason": reason,
                           "answer": (answer or "")[:300],
                           "ts": round(time.time() - self.started, 3)})

    def to_json(self, path=None):
        data = {"goal": self.goal, "total_seconds": round(time.time() - self.started, 3),
                "steps": self.steps}
        text = json.dumps(data, ensure_ascii=False, indent=2)
        if path:
            with open(path, "w", encoding="utf-8") as f:
                f.write(text)
        return text

    def pretty(self) -> str:
        """人眼快速扫一遍用的紧凑视图。"""
        lines = ["目标：%s" % self.goal]
        for s in self.steps:
            if s["phase"] == "think":
                tools = "、".join(t["name"] for t in s["tools"]) or "（不再调用工具）"
                lines.append("  [%5.2fs] 第%s轮 想 → %s" % (s["ts"], s["step"], tools))
            elif s["phase"] == "act":
                lines.append("  [%5.2fs] 第%s轮 做 → %s(%s) 用时 %dms"
                             % (s["ts"], s["step"], s["tool"], s["args"], s["cost_ms"]))
                lines.append("           看 → %s" % s["result"])
            else:
                lines.append("  [%5.2fs] 结束：%s" % (s["ts"], s["reason"]))
        return "\n".join(lines)

    def diagnose(self) -> list:
        """从轨迹里自动挑出三种最常见的异常形态。"""
        problems = []
        acts = [s for s in self.steps if s["phase"] == "act"]
        seen = {}
        for a in acts:
            key = "%s|%s" % (a["tool"], a["args"])
            seen[key] = seen.get(key, 0) + 1
        for key, n in seen.items():
            if n >= 3:
                problems.append("同一动作重复 %d 次，疑似原地打转：%s" % (n, key))
        slow = [a for a in acts if a["cost_ms"] > 5000]
        if slow:
            problems.append("有 %d 次工具调用超过 5 秒，循环整体延迟会被放大" % len(slow))
        empty = [a for a in acts if not a["result"].strip()]
        if empty:
            problems.append("有 %d 次工具返回为空，模型很可能据此编造内容" % len(empty))
        return problems


if __name__ == "__main__":
    tr = Trace("对比两个城市人口")
    tr.log_act(1, "lookup", '{"city":"杭州"}', "1252 万", 120)
    tr.log_act(2, "lookup", '{"city":"杭州"}', "1252 万", 118)
    tr.log_act(3, "lookup", '{"city":"杭州"}', "1252 万", 130)
    tr.log_stop("检测到重复动作")
    print(tr.pretty())
    print("诊断：", tr.diagnose())

"""护栏：把 Agent 的四种典型失败模式，一条条堵在执行工具之前。

这些检查都发生在「模型已经写好动作、你还没真正执行」的那个缝隙里。
铁律在这里体现得最直白：循环的执行权在你的代码手上，所以拦得住。
"""
import json
import re


class GuardError(Exception):
    """护栏拒绝执行；把它转成 observation 交回模型，而不是让循环崩掉。"""


ALLOWED_TOOLS = {"search_docs", "lookup_population", "calculate", "send_email"}
HIGH_RISK = {"send_email", "delete_record", "transfer_money"}
SQL_WRITE = re.compile(r"\b(delete|drop|update|insert|alter|truncate)\b", re.I)


def check_known_tool(name: str):
    """幻觉调用：模型会凭空发明一个不存在的工具名，还填得有模有样。"""
    if name not in ALLOWED_TOOLS:
        raise GuardError("工具 %s 不在清单内，可选：%s"
                         % (name, "、".join(sorted(ALLOWED_TOOLS))))


def check_args(name: str, arguments: str) -> dict:
    """参数幻觉：字段名拼错、该填数字填了整句话，都要在执行前拦下。"""
    try:
        args = json.loads(arguments)
    except json.JSONDecodeError as exc:
        raise GuardError("参数不是合法 JSON：%s" % exc)
    if not isinstance(args, dict):
        raise GuardError("参数必须是一个对象")
    if name == "calculate":
        expr = str(args.get("expression", ""))
        if not re.fullmatch(r"[\d\s\.\+\-\*/\(\)%]+", expr):
            raise GuardError("计算表达式只允许数字与运算符，收到：%s" % expr)
    if name == "search_docs" and SQL_WRITE.search(str(args.get("keyword", ""))):
        raise GuardError("检索关键词里出现写操作语句，已拒绝")
    return args


def check_budget(step: int, max_steps: int, tokens: int, max_tokens: int):
    """步数爆炸与成本失控：这两件事必须同时看。

    只限步数挡不住成本——一次塞进去几万 token 的长上下文，
    三轮就能把预算烧完。
    """
    if step > max_steps:
        raise GuardError("已达最大步数 %d" % max_steps)
    if tokens > max_tokens:
        raise GuardError("已超 token 预算 %d" % max_tokens)


def needs_human(name: str, args: dict) -> bool:
    """高风险动作不由模型拍板：发邮件、删数据、转账一律先问人。

    这不是保守，而是循环的必然要求：Agent 可以跑很多轮，
    错一次就可能对外部世界造成不可撤销的改动。
    """
    if name not in HIGH_RISK:
        return False
    if name == "send_email" and args.get("to", "").endswith("@example.com"):
        return False        # 内部测试地址可以放行
    return True


def guard(name: str, arguments: str, step: int, tokens: int,
          max_steps=8, max_tokens=20000, approver=None) -> dict:
    """统一入口：过了这一关才允许真正执行工具。"""
    check_budget(step, max_steps, tokens, max_tokens)
    check_known_tool(name)
    args = check_args(name, arguments)
    if needs_human(name, args):
        if approver is None or not approver(name, args):
            raise GuardError("高风险动作 %s 未获批准，已取消" % name)
    return args


if __name__ == "__main__":
    cases = [("查天气", "{}"),
             ("calculate", '{"expression": "import os"}'),
             ("send_email", '{"to": "[email protected]", "body": "涨薪申请"}'),
             ("calculate", '{"expression": "(1252-954)/954*100"}')]
    for n, a in cases:
        try:
            print("放行：", n, guard(n, a, step=1, tokens=10,
                                   approver=lambda *_: False))
        except GuardError as exc:
            print("拦下：", n, "→", exc)

"""成本推演：为什么一条链的花费可以估准，而 Agent 只能估一个区间。

链的调用次数是常数，Agent 的调用次数是随机变量。更麻烦的是，
每一轮都要把之前所有轮次的消息重新发一遍，上下文是累加的——
所以总开销大致按轮数的平方增长，而不是线性。这段代码把这件事算给你看。
"""


def chain_cost(steps: int, prompt_tokens: int, output_tokens: int) -> int:
    """链：每一步各发各的提示词，互不叠加。"""
    return steps * (prompt_tokens + output_tokens)


def agent_cost(rounds: int, system_tokens: int, question_tokens: int,
               think_tokens: int, observation_tokens: int) -> int:
    """循环：第 n 轮要重发前 n-1 轮的全部内容。

    每轮的输入 = 系统提示 + 问题 + 已累积的（思考 + 观察）。
    """
    total = 0
    carried = system_tokens + question_tokens
    for _ in range(rounds):
        total += carried + think_tokens          # 这一轮的输入 + 输出
        carried += think_tokens + observation_tokens
    return total


def compare(max_rounds: int = 8):
    sys_t, q_t, think_t, obs_t = 300, 60, 120, 400
    print("轮数  累计 token   相对 2 轮的倍数")
    base = agent_cost(2, sys_t, q_t, think_t, obs_t)
    for r in range(2, max_rounds + 1):
        c = agent_cost(r, sys_t, q_t, think_t, obs_t)
        print("%3d   %9d   %6.1f 倍" % (r, c, c / base))
    print("\n同样办完这件事，写死的四步链：%d token"
          % chain_cost(4, sys_t + q_t, think_t))


def budget_rounds(max_tokens: int, system_tokens=300, question_tokens=60,
                  think_tokens=120, observation_tokens=400) -> int:
    """反过来问：给定预算，最多允许转几轮？把结果配成 max_steps。"""
    r = 0
    while agent_cost(r + 1, system_tokens, question_tokens,
                     think_tokens, observation_tokens) <= max_tokens:
        r += 1
    return r


if __name__ == "__main__":
    compare()
    for budget in (5000, 20000, 100000):
        print("预算 %6d token → 最多约 %d 轮" % (budget, budget_rounds(budget)))
    # 结论：观察结果的长度是成本的主要杠杆。
    # 把一次检索返回的几千字符截断到几百字符，往往比把 max_steps 从 8 调到 6
    # 省得多，而且不损失解决问题的能力。