如何在LLamaIndex Workflow中实现带短期会话记忆的RAG聊天机器人

核心实现思路 #

不用搭建复杂的外部数据库+人工循环逻辑，LLamaIndex内置的记忆组件可直接结合Workflow实现需求，具体步骤如下：

• 复用Chat Engine记忆组件：Chat Engine中的ChatMemoryBuffer等记忆类可直接在Workflow节点中调用，无需从零构建。初始化时设置记忆的最大token数或消息条数，组件会自动处理短期记忆的截断逻辑。
• Workflow内传递记忆实例：将记忆对象作为Workflow的全局状态或节点间参数传递，每次会话请求时，先从记忆中加载历史对话，再与当前提问一同传入RAG检索+生成节点。
• 简化超时清理逻辑：无需手动实现循环等待，通过内存存储+定时任务处理会话超时。给每个会话分配唯一ID并绑定记忆实例，定时扫描超过阈值的会话ID并删除对应记忆即可。

代码示例 #

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.memory import ChatMemoryBuffer
from llama_index.core.workflow import Workflow, step
from llama_index.llms.openai import OpenAI
import time
import threading

# 初始化RAG索引
docs = SimpleDirectoryReader("data").load_data()
index = VectorStoreIndex.from_documents(docs)
llm = OpenAI(model="gpt-3.5-turbo")

# 会话存储：key=会话ID，value=(记忆实例, 最后活跃时间)
session_store = {}

# 定时清理超时会话（30分钟无活动则清除）
def clean_expired_sessions():
    while True:
        current_time = time.time()
        expired_ids = [sid for sid, (mem, last_time) in session_store.items() if current_time - last_time > 1800]
        for sid in expired_ids:
            del session_store[sid]
        time.sleep(300)  # 每5分钟检查一次

# 启动清理后台线程
threading.Thread(target=clean_expired_sessions, daemon=True).start()

class RAGChatWorkflow(Workflow):
    @step
    async def load_memory(self, session_id: str, user_query: str):
        # 加载或初始化会话记忆
        if session_id not in session_store:
            memory = ChatMemoryBuffer.from_defaults(token_limit=10000, llm=llm)
            session_store[session_id] = (memory, time.time())
        else:
            memory, _ = session_store[session_id]
            # 更新会话最后活跃时间
            session_store[session_id] = (memory, time.time())
        # 获取历史对话记录
        chat_history = memory.get_all()
        return chat_history, memory, user_query

    @step
    async def rag_response(self, chat_history: list, memory: ChatMemoryBuffer, user_query: str):
        # 构建带历史对话的RAG查询引擎
        query_engine = index.as_chat_engine(
            chat_mode="context",
            memory=memory,
            llm=llm,
            system_prompt="你是基于给定文档的RAG聊天机器人，需结合历史对话内容回答用户问题"
        )
        response = await query_engine.achat(user_query)
        # 记忆组件会自动保存当前对话，无需手动处理
        return response.response

# 使用示例
async def main():
    workflow = RAGChatWorkflow()
    session_id = "user_123"
    # 第一轮提问
    resp1 = await workflow.run(session_id=session_id, user_query="文档里关于XX的内容是什么？")
    print(resp1)
    # 跟进提问
    resp2 = await workflow.run(session_id=session_id, user_query="那XX和YY的区别是什么？")
    print(resp2)

if __name__ == "__main__":
    import asyncio
    asyncio.run(main())

关键说明 #

• 记忆组件复用：ChatMemoryBuffer会自动管理对话历史长度，超过设定的token限制时会自动截断最早的消息，无需手动处理。
• 会话存储轻量化：用内存字典+后台定时线程实现会话超时清理，比外部数据库更轻量；若需持久化，可直接替换为Redis等键值数据库。
• Workflow节点解耦：将加载记忆和生成RAG响应拆分为独立节点，逻辑清晰，便于后续扩展意图识别、权限校验等功能。

内容的提问来源于stack exchange，提问作者hristogg