在上一节内容中详细介绍了短期记忆管理的3种方式，删除、修剪和总结。在本节内容中，我们将基于第4.10节的内容实现一个类似 ChatGPT 的 RAG 问答工具，做到每次回答完用户问题后都判断是否需要对历史消息进行一次总结，并同步将会话消息实时保存到数据库，使得每次启动程序时都能恢复历史会话。
5.4 拥有短期记忆的 RAG Agent
在上一节内容中详细介绍了短期记忆管理的3种方式，删除、修剪和总结。在本节内容中，我们将基于第4.10节的内容实现一个类似 ChatGPT 的 RAG 问答工具，做到每次回答完用户问题后都判断是否需要对历史消息进行一次总结，并同步将会话消息实时保存到数据库，使得每次启动程序时都能恢复历史会话。
5.4.1 流程构建思路
为了实现上述目的，需要在图4-13的图结构之上再添加一个条件判断边和一个节点，来判断是否需要对历史记忆进行总结以及如何总结，整个流程如图5-6所示。
图 5-6 带记忆能力的 RAG Agent 图结构可视化
从图5-6可以看出，在每一轮回答结束以后两个分支都会统一汇集到 check_messages 节点，然后再通过条件分支判断是否需要进行历史消息总结。下面逐一对这部分内容进行介绍。
5.4.2 分支汇集与条件判断
由于 decide 这条边已经用 tools_condition 在做"是否调用工具"的分流，即 add_conditional_edges("decide",tools_condition, {"tools": "my_retrieve", END: END})，而 tools_condition 的映射是静态的，没法在左侧的 END 分支里再动态判断一次是否应该进行记忆总结，所以需要添加一个空节点来做汇集，即图5-6中的 check_messages，其对应的实现为：
1 def passthrough(state: MyState) -> MyState:
2 return {}在上述代码中，第2行 return {} 表示更新内容为空，即不对传入的 state 做任何修改。
进一步，实现是否需要进行总结的条件判断，示例代码如下：
1 def count_messages(state: MyState) -> Literal[END, "to_summarize"]:
2 messages = state["messages"]
3 threshold = 5
4 if len(messages) > threshold:
5 return "to_summarize"
6 else:
7 return END在上述代码中，我们简单地使用消息条数来进行判断，当然也可以是其它策略。第4~7行是根据条件判断流向 to_summarize 节点或直接结束。
5.4.3 历史记忆总结
接着，在第5.2.4节实现的 summarize_conversation 基础之上，再来对内部进行一些改造，使得能够更加灵活的总结和删除对应的历史记忆，示例代码如下：
1 def summarize_conversation(state: MyState) -> MyState:
2 ......
3 trimmed = trim_messages(state["messages"], max_tokens=5, strategy="last",
4 token_counter=len, start_on="human", end_on=("human", "ai"))
5 trimmed_msg_ids = [msg.id for msg in trimmed] # 保留下来的消息
6 delete_messages, summary_messages = [], []
7 for msg in state["messages"]:
8 if msg.id not in trimmed_msg_ids:
9 delete_messages.append(RemoveMessage(id=msg.id))
10 if msg.type != 'tool':
11 summary_messages.append(msg)
12 summary_messages += [HumanMessage(content=summary_message)]
13 response = get_llm_model().invoke(summary_messages)
14 return {"summary": response.content, "messages": delete_messages}在上述代码中，第3~4行中先使用 trim_messages 函数来对需要保留的消息记录做一次条件筛选，最多保留最后的5条消息，这样能够使得保留下来的消息更上下文更加完整。第5行是得到待保留消息的 ID，第7~11行是遍历原始状态中的每一条消息，然后标记待删除的消息，以及根据消息类型筛选待总结的消息，因为工具返回的内容篇幅通常很长，而且不一定都有用，所以去掉这部分内容。
5.4.4 搭建完整流程
在完成上述节点定义以后，最后来完成整个图的构建过程，示例代码如下：
1 def get_workflow(checkpointer: PostgresSaver):
2 workflow = StateGraph(MyState)
3 ...
4 workflow.add_node("check_messages", passthrough)
5 workflow.add_node("summarize", summarize_conversation)
6 workflow.add_edge(START, "decide")
7 workflow.add_conditional_edges("decide",tools_condition,
8 {"tools": "my_retrieve", END: "check_messages"})
9 workflow.add_conditional_edges("my_retrieve", grade_documents)
10 workflow.add_edge("generate_answer", "check_messages")
11 workflow.add_conditional_edges("check_messages", count_messages,
12 {"to_summarize": "summarize", END: END})
13 workflow.add_edge("summarize", END)
14 workflow.add_edge("rewrite_question", "decide")
15 graph = workflow.compile(checkpointer=checkpointer)
16 return graph在上述代码中......