LangChain Tools 与 Agent 实战

"""内置工具与自定义工具的分界线在哪。

社区里已经有大量现成工具（搜索、维基、文件操作、Shell、SQL……），
装上对应的集成包就能直接塞进 tools 列表，不用自己写。

判断标准很简单：
    这件事任何人做都一样      → 用现成的（搜索、爬网页、读文件）
    这件事只有你们公司这么做  → 自己写（查你们的库、算你们的价、下你们的单）
"""

import os

from langchain.tools import tool


def load_search_tool():
    """搜索类工具的典型接法：装包 → 配 key → 实例化 → 塞进 tools。

    pip install langchain-tavily
    export TAVILY_API_KEY=...
    """
    if not os.environ.get("TAVILY_API_KEY"):
        return None
    from langchain_tavily import TavilySearch

    return TavilySearch(max_results=3)


@tool
def get_internal_headcount(dept: str) -> str:
    """查询公司某个部门的在编人数。仅限内部部门名，例如 研发部、售后部。"""
    table = {"研发部": 128, "售后部": 46, "市场部": 31}
    if dept not in table:
        return f"没有 {dept} 这个部门"
    return f"{dept} 在编 {table[dept]} 人"


def build_toolbox() -> list:
    """把现成工具与自定义工具混在一个列表里——对 Agent 来说它们没有区别。"""
    box = [get_internal_headcount]
    search = load_search_tool()
    if search is not None:
        box.insert(0, search)
    else:
        print("未配置 TAVILY_API_KEY，这次只带自定义工具")
    return box


if __name__ == "__main__":
    tools = build_toolbox()
    for t in tools:
        print("-", t.name, "|", t.description.splitlines()[0][:50])

    # 工具箱不是越大越好：
    # 工具越多，每次调模型的上下文越长、选错的概率也越高。
    # 超过十来个就该按场景分组，或者用动态工具选择只挂当前用得上的那几个。
    print("当前工具数：", len(tools))

"""Tool 的四要素：名称、功能描述、参数结构、要调用的函数。

用 @tool 装饰器把一个普通函数升级成 Tool，然后把这四样东西
逐个打印出来看——模型收到的就是这些信息，一个字都不多。

运行：python tool_basics.py
依赖：pip install "langchain>=1.0"
"""

import json

from langchain.tools import tool


@tool
def add_number(a: int, b: int) -> int:
    """两个整数相加，返回它们的和。"""
    return a + b


# name / description / return_direct 都可以在装饰器里改写。
# 注意第一个位置参数就是工具名，传字符串即可覆盖函数名。
@tool("add_two_number", description="计算两个整数的和", return_direct=True)
def add_number_v2(a: int, b: int) -> int:
    """两个整数相加。"""
    return a + b


def show(t) -> None:
    """把一个 Tool 的四要素摊开打印。"""
    print("① 名称      name         =", t.name)
    print("② 功能描述  description  =", t.description)
    print("③ 参数结构  args         =", json.dumps(t.args, ensure_ascii=False))
    print("   完整入参 schema        =",
          json.dumps(t.args_schema.model_json_schema(), ensure_ascii=False)
          if hasattr(t.args_schema, "model_json_schema") else t.args_schema)
    print("④ 要调用的函数            =", t.func.__name__ if t.func else "(异步实现)")
    print("   return_direct         =", t.return_direct)
    print("-" * 60)


if __name__ == "__main__":
    show(add_number)
    show(add_number_v2)

    # 工具自己也是 Runnable：invoke 传一个参数字典，键名即形参名。
    print("add_number.invoke        ->", add_number.invoke({"a": 10, "b": 20}))
    print("add_two_number.invoke    ->", add_number_v2.invoke({"a": 10, "b": 20}))

    # 类型注解不是装饰用的：它决定了 args 里每个参数的 type。
    # 去掉 a: int 之后，参数类型会退化，模型更容易填错。

"""description 写得好不好，差别有多大？把两版工具摆在一起看。

同一个功能写两份描述：一份含糊、一份说清「什么时候用、输入是什么、
产出是什么、什么时候别用」。把两份分别交给同一个模型、问同一句话，
看它挑不挑得中。

运行前先配好模型访问信息（以 OpenAI 兼容网关为例）：
    export OPENAI_API_KEY=...
    export OPENAI_BASE_URL=...
"""

import os

from langchain.chat_models import init_chat_model
from langchain.tools import tool


@tool("query_order", description="查订单")
def bad_version(order_id: str) -> str:
    """查订单。"""
    return "{\"order_id\": \"%s\", \"status\": \"已发货\"}" % order_id


@tool(
    "query_order",
    description=(
        "按订单号查询一笔订单的状态、金额与物流单号。"
        "当用户提到订单号、问「我的单子到哪了」「发货没有」时使用。"
        "入参 order_id 是形如 SO20251208001 的订单编号；"
        "用户只说了商品名而没给订单号时，先反问订单号，不要猜。"
    ),
)
def good_version(order_id: str) -> str:
    """按订单号查询订单状态。"""
    return "{\"order_id\": \"%s\", \"status\": \"已发货\"}" % order_id


QUESTION = "帮我看看 SO20251208001 这一单发货了吗"


def ask(tool_obj) -> None:
    """把单个工具绑定到模型上，看模型会不会主动选它。"""
    if not os.environ.get("OPENAI_API_KEY"):
        print("未配置 OPENAI_API_KEY，跳过实际请求")
        return
    model = init_chat_model("openai:gpt-4o-mini", temperature=0)
    msg = model.bind_tools([tool_obj]).invoke(QUESTION)
    if msg.tool_calls:
        print("选中工具:", msg.tool_calls[0]["name"], "参数:", msg.tool_calls[0]["args"])
    else:
        print("没有选工具，模型直接回了文本:", msg.text())


if __name__ == "__main__":
    print("【含糊版 description】", bad_version.description)
    ask(bad_version)
    print("-" * 60)
    print("【说清楚版 description】", good_version.description)
    ask(good_version)

    # 一句话记住：description 是给模型的选择依据，不是给同事的注释。
    # 它要回答四个问题：干什么、什么时候用、参数怎么填、什么时候别用。

"""用 Pydantic 给工具参数写一份正经的 schema。

函数签名只能表达「参数叫什么、是什么类型」，
args_schema 还能表达「这个参数是什么意思、有哪些合法取值、默认值是多少」。
这些信息会随工具一起送进模型的上下文，直接决定模型填参数的准确率。
"""

import json
from typing import Literal

from langchain.tools import tool
from pydantic import BaseModel, Field


class WeatherInput(BaseModel):
    """查询天气所需要的入参。"""

    city: str = Field(description="城市名称，使用中文，例如 北京、杭州")
    unit: Literal["摄氏度", "华氏度"] = Field(
        default="摄氏度",
        description="温度单位，默认摄氏度",
    )
    days: int = Field(
        default=1,
        ge=1,
        le=7,
        description="要查询的天数，1 表示只看今天，最多 7 天",
    )


@tool(args_schema=WeatherInput)
def get_weather(city: str, unit: str = "摄氏度", days: int = 1) -> str:
    """查询指定城市未来若干天的天气。需要实时天气、气温、下不下雨时使用。"""
    sample = {"北京": (18, 32, "晴"), "杭州": (21, 29, "多云")}
    low, high, text = sample.get(city, (20, 28, "晴"))
    if unit == "华氏度":
        low, high = round(low * 9 / 5 + 32), round(high * 9 / 5 + 32)
    return f"{city} 未来 {days} 天：{text}，{low}~{high}{unit}"


if __name__ == "__main__":
    print("name        =", get_weather.name)
    print("description =", get_weather.description)
    print("args        =", json.dumps(get_weather.args, ensure_ascii=False, indent=2))

    print(get_weather.invoke({"city": "北京"}))
    print(get_weather.invoke({"city": "杭州", "unit": "华氏度", "days": 3}))

    # 越界的参数会被 Pydantic 直接拦下，轮不到你的业务代码处理。
    try:
        get_weather.invoke({"city": "北京", "days": 99})
    except Exception as exc:  # noqa: BLE001  演示校验失败的现象
        print("参数校验失败：", type(exc).__name__)

"""StructuredTool.from_function：把「已经存在的函数」包成 Tool。

@tool 适合新写的函数——你能直接在函数上加装饰器。
但很多时候函数来自别人的模块、来自 SDK，你不想（也不该）去改它，
这时候用 StructuredTool.from_function 在外面包一层。
它还能同时挂上同步实现 func 与异步实现 coroutine。
"""

import asyncio
import json

from langchain_core.tools import StructuredTool
from pydantic import BaseModel, Field


# ---- 假设下面两个函数来自已有业务模块，一个字都不想改 ----------------
def search_docs(query: str, top_k: int = 3) -> str:
    hits = [f"《{query}》相关文档 {i + 1}" for i in range(top_k)]
    return "；".join(hits)


async def asearch_docs(query: str, top_k: int = 3) -> str:
    await asyncio.sleep(0)
    return search_docs(query, top_k)


# ---- 在外面包一层，把四要素补齐 ------------------------------------
class SearchInput(BaseModel):
    query: str = Field(description="检索关键词，写用户真正想找的那个词")
    top_k: int = Field(default=3, ge=1, le=10, description="返回的文档条数")


search_tool = StructuredTool.from_function(
    func=search_docs,                 # ④ 要调用的函数（同步）
    coroutine=asearch_docs,           #    异步实现，ainvoke 时自动走这条
    name="search_internal_docs",      # ① 名称：小写加下划线，别带空格
    description=(                     # ② 功能描述：写给模型看的
        "在公司内部知识库里检索文档。"
        "当用户问到内部流程、规范、历史方案时使用；"
        "不要用它查实时新闻或计算数值。"
    ),
    args_schema=SearchInput,          # ③ 参数结构
    return_direct=False,
)


if __name__ == "__main__":
    print("name        =", search_tool.name)
    print("description =", search_tool.description)
    print("args        =", json.dumps(search_tool.args, ensure_ascii=False))

    print(search_tool.invoke({"query": "报销流程", "top_k": 2}))
    print(asyncio.run(search_tool.ainvoke({"query": "值班制度"})))

"""三代写法的演进：同一个需求，三种年代的代码长相。

同一件事——「给模型一个搜索工具，让它自己决定要不要用」——
在 LangChain 的不同阶段有三种写法。前两种在老项目里还大量存在，
你要认得出来，也要知道搬到当代写法之后对应哪一行。

本文件里只有第三代是可运行的；前两代以字符串形式保留长相，
避免导入已经不在 langchain 主包里的名字。
"""

ERA_1 = '''
# 第一代：一体化创建，三行出一个执行器
from langchain.agents import initialize_agent, AgentType

agent_executor = initialize_agent(
    tools=[search_tool],
    llm=llm,
    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,   # 用枚举挑一套内置提示词
    verbose=True,
)
agent_executor.invoke("今天北京的天气怎么样？")

# 上手快，但提示词是内置的、看不见也改不动；
# 想换一句措辞就只能换一个 AgentType 碰运气。
'''

ERA_2 = '''
# 第二代：Agent 与运行时分离
from langchain.agents import AgentExecutor, create_react_agent
from langchain import hub

prompt = hub.pull("hwchase17/react")           # 提示词终于拿得到了
agent = create_react_agent(llm=llm, tools=tools, prompt=prompt)
agent_executor = AgentExecutor(agent=agent, tools=tools)
agent_executor.invoke({"input": "今天北京的天气怎么样？"})

# 决策逻辑（agent）与执行循环（AgentExecutor）分开了，提示词可自定义；
# 代价是代码变长，而且要在循环里插一段自己的逻辑仍然很别扭。
'''

ERA_3 = '''
# 第三代：统一入口 create_agent
from langchain.agents import create_agent

agent = create_agent(
    model="openai:gpt-4o-mini",
    tools=tools,
    system_prompt="今天北京的天气怎么样？这类实时问题必须调用工具。",
    middleware=[...],          # 想插东西就插 middleware，不用改循环
)
agent.invoke({"messages": [{"role": "user", "content": "今天北京的天气怎么样？"}]})
'''

# 逐项的搬迁关系，照着改即可
MIGRATION = [
    ("导入路径", "langgraph.prebuilt", "langchain.agents"),
    ("工厂函数", "create_react_agent", "create_agent"),
    ("提示词参数", "prompt=", "system_prompt=（传字符串）"),
    ("动态提示词", "自己拼 prompt", "@dynamic_prompt middleware"),
    ("调用前后插逻辑", "pre_model_hook / post_model_hook / state_modifier",
     "before_model / after_model / @wrap_model_call / @wrap_tool_call"),
    ("结构化输出", "response_format=（含 prompted 形式）",
     "response_format=ToolStrategy(...) 或 ProviderStrategy(...)"),
    ("流式节点名", '"agent"', '"model"'),
    ("运行时上下文", 'config["configurable"] 里塞', "invoke(context=...) + context_schema="),
    ("入参形态", '{"input": "..."}', '{"messages": [{"role": "user", "content": "..."}]}'),
]


if __name__ == "__main__":
    for title, block in (("第一代", ERA_1), ("第二代", ERA_2), ("第三代", ERA_3)):
        print("=" * 66)
        print(title, block)

    print("=" * 66)
    print("%-16s %-42s %s" % ("变化点", "旧写法", "当代写法"))
    for row in MIGRATION:
        print("%-16s %-42s %s" % row)

"""最小可跑的 Agent：一个工具 + 一句 system_prompt + 一次 invoke。

这是 LangChain 1.x 里构建 Agent 的当代写法：
    from langchain.agents import create_agent

四步：定义工具 → 建 Agent → invoke → 取最后一条消息。
循环（要不要再调一次工具）由 create_agent 内部替你转，你不用手写 while。

环境：
    pip install "langchain>=1.0" langchain-openai
    export OPENAI_API_KEY=...
    export OPENAI_BASE_URL=...        # 用官方接口时可以不设
"""

import os

from langchain.agents import create_agent
from langchain.tools import tool


@tool
def get_city_population(city: str) -> str:
    """查询一个城市的常住人口数量（单位：万人）。需要人口数据时使用。"""
    table = {"北京": 2183, "上海": 2487, "郑州": 1300, "杭州": 1252}
    if city not in table:
        return f"没有 {city} 的人口数据"
    return f"{city} 常住人口约 {table[city]} 万人"


def build_agent():
    """把模型、工具、system_prompt 装配成一个 Agent。"""
    return create_agent(
        model="openai:gpt-4o-mini",          # 也可以传已经实例化好的模型对象
        tools=[get_city_population],         # 工具列表，就是那个「工具箱」
        system_prompt=(                      # 传字符串，不要传 SystemMessage 对象
            "你是一个数据助手。需要具体数字时必须调用工具，"
            "工具没有数据就如实说没有，不要自己编。"
        ),
    )


def main() -> None:
    if not os.environ.get("OPENAI_API_KEY"):
        print("先配置 OPENAI_API_KEY 再运行")
        return

    agent = build_agent()
    # 入参固定是一个字典，键 messages 装一个消息列表
    result = agent.invoke(
        {"messages": [{"role": "user", "content": "北京和杭州哪个人多？多多少万？"}]}
    )

    # 返回值也是状态字典，messages 里是这一轮跑出来的完整消息链
    for m in result["messages"]:
        print(f"[{m.type:9s}]", (m.text() or "").strip()[:80],
              "| tool_calls:", getattr(m, "tool_calls", None))

    print("\n最终答复：", result["messages"][-1].text())


if __name__ == "__main__":
    main()

"""门店助手的工具箱：三件真正能跑的工具。

- query_stock   查某个门店某个商品还有多少件（查询类）
- quote_price   按数量与折扣算总价（计算类）
- create_order  下一张订单（写操作，带安全边界）

工具的边界原则：
读操作放开；写操作必须自己校验参数、自己限额，
不要指望 system_prompt 能拦住一个填错参数的模型。
"""

from datetime import datetime

from langchain.tools import tool
from pydantic import BaseModel, Field

# 演示用的内存数据，真项目换成数据库查询即可
STOCK = {
    ("北京朝阳店", "拿铁豆"): 120,
    ("北京朝阳店", "冷萃瓶"): 8,
    ("杭州西湖店", "拿铁豆"): 40,
    ("杭州西湖店", "冷萃瓶"): 0,
}
PRICE = {"拿铁豆": 68.0, "冷萃瓶": 25.0}
ORDERS: list[dict] = []
MAX_QTY_PER_ORDER = 500


class StockInput(BaseModel):
    shop: str = Field(description="门店名称，例如 北京朝阳店、杭州西湖店")
    item: str = Field(description="商品名称，例如 拿铁豆、冷萃瓶")


@tool(args_schema=StockInput)
def query_stock(shop: str, item: str) -> str:
    """查询某个门店某个商品的当前库存件数。
    当用户问「还有没有货」「库存多少」「够不够」时使用。
    入参是门店名与商品名，两者都必须由用户给出，不要替用户猜门店。"""
    if (shop, item) not in STOCK:
        return f"查无记录：{shop} 没有登记过 {item}"
    return f"{shop} 的 {item} 当前库存 {STOCK[(shop, item)]} 件"


class QuoteInput(BaseModel):
    item: str = Field(description="商品名称")
    qty: int = Field(ge=1, le=MAX_QTY_PER_ORDER, description="购买数量，正整数")
    discount: float = Field(default=1.0, gt=0, le=1,
                            description="折扣系数，0.9 表示九折，默认不打折")


@tool(args_schema=QuoteInput)
def quote_price(item: str, qty: int, discount: float = 1.0) -> str:
    """按商品单价、数量和折扣计算订单总价。
    当用户问「多少钱」「一共几钱」「打完折是多少」时使用。
    这个工具只算钱，不查库存，也不会真的下单。"""
    if item not in PRICE:
        return f"没有 {item} 的价格"
    total = PRICE[item] * qty * discount
    return (f"{item} 单价 {PRICE[item]:.2f} 元 × {qty} 件 × 折扣 {discount} "
            f"= {total:.2f} 元")


class OrderInput(BaseModel):
    shop: str = Field(description="下单门店")
    item: str = Field(description="商品名称")
    qty: int = Field(ge=1, le=MAX_QTY_PER_ORDER, description="下单数量")


@tool(args_schema=OrderInput)
def create_order(shop: str, item: str, qty: int) -> str:
    """为指定门店创建一张采购订单。这是写操作，会真的产生一条订单记录。
    只有在用户明确说了「下单」「就要这些」之后才使用；
    库存不足时不要下单，先把缺口告诉用户。"""
    have = STOCK.get((shop, item))
    if have is None:
        return f"下单失败：{shop} 没有 {item} 这个商品"
    if qty > have:
        return f"下单失败：{shop} 的 {item} 只剩 {have} 件，不足 {qty} 件"
    STOCK[(shop, item)] = have - qty
    oid = f"SO{datetime.now():%Y%m%d}{len(ORDERS) + 1:03d}"
    ORDERS.append({"id": oid, "shop": shop, "item": item, "qty": qty})
    return f"下单成功，订单号 {oid}，{shop} 的 {item} 剩余 {STOCK[(shop, item)]} 件"


SHOP_TOOLS = [query_stock, quote_price, create_order]


if __name__ == "__main__":
    print(query_stock.invoke({"shop": "北京朝阳店", "item": "冷萃瓶"}))
    print(quote_price.invoke({"item": "拿铁豆", "qty": 30, "discount": 0.9}))
    print(create_order.invoke({"shop": "北京朝阳店", "item": "冷萃瓶", "qty": 20}))

"""门店助手完整案例：三个自定义工具串成一个能自主决策的 Agent。

用户一句「杭州西湖店的拿铁豆还有多少？按九折买 30 袋要多少钱？」
里面藏着两件事：先查库存、再算价钱。
你不需要写「先调哪个再调哪个」——把工具交出去，顺序由模型自己定。

环境：
    pip install "langchain>=1.0" langchain-openai
    export OPENAI_API_KEY=...
    export OPENAI_BASE_URL=...
"""

import os

from langchain.agents import create_agent

from shop_tools import SHOP_TOOLS

SYSTEM_PROMPT = """你是一家连锁咖啡店的门店助手。

规则：
1. 涉及库存、价格、下单的问题一律调用工具，不要凭印象回答。
2. 用户没说门店时，先问清楚是哪家店，不要自己挑一家。
3. 下单前必须先查库存；库存不足就把缺口说明白，不要下单。
4. 工具返回「查无记录」时如实转述，不要编造数字。
"""


def build_agent():
    return create_agent(
        model="openai:gpt-4o-mini",
        tools=SHOP_TOOLS,
        system_prompt=SYSTEM_PROMPT,
    )


def run(agent, question: str) -> None:
    print("=" * 66)
    print("用户：", question)
    result = agent.invoke({"messages": [{"role": "user", "content": question}]})

    for m in result["messages"]:
        if m.type == "ai" and getattr(m, "tool_calls", None):
            for call in m.tool_calls:
                print(f"  → 决定调用 {call['name']}，参数 {call['args']}")
        elif m.type == "tool":
            print(f"  ← {m.name} 返回：{m.content}")
    print("助手：", result["messages"][-1].text())


def main() -> None:
    if not os.environ.get("OPENAI_API_KEY"):
        print("先配置 OPENAI_API_KEY 再运行")
        return
    agent = build_agent()

    # 一轮里要连着用两个工具：查库存 → 算价钱
    run(agent, "杭州西湖店的拿铁豆还有多少？按九折买 30 袋要多少钱？")

    # 库存为 0 的情况，看它会不会老老实实说没货
    run(agent, "杭州西湖店的冷萃瓶帮我下 10 件")

    # 缺少门店信息，看它会不会反问而不是瞎选一家
    run(agent, "拿铁豆还有货吗")


if __name__ == "__main__":
    main()

"""把 Agent 的执行循环看清楚：用 stream 逐节点观察。

invoke 只给你最后的结果，出了问题两眼一抹黑。
stream 会在每个节点跑完后吐一块更新出来，你能看见
模型什么时候决定调工具、工具返回了什么、它又转回模型第几次。

节点名在 LangChain 1.x 里是 "model" 与 "tools"。
（旧版 langgraph.prebuilt.create_react_agent 里模型节点叫 "agent"，
 迁移到 create_agent 之后这个名字变了，按 "agent" 取值会一直取到空。）
"""

import os

from langchain.agents import create_agent

from shop_tools import SHOP_TOOLS


def watch(agent, question: str) -> None:
    print("用户：", question)
    step = 0
    for chunk in agent.stream(
        {"messages": [{"role": "user", "content": question}]},
        stream_mode="updates",
    ):
        step += 1
        for node, update in chunk.items():
            msgs = update.get("messages", []) if isinstance(update, dict) else []
            for m in msgs:
                if node == "model":
                    calls = getattr(m, "tool_calls", None)
                    if calls:
                        names = "、".join(c["name"] for c in calls)
                        print(f"第 {step} 步 [model] 决定调用：{names}")
                    else:
                        print(f"第 {step} 步 [model] 给出最终答复：{m.text()[:60]}")
                elif node == "tools":
                    print(f"第 {step} 步 [tools] {m.name} → {m.content[:60]}")
    print("-" * 66)


def token_stream(agent, question: str) -> None:
    """只想看最终答复一个字一个字吐出来，用 stream_mode="messages"。"""
    for token, meta in agent.stream(
        {"messages": [{"role": "user", "content": question}]},
        stream_mode="messages",
    ):
        if meta.get("langgraph_node") == "model" and token.text():
            print(token.text(), end="", flush=True)
    print()


def main() -> None:
    if not os.environ.get("OPENAI_API_KEY"):
        print("先配置 OPENAI_API_KEY 再运行")
        return
    agent = create_agent(
        model="openai:gpt-4o-mini",
        tools=SHOP_TOOLS,
        system_prompt="你是门店助手，涉及库存与价格一律调用工具。",
    )
    watch(agent, "北京朝阳店的拿铁豆还有多少？买 20 袋原价多少钱？")
    token_stream(agent, "再帮我算一下 50 袋打八折多少钱")


if __name__ == "__main__":
    main()

"""让 Agent 记住上一轮：checkpointer + thread_id。

Agent 的每一轮 invoke 默认是互相独立的——上一轮查过哪家门店，
这一轮它并不知道。给 create_agent 传一个 checkpointer，
再在调用时带上 thread_id，状态就会被持久化并按会话隔离。

记忆本身的原理（模型没有记忆、历史要一轮轮重新塞回去）在会话记忆那一讲，
这里只讲 Agent 侧这条路子怎么接。
"""

import os

from langchain.agents import create_agent
from langgraph.checkpoint.memory import InMemorySaver

from shop_tools import SHOP_TOOLS


def build_agent():
    return create_agent(
        model="openai:gpt-4o-mini",
        tools=SHOP_TOOLS,
        system_prompt="你是门店助手，涉及库存与价格一律调用工具。",
        # 演示用内存实现；进程一退就没了。
        # 要跨进程留存，换成 SQLite / Postgres 的 saver 实现。
        checkpointer=InMemorySaver(),
    )


def ask(agent, text: str, thread_id: str) -> str:
    cfg = {"configurable": {"thread_id": thread_id}}
    result = agent.invoke({"messages": [{"role": "user", "content": text}]}, cfg)
    answer = result["messages"][-1].text()
    print(f"[{thread_id}] 用户：{text}")
    print(f"[{thread_id}] 助手：{answer}\n")
    return answer


def main() -> None:
    if not os.environ.get("OPENAI_API_KEY"):
        print("先配置 OPENAI_API_KEY 再运行")
        return
    agent = build_agent()

    # 同一个 thread_id：第二句里的「这家店」指代第一句的门店
    ask(agent, "北京朝阳店的拿铁豆还有多少？", thread_id="shop-001")
    ask(agent, "那这家店的冷萃瓶呢？", thread_id="shop-001")

    # 换一个 thread_id：互不相通，它不知道刚才聊的是哪家店
    ask(agent, "那这家店的冷萃瓶呢？", thread_id="shop-999")

    # 想看某个会话当前存了什么，直接读状态
    state = agent.get_state({"configurable": {"thread_id": "shop-001"}})
    print("shop-001 累计消息条数:", len(state.values["messages"]))


if __name__ == "__main__":
    main()

"""让 Agent 的最终答复是一个结构化对象，而不是一段人话。

Agent 跑完循环之后通常吐一段自然语言，好看但没法直接入库。
给 create_agent 传 response_format，最终结果里会多出 structured_response，
它是一个通过校验的 Pydantic 对象，字段名、类型都由你定。

ToolStrategy 的做法是把这个 schema 也当作一个「工具」交给模型，
让它在结束前按 schema 填一次——所以对支持工具调用的模型都能用。
"""

import os

from langchain.agents import create_agent
from langchain.agents.structured_output import ToolStrategy
from pydantic import BaseModel, Field

from shop_tools import SHOP_TOOLS


class StockAnswer(BaseModel):
    """一次库存问询的结构化结论。"""

    shop: str = Field(description="门店名称")
    item: str = Field(description="商品名称")
    qty: int = Field(description="当前库存件数，查不到写 -1")
    enough: bool = Field(description="是否满足用户要的数量")
    note: str = Field(description="一句话说明，写给人看")


def build_agent():
    return create_agent(
        model="openai:gpt-4o-mini",
        tools=SHOP_TOOLS,
        system_prompt="你是门店助手，涉及库存与价格一律调用工具。",
        response_format=ToolStrategy(StockAnswer),
    )


def main() -> None:
    if not os.environ.get("OPENAI_API_KEY"):
        print("先配置 OPENAI_API_KEY 再运行")
        return

    agent = build_agent()
    result = agent.invoke(
        {"messages": [{"role": "user", "content": "北京朝阳店的冷萃瓶够不够 20 件？"}]}
    )

    answer: StockAnswer = result["structured_response"]
    print("类型      :", type(answer).__name__)
    print("门店      :", answer.shop)
    print("商品      :", answer.item)
    print("库存      :", answer.qty)
    print("是否够用  :", answer.enough)
    print("说明      :", answer.note)

    # 拿到的是对象，直接进数据库或者拼下一段逻辑，不用再去正则抠数字
    print("入库用的字典:", answer.model_dump())


if __name__ == "__main__":
    main()

"""给 Agent 的循环加护栏：middleware 的四个常用钩子。

create_agent 内部的循环是「调模型 → 执行工具 → 回填 → 再调模型」，
middleware 就是在这条循环的缝隙里插自己的代码：

    before_agent      整个任务开始前跑一次
    before_model      每次调模型之前
    wrap_model_call   把调模型这一步整个包住，可以改请求、可以重试
    wrap_tool_call    把执行工具这一步整个包住，异常在这里兜住
    after_model       每次模型返回之后
    after_agent       整个任务结束后跑一次

下面三段分别解决三个真会踩到的问题：工具报错、步数失控、写操作没人复核。
"""

from collections.abc import Callable

from langchain.agents import create_agent
from langchain.agents.middleware import (
    HumanInTheLoopMiddleware,
    SummarizationMiddleware,
    before_model,
    wrap_tool_call,
)
from langchain.messages import ToolMessage
from langchain.tools.tool_node import ToolCallRequest

MAX_MODEL_CALLS = 8


@wrap_tool_call
def catch_tool_errors(
    request: ToolCallRequest,
    handler: Callable[[ToolCallRequest], ToolMessage],
) -> ToolMessage:
    """工具抛异常时，把异常转成一条 ToolMessage 还给模型。

    不加这一层，工具里一个 KeyError 就会把整个 Agent 打断；
    加了之后模型看得见错误内容，通常会改参数重试一次。
    """
    try:
        return handler(request)
    except Exception as exc:  # noqa: BLE001  这里就是要兜住所有异常
        return ToolMessage(
            content=f"工具执行失败：{exc}。请检查参数后重试，或者告诉用户查不到。",
            tool_call_id=request.tool_call["id"],
        )


@before_model
def limit_steps(state) -> dict | None:
    """步数刹车：模型被调用超过上限就强行收尾。

    模型偶尔会在两个工具之间来回打转。没有上限的循环
    既烧钱又不会自己停，生产环境必须有这一层。
    """
    used = sum(1 for m in state["messages"] if m.type == "ai")
    if used >= MAX_MODEL_CALLS:
        return {
            "messages": [{"role": "assistant",
                          "content": "这个问题我尝试了多轮仍未解决，先交给人工处理。"}],
            "jump_to": "end",
        }
    return None


def build_guarded_agent(tools):
    """把三层护栏一起装上。"""
    return create_agent(
        model="openai:gpt-4o-mini",
        tools=tools,
        system_prompt="你是门店助手，涉及库存与价格一律调用工具。",
        middleware=[
            limit_steps,            # 步数上限
            catch_tool_errors,      # 工具异常兜底
            # 写操作在执行前暂停，等人点头；把工具名列进 interrupt_on
            HumanInTheLoopMiddleware(
                interrupt_on={"create_order": True},
            ),
            # 历史太长时自动压缩，避免把上下文窗口撑爆
            SummarizationMiddleware(
                model="openai:gpt-4o-mini",
                max_tokens_before_summary=4000,
            ),
        ],
    )


if __name__ == "__main__":
    from shop_tools import SHOP_TOOLS

    agent = build_guarded_agent(SHOP_TOOLS)
    print("已装配的护栏：步数上限 / 工具异常兜底 / 写操作人工确认 / 历史摘要压缩")
    print("提示：装了 HumanInTheLoopMiddleware 之后必须配 checkpointer，")
    print("      否则中断之后没有地方存状态，恢复不回来。")

"""写操作前停一下：HumanInTheLoopMiddleware 的完整跑法。

查库存查错了最多是答案不准；下错单是真赔钱。
所以凡是写操作，正确的姿势是让 Agent 在执行工具之前停住，
把「它打算调哪个工具、参数填了什么」交给人看一眼。

中断要能恢复，就必须有 checkpointer 存状态，还要有 thread_id 找回来。
"""

import os

from langchain.agents import create_agent
from langchain.agents.middleware import HumanInTheLoopMiddleware
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.types import Command

from shop_tools import SHOP_TOOLS


def build_agent():
    return create_agent(
        model="openai:gpt-4o-mini",
        tools=SHOP_TOOLS,
        system_prompt="你是门店助手。下单前先查库存，库存不足不要下单。",
        middleware=[
            HumanInTheLoopMiddleware(
                # 只拦写操作；查询类工具照常自动执行
                interrupt_on={"create_order": True},
                description_prefix="以下写操作需要人工确认：",
            )
        ],
        checkpointer=InMemorySaver(),
    )


def main() -> None:
    if not os.environ.get("OPENAI_API_KEY"):
        print("先配置 OPENAI_API_KEY 再运行")
        return

    agent = build_agent()
    cfg = {"configurable": {"thread_id": "hitl-001"}}

    result = agent.invoke(
        {"messages": [{"role": "user",
                       "content": "北京朝阳店的拿铁豆下单 10 袋"}]},
        cfg,
    )

    # 被中断时，结果里会带上 __interrupt__，里面是等待确认的工具调用
    interrupts = result.get("__interrupt__")
    if not interrupts:
        print("这一轮没有触发确认：", result["messages"][-1].text())
        return

    print("等待确认的动作：", interrupts[0].value)

    # 人点头 → accept；也可以 edit 改参数，或者 reject 附上拒绝原因
    approved = agent.invoke(
        Command(resume={"decisions": [{"type": "accept"}]}),
        cfg,
    )
    print("确认之后：", approved["messages"][-1].text())

    # 拒绝的写法（换成这一段即可）：
    # agent.invoke(
    #     Command(resume={"decisions": [
    #         {"type": "reject", "message": "本月采购额度已用完，先不要下单"}
    #     ]}),
    #     cfg,
    # )


if __name__ == "__main__":
    main()

"""可复用骨架：一个带护栏的 Agent，改五处 TODO 就能上手。

结构固定为四块：
    ① 工具区    你的业务能力，一个工具做一件事
    ② 护栏区    工具异常兜底 + 步数上限 + 写操作确认
    ③ 装配区    create_agent 把模型、工具、提示词、护栏拼起来
    ④ 调用区    invoke 拿结果，stream 看过程

环境：
    pip install "langchain>=1.0" langchain-openai
    export OPENAI_API_KEY=...
    export OPENAI_BASE_URL=...
"""

import os
from collections.abc import Callable

from langchain.agents import create_agent
from langchain.agents.middleware import before_model, wrap_tool_call
from langchain.messages import ToolMessage
from langchain.tools import tool
from langchain.tools.tool_node import ToolCallRequest
from langgraph.checkpoint.memory import InMemorySaver
from pydantic import BaseModel, Field

MAX_MODEL_CALLS = 8


# ---------- ① 工具区：TODO 1 换成你自己的工具 ------------------------
class LookupInput(BaseModel):
    keyword: str = Field(description="TODO 2：说清这个参数是什么、怎么填、有什么格式要求")


@tool(args_schema=LookupInput)
def lookup_something(keyword: str) -> str:
    """TODO 3：这段描述是写给模型看的，决定它选不选这个工具。
    要回答四件事：这个工具干什么、用户说什么话时该用它、
    参数从哪里取、什么情况下不要用它。"""
    return f"这里返回 {keyword} 的查询结果"


TOOLS = [lookup_something]


# ---------- ② 护栏区：一般不用改 --------------------------------------
@wrap_tool_call
def catch_tool_errors(
    request: ToolCallRequest,
    handler: Callable[[ToolCallRequest], ToolMessage],
) -> ToolMessage:
    """工具抛异常时转成消息还给模型，别让整个任务直接崩掉。"""
    try:
        return handler(request)
    except Exception as exc:  # noqa: BLE001
        return ToolMessage(
            content=f"工具执行失败：{exc}。请检查参数后重试。",
            tool_call_id=request.tool_call["id"],
        )


@before_model
def limit_steps(state) -> dict | None:
    """步数刹车，避免模型在工具之间无限打转。"""
    used = sum(1 for m in state["messages"] if m.type == "ai")
    if used >= MAX_MODEL_CALLS:
        return {
            "messages": [{"role": "assistant", "content": "多轮尝试未果，转人工处理。"}],
            "jump_to": "end",
        }
    return None


# ---------- ③ 装配区：TODO 4 换模型、TODO 5 换人设 --------------------
def build_agent():
    return create_agent(
        model="openai:gpt-4o-mini",                      # TODO 4
        tools=TOOLS,
        system_prompt=(                                  # TODO 5
            "你是一个业务助手。需要具体数据时一律调用工具，"
            "工具查不到就如实说查不到，不要编造。"
        ),
        middleware=[limit_steps, catch_tool_errors],
        checkpointer=InMemorySaver(),
    )


# ---------- ④ 调用区 --------------------------------------------------
def ask(agent, text: str, thread_id: str = "default") -> str:
    cfg = {"configurable": {"thread_id": thread_id}}
    result = agent.invoke({"messages": [{"role": "user", "content": text}]}, cfg)
    return result["messages"][-1].text()


def watch(agent, text: str, thread_id: str = "default") -> None:
    cfg = {"configurable": {"thread_id": thread_id}}
    for chunk in agent.stream(
        {"messages": [{"role": "user", "content": text}]}, cfg, stream_mode="updates"
    ):
        for node, update in chunk.items():
            print("[%s]" % node, update)


if __name__ == "__main__":
    if not os.environ.get("OPENAI_API_KEY"):
        print("先配置 OPENAI_API_KEY 再运行")
    else:
        bot = build_agent()
        print(ask(bot, "帮我查一下 报销流程"))

"""可复用骨架：一个工具该长什么样。

三类工具各有各的写法边界，这份模板把三类都摆出来：
    读操作    放开用，返回值尽量短、尽量结构化
    算操作    参数在 schema 层就校验掉，别把脏数据放进计算
    写操作    自己限额、自己幂等，再配 HumanInTheLoopMiddleware 复核

改 TODO 的地方即可。
"""

from datetime import datetime

from langchain.tools import tool
from pydantic import BaseModel, Field

MAX_ROWS = 50              # TODO 1：读操作的返回上限
MAX_AMOUNT = 10000.0       # TODO 2：写操作的单次额度上限
_DONE: dict[str, str] = {}  # 幂等表：请求号 → 已产生的结果


# ---------- 读操作 ----------------------------------------------------
class SearchInput(BaseModel):
    keyword: str = Field(description="TODO 3：检索关键词，说明从用户的哪句话里取")
    limit: int = Field(default=10, ge=1, le=MAX_ROWS,
                       description=f"返回条数，最多 {MAX_ROWS} 条")


@tool(args_schema=SearchInput)
def search_records(keyword: str, limit: int = 10) -> str:
    """TODO 4：一句话说清这个工具查什么表、什么时候该用、什么时候不该用。"""
    rows = [f"{keyword}-记录{i + 1}" for i in range(min(limit, MAX_ROWS))]
    # 返回值越短，模型越不容易被淹没；超长结果先截断再给它。
    return "；".join(rows) if rows else "没有查到匹配的记录"


# ---------- 算操作 ----------------------------------------------------
class CalcInput(BaseModel):
    amount: float = Field(gt=0, description="金额，必须为正数")
    rate: float = Field(gt=0, le=1, description="折扣系数，0.85 表示八五折")


@tool(args_schema=CalcInput)
def calc_amount(amount: float, rate: float) -> str:
    """按折扣系数计算实付金额。只负责算钱，不查数据也不写库。"""
    # 校验交给 Pydantic，这里只写业务；
    # 千万别用 eval 跑模型给的表达式，那等于把执行权交出去。
    return f"{amount:.2f} × {rate} = {amount * rate:.2f}"


# ---------- 写操作 ----------------------------------------------------
class WriteInput(BaseModel):
    request_id: str = Field(description="业务请求号，同一笔业务重试时必须一致")
    amount: float = Field(gt=0, le=MAX_AMOUNT,
                          description=f"金额，单次不得超过 {MAX_AMOUNT}")
    memo: str = Field(default="", description="备注，可留空")


@tool(args_schema=WriteInput)
def submit_request(request_id: str, amount: float, memo: str = "") -> str:
    """提交一笔付款申请。这是写操作，会真的产生一条记录。
    只有在用户明确确认之后才使用；金额超限时不要拆单绕过上限。"""
    if request_id in _DONE:            # 幂等：同一个请求号只生效一次
        return f"该请求已处理过：{_DONE[request_id]}"
    if amount > MAX_AMOUNT:            # schema 之外再兜一层，防止绕过
        return f"金额 {amount} 超过单次上限 {MAX_AMOUNT}，已拒绝"
    result = f"受理成功，单号 RQ{datetime.now():%Y%m%d%H%M%S}，金额 {amount:.2f}"
    _DONE[request_id] = result
    return result


TOOLS = [search_records, calc_amount, submit_request]


if __name__ == "__main__":
    print(search_records.invoke({"keyword": "报销", "limit": 3}))
    print(calc_amount.invoke({"amount": 268, "rate": 0.85}))
    print(submit_request.invoke({"request_id": "R-001", "amount": 500}))
    print(submit_request.invoke({"request_id": "R-001", "amount": 500}))