1
import os
2
import time
3

4
from dotenv import load_dotenv
5
from langchain.agents import create_agent
6
from langchain.agents.middleware import (
7
    AgentMiddleware,
8
    ContextEditingMiddleware,
9
    ClearToolUsesEdit,
10
    ModelCallLimitMiddleware,
11
    PIIMiddleware,
12
    ToolRetryMiddleware,
13
)
14
from langchain.chat_models import init_chat_model
15
from langchain.tools import tool
16

17
load_dotenv(override=True)
18

19
model = init_chat_model(
20
    model="deepseek-v4-flash",
21
    model_provider="openai",
22
    api_key=os.getenv("DEEPSEEK_API_KEY"),
23
    base_url=os.getenv("DEEPSEEK_BASE_URL"),
24
)
25

26
# ---------- 2-1 限制模型调用次数 ----------
27
@tool
28
def add(a: int, b: int) -> int:
29
    """两数相加
30

31
    Args:
32
        a: 第一个数
33
        b: 第二个数
34
    """
35
    return a + b
36

37
agent_limit = create_agent(
38
    model=model,
39
    tools=[add],
40
    middleware=[ModelCallLimitMiddleware(run_limit=2, exit_behavior="end")],
41
)
42
r = agent_limit.invoke({"messages": ["帮我连续算三次：(1+2)、(3+4)、(5+6)，最后汇总"]})
43
print("消息条数:", len(r["messages"]))
44

45
# ---------- 2-2 工具失败自动重试 ----------
46
attempts = {"n": 0}
47

48
@tool
49
def flaky_tool(query: str) -> str:
50
    """一个前两次必失败、第三次成功的工具
51

52
    Args:
53
        query: 查询内容
54
    """
55
    attempts["n"] += 1
56
    print(f"   第 {attempts['n']} 次调用，时间 {time.strftime('%H:%M:%S')}")
57
    if attempts["n"] < 3:
58
        raise TimeoutError("模拟网络超时")
59
    return f"{query} 的查询结果：成功"
60

61
agent_retry = create_agent(
62
    model=model,
63
    tools=[flaky_tool],
64
    middleware=[
65
        ToolRetryMiddleware(
66
            max_retries=3,
67
            backoff_factor=2.0,
68
            initial_delay=0.5,
69
            max_delay=5.0,
70
            jitter=True,
71
            retry_on=(TimeoutError,),
72
            on_failure="continue",
73
        )
74
    ],
75
)
76
r2 = agent_retry.invoke({"messages": ["用 flaky_tool 查一下北京天气"]})
77
print(r2["messages"][-1].content)
78

79
# ---------- 2-3 中间件顺序 ----------
80
class A(AgentMiddleware):
81
    def before_model(self, state, runtime):
82
        print("  [A] before_model")
83
        return None
84

85
    def after_model(self, state, runtime):
86
        print("  [A] after_model")
87
        return None
88

89
class B(AgentMiddleware):
90
    def before_model(self, state, runtime):
91
        print("  [B] before_model")
92
        return None
93

94
    def after_model(self, state, runtime):
95
        print("  [B] after_model")
96
        return None
97

98
agent_order = create_agent(model=model, tools=[], middleware=[A(), B()])
99
agent_order.invoke({"messages": ["测试"]})
100
# 输出：A before → B before → B after → A after（洋葱模型）
101

102
# ---------- 3-1 组合一套"生产级"中间件 ----------
103
agent_prod = create_agent(
104
    model=model,
105
    tools=[flaky_tool],
106
    middleware=[
107
        # 1. 最先：敏感信息脱敏（越早拦越好，避免泄露给模型）
108
        PIIMiddleware("email", strategy="redact", apply_to_input=True),
109
        # 2. 其次：成本硬约束，防止无限调用
110
        ModelCallLimitMiddleware(run_limit=10, exit_behavior="end"),
111
        # 3. 再次：控制上下文体积（只影响发给模型的消息）
112
        ContextEditingMiddleware(edits=[ClearToolUsesEdit(trigger=100000, keep=3)]),
113
        # 4. 最后：失败后的兜底重试
114
        ToolRetryMiddleware(max_retries=3, initial_delay=0.5, backoff_factor=2.0, on_failure="continue"),
115
    ],
116
)
117
r3 = agent_prod.invoke({"messages": ["我的邮箱是 a@b.com，用它查一下北京天气"]})
118
print(r3["messages"][-1].content)