AI入门实战|Agent骨架与工具调用|Memory记忆模块

向量知识库/长期记忆

对话短期记忆仅保存在内存，程序重启即丢失；向量数据库将文档、问答、知识文本向量化持久化存储，实现长期知识库、离线检索、文档问答

流程：文档加载 → 文本分割 → 向量化 → 存入向量库 → 相似检索 → 结合大模型问答

源码(sample-04-demo)

https://gitee.com/kcnf-python/ai-in-action

关键代码解析

模块分工

TextLoader：加载本地 txt 文档
RecursiveCharacterTextSplitter：长文本分片，避免单次内容过长
OpenAIEmbeddings：将文本转为向量（DeepSeek 兼容 Embedding 接口）
Chroma：向量存储 + 相似检索
RetrievalQA：检索 + 问答链路，检索知识库内容 + 交给大模型回答

持久化原理

persist_directory 指定本地文件夹，向量、元数据全部落地磁盘
程序重启不再重新导入文档，直接加载已有向量库 → 实现长期记忆

为什么能当 “记忆”

普通对话记忆只存对话文本；向量库可以
存储历史知识、业务文档、手册
根据问题语义相似检索，而非关键词匹配
数据永久保存，多轮会话、重启不丢失

代码分析

import os
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_community.document_loaders import TextLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough

# 加载环境变量
load_dotenv()

# ==================== 配置项 ====================
CHROMA_PERSIST_DIR = "./chroma_db"
DOC_PATH = "./knowledge.txt"

# 初始化大模型
llm = ChatOpenAI(
    api_key=os.getenv("DEEPSEEK_API_KEY"),
    base_url=os.getenv("DEEPSEEK_BASE_URL"),
    model="deepseek-chat",
    temperature=0
)

# 初始化向量化模型
embeddings = OpenAIEmbeddings(
    api_key=os.getenv("DEEPSEEK_API_KEY"),
    base_url=os.getenv("DEEPSEEK_BASE_URL")
)

# ==================== 初始化向量库（持久化/长期记忆） ====================
def get_chroma_vector_store():
    if os.path.exists(CHROMA_PERSIST_DIR):
        print("加载已存在的向量知识库...")
        vector_store = Chroma(
            persist_directory=CHROMA_PERSIST_DIR,
            embedding_function=embeddings
        )
    else:
        print("首次运行：加载文档并构建向量库...")
        # 加载文档
        loader = TextLoader(DOC_PATH, encoding="utf-8")
        docs = loader.load()
        # 文本切片
        splitter = RecursiveCharacterTextSplitter(
            chunk_size=200,
            chunk_overlap=20
        )
        split_docs = splitter.split_documents(docs)
        # 存入并持久化
        vector_store = Chroma.from_documents(
            documents=split_docs,
            embedding=embeddings,
            persist_directory=CHROMA_PERSIST_DIR
        )
        vector_store.persist()
        print("文档入库 & 持久化完成！")
    return vector_store

# ==================== 构建检索问答链路（LCEL 新版写法） ====================
def build_qa_chain(vector_store):
    # 转为检索器，取相似度前2条
    retriever = vector_store.as_retriever(k=2)

    # 自定义提示词模板
    prompt = ChatPromptTemplate.from_template("""
请根据下面的知识库内容回答用户问题。
如果知识库没有相关信息，请如实说明。

知识库内容：{context}
用户问题：{question}
""")

    # LCEL 链式组合：检索 -> 填充提示词 -> 模型推理 -> 解析输出
    qa_chain = (
        {"context": retriever, "question": RunnablePassthrough()}
        | prompt
        | llm
        | StrOutputParser()
    )
    return qa_chain

# ==================== 交互入口 ====================
if __name__ == "__main__":
    store = get_chroma_vector_store()
    qa_chain = build_qa_chain(store)

    print("\n===== ChromaDB 知识库问答（输入 exit 退出）=====")
    while True:
        user_q = input("你：").strip()
        if user_q.lower() == "exit":
            print("对话结束")
            break
        answer = qa_chain.invoke(user_q)
        print("助手：", answer, "\n")