AI 工具开发实战（2）：开发一个本地 RAG 知识库——丢一个文件夹进去，直接问答

localrag/
├── localrag.py          # CLI 主程序
├── indexer.py           # 文档索引
├── retriever.py         # 检索引擎
├── requirements.txt
└── docs/                # 放文档的文件夹

# requirements.txt
click==8.1.7
openai==1.50.0
python-dotenv==1.0.0
sentence-transformers==3.1.0
PyMuPDF==1.24.0
numpy==1.26.0

# indexer.py
import os
import json
import hashlib
from pathlib import Path
import fitz  # PyMuPDF

class DocumentIndexer:
    """文档索引：读取文件夹、切分文本、向量化。"""

    def __init__(self, chunk_size=512, chunk_overlap=128):
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap

    def load_folder(self, folder_path: str) -> list:
        """加载文件夹中的所有文档。"""
        docs = []
        for f in sorted(Path(folder_path).iterdir()):
            if f.suffix.lower() == '.pdf':
                text = self._parse_pdf(str(f))
            elif f.suffix.lower() in ('.txt', '.md'):
                text = f.read_text(encoding='utf-8', errors='ignore')
            else:
                continue
            if text.strip():
                docs.append({"source": f.name, "content": text})
        return docs

    def _parse_pdf(self, path):
        doc = fitz.open(path)
        text = '\n'.join(p.get_text() for p in doc)
        doc.close()
        return text

    def chunk_documents(self, docs: list) -> list:
        """将文档切分为重叠的 chunk。"""
        chunks = []
        for doc in docs:
            text = doc["content"]
            # 按段落切
            paragraphs = text.split('\n\n')
            buffer = ""
            for para in paragraphs:
                para = para.strip()
                if not para:
                    continue
                if len(buffer) + len(para)  list:
        """检索最相关的 chunk。"""
        q_vec = self.model.encode([query], normalize_embeddings=True)[0]
        scores = np.dot(self.vectors, q_vec)  # 余弦相似度
        top_indices = np.argsort(scores)[-top_k:][::-1]

        results = []
        for idx in top_indices:
            if scores[idx] `")
        return

    # 检索
    results = retriever.search(question, top_k=top)
    if not results:
        click.echo("❌ 没有找到相关内容")
        return

    # 构建 Prompt
    context = "\n\n".join(
        f"[{r['source']}] {r['text'][:500]}" for r in results
    )

    # 调用 LLM
    response = client.chat.completions.create(
        model="deepseek-chat",
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT.format(context=context)},
            {"role": "user", "content": question},
        ],
        temperature=0.3,
        stream=True,
    )

    # 流式输出
    for chunk in response:
        if chunk.choices[0].delta.content:
            click.echo(chunk.choices[0].delta.content, nl=False)
    click.echo()

    # 显示来源
    if show_sources:
        click.echo(f"\n📖 引用来源：")
        for i, r in enumerate(results):
            click.echo(f"  [{i+1}] {r['source']} (相关度: {r['score']:.2f})")


if __name__ == "__main__":
    cli()

# 1. 索引文档
python localrag.py index ./docs
# ✅ 索引完成：5 个文档 → 42 个片段

# 2. 提问
python localrag.py ask "这个项目的架构是什么？"
# 这个项目采用微服务架构，包含 API 服务、数据库、缓存三层...

# 3. 带来源引用
python localrag.py ask "部署流程是怎样的？" --show-sources
# ...（回答内容）...
# 📖 引用来源：
#   [1] deploy.md (相关度: 0.87)
#   [2] architecture.md (相关度: 0.72)