模块8: RAG — 检索增强生成

LLM有一个知识截止日期 — 它们只知道训练数据中包含的内容。它们无法回答关于截止日期之后的事件、您公司的内部文档或任何从未训练过的私有数据的问题。它们在不确定的时候还会自信地产生幻觉。

检索增强生成（RAG）通过在查询时检索相关文档并将其注入LLM的上下文来解决这个问题。模型不是依赖记忆中的知识，而是阅读实际的源材料并生成基于该材料的答案。

RAG是生产AI系统中部署最广泛的模式，因为它比微调更简单、更便宜，可以即时更新（只需添加或删除文档），并且提供回溯到源材料的引用。

RAG系统有两个阶段：索引阶段（离线，运行一次或定期运行）和查询阶段（在线，每个用户问题运行一次）。理解这两个阶段对于构建有效的管道至关重要。

索引阶段（离线）

1. 加载 — 从文件、数据库、API或网页读取文档
2. 分块 — 将文档分割成更小的、有意义的片段
3. 嵌入 — 使用嵌入模型将每个片段转换为向量（一组数字）
4. 存储 — 将向量和原始文本保存到向量数据库中

查询阶段（在线）

1. 嵌入查询 — 使用相同的嵌入模型将用户的问题转换为向量
2. 检索 — 使用向量相似性搜索找到最相似的文档片段
3. 增强 — 将检索到的片段作为上下文插入LLM提示中
4. 生成 — LLM阅读上下文并产生基于事实的答案

分块可以说是RAG管道中最重要的步骤。如果您的片段太大，它们会浪费上下文token并稀释相关性。如果太小，它们会丢失上下文，模型无法综合出连贯的答案。

目标是创建自包含的（独立就有意义）、专注的（关于一个主题）且大小适当的（通常200-1000个token）片段。

带重叠的分块

def chunk_text(text: str, chunk_size: int = 500, overlap: int = 50) -> list[str]:
    """Split text into overlapping chunks by character count."""
    chunks = []
    start = 0
    while start < len(text):
        end = start + chunk_size
        chunk = text[start:end]
        # Try to end at a sentence boundary
        last_period = chunk.rfind(". ")
        if last_period > chunk_size * 0.5:
            end = start + last_period + 1
            chunk = text[start:end]
        chunks.append(chunk.strip())
        start = end - overlap  # overlap for context continuity
    return chunks

# Example
text = open("company_handbook.txt").read()
chunks = chunk_text(text, chunk_size=800, overlap=100)
print(f"Created {len(chunks)} chunks from {len(text)} characters")

生成嵌入

# Using OpenAI embeddings (most popular choice)
import openai

client = openai.OpenAI()

def get_embeddings(texts: list[str]) -> list[list[float]]:
    """Generate embeddings for a batch of texts."""
    response = client.embeddings.create(
        model="text-embedding-3-small",  # 1536 dimensions, cheap
        input=texts
    )
    return [item.embedding for item in response.data]

# Embed a single query
query_embedding = get_embeddings(["What is the refund policy?"])[0]

# Embed all chunks
chunk_embeddings = get_embeddings(chunks)

向量数据库专为存储和搜索高维向量而构建。它们使用专门的索引算法（如HNSW或IVF），即使在数百万个向量上也能快速进行相似性搜索。

设置ChromaDB

import chromadb

# Persistent storage (survives restarts)
client = chromadb.PersistentClient(path="./chroma_db")

# Create a collection
collection = client.get_or_create_collection(
    name="company_docs",
    metadata={"hnsw:space": "cosine"}  # cosine similarity
)

# Index documents
collection.add(
    documents=chunks,                          # raw text
    ids=[f"chunk_{i}" for i in range(len(chunks))],  # unique IDs
    metadatas=[{"source": "handbook.pdf", "page": i // 5}
               for i in range(len(chunks))]     # metadata for filtering
)

# Query
results = collection.query(
    query_texts=["What is the vacation policy?"],
    n_results=5,
    where={"source": "handbook.pdf"}  # optional metadata filter
)

for doc, score in zip(results["documents"][0], results["distances"][0]):
    print(f"[{score:.3f}] {doc[:100]}...")

让我们构建一个生产质量的RAG管道，索引一组文档并从中回答问题。此示例使用ChromaDB进行存储，Anthropic的Claude进行生成。

import anthropic
import chromadb
import os
import json

class RAGPipeline:
    """Complete RAG pipeline: index, retrieve, generate."""

    def __init__(self, collection_name: str = "knowledge_base"):
        self.ai = anthropic.Anthropic()
        self.db = chromadb.PersistentClient(path="./rag_db")
        self.collection = self.db.get_or_create_collection(collection_name)

    def index_document(self, text: str, source: str, chunk_size: int = 800):
        """Chunk and index a document."""
        chunks = chunk_text(text, chunk_size=chunk_size, overlap=100)
        self.collection.add(
            documents=chunks,
            ids=[f"{source}_{i}" for i in range(len(chunks))],
            metadatas=[{"source": source, "chunk_index": i}
                       for i in range(len(chunks))]
        )
        print(f"Indexed {len(chunks)} chunks from '{source}'")

    def index_directory(self, dir_path: str):
        """Index all .txt and .md files in a directory."""
        for fname in os.listdir(dir_path):
            if fname.endswith((".txt", ".md")):
                path = os.path.join(dir_path, fname)
                text = open(path, "r", encoding="utf-8").read()
                self.index_document(text, source=fname)

    def retrieve(self, query: str, n_results: int = 5) -> list[dict]:
        """Retrieve relevant chunks for a query."""
        if self.collection.count() == 0:
            return []
        results = self.collection.query(
            query_texts=[query],
            n_results=min(n_results, self.collection.count())
        )
        return [
            {"text": doc, "source": meta["source"], "score": dist}
            for doc, meta, dist in zip(
                results["documents"][0],
                results["metadatas"][0],
                results["distances"][0]
            )
        ]

    def query(self, question: str, n_results: int = 5) -> dict:
        """Full RAG: retrieve context, then generate an answer."""
        chunks = self.retrieve(question, n_results)

        if not chunks:
            return {"answer": "No documents indexed yet.", "sources": []}

        # Build context string with source attribution
        context_parts = []
        for i, c in enumerate(chunks, 1):
            context_parts.append(f"[Source {i}: {c['source']}]\n{c['text']}")
        context = "\n\n---\n\n".join(context_parts)

        system_prompt = f"""You are a helpful assistant that answers questions
based ONLY on the provided context. If the context does not contain
the answer, say "I don't have enough information to answer that."
Always cite which source(s) you used, e.g., [Source 1].

Context:
{context}"""

        response = self.ai.messages.create(
            model="claude-sonnet-4-20250514",
            max_tokens=1024,
            system=system_prompt,
            messages=[{"role": "user", "content": question}]
        )

        return {
            "answer": response.content[0].text,
            "sources": [c["source"] for c in chunks],
            "chunks_used": len(chunks)
        }

# Usage
rag = RAGPipeline()
rag.index_directory("./company_docs")
result = rag.query("What is our remote work policy?")
print(result["answer"])
print("Sources:", result["sources"])

衡量RAG质量需要分别评估检索步骤和生成步骤。一个常见的错误是只评估最终答案，而不了解失败是来自糟糕的检索还是糟糕的生成。

简单评估脚本

def evaluate_rag(pipeline: RAGPipeline, test_cases: list[dict]) -> dict:
    """Evaluate RAG pipeline on test cases.
    Each test case: {"question": str, "expected_source": str, "expected_keywords": [str]}
    """
    results = {"retrieval_hits": 0, "answer_relevant": 0, "total": len(test_cases)}

    for tc in test_cases:
        result = pipeline.query(tc["question"])

        # Check retrieval: did we find the right source?
        if tc["expected_source"] in result["sources"]:
            results["retrieval_hits"] += 1

        # Check answer: does it contain expected keywords?
        answer_lower = result["answer"].lower()
        if all(kw.lower() in answer_lower for kw in tc["expected_keywords"]):
            results["answer_relevant"] += 1

    results["retrieval_accuracy"] = results["retrieval_hits"] / results["total"]
    results["answer_accuracy"] = results["answer_relevant"] / results["total"]
    return results

# Example test cases
test_cases = [
    {"question": "What is the vacation policy?",
     "expected_source": "hr_handbook.txt",
     "expected_keywords": ["days", "annual"]},
    {"question": "How do I submit expenses?",
     "expected_source": "finance_guide.txt",
     "expected_keywords": ["receipt", "submit"]},
]
metrics = evaluate_rag(rag, test_cases)
print(f"Retrieval accuracy: {metrics['retrieval_accuracy']:.0%}")
print(f"Answer accuracy: {metrics['answer_accuracy']:.0%}")