主题
🎛️ 9.11 自定义中间件与生命周期钩子
“预置中间件是别人造好的轮子;自定义中间件,是你亲手给 Agent 装上的神经与韧带。”
9.7 讲的是官方预置中间件“即插即用”,本节我们下沉到源码级:搞懂 Agent 循环的每一步,并亲手在关键节点上挂载自己的钩子(Hooks),实现日志、限流、重试、计数等任意横切能力。
💡 先看一个直观例子:流水线质检员与包裹代收点
1. 中间件到底插在哪?
回顾 Agent 的核心循环:调用模型 → 模型决定调工具 → 执行工具 → 把结果喂回模型 → … → 结束。中间件就是插在这条循环上的“钩子”。
2. 两种钩子风格(官方定义)
| 风格 | 作用时机 | 一句话比喻 |
|---|---|---|
| Node-style 钩子 | 在固定节点依次经过 | 流水线上的质检员:到点必查,查完放行 |
| Wrap-style 钩子 | 包裹每次模型/工具调用 | 快递驿站的代收点:可决定包裹是否发出、发几次、甚至拦截 |
Node-style 钩子(4 个):
| 钩子 | 触发时机 |
|---|---|
before_agent | Agent 启动前(每次调用仅一次) |
before_model | 每次模型调用前 |
after_model | 每次模型响应后 |
after_agent | Agent 结束后(每次调用仅一次) |
Wrap-style 钩子(2 个):
| 钩子 | 触发时机 |
|---|---|
wrap_model_call | 包裹每次模型调用 |
wrap_tool_call | 包裹每次工具调用 |
💻 核心实操:从零手写中间件
完整可运行脚本见 code/s11_custom_middleware.py。中间件有两种创建姿势——装饰器式(单钩子、快速原型)与类式(多钩子、可配置、可复用)。
1. Node-style:消息上限熔断 + 响应日志
python
# code/s11_custom_middleware.py
from langchain.agents import create_agent
from langchain.agents.middleware import before_model, after_model, AgentState
from langchain.messages import AIMessage
from langgraph.runtime import Runtime
from typing import Any
@before_model(can_jump_to=["end"]) # 声明本钩子有权"跳到 end"提前结束
def check_message_limit(state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
if len(state["messages"]) >= 50:
return { # 返回 dict 即持久更新状态 + 跳转
"messages": [AIMessage("对话已达上限,请开启新会话。")],
"jump_to": "end",
}
return None # 返回 None 表示不干预
@after_model
def log_response(state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
print(f"📝 模型返回:{state['messages'][-1].content}")
return None
agent = create_agent(model=llm, tools=tools, middleware=[check_message_limit, log_response])2. Wrap-style:模型调用自动重试
wrap_model_call 能控制 handler 被调用 0 次 / 1 次 / 多次——重试、短路、缓存都靠它:
python
from langchain.agents.middleware import wrap_model_call, ModelRequest, ModelResponse
from typing import Callable
@wrap_model_call
def retry_model(request: ModelRequest, handler: Callable[[ModelRequest], ModelResponse]) -> ModelResponse:
for attempt in range(3):
try:
return handler(request)
except Exception as e:
if attempt == 2:
raise
print(f"🔁 第 {attempt + 1}/3 次重试,错误:{e}")3. 类式中间件:同步 + 异步双实现
类式中间件继承 AgentMiddleware,可同时定义 sync / async 钩子,并携带初始化配置:
python
from langchain.agents.middleware import AgentMiddleware, AgentState
from langgraph.runtime import Runtime
from typing import Any
class LoggingMiddleware(AgentMiddleware):
def before_model(self, state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
print(f"🔍 即将调用模型,当前消息数:{len(state['messages'])}")
return None
def after_model(self, state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
print(f"✅ 模型返回:{state['messages'][-1].content}")
return None
# 异步版本(FastAPI / astream 场景自动走这里)
async def abefore_model(self, state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
return self.before_model(state, runtime)4. 自定义状态 Schema:让中间件拥有“记忆”
中间件可通过 state_schema 扩展 Agent 状态,在多个钩子之间共享计数器/标志位:
python
from langchain.agents.middleware import AgentMiddleware, AgentState
from typing_extensions import NotRequired
from typing import Any
class CounterState(AgentState):
model_call_count: NotRequired[int]
class CallCounterMiddleware(AgentMiddleware):
state_schema = CounterState # 挂载自定义状态
def before_model(self, state: CounterState, runtime) -> dict[str, Any] | None:
if state.get("model_call_count", 0) > 10:
return {"jump_to": "end"} # 超限直接收尾,防止费用失控
return None
def after_model(self, state: CounterState, runtime) -> dict[str, Any] | None:
return {"model_call_count": state.get("model_call_count", 0) + 1}5. Wrap-style 的状态更新:ExtendedModelResponse + Command
Wrap-style 钩子想更新状态,不能直接返回 dict,要返回 ExtendedModelResponse 包裹 Command:
python
from langchain.agents.middleware import wrap_model_call, ModelRequest, ModelResponse, ExtendedModelResponse
from langgraph.types import Command
@wrap_model_call
def track_usage(request: ModelRequest, handler) -> ExtendedModelResponse:
response = handler(request)
return ExtendedModelResponse(
model_response=response,
command=Command(update={"last_tokens": 150}), # 经 graph reducer 合并进状态
)🔑 记忆口诀:Node-style 钩子返回
dict直接改状态;Wrap-style 钩子返回ExtendedModelResponse(包Command)改状态。多条中间件叠加时,messages走累加 reducer,其余字段外层覆盖内层。
6. 快捷姿势一:@dynamic_prompt 装饰器——动态系统提示的"一行式"
9.10 里"每次模型调用前动态改系统提示"要用 wrap_model_call 手写十几行;官方还提供了一行装饰器的快捷姿势(底层就是替你包好了 wrap_model_call,1.3 后期版本已可用,1.4 环境实测在列)。好比酒店前台不用每次都找经理审批换房卡,出示工牌(装饰器)直接办:
python
# code/s11_custom_middleware.py 可直接扩展:动态系统提示一行搞定
from langchain.agents.middleware import dynamic_prompt, ModelRequest
@dynamic_prompt
def personalized_prompt(request: ModelRequest) -> str:
user_tier = request.runtime.context.user_tier # 读 9.10 的 Runtime Context
base = request.system_message.text if request.system_message else ""
return f"{base}\n\n当前用户等级:{user_tier},回答风格请相应调整。"与手写 wrap_model_call 的关系:能力等价,场景不同——只改提示词用 @dynamic_prompt;还要改工具、模型参数或多轮 handler 才需要完整的 wrap_model_call。
7. 快捷姿势二:TracePolicy——给中间件链路"打码"的可观测开关
中间件的钩子执行也会被追踪(LangSmith 等平台能看到每次钩子收到的输入/输出)。每个中间件可声明 trace_policy(由 LangGraph 提供,1.3 后期版本已可用,1.4 环境实测在列),控制自己这一跳记录什么——好比快递面单上只写"日用品",不写具体型号,防窥探又不影响配送:
python
from langgraph.types import TracePolicy, omit_payload
from langchain.agents.middleware import AgentMiddleware
class SensitiveAuditMiddleware(AgentMiddleware):
# 钩子的输入照记(要看消息数),输出不记(防止敏感审查结论进日志)
trace_policy = TracePolicy(process_outputs=omit_payload)
def after_model(self, state, runtime):
...两个工程要点:
process_inputs/process_outputs分别对本节点的入参/返回值做变换,omit_payload是官方提供的"整段丢弃"助手;- 它只改追踪记录,不是脱敏正道——真要全链路(含子运行)脱敏,官方明确建议用 LangSmith 客户端的
hide_inputs/anonymizer。另有configure_trace_policy(...)可在启动时设置进程级默认策略,单个中间件的trace_policy属性优先级更高。
🖥️ 实验台真实调用演示(9.11 页可交互)
code/s11_custom_middleware.py 的演示已从「仅装配」升级为真实执行,点实验台按钮即可观察:
| 演示 | 观察点 |
|---|---|
| 🚦 Node-style 钩子 | 场景 A 正常放行(after_model 打印响应日志);场景 B 伪造 50 条历史 → before_model 熔断 jump_to='end',模型零调用、零 Token |
| 🔁 Wrap-style 重试 | 用「前 2 次必失败」的 FakeChatModel 模拟网络抖动(零 Token 可复现),终端实时打印 🔁 第 1/3 次重试,第 3 次自愈 |
| 🧱 类式中间件 | 三轮对话日志 + model_call_count 真实累计;state_schema 让中间件跨轮「记忆」,超限后 before_model 直接收尾 |
| 🕶️ TracePolicy | 离线验证 TracePolicy(process_outputs=omit_payload) 构造与挂载:钩子照常执行、追踪输出整段丢弃;configure_trace_policy 进程级默认设置与清理全程可见(Fake 模型零 Token) |
💡 教学设计:重试演示故意不用真实网络故障(不可复现),而是注入一个确定失败两次的假模型—— 既能 100% 触发重试路径,又不烧一分钱 Token。
📚 权威官方资料直达
- 🔗 自定义中间件官方指南(本节核心):Custom Middleware
- 🔗 中间件总览与 Agent 循环:Middleware Overview
- 🔗 预置中间件全家桶:Prebuilt Middleware
- 🔗 中间件 API 参考:Middleware API Reference---
中间件组合要防止重复执行
中间件的包裹顺序会影响鉴权、重试、计费和脱敏。如果外层与内层都重试三次,一次失败最多可能触发九次调用;如果状态更新发生在重试内部,还可能重复累计。组合前应画出调用顺序,并为每层规定职责。
需要副作用的中间件应使用幂等键,确保同一次运行重复进入时不会重复写入。安全拦截失败通常要终止请求,观测上报失败则可以降级继续。把顺序、失败策略和状态变化写进测试,才能避免升级后行为悄悄改变。
🎯 本节小结与思考
- 核心收获:掌握了中间件的两大钩子风格(Node-style 4 个 + Wrap-style 2 个)、装饰器式与类式两种创建姿势、
state_schema自定义状态、以及两种状态更新机制(返回dict/ExtendedModelResponse + Command);并上手了两个官方快捷姿势——@dynamic_prompt一行动态系统提示、TracePolicy中间件链路追踪脱敏。 - 下一步探索:中间件既能做“增强”,也能做“防守”。如何用它构建护栏(Prompt 注入防护、PII 检测、内容审查)并做测试评估?下一节我们学习 9.12 生产级防护:护栏安全与测试评估。