"""
LangChain 1.0 - Context Management (上下文管理)
==============================================
本模块重点讲解：
1. SummarizationMiddleware - 自动摘要中间件（LangChain 1.0 新增）
2. trim_messages - 消息修剪工具
3. 管理对话长度，避免超 token
4. 中间件的使用
"""
from langchain.agents import create_agent
from langchain.agents.middleware import SummarizationMiddleware
from langgraph.checkpoint.memory import InMemorySaver
from langchain_core.messages import trim_messages
# 模拟一个长对话历史
from langchain_core.messages import HumanMessage, AIMessage
import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from langchain.tools import tool
from init_model import get_chat_model
chat_model = get_chat_model()
@tool
def calculator(operation: str, a: float, b: float) -> str:
"""执行数学计算"""
ops = {
"add": lambda x, y: x + y,
"multiply": lambda x, y: x * y,
}
result = ops.get(operation, lambda x, y: 0)(a, b)
return f"{a} {operation} {b} = {result}"
# 问题演示 - 对话历史无限增长
# ============================================================================
# 示例 1：解决方案 1 - SummarizationMiddleware（推荐）
# ============================================================================
def create_agent_with_summarization():
"""
示例2：使用 SummarizationMiddleware 自动摘要
关键：LangChain 1.0 新增的中间件
当消息数超过阈值时，自动摘要旧消息
"""
agent = create_agent(
model=chat_model,
tools=[calculator],
checkpointer=InMemorySaver(),
middleware=[
SummarizationMiddleware(model=chat_model, max_messages=2, trigger_tokens=30)
], # <--- 关键：设为一个小值，强制早期触发)],
)
config = {"configurable": {"thread_id": "with_summary2"}}
conversations = [
"计算 2 + 3",
"结果再加10",
"结果再加15",
"结果再加20",
"结果再加20",
"结果再加20",
"结果再加20",
"结果再加20",
"结果再加20",
"请总结一下计算过程",
]
for msg in conversations:
print(f"\n用户: {msg}")
response = agent.invoke(
{"messages": [{"role": "user", "content": msg}]}, config=config
)
print(response["messages"][-1].content)
# 结果：
# 用户: 计算 2 + 3
# 2 + 3 = 5
# 用户: 结果再加10
# 5 + 10 = 15
# 用户: 结果再加15
# 15 + 15 = 30
# 用户: 结果再加20
# 30 + 20 = 50
# 用户: 结果再加20
# 50 + 20 = 70
# 用户: 结果再加20
# 70 + 20 = 90
# 用户: 结果再加20
# 90 + 20 = 110
# 用户: 结果再加20
# 110 + 20 = 130
# 用户: 结果再加20
# 130 + 20 = 150
# 用户: 请总结一下计算过程
# ## 计算过程总结
# ### 初始计算
# - **第一步**：2 + 3 = **5**
# ### 后续累加
# | 步骤 | 运算 | 结果 |
# |------|------|------|
# | 第二步 | 5 + 10 | 15 |
# | 第三步 | 15 + 15 | 30 |
# | 第四步 | 30 + 20 | 50 |
# | 第五步 | 50 + 20 | 70 |
# | 第六步 | 70 + 20 | 90 |
# | 第七步 | 90 + 20 | 110 |
# | 第八步 | 110 + 20 | 130 |
# | 第九步 | 130 + 20 | **150** |
# ### 总计
# - **初始值**：2 + 3 = 5
# - **累计增加**：10 + 15 + 20×6 = 145
# - **最终结果**：**150**
# 共进行了 **9 次加法运算**，从最初的 5 逐步增加到最终的 150。
# 太棒了！🎉 SummarizationMiddleware 已经完美生效了！
# ✅ 为什么说它生效了？
# 你提供的输出结果是最好的证明：
# 1.上下文长度限制：你设置了 max_messages=2。这意味着在最后一轮对话时，Agent 的“短期记忆”里实际上只保留了最后两条原始消息（即第9轮的“130+20”和第10轮的提问）。
# 2.信息完整性：如果没有摘要机制，Agent 应该完全忘记第1轮的 2+3=5 以及中间的累加过程，只能回答“我只记得最近加了20”。
# 3.最终表现：Agent 却能够精确地列出从第一步到第九步的所有细节，甚至算出了累计增加量（145）。
# 这说明：之前的那些被“移除”出短期记忆的对话，已经被中间件自动压缩成了一条高密度的摘要（Summary），并作为系统提示词传给了模型。
# 模型正是读取了这条摘要，才“回忆”起了整个计算链条。
# ============================================================================
# 示例 2：理解 SummarizationMiddleware 参数
# ============================================================================
# def demo_summarization_middleware():
# SummarizationMiddleware(
# model: BaseChatModel,
# max_messages: int = 20,
# trigger_tokens: int | None = None,
# summary_prompt: ChatPromptTemplate | None = None,
# )
# model BaseChatModel 必填 (无默认) 指定执行摘要任务的模型实例。 • 省钱策略：主 Agent 用 gpt-4o，摘要用 gpt-4o-mini。
# • 一致性：通常直接使用 Agent 的主模型。
# max_messages int 20 (视版本而定) 触发总结时，保留最近多少条原始消息不压缩。 • 强逻辑/数学：设为 2 ~ `4`。保留最近几步推理，防止逻辑断层。<br>• <v>**<v>创作/聊天<v>**<v>：设为 `10` ~ 20。保留更多近期上下文以保持连贯性。
# trigger_tokens int / None None (通常为窗口 75%) 上下文 Token 数超过此值时，强制触发摘要。 • 调试/测试：设为 100 ~ `200`。强制在短对话中触发以便观察效果。<br>• <v>**<v>生产环境<v>**<v>：设为模型上限的 `50%` ~ 70%。平衡性能与成本。
# summary_prompt ChatPromptTemplate 内置默认模板 自定义如何生成摘要的指令模板。 • 特定领域：若需提取 JSON、代码或特定实体，需自定义此 Prompt。
# • 默认行为：通常只需总结关键事实、待办和实体，忽略闲聊。
# ============================================================================
# 示例 3：手动消息修剪（trim_messages）
# ============================================================================
from langchain_core.messages import (
HumanMessage,
AIMessage,
SystemMessage,
trim_messages,
)
def demo1_trim_messages():
print("=== 开始演示上下文裁剪 ===")
# 1. 初始化 Agent (保持你的配置)
# 注意：如果你手动 trim 了，Middleware 可能不会触发，这是正常的。
agent = create_agent(
model=chat_model,
checkpointer=InMemorySaver(),
middleware=[
SummarizationMiddleware(model=chat_model, max_messages=2, trigger_tokens=30)
],
)
config = {"configurable": {"thread_id": "test_trim_final"}}
# 2. 构造长历史消息
messages = [
HumanMessage(content="消息 1：今天天气不错。"),
AIMessage(content="回复 1：是的，适合出门。"),
HumanMessage(content="消息 2：我想去公园。"),
AIMessage(content="回复 2：公园是个好主意。"),
HumanMessage(content="消息 3：但是公园有点远。"),
AIMessage(content="回复 3：那我们可以去近一点的地方。"),
HumanMessage(content="消息 4：附近有什么推荐吗？"),
AIMessage(content="回复 4：附近的植物园很不错。"),
]
print(f"原始消息数量: {len(messages)}")
# 3. 执行裁剪
# 目标：保留最后的部分，但限制 Token 数量
trimmed = trim_messages(
messages,
max_tokens=30, # 设一个小值，强制它丢弃前面的“消息 1/2”
token_counter="approximate",
strategy="last",
start_on="human", # 确保第一条是用户说的
)
print(f"裁剪后消息数量: {len(trimmed)}")
print("--- 裁剪后的上下文内容 ---")
for i, msg in enumerate(trimmed):
role = "User" if isinstance(msg, HumanMessage) else "AI"
print(f"[{i}] {role}: {msg.content}")
print("--------------------------")
# 4. 【关键】构建输入
# 必须把 trimmed 列表作为一个整体历史传进去，而不是循环单条发送
# 我们再加一个新的问题，测试 Agent 是否记得前面的内容
new_question = HumanMessage(content="结合刚才的对话，你觉得去植物园好还是公园好？")
# 合并历史和新问题
input_payload = {"messages": trimmed + [new_question]}
print("\n正在调用 Agent (带着裁剪后的历史 + 新问题)...")
try:
response = agent.invoke(input_payload, config=config)
# 获取最后一条消息（AI 的回答）
final_answer = response["messages"][-1]
print("\n=== Agent 的最终回答 ===")
print(final_answer.content)
print("========================")
# 验证：如果回答里提到了“植物园”或“公园”，说明上下文保留成功了
if "植物园" in final_answer.content or "公园" in final_answer.content:
print("✅ 成功：Agent 记住了裁剪后的上下文！")
else:
print("⚠️ 警告：Agent 的回答似乎没有包含上下文信息。")
except Exception as e:
print(f"❌ 调用失败: {e}")
import traceback
traceback.print_exc()
##结果
# === 开始演示上下文裁剪 ===
# 原始消息数量: 8
# 裁剪后消息数量: 2
# --- 裁剪后的上下文内容 ---
# [0] User: 消息 4：附近有什么推荐吗？
# [1] AI: 回复 4：附近的植物园很不错。
# --------------------------
# 正在调用 Agent (带着裁剪后的历史 + 新问题)...
# === Agent 的最终回答 ===
# 结合刚才的对话，既然我之前特意推荐了植物园，那我依然觉得**植物园**会更好一些。
# 相比普通公园，植物园通常有以下优势：
# 1. **景观更丰富**：植物种类更多，四季景色变化更明显，观赏性强。
# 2. **环境更清幽**：一般比普通公园更适合静心散步、呼吸新鲜空气。
# 3. **体验更独特**：之前提到它“很不错”，说明它有独特的亮点，不仅仅是普 通的绿地。
# 当然，如果你只是想进行简单的运动或带孩子玩耍，普通公园也很方便；但如果是为了放松身心、欣赏自然，植物园会是更优选。
# ========================
# ✅ 成功：Agent 记住了裁剪后的上下文！
# ============================================================================
# 示例 4：对比不同策略
# ============================================================================
def example_4_comparison():
"""
策略对比：
1. 不做处理（默认）
优点：保留完整历史
缺点：会超 token、成本高
适用：短对话
2. SummarizationMiddleware（推荐）
优点：
- 自动化，无需手动管理
- 保留重要信息（通过摘要）
- 平滑过渡
缺点：
- 摘要可能丢失细节
- 额外的摘要成本
适用：长对话、需要保留上下文
3. trim_messages（手动修剪）
优点：
- 精确控制
- 简单直接
- 无额外成本
缺点：
- 旧消息完全丢失
- 可能断开上下文
适用：只需要最近 N 轮
4. 滑动窗口（自定义）
优点：
- 保留系统消息 + 最近消息
- 可控成本
缺点：
- 需要自己实现
适用：有明确规则的场景
推荐方案：
- 短对话（<10轮）：不处理
- 中长对话：SummarizationMiddleware
- 只要最近几轮：trim_messages
"""
# ============================================================================
# 示例 5：实际应用 - 客服机器人
# ============================================================================
def example_5_customer_service():
"""
客服机器人场景：
1. 问题：“我想知道北京天气”
2. 客服：“好的，我会查询北京的天气。”
3. 问题：“我想知道北京的温度”
4. 客服：“北京当前天气晴朗，温度20°C，空气质量良好。”
5. 问题：“我想知道北京的湿度”
6. 客服：“我没有查询到北京的湿度信息。”
场景：客服对话可能很长，需要管理上下文
"""
agent = create_agent(
model=chat_model,
tools=[calculator],
system_prompt="""你是客服助手。
1. 你是一个客服助手，负责回答用户的问题。
2. 你只能使用提供的工具（如计算器）来回答问题。
3. 你不能编造信息，只能基于提供的上下文回答。
4. 你不能回答与客服助手无关的问题。
""",
checkpointer=InMemorySaver(),
middleware=[
SummarizationMiddleware(
model=chat_model,
trim_tokens_to_summarize=800,
)
],
)
config = {"configurable": {"thread_id": "customer_123"}}
# 模拟客服对话
conversations = [
"你好，我想咨询订单",
"我的订单号是 12345",
"帮我算一下 100 乘以 2 的优惠价",
"谢谢",
]
for msg in conversations:
print(f"\n客户: {msg}")
response = agent.invoke(
{"messages": [{"role": "user", "content": msg}]}, config=config
)
print(f"客服: {response['messages'][-1].content}")
print(f"\n总消息数: {len(response['messages'])}")
print("\n关键点：")
print(" - 自动管理对话长度")
print(" - 重要信息（订单号）通过摘要保留")
print(" - 适合生产环境")
if __name__ == "__main__":
# create_agent_with_summarization()
# demo1_trim_messages()
example_5_customer_service()