构建本地 RAG 系统

本地 RAG 原理

RAG（Retrieval-Augmented Generation，检索增强生成）解决了大模型的两大局限：知识截止日期和无法访问私有数据。本地 RAG 系统完全运行在本机，无任何数据外泄风险：

环境准备：安装依赖与嵌入模型

# 安装 Python 依赖
pip install chromadb ollama pypdf markdown

# 拉取嵌入模型（必须）
ollama pull nomic-embed-text    # 推荐：体积小(274MB)，速度快，中英双语
ollama pull mxbai-embed-large   # 备选：更高质量，1024维，中文稍弱

# 拉取生成模型
ollama pull qwen2.5:7b   # 中文 RAG 推荐

# 验证嵌入模型可用
curl http://localhost:11434/api/embeddings \
  -d '{"model": "nomic-embed-text", "prompt": "测试"}' | \
  python3 -c "import sys,json; d=json.load(sys.stdin); print(f'维度: {len(d[\"embedding\"])}')"

完整 RAG 系统实现

"""
local_rag.py — 完整的本地 RAG 知识库问答系统
依赖：pip install chromadb ollama pypdf
"""
import ollama
import chromadb
from pathlib import Path
from typing import Optional
import hashlib

# ── 配置 ──────────────────────────────────────────────
EMBED_MODEL = "nomic-embed-text"
CHAT_MODEL  = "qwen2.5:7b"
CHUNK_SIZE  = 500   # 每块字符数
CHUNK_OVERLAP = 50  # 相邻块重叠字符数
TOP_K = 5           # 检索返回最相关的 k 个块

# ── 初始化 ChromaDB（持久化到本地磁盘）─────────────────
db_client = chromadb.PersistentClient(path="./rag_db")
collection = db_client.get_or_create_collection(
    name="knowledge_base",
    metadata={"hnsw:space": "cosine"}  # 使用余弦相似度
)

# ── 步骤1：文档加载 ──────────────────────────────────
def load_document(file_path: str) -> str:
    """支持 txt、md、pdf 格式的文档加载。"""
    path = Path(file_path)
    if path.suffix == ".txt":
        return path.read_text(encoding="utf-8")
    elif path.suffix == ".md":
        return path.read_text(encoding="utf-8")
    elif path.suffix == ".pdf":
        import pypdf
        reader = pypdf.PdfReader(file_path)
        return " ".join(page.extract_text() for page in reader.pages)
    else:
        raise ValueError(f"不支持的文件格式: {path.suffix}")

# ── 步骤2：文本分块 ──────────────────────────────────
def chunk_text(text: str, size: int = CHUNK_SIZE, overlap: int = CHUNK_OVERLAP) -> list[str]:
    """按字符数分块，相邻块有重叠。"""
    chunks = []
    start = 0
    while start < len(text):
        end = min(start + size, len(text))
        chunk = text[start:end].strip()
        if chunk:
            chunks.append(chunk)
        start += size - overlap
    return chunks

# ── 步骤3：嵌入并存入向量库 ─────────────────────────
def add_document(file_path: str):
    """加载文档，分块，嵌入，存入 ChromaDB。"""
    print(f"正在处理：{file_path}")
    text = load_document(file_path)
    chunks = chunk_text(text)
    print(f"  分块数量：{len(chunks)}")

    for i, chunk in enumerate(chunks):
        # 生成向量嵌入
        resp = ollama.embeddings(model=EMBED_MODEL, prompt=chunk)
        embedding = resp["embedding"]

        # 生成唯一 ID（文件名 + 块序号）
        doc_id = f"{hashlib.md5(file_path.encode()).hexdigest()[:8]}_chunk_{i}"

        # 存入向量库
        collection.upsert(
            ids=[doc_id],
            embeddings=[embedding],
            documents=[chunk],
            metadatas=[{"source": file_path, "chunk_index": i}]
        )

    print(f"  ✓ 已添加 {len(chunks)} 个块到知识库")

# ── 步骤4：检索相关块 ───────────────────────────────
def retrieve(query: str, k: int = TOP_K) -> list[str]:
    """将查询嵌入，在向量库中检索最相关的 k 个文本块。"""
    query_emb = ollama.embeddings(model=EMBED_MODEL, prompt=query)["embedding"]
    results = collection.query(
        query_embeddings=[query_emb],
        n_results=k,
        include=["documents", "distances", "metadatas"]
    )
    return results["documents"][0]  # 返回文本块列表

# ── 步骤5：RAG 问答 ──────────────────────────────────
def rag_query(question: str) -> str:
    """检索相关文档，生成基于上下文的答案。"""
    # 检索相关文本块
    relevant_chunks = retrieve(question)
    context = "\n\n---\n\n".join(relevant_chunks)

    # 构建 RAG Prompt
    prompt = f"""请基于以下文档内容回答用户的问题。
如果文档中没有相关信息，请明确说明"文档中未找到相关信息"，不要凭空捏造。

【文档内容】
{context}

【用户问题】
{question}

【回答】"""

    response = ollama.generate(
        model=CHAT_MODEL,
        prompt=prompt,
        options={"temperature": 0.3}
    )
    return response["response"]

# ── 主程序 ──────────────────────────────────────────
if __name__ == "__main__":
    # 添加文档到知识库
    add_document("./docs/company_handbook.pdf")
    add_document("./docs/product_manual.md")

    # 查询
    questions = [
        "公司的休假政策是什么？",
        "产品如何重置出厂设置？",
    ]
    for q in questions:
        print(f"\n❓ {q}")
        print(f"💡 {rag_query(q)}\n")

LangChain 与 LlamaIndex 集成

流行的 RAG 框架都内置了对 Ollama 的支持，几行代码即可集成：

# ── LangChain + Ollama ──────────────────────────────
# pip install langchain langchain-ollama langchain-chroma
from langchain_ollama import OllamaLLM, OllamaEmbeddings
from langchain_chroma import Chroma
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.chains import RetrievalQA

# 初始化模型
llm = OllamaLLM(model="qwen2.5:7b")
embeddings = OllamaEmbeddings(model="nomic-embed-text")

# 文档分块
splitter = RecursiveCharacterTextSplitter(
    chunk_size=500, chunk_overlap=50, separators=["\n\n", "\n", "。", " "]
)
with open("./docs/handbook.txt") as f:
    docs = splitter.create_documents([f.read()])

# 创建向量库并存入
vectorstore = Chroma.from_documents(
    docs, embeddings, persist_directory="./lc_rag_db"
)

# 创建 RAG Chain
qa = RetrievalQA.from_chain_type(
    llm=llm,
    retriever=vectorstore.as_retriever(search_kwargs={"k": 5})
)

result = qa.invoke({"query": "员工福利有哪些？"})
print(result["result"])