200行代码构建私有知识库：RAG + LLM 实战指南

背景：为什么你需要一套会"说话"的知识库

你是否遇到过这些场景：

收藏了上百篇技术文章，需要时却怎么也找不到

公司内部文档越堆越多，新人入职全靠口口相传

自己的笔记散落各处，按关键词搜索全凭运气

传统搜索的问题是匹配关键词而非语义。当你问"如何优化慢查询"时，它不会意识到这和"SQL 性能调优"是同一件事。LLM 时代的解决方案是 RAG（检索增强生成）：先用向量相似度找到最相关的内容，再让 LLM 基于这些内容回答问题。结果远胜于关键词搜索，且完全可控——幻觉问题被大幅减少。本文要构建的正是这样一套系统。

Document Input (PDF / TXT / MD)
       v
   Text Chunking
       v
   Vectorization (Embedding API)
       v
   Store in Vector DB (ChromaDB)
       v
User Query -> Retrieve relevant chunks -> Assemble Prompt -> LLM generates answer

向量化 & 问答模型：通过 WRouter 调用 text-embedding-3-small + gpt-4o

向量数据库：ChromaDB（本地部署，零配置）

文档解析：LangChain Document Loaders

pip install openai chromadb langchain langchain-community tiktoken pypdf

WRouter 兼容 OpenAI SDK——只需替换 base_url，其他代码无需改动：

from openai import OpenAI

client = OpenAI(api_key="Your WRouter API Key", base_url="https://www.wrouter.ai/v1")

WRouter 支持 GPT-4o、Claude 3.5、Gemini 等主流模型，一个 Key 统一管理——无需在各提供商处分别注册账号。注册地址：www.wrouter.ai

第一步：文档加载与分块

from langchain_community.document_loaders import PyPDFLoader, TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
import os

def load_documents(file_path: str):
    """支持 PDF 和纯文本"""
    ext = os.path.splitext(file_path)[-1].lower()
    if ext == ".pdf":
        loader = PyPDFLoader(file_path)
    else:
        loader = TextLoader(file_path, encoding="utf-8")
    return loader.load()

def split_documents(docs, chunk_size=500, chunk_overlap=50):
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        separators=["\n\n", "\n", "。", "！", "？", " ", ""]
    )
    return splitter.split_documents(docs)

如何选择 chunk_size？

太小（< 200 字符）：上下文不足，LLM 无法回答完整问题

太大（> 1000 字符）：检索精度下降，且消耗更多 token

推荐 400–600 字符，保留约 10% 重叠以防止语义断裂

第二步：向量化并存入 ChromaDB

import chromadb
from openai import OpenAI

client = OpenAI(api_key="Your WRouter API Key", base_url="https://www.wrouter.ai/v1")

# 初始化 ChromaDB（数据本地持久化）
chroma_client = chromadb.PersistentClient(path="./chroma_db")
collection = chroma_client.get_or_create_collection(
    name="knowledge_base",
    metadata={"hnsw:space": "cosine"}  # 使用余弦相似度
)

def embed_texts(texts: list[str]) -> list[list[float]]:
    """批量向量化，每次最多 2048 条"""
    response = client.embeddings.create(
        model="text-embedding-3-small",
        input=texts
    )
    return [item.embedding for item in response.data]

def index_documents(chunks):
    """将文档块向量化并存入数据库"""
    batch_size = 100  # 每批处理 100 条
    for i in range(0, len(chunks), batch_size):
        batch = chunks[i:i + batch_size]
        texts = [chunk.page_content for chunk in batch]
        metadatas = [chunk.metadata for chunk in batch]
        ids = [f"chunk_{i + j}" for j in range(len(batch))]

        embeddings = embed_texts(texts)

        collection.add(
            embeddings=embeddings,
            documents=texts,
            metadatas=metadatas,
            ids=ids
        )
        print(f"已处理 {min(i + batch_size, len(chunks))}/{len(chunks)} 个 chunk")

# 使用示例
docs = load_documents("your_document.pdf")
chunks = split_documents(docs)
index_documents(chunks)
print(f"索引构建完成，共 {len(chunks)} 个文档块")

第三步：检索 + 问答

def retrieve(query: str, top_k: int = 5) -> list[str]:
    """检索与问题最相关的文档块"""
    query_embedding = embed_texts([query])[0]

    results = collection.query(
        query_embeddings=[query_embedding],
        n_results=top_k,
        include=["documents", "metadatas", "distances"]
    )
    # 过滤掉相似度过低的结果（余弦距离 > 0.5 表示弱相关）
    docs_with_scores = zip(
        results["documents"][0],
        results["distances"][0]
    )
    return [doc for doc, dist in docs_with_scores if dist < 0.5]

def answer(query: str) -> str:
    """核心 RAG 问答函数"""
    # 1. 检索相关文档
    relevant_chunks = retrieve(query)
    if not relevant_chunks:
        return "抱歉，在知识库中未找到与该问题相关的内容。"

    # 2. 组装上下文
    context = "\n\n---\n\n".join(relevant_chunks)

    # 3. 构建 Prompt
    system_prompt = """你是一个基于知识库回答问题的助手。请严格基于下方提供的上下文进行回答，不要编造信息。如果上下文包含的信息不足，请明确告知用户。"""

    user_prompt = f"""Context: {context}
Question: {query}
请基于以上内容回答。"""

    # 4. 调用 LLM
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ],
        temperature=0.3  # 知识库问答建议使用低温度
    )
    return response.choices[0].message.content

第四步：添加简易命令行界面

def main():
    print("📚 个人知识库已就绪。输入问题进行查询（输入 quit 退出）\n")
    while True:
        query = input("You: ").strip()
        if query.lower() in ("quit", "exit", "q"):
            break
        if not query:
            continue
        print("\nAI: ", end="", flush=True)
        response = answer(query)
        print(response)
        print()

if __name__ == "__main__":
    # 首次运行时构建索引：
    # docs = load_documents("your_document.pdf")
    # chunks = split_documents(docs)
    # index_documents(chunks)

    main()

以一份 50 页的技术文档为例，索引构建约需 8 秒，向量化成本约 ¥0.02（使用 WRouter 调用 text-embedding-3-small）。问答示例：

You: 这份文档提到了哪些性能优化方案？

AI: 根据文档内容，涉及以下性能优化方案：
1. 数据库层：建议在高频查询字段上建立复合索引，并开启查询缓存……
2. 应用层：使用异步处理减少主线程阻塞……
3. 网络层：启用 HTTP/2 并配置合理的 CDN 策略……
（来源：第 23 页、第 31 页）

多文档管理：为每个文档创建独立的 Collection，支持跨库检索

混合检索：结合向量检索 + BM25 关键词检索，优势互补

对话记忆：维护多轮对话历史，支持追问

切换更强模型：将 gpt-4o 换为 claude-3-5-sonnet 处理复杂推理任务——通过 WRouter 改一行配置即可

整个系统的核心代码不到 150 行，却实现了：

自动文档解析与分块

语义向量化与持久化存储

相似度检索 + LLM 问答的完整闭环

对于 API 访问，推荐使用 WRouter（www.wrouter.ai）、Openrouter（www.openrouter.ai）这类 API Gateway。兼容 OpenAI SDK，支持统一调用 GPT-4o、Claude、Gemini 等模型——非常适合个人开发者和小团队快速上手，无需在各模型提供商处分别注册账号和管理 Key。