Langgraph 实现类似Langchain记忆上下文过长裁剪

锋行天下2026-10-30 15:03
python 复制代码
from typing import TypedDict, Any
from langchain_openai import ChatOpenAI
from langchain_core.messages import AnyMessage
from langgraph.graph import StateGraph, START, END, MessagesState
from langgraph.prebuilt import ToolNode
from langgraph.checkpoint.memory import InMemorySaver
from langmem.short_term import SummarizationNode, RunningSummary

# 1. 定义状态：MessagesState 基础上增加 context 存放运行中的摘要
class State(MessagesState):
    context: dict[str, RunningSummary]  # 存放摘要状态，避免重复总结旧消息

class LLMInputState(TypedDict):
    summarized_messages: list[AnyMessage]

# 2. 初始化模型
model = ChatOpenAI(model="gpt-4o")
summarization_model = model.bind(max_tokens=128)  # 专门用来生成摘要的模型，限制输出长度

# 3. 创建摘要节点
summarization_node = SummarizationNode(
    token_counter=model.get_num_tokens_from_messages,  # 精确计算 token
    model=summarization_model,
    max_tokens=256,                    # 压缩后允许的 token 总量
    max_tokens_before_summary=1024,    # 超过这个阈值才触发摘要
    max_summary_tokens=128,            # 摘要本身的最大 token
)

# 4. 调用 LLM 的节点（注意：输入用 summarized_messages，不是 messages）
def call_model(state: LLMInputState):
    response = model.bind_tools(tools).invoke(state["summarized_messages"])
    return {"messages": [response]}

# 5. 构建图
builder = StateGraph(State)
builder.add_node("summarize", summarization_node)
builder.add_node("call_model", call_model)
builder.add_node("tools", ToolNode(tools))

builder.set_entry_point("summarize")
builder.add_edge("summarize", "call_model")
# ... 后续 router 和 tools 的边按需添加

graph = builder.compile(checkpointer=InMemorySaver())