RAG实战 || 代码实现流程 || 简洁明了
requirements.txt
langchain
faiss-cpu
sentence-transformers
pypdf
langchain-openai
langchain-community
langchain-huggingface
我们的RAG系统主要分为两个阶段:
阶段一:数据索引 (通过 ingest.py
实现)
此阶段的目标是将非结构化的文档(如PDF, TXT)处理成一个可供快速检索的向量数据库。这是一个离线过程。
工作流程图:
[原始文档 (.txt, .pdf)] -> [1. 文档加载] -> [2. 文本分割] -> [3. 文本嵌入] -> [4. 存入向量数据库]
-
文档加载 (Document Loading):
- 使用
PyPDFLoader
和TextLoader
等工具,从data/
目录加载原始文件内容。
- 使用
-
文本分割 (Text Splitting):
- 由于LLM的上下文窗口长度有限,且为了检索更精确的片段,我们使用
RecursiveCharacterTextSplitter
将长文档切分成更小的文本块(Chunks)。这可以确保每个块都包含有意义的、连贯的信息。
- 由于LLM的上下文窗口长度有限,且为了检索更精确的片段,我们使用
-
文本嵌入 (Text Embedding):
- 这是RAG的核心步骤之一。我们使用
HuggingFaceEmbeddings
加载一个预训练的句子转换模型(sentence-transformers/all-MiniLM-L6-v2
)。 - 该模型将每个文本块转换为一个高维度的数学向量(Embedding)。这些向量能够捕捉文本的语义含义。在向量空间中,意思相近的文本块,其向量距离也更近。
- 这是RAG的核心步骤之一。我们使用
-
存入向量数据库 (Vector Store):
- 我们使用
FAISS
(Facebook AI Similarity Search) 作为向量数据库。 - 它专门为高效的向量相似度搜索而设计。所有文本块的嵌入向量都被存储在这个数据库中,并构建索引以加速后续的检索过程。
- 我们使用
# -*- coding: utf-8 -*-
import os
import sys
from langchain_community.document_loaders import PyPDFLoader, TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISSDATA_PATH = "data/"
DB_FAISS_PATH = "vectorstore/db_faiss"def create_vector_db():"""Creates a vector database from documents in the data path."""# 1. Load documentsdocuments = []for f in os.listdir(DATA_PATH):file_path = os.path.join(DATA_PATH, f)if f.endswith(".pdf"):loader = PyPDFLoader(file_path)documents.extend(loader.load())elif f.endswith(".txt"):loader = TextLoader(file_path, encoding='utf-8')documents.extend(loader.load())if not documents:print(f"Error: No documents found in '{DATA_PATH}'.")print("Please add some .pdf or .txt files to the 'data' directory and try again.")sys.exit(1)print(f"Loaded {len(documents)} document(s).")# 2. Split documents into chunkstext_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)texts = text_splitter.split_documents(documents)if not texts:print("Error: Could not split documents into chunks. Please check the content of your files.")sys.exit(1)print(f"Split documents into {len(texts)} chunks.")# 3. Create embeddings and vector storeprint("Loading embedding model...")embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2",model_kwargs={"device": "cpu"},)print("Creating vector database... This may take a moment.")db = FAISS.from_documents(texts, embeddings)db.save_local(DB_FAISS_PATH)print(f"Vector database created and saved to '{DB_FAISS_PATH}'.")if __name__ == "__main__":create_vector_db()
阶段二:问答与推理 (通过 main.py
实现)
此阶段是用户与系统交互的实时过程。
工作流程图:
[用户问题] -> [1. 问题嵌入] -> [2. 向量检索] -> [3. 上下文增强] -> [4. LLM生成答案] -> [返回给用户]
-
问题嵌入 (Query Embedding):
- 当用户输入一个问题时,系统使用与索引阶段相同的嵌入模型将问题也转换成一个向量。
-
向量检索 (Vector Retrieval):
- 系统拿着问题的向量,在FAISS向量数据库中执行相似度搜索。
- 它的目标是找到与问题向量最接近的N个文本块向量(在我们的代码中,N=2)。这些最相似的文本块被认为是与用户问题最相关的“上下文”。
-
上下文增强 (Context Augmentation):
- 将上一步检索到的文本块内容与用户的原始问题,按照一个预设的
PromptTemplate
格式进行组合。 - 我们的模板大致如下:
“请根据以下上下文回答问题。上下文:{检索到的文本块}。问题:{用户的问题}。答案:”
- 将上一步检索到的文本块内容与用户的原始问题,按照一个预设的
-
LLM生成答案 (Answer Generation):
- 将这个增强后的完整Prompt发送给我们自定义的LLM(
gemini-2.5-pro
)。 - LLM现在不再依赖其内部的、可能过时的知识,而是被强制要求基于我们提供的上下文来生成答案。这极大地提高了答案的准确性,并有效抑制了幻觉。
- 将这个增强后的完整Prompt发送给我们自定义的LLM(
# -*- coding: utf-8 -*-
import os
import sys
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_openai import ChatOpenAI
from langchain.prompts import PromptTemplate
from langchain.chains import RetrievalQA# --- Configuration ---
DB_FAISS_PATH = "vectorstore/db_faiss"
LLM_API_BASE = " #需要自己配置
LLM_API_KEY = ""
LLM_MODEL_NAME = "gemini-2.5-pro"
EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
PROMPT_TEMPLATE = """Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer.{context}Question: {question}
Answer:"""def set_prompt():"""Prompt template for QA retrieval for each vectorstore"""prompt = PromptTemplate(template=PROMPT_TEMPLATE, input_variables=["context", "question"])return promptdef retrieval_qa_chain(llm, prompt, db):"""Creates a retrieval QA chain"""chain = RetrievalQA.from_chain_type(llm=llm,chain_type="stuff",retriever=db.as_retriever(search_kwargs={"k": 2}),return_source_documents=True,chain_type_kwargs={"prompt": prompt},)return chaindef load_llm():"""Loads the LLM"""llm = ChatOpenAI(model=LLM_MODEL_NAME,base_url=LLM_API_BASE,api_key=LLM_API_KEY,temperature=0.7,)return llmdef qa_bot():"""Initializes the QA bot"""if not os.path.exists(DB_FAISS_PATH):print(f"Error: Vector database not found at '{DB_FAISS_PATH}'.")print("Please run 'python ingest.py' first to create the database.")sys.exit(1)embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL,model_kwargs={"device": "cpu"},)# The allow_dangerous_deserialization argument is needed for FAISS.load_localdb = FAISS.load_local(DB_FAISS_PATH, embeddings, allow_dangerous_deserialization=True)llm = load_llm()qa_prompt = set_prompt()qa = retrieval_qa_chain(llm, qa_prompt, db)return qadef main():"""Main function to run the bot"""chain = qa_bot()print("你好!我是你的RAG问答机器人。输入 'exit' 来退出。")while True:query = input("请输入你的问题: ")if query.lower() == "exit":breakres = chain.invoke({"query": query})answer, docs = res["result"], res["source_documents"]print("--- 答案 ---")print(answer)if docs:print("--- 参考来源 ---")for document in docs:print(f"来源: {document.metadata.get('source', 'N/A')}")if __name__ == "__main__":main()