完整本地 RAG 检索增强生成实战代码

技术民工之路

235人浏览 · 2026-07-12 11:16:47

技术民工之路 · 2026-07-12 11:16:47 发布

方案说明

依赖：langchain + Ollama本地大模型 + Chroma向量库 + sentence-transformers本地向量化
无需付费API，全部本地运行，私有化知识库问答
流程：文档加载 → 文本分割 → 向量入库 → 相似度检索 → 大模型结合上下文回答
支持TXT文档，可扩展PDF/Word

一、环境安装

# 安装依赖包
pip install langchain langchain-community chromadb sentence-transformers ollama
pip install pypdf  # 如需读取PDF额外安装

前置准备

下载安装Ollama
终端拉取轻量大模型（任选其一）

ollama pull qwen:7b    # 通义千问7B中文推荐
ollama pull llama3:8b

二、完整RAG可运行代码

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import SentenceTransformerEmbeddings
from langchain_community.llms import Ollama
from langchain.chains import RetrievalQA

# ====================== 1. 全局配置 ======================
# 本地向量模型（轻量中文嵌入模型）
embedding_model = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
# 本地大模型名称（和ollama pull对应）
llm_model_name = "qwen:7b"
# 向量数据库持久化存储路径
chroma_persist_path = "./chroma_db"
# 知识库文档路径，新建test.txt放入你的资料
doc_path = "./test.txt"

# ====================== 2. 加载并切割文档 ======================
def load_and_split_docs(file_path):
    # 加载txt文档
    loader = TextLoader(file_path, encoding="utf-8")
    documents = loader.load()

    # 文本分割器，防止上下文过长
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=300,        # 单段文本长度
        chunk_overlap=50,      # 分段重叠，保留上下文关联
        separators=["\n\n", "\n", "。", "，", " "]
    )
    split_docs = text_splitter.split_documents(documents)
    print(f"文档切割完成，共 {len(split_docs)} 个文本片段")
    return split_docs

# ====================== 3. 构建/加载向量数据库 ======================
def create_chroma_vector_db(docs):
    # 创建向量库并持久化保存
    vector_db = Chroma.from_documents(
        documents=docs,
        embedding=embedding_model,
        persist_directory=chroma_persist_path
    )
    vector_db.persist()
    print("向量数据库创建完成并保存到本地")
    return vector_db

def load_exist_vector_db():
    # 读取已存在的向量库，无需重复向量化文档
    vector_db = Chroma(
        persist_directory=chroma_persist_path,
        embedding_function=embedding_model
    )
    print("读取本地已有向量数据库")
    return vector_db

# ====================== 4. 初始化RAG问答链 ======================
def build_rag_chain(vector_db):
    # 初始化本地Ollama大模型
    llm = Ollama(model=llm_model_name, temperature=0.1)  # temperature越低回答越严谨

    # 检索器：查询时返回top3最相似文档片段
    retriever = vector_db.as_retriever(search_kwargs={"k": 3})

    # 构建检索问答链路
    rag_qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",  # 简单模式：把检索到的上下文全部传入prompt
        retriever=retriever,
        return_source_documents=True  # 返回参考原文片段，方便溯源
    )
    return rag_qa_chain

# ====================== 5. 主问答交互逻辑 ======================
def main():
    # 第一步：加载文档、切片、生成向量库（首次运行执行，后续注释掉）
    docs = load_and_split_docs(doc_path)
    vector_db = create_chroma_vector_db(docs)

    # 【后续运行直接读取库，注释上面两行，启用下面一行】
    # vector_db = load_exist_vector_db()

    # 初始化问答链
    qa_chain = build_rag_chain(vector_db)

    print("===== 本地RAG知识库问答机器人启动完成 =====")
    print("输入exit退出对话\n")

    # 循环问答
    while True:
        user_query = input("请输入你的问题：")
        if user_query.lower() == "exit":
            print("程序退出")
            break

        # 调用RAG链路
        result = qa_chain.invoke({"query": user_query})
        answer = result["result"]
        source_docs = result["source_documents"]

        print("\n【AI回答】")
        print(answer)

        # 打印引用的知识库原文片段
        print("\n【参考文档片段】")
        for idx, src in enumerate(source_docs):
            print(f"{idx+1}. {src.page_content}\n")

if __name__ == "__main__":
    main()

三、使用步骤

在代码同目录新建 test.txt，写入你的知识库内容（小说、笔记、产品手册、教程等）
打开终端，先启动Ollama（后台保持运行）
首次运行代码：自动加载文档、切割、生成向量库保存到 chroma_db 文件夹
第二次及以后运行：注释文档加载创建库代码，启用 load_exist_vector_db 加速启动
输入问题即可基于你的文档回答，大模型不会胡说八道

四、扩展升级功能

1. 支持PDF文件读取

替换 load_and_split_docs 函数加载器

from langchain_community.document_loaders import PyPDFLoader
def load_pdf(file_path):
    loader = PyPDFLoader(file_path)
    docs = loader.load()
    return docs

2. 自定义Prompt（约束大模型只基于文档回答）

from langchain.prompts import PromptTemplate

# 自定义提示词模板
prompt_template = """
你是专业知识库助手，仅能使用下方提供的参考文档内容回答用户问题。
如果文档中没有相关信息，直接回答：知识库暂无相关内容，不要编造信息。

参考文档：
{context}

用户问题：{question}
"""
prompt = PromptTemplate(input_variables=["context", "question"], template=prompt_template)

# 构建链时传入自定义prompt
rag_qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff",
    retriever=retriever,
    return_source_documents=True,
    chain_type_kwargs={"prompt": prompt}
)

3. 网页端Web服务（FastAPI封装RAG接口）

pip install fastapi uvicorn

新增接口文件 rag_api.py，可前端对接调用：

from fastapi import FastAPI
import uvicorn
app = FastAPI()

# 复用上面代码中初始化好的qa_chain
@app.post("/chat")
def chat(query: str):
    res = qa_chain.invoke({"query": query})
    return {
        "answer": res["result"],
        "reference": [doc.page_content for doc in res["source_documents"]]
    }

if __name__ == "__main__":
    uvicorn.run("rag_api:app", host="0.0.0.0", port=8000)

启动命令：uvicorn rag_api:app --reload

五、常见问题解决

Ollama连接失败：确认终端执行ollama serve开启服务
模型内存不足：更换更小模型 ollama pull qwen:4b
中文乱码：文件保存为UTF-8编码
回答偏离文档：降低temperature=0，使用自定义Prompt限制幻觉# 完整本地RAG检索增强生成实战代码

方案说明

依赖：langchain + Ollama本地大模型 + Chroma向量库 + sentence-transformers本地向量化
无需付费API，全部本地运行，私有化知识库问答
流程：文档加载 → 文本分割 → 向量入库 → 相似度检索 → 大模型结合上下文回答
支持TXT文档，可扩展PDF/Word

一、环境安装

# 安装依赖包
pip install langchain langchain-community chromadb sentence-transformers ollama
pip install pypdf  # 如需读取PDF额外安装

前置准备

下载安装Ollama：https://ollama.com/
终端拉取轻量大模型（任选其一）

ollama pull qwen:7b    # 通义千问7B中文推荐
ollama pull llama3:8b

二、完整RAG可运行代码

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import SentenceTransformerEmbeddings
from langchain_community.llms import Ollama
from langchain.chains import RetrievalQA

# ====================== 1. 全局配置 ======================
# 本地向量模型（轻量中文嵌入模型）
embedding_model = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
# 本地大模型名称（和ollama pull对应）
llm_model_name = "qwen:7b"
# 向量数据库持久化存储路径
chroma_persist_path = "./chroma_db"
# 知识库文档路径，新建test.txt放入你的资料
doc_path = "./test.txt"

# ====================== 2. 加载并切割文档 ======================
def load_and_split_docs(file_path):
    # 加载txt文档
    loader = TextLoader(file_path, encoding="utf-8")
    documents = loader.load()

    # 文本分割器，防止上下文过长
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=300,        # 单段文本长度
        chunk_overlap=50,      # 分段重叠，保留上下文关联
        separators=["\n\n", "\n", "。", "，", " "]
    )
    split_docs = text_splitter.split_documents(documents)
    print(f"文档切割完成，共 {len(split_docs)} 个文本片段")
    return split_docs

# ====================== 3. 构建/加载向量数据库 ======================
def create_chroma_vector_db(docs):
    # 创建向量库并持久化保存
    vector_db = Chroma.from_documents(
        documents=docs,
        embedding=embedding_model,
        persist_directory=chroma_persist_path
    )
    vector_db.persist()
    print("向量数据库创建完成并保存到本地")
    return vector_db

def load_exist_vector_db():
    # 读取已存在的向量库，无需重复向量化文档
    vector_db = Chroma(
        persist_directory=chroma_persist_path,
        embedding_function=embedding_model
    )
    print("读取本地已有向量数据库")
    return vector_db

# ====================== 4. 初始化RAG问答链 ======================
def build_rag_chain(vector_db):
    # 初始化本地Ollama大模型
    llm = Ollama(model=llm_model_name, temperature=0.1)  # temperature越低回答越严谨

    # 检索器：查询时返回top3最相似文档片段
    retriever = vector_db.as_retriever(search_kwargs={"k": 3})

    # 构建检索问答链路
    rag_qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",  # 简单模式：把检索到的上下文全部传入prompt
        retriever=retriever,
        return_source_documents=True  # 返回参考原文片段，方便溯源
    )
    return rag_qa_chain

# ====================== 5. 主问答交互逻辑 ======================
def main():
    # 第一步：加载文档、切片、生成向量库（首次运行执行，后续注释掉）
    docs = load_and_split_docs(doc_path)
    vector_db = create_chroma_vector_db(docs)

    # 【后续运行直接读取库，注释上面两行，启用下面一行】
    # vector_db = load_exist_vector_db()

    # 初始化问答链
    qa_chain = build_rag_chain(vector_db)

    print("===== 本地RAG知识库问答机器人启动完成 =====")
    print("输入exit退出对话\n")

    # 循环问答
    while True:
        user_query = input("请输入你的问题：")
        if user_query.lower() == "exit":
            print("程序退出")
            break

        # 调用RAG链路
        result = qa_chain.invoke({"query": user_query})
        answer = result["result"]
        source_docs = result["source_documents"]

        print("\n【AI回答】")
        print(answer)

        # 打印引用的知识库原文片段
        print("\n【参考文档片段】")
        for idx, src in enumerate(source_docs):
            print(f"{idx+1}. {src.page_content}\n")

if __name__ == "__main__":
    main()

三、使用步骤

在代码同目录新建 test.txt，写入你的知识库内容（小说、笔记、产品手册、教程等）
打开终端，先启动Ollama（后台保持运行）
首次运行代码：自动加载文档、切割、生成向量库保存到 chroma_db 文件夹
第二次及以后运行：注释文档加载创建库代码，启用 load_exist_vector_db 加速启动
输入问题即可基于你的文档回答，大模型不会胡说八道

四、扩展升级功能

1. 支持PDF文件读取

替换 load_and_split_docs 函数加载器

from langchain_community.document_loaders import PyPDFLoader
def load_pdf(file_path):
    loader = PyPDFLoader(file_path)
    docs = loader.load()
    return docs

2. 自定义Prompt（约束大模型只基于文档回答）

from langchain.prompts import PromptTemplate

# 自定义提示词模板
prompt_template = """
你是专业知识库助手，仅能使用下方提供的参考文档内容回答用户问题。
如果文档中没有相关信息，直接回答：知识库暂无相关内容，不要编造信息。

参考文档：
{context}

用户问题：{question}
"""
prompt = PromptTemplate(input_variables=["context", "question"], template=prompt_template)

# 构建链时传入自定义prompt
rag_qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff",
    retriever=retriever,
    return_source_documents=True,
    chain_type_kwargs={"prompt": prompt}
)

3. 网页端Web服务（FastAPI封装RAG接口）

pip install fastapi uvicorn

新增接口文件 rag_api.py，可前端对接调用：

from fastapi import FastAPI
import uvicorn
app = FastAPI()

# 复用上面代码中初始化好的qa_chain
@app.post("/chat")
def chat(query: str):
    res = qa_chain.invoke({"query": query})
    return {
        "answer": res["result"],
        "reference": [doc.page_content for doc in res["source_documents"]]
    }

if __name__ == "__main__":
    uvicorn.run("rag_api:app", host="0.0.0.0", port=8000)