01. 简单 RAG

让我们从最基础的 RAG 实现开始，了解 RAG 系统的核心组件和工作流程。

学习目标

通过本章学习，你将掌握：

• RAG 系统的基本架构
• 文档处理和文本分块
• 向量嵌入的创建和使用
• 语义搜索的实现
• 基于检索结果的问答生成

系统架构

graph TD
    A[PDF文档] --> B[文本提取]
    B --> C[文本分块]
    C --> D[创建嵌入向量]
    D --> E[向量数据库]
    F[用户问题] --> G[问题向量化]
    G --> H[语义搜索]
    E --> H
    H --> I[检索相关文档]
    I --> J[构建提示]
    J --> K[LLM生成回答]

核心组件

1. 文档处理

PDF 文本提取

extract_text_from_pdf.py

import fitz  # PyMuPDF

def extract_text_from_pdf(pdf_path):
    """
    从PDF文件中提取文本内容
    """
    mypdf = fitz.open(pdf_path)
    all_text = ""

    # 遍历PDF的每一页
    for page_num in range(mypdf.page_count):
        page = mypdf[page_num]
        text = page.get_text("text")
        all_text += text

    return all_text

# 使用示例
pdf_path = "data/AI_Information.pdf"
extracted_text = extract_text_from_pdf(pdf_path)
print(f"提取的文本长度: {len(extracted_text)} 字符")

关键点：

• 使用 PyMuPDF 库处理 PDF 文件
• 逐页提取文本并合并
• 支持多种 PDF 格式

文本分块

chunk_text.py

def chunk_text(text, n, overlap):
    """
    将文本分割成固定大小的块，支持重叠

    Args:
        text: 原始文本
        n: 每块的字符数
        overlap: 重叠的字符数
    """
    chunks = []

    # 步长 = 块大小 - 重叠大小
    step = n - overlap

    for i in range(0, len(text), step):
        chunk = text[i:i + n]
        chunks.append(chunk)

    return chunks

# 使用示例
text_chunks = chunk_text(extracted_text, 1000, 200)
print(f"创建了 {len(text_chunks)} 个文本块")
print(f"第一个文本块:\n{text_chunks[0]}")

重叠的重要性：

• 避免句子或概念被截断
• 提供上下文连续性
• 提高检索准确性

2. 向量嵌入

create_embeddings.py

from llama_index.embeddings.huggingface import HuggingFaceEmbedding

def create_embeddings(text, model="BAAI/bge-base-en-v1.5"):
    """
    为文本创建向量嵌入
    """
    embedding_model = HuggingFaceEmbedding(model_name=model)

    if isinstance(text, list):
        # 批量处理文本列表
        response = embedding_model.get_text_embedding_batch(text)
    else:
        # 处理单个文本
        response = embedding_model.get_text_embedding(text)

    return response

# 创建所有文本块的嵌入
chunks_embeddings = create_embeddings(text_chunks)
print(f"嵌入向量维度: {len(chunks_embeddings[0])}")

模型选择：

• BAAI/bge-base-en-v1.5: 英文文本，平衡性能和质量
• BAAI/bge-large-zh-v1.5: 中文文本
• sentence-transformers/all-MiniLM-L6-v2: 轻量级选择

3. 语义搜索

semantic_search.py

import numpy as np

def cosine_similarity(vec1, vec2):
    """
    计算两个向量的余弦相似度
    """
    return np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))

def semantic_search(query, text_chunks, embeddings, k=5):
    """
    基于语义相似度搜索最相关的文档块
    """
    # 为查询创建嵌入向量
    query_embedding = create_embeddings(query)
    similarity_scores = []

    # 计算查询与每个文档块的相似度
    for i, chunk_embedding in enumerate(embeddings):
        similarity = cosine_similarity(
            np.array(query_embedding),
            np.array(chunk_embedding)
        )
        similarity_scores.append((i, similarity))

    # 按相似度排序并返回top-k结果
    similarity_scores.sort(key=lambda x: x[1], reverse=True)
    top_indices = [index for index, _ in similarity_scores[:k]]

    return [text_chunks[index] for index in top_indices]

# 示例搜索
query = "What is artificial intelligence?"
top_chunks = semantic_search(query, text_chunks, chunks_embeddings, k=3)

print(f"查询: {query}")
for i, chunk in enumerate(top_chunks):
    print(f"结果 {i+1}:\n{chunk}\n" + "="*50)

4. 答案生成

generate_response.py

from openai import OpenAI
import os

# 初始化OpenAI客户端
client = OpenAI(
    base_url="https://openrouter.ai/api/v1",
    api_key=os.getenv("OPENROUTER_API_KEY")
)

def generate_response(system_prompt, user_message, model="meta-llama/llama-3.2-3b-instruct:free"):
    """
    基于检索到的上下文生成回答
    """
    response = client.chat.completions.create(
        model=model,
        temperature=0,  # 设置为0以获得确定性回答
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_message}
        ]
    )
    return response

# 系统提示词
system_prompt = """你是一个AI助手，严格基于给定的上下文回答问题。
如果答案无法从提供的上下文中直接得出，请回答："我没有足够的信息来回答这个问题。" """

# 构建用户提示
context = "\n".join([f"上下文 {i+1}:\n{chunk}" for i, chunk in enumerate(top_chunks)])
user_prompt = f"{context}\n\n问题: {query}"

# 生成回答
ai_response = generate_response(system_prompt, user_prompt)
print(f"AI回答: {ai_response.choices[0].message.content}")

完整示例

将所有组件组合成一个完整的 RAG 系统：

simple_rag_complete.py

import fitz
import os
import numpy as np
import json
from openai import OpenAI
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from dotenv import load_dotenv

load_dotenv()

class SimpleRAG:
    def __init__(self, model_name="BAAI/bge-base-en-v1.5"):
        self.embedding_model = HuggingFaceEmbedding(model_name=model_name)
        self.client = OpenAI(
            base_url="https://openrouter.ai/api/v1",
            api_key=os.getenv("OPENROUTER_API_KEY")
        )
        self.text_chunks = []
        self.embeddings = []

    def load_document(self, pdf_path):
        """加载并处理PDF文档"""
        # 提取文本
        text = self.extract_text_from_pdf(pdf_path)

        # 分块
        self.text_chunks = self.chunk_text(text, 1000, 200)

        # 创建嵌入
        self.embeddings = self.create_embeddings(self.text_chunks)

        print(f"已加载文档，创建了 {len(self.text_chunks)} 个文本块")

    def extract_text_from_pdf(self, pdf_path):
        mypdf = fitz.open(pdf_path)
        all_text = ""
        for page_num in range(mypdf.page_count):
            page = mypdf[page_num]
            all_text += page.get_text("text")
        return all_text

    def chunk_text(self, text, n, overlap):
        chunks = []
        for i in range(0, len(text), n - overlap):
            chunks.append(text[i:i + n])
        return chunks

    def create_embeddings(self, texts):
        return self.embedding_model.get_text_embedding_batch(texts)

    def cosine_similarity(self, vec1, vec2):
        return np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))

    def search(self, query, k=3):
        """搜索相关文档"""
        query_embedding = self.embedding_model.get_text_embedding(query)
        similarities = []

        for i, chunk_embedding in enumerate(self.embeddings):
            similarity = self.cosine_similarity(
                np.array(query_embedding),
                np.array(chunk_embedding)
            )
            similarities.append((i, similarity))

        similarities.sort(key=lambda x: x[1], reverse=True)
        top_indices = [i for i, _ in similarities[:k]]

        return [self.text_chunks[i] for i in top_indices]

    def answer(self, query):
        """回答问题"""
        # 检索相关文档
        relevant_chunks = self.search(query, k=3)

        # 构建提示
        context = "\n".join([f"上下文 {i+1}:\n{chunk}"
                           for i, chunk in enumerate(relevant_chunks)])

        system_prompt = """你是一个AI助手，严格基于给定的上下文回答问题。
如果答案无法从提供的上下文中直接得出，请回答："我没有足够的信息来回答这个问题。" """

        user_prompt = f"{context}\n\n问题: {query}"

        # 生成回答
        response = self.client.chat.completions.create(
            model="meta-llama/llama-3.2-3b-instruct:free",
            temperature=0,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt}
            ]
        )

        return response.choices[0].message.content

# 使用示例
if __name__ == "__main__":
    # 创建RAG实例
    rag = SimpleRAG()

    # 加载文档
    rag.load_document("data/AI_Information.pdf")

    # 提问
    question = "What is artificial intelligence?"
    answer = rag.answer(question)

    print(f"问题: {question}")
    print(f"回答: {answer}")

性能优化建议

1. 文本分块优化

• 固定大小 vs 语义分块: 考虑基于句子或段落的分块
• 重叠大小: 通常设置为块大小的 10-20%
• 块大小: 根据模型上下文窗口调整（512-2048 字符）

2. 嵌入模型选择

• 速度 vs 质量: 平衡推理速度和嵌入质量
• 多语言支持: 选择支持目标语言的模型
• 领域适应: 考虑领域特定的嵌入模型

3. 检索优化

• Top-K 选择: 根据问题复杂度调整检索数量
• 阈值过滤: 设置相似度阈值过滤低质量结果
• 多样性: 在相似度和多样性之间平衡

常见问题

Q: 为什么有时候检索结果不准确？

A: 可能的原因：

• 文本分块策略不当，破坏了语义完整性
• 查询与文档的表达方式差异较大
• 嵌入模型不适合当前领域

Q: 如何提高回答质量？

A: 几个建议：

• 改进系统提示词，提供更清晰的指导
• 增加检索的文档数量（Top-K）
• 使用更强大的生成模型
• 对检索结果进行重排序

Q: 系统运行速度慢怎么办？

A: 优化方案：

• 使用更轻量的嵌入模型
• 实现向量索引（如 FAISS）
• 缓存常见查询的结果
• 并行处理文档分块

下一步

恭喜你完成了第一个 RAG 系统！虽然这个版本比较基础，但它包含了 RAG 的所有核心组件。

在下一章 语义分块 中，我们将学习如何改进文档分块策略，提高检索质量。