书接上回，前一篇文章（点击直达）简单说了下向量库相关知识，这边文章主要讲解如何具体实现RAG，简单实现，原理都一样！！！

一、环境准备

1. 安装依赖库

创建 Python 环境，执行以下命令安装核心依赖：

pip install langchain langchain_chroma langchain_community langchain_core
pip install transformers torch sentence-transformers

2. 准备 API Key

前往阿里通义千问开放平台，获取你的DASHSCOPE_API_KEY。

二、核心代码实现

完整源码（更改通过上篇内容代码保存的向量库位置）

import os

from langchain_chroma import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.llms import Tongyi
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser

# ---------------------- 配置 ----------------------
DB_DIR = "./chroma_db/chroma_db_txt"

os.environ["DASHSCOPE_API_KEY"] = "你自己的KEY"


print("加载向量库...")
# ---------------------- 向量库 & 模型 ----------------------
embeddings = HuggingFaceEmbeddings(
    model_name="BAAI/bge-large-zh-v1.5",
    model_kwargs={"device": "cuda"},
    encode_kwargs={"normalize_embeddings": True}
)
db = Chroma(persist_directory=DB_DIR, embedding_function=embeddings)

retriever = db.as_retriever(search_kwargs={"k": 5})

llm = Tongyi(model_name="qwen-turbo", temperature=0.1, max_tokens=1024)

# ---------------------- 提示词 ----------------------
template = """
你是一个专业的问答助手，请根据下面的参考资料回答问题。
如果参考资料中没有答案，请直接说“没有找到相关信息”。

参考资料：
{context}

问题：{question}

请回答，并在最后列出你参考了哪些文件。
回答格式要求：
【回答】
xxx

【参考文件】
xxx
"""

prompt = PromptTemplate.from_template(template)

# ====================== ✅ 修复：把内容 + 文件名一起传给模型 ======================
def format_docs(docs):
    formatted = []
    for doc in docs:
        content = doc.page_content
        filename = doc.metadata.get("filename", "未知文件")
        formatted.append(f"【内容】：{content}\n【来源文件】：{filename}")
    return "\n\n------------------------\n\n".join(formatted)

# 调试检索
def debug_retriever(query):
    print("\n====== 检索到的知识库内容 ======")
    docs = retriever.invoke(query)
    for i, doc in enumerate(docs):
        print(f"[{i+1}] {doc.metadata.get('filename', '未知文件')} | {doc.page_content}")
    print("==================================\n")

# RAG 链
rag_chain = (
    {"context": retriever | format_docs, "question": RunnablePassthrough()}
    | prompt
    | llm
    | StrOutputParser()
)

# ---------------------- 启动 ----------------------
# 增加知识库向量标签，补充回答文本来源
if __name__ == "__main__":
    print("\n 知识库启动成功！输入 exit 退出")
    while True:
        query = input("\n输入：")
        if query.lower() == "exit":
            print("结束")
            break
        debug_retriever(query)
        answer = rag_chain.invoke(query)
        print("\nAI：", answer)

三、代码核心解析

1. Embedding 模型配置

使用BAAI/bge-large-zh-v1.5中文 embedding 模型，支持 GPU 加速，生成的向量相似度更高，检索效果更好。

2. 向量库加载

直接从本地持久化目录加载已构建好的 Chroma 向量库，无需重复构建，高效便捷。

3. 检索器优化

设置k=5，每次检索5 条最相关的知识库文本，平衡精准度与响应速度。

4. 文档格式化（核心亮点）

自定义format_docs函数，将文本内容 + 来源文件名一起传给大模型，实现答案溯源，彻底解决大模型 “不知道自己参考了什么” 的问题！

5. RAG 链构建

通过 LangChain 的链式调用，串联检索→格式化→提示词→大模型→解析，代码简洁、扩展性极强。

6. 调试功能

内置debug_retriever函数，运行时会打印检索到的文本和来源文件，方便排查检索效果。

四、常见问题解决

1. GPU 报错（CUDA 不可用）

解决方案：将 embedding 模型的device改为cpu：

model_kwargs={"device": "cpu"}