06_记忆和知识检索

为AI注入长期记忆！本节课学习如何实现对话记忆管理和知识库检索（RAG）。掌握向量数据库的基本使用，让AI能够记住历史对话，并从海量文档中精准检索相关信息，构建真正智能的问答系统。

---

🎯 学习目标

1. 理解对话记忆的重要性和实现方式
2. 掌握RAG（检索增强生成）的基本原理
3. 使用向量数据库存储和检索知识
4. 构建带知识库的智能问答系统
5. 优化检索效果

---

1. 对话记忆管理

1.1 为什么需要记忆？

没有记忆：

用户：我叫小明
AI：你好，小明！

用户：我叫什么名字？
AI：抱歉，我不知道你的名字。  ❌

有记忆：

用户：我叫小明
AI：你好，小明！

用户：我叫什么名字？
AI：你叫小明。  ✅

1.2 记忆类型

短期记忆（Session Memory）

• 当前对话的上下文
• 保存在State中
• 会话结束后清除

长期记忆（Persistent Memory）

• 跨会话保存
• 存储在数据库中
• 可以长期查询

---

2. RAG（检索增强生成）

2.1 什么是RAG？

问题： LLM的知识是固定的，无法回答最新信息或专有知识

解决方案： RAG = Retrieval (检索) + Augmented (增强) + Generation (生成)

流程：

用户提问 → 检索知识库 → 找到相关文档 → 拼接到提示词 → LLM生成回答

2.2 核心组件

1. 文档切分：把长文档切成小块
2. 向量化：把文本转成向量（数字数组）
3. 存储：保存到向量数据库
4. 检索：找最相似的文档块
5. 生成：基于检索结果回答问题

---

3. 实战：简单的RAG系统

创建 01_simple_rag.py：

"""
简单的RAG问答系统
学习目标：理解RAG的基本流程
"""
import os
from typing import TypedDict, List
from typing_extensions import NotRequired
from langgraph.graph import StateGraph, END
from langchain_community.llms import Tongyi
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import DashScopeEmbeddings
from dotenv import load_dotenv

load_dotenv()


# 知识库文档
KNOWLEDGE_BASE = """
公司政策文档：

1. 请假制度
- 年假：每年10天带薪年假
- 病假：需提供医院证明
- 事假：需提前3天申请

2. 报销制度  
- 差旅费：需提供发票和审批单
- 餐饮费：每天限额200元
- 交通费：公共交通实报实销

3. 工作时间
- 上班时间：9:00-18:00
- 午休时间：12:00-13:00
- 弹性工作：可申请远程办公
"""


# 定义状态
class State(TypedDict):
    question: str
    retrieved_docs: NotRequired[List[str]]
    answer: NotRequired[str]


# 初始化向量数据库
def init_vectorstore():
    """初始化向量数据库"""
    # 文档切分
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=200,
        chunk_overlap=50
    )
    docs = splitter.create_documents([KNOWLEDGE_BASE])
    
    # 创建向量数据库
    embeddings = DashScopeEmbeddings(
        model="text-embedding-v1",
        dashscope_api_key=os.getenv("DASHSCOPE_API_KEY")
    )
    
    vectorstore = FAISS.from_documents(docs, embeddings)
    return vectorstore


# 全局向量数据库
vectorstore = init_vectorstore()


# 节点1：检索相关文档
def retrieve_node(state: State) -> State:
    """从知识库检索相关文档"""
    question = state["question"]
    
    print(f"[检索] 查询问题: {question}")
    
    # 检索最相关的2个文档
    docs = vectorstore.similarity_search(question, k=2)
    retrieved = [doc.page_content for doc in docs]
    
    state["retrieved_docs"] = retrieved
    print(f"[检索] 找到{len(retrieved)}个相关文档")
    
    return state


# 节点2：生成回答
def generate_node(state: State) -> State:
    """基于检索结果生成回答"""
    llm = Tongyi(
        model="qwen-turbo",
        api_key=os.getenv("DASHSCOPE_API_KEY")
    )
    
    # 构造提示词
    context = "\n\n".join(state["retrieved_docs"])
    prompt = f"""
基于以下知识库内容回答问题：

知识库：
{context}

问题：{state["question"]}

要求：
1. 只根据知识库内容回答
2. 如果知识库中没有相关信息，明确说明
3. 回答要简洁明了

回答：
"""
    
    print("[生成] AI生成回答中...")
    answer = llm.invoke(prompt)
    
    state["answer"] = answer
    return state


# 构建图
def create_graph():
    """创建RAG图"""
    graph = StateGraph(State)
    
    graph.add_node("retrieve", retrieve_node)
    graph.add_node("generate", generate_node)
    
    graph.set_entry_point("retrieve")
    graph.add_edge("retrieve", "generate")
    graph.add_edge("generate", END)
    
    return graph.compile()


def main():
    """测试RAG系统"""
    app = create_graph()
    
    questions = [
        "公司的年假有多少天？",
        "报销餐饮费有什么限制？",
        "上班时间是几点到几点？",
        "公司的股票期权政策是什么？",  # 知识库中没有
    ]
    
    print("="*60)
    print("RAG问答系统")
    print("="*60)
    
    for question in questions:
        print(f"\n{'='*60}")
        print(f"问题: {question}")
        print(f"{'='*60}")
        
        result = app.invoke({"question": question})
        
        print(f"\n回答: {result['answer']}\n")


if __name__ == "__main__":
    main()