Python RAG 文件问答 Agent

用 Python 实现 RAG（检索增强生成），理解 chunking、embedding、向量检索的完整流程。

#type / howto #status / growing #tech / lang / python #resource / python #tech / ai #resource / agent

[!info] related notes

前置笔记: Python Tool Calling Agent

所属 MOC: 学习 AI Agent MOC

相关概念: Augmented LLM, Embedding 模型, Reranker 模型

相关资源: OpenRouter

Python RAG 文件问答 Agent

目标

实现 RAG（Retrieval Augmented Generation）— 先从文档里搜相关内容，再让模型基于真实内容回答。

和前两天的区别

	Day 2（Tool Agent）	Day 3（RAG Agent）
模型调用	工具返回精确结果（天气、计算）	检索到的文档段落，语义相关但不精确
适用场景	结构化操作	知识问答、文档查询
核心能力	function calling	embedding + 向量搜索

前置条件

uv add sentence-transformers chromadb requests python-dotenv
uv run python main.py

项目结构：

day3-rag-agent/
├── main.py        # Agent 主程序
├── rag.py         # RAG 模块（切块、向量化、存储、检索）
├── data/
│   └── python-basics.md   # 测试文档
└── .env

核心理解：3 个概念

1. Chunking（切块）

一整本书 → 按段落切成 ~300 字的小块。

为什么？LLM 有上下文限制，而且”找到相关的一小段”比”塞进整本书”效果好。

2. Embedding（向量化）

"装饰器是一个函数" → [0.12, -0.34, 0.56, ...]  (384 个数字)

意思相近的句子，生成的数字也相近。这是”语义搜索”的基础。

为什么是 384 维？ 由模型决定。all-MiniLM-L6-v2 输出 384 维，all-mpnet-base-v2 输出 768 维，OpenAI text-embedding-3-small 输出 1536 维。维度越多越精确，但也越慢越大。384 是速度和效果的平衡点。

3. 检索 + 生成（RAG 流程）

用户: "什么是装饰器？"
  → 向量搜索 → 找到最相关的 3 个段落
  → 把段落塞进 prompt: "根据以下资料回答：[段落1][段落2][段落3]"
  → LLM 基于真实资料回答

完整代码

rag.py — RAG 模块

[!note]- 展开查看 rag.py

"""
RAG 模块 — 检索增强生成

负责：
1. 读取文件
2. 切分文本（chunking）
3. 向量化（embedding）
4. 存入向量数据库
5. 检索相关内容

原理：
  "向量"就是一组数字，用来表示文本的"含义"。
  意思相近的文本，向量也相近。
  所以找相关内容 = 找向量最相似的文本块。
"""

import os
from pathlib import Path

# sentence-transformers: 本地运行的 embedding 模型
# 第一次运行会下载模型（约 80MB），之后从缓存加载
from sentence_transformers import SentenceTransformer

# chromadb: 本地向量数据库，不需要启动服务器
import chromadb


# ============================================================
# 1. 文本切分 (Chunking)
# ============================================================

def chunk_text(text: str, chunk_size: int = 300, overlap: int = 50) -> list[str]:
    """
    把长文本切成小块。

    为什么要切？
      - LLM 的 context window 有限，不能把整本书塞进去
      - 太长的文本 embedding 效果差
      - 切块后可以精确定位哪一段和问题相关

    参数:
        text: 原始文本
        chunk_size: 每块的目标字符数
        overlap: 块之间的重叠字符数（防止重要信息被切断）

    返回:
        文本块列表
    """
    # 按段落先拆分
    paragraphs = text.split("\n\n")

    chunks = []
    current_chunk = ""

    for para in paragraphs:
        para = para.strip()
        if not para:
            continue

        # 如果当前块加上新段落不超过限制，就合并
        if len(current_chunk) + len(para) < chunk_size:
            current_chunk += ("\n\n" if current_chunk else "") + para
        else:
            # 当前块已满，保存它
            if current_chunk:
                chunks.append(current_chunk)
            # 新段落太长的话，直接作为一块（或进一步切分）
            if len(para) > chunk_size:
                # 按句子切分长段落
                sentences = para.replace("。", "。\n").replace(". ", ".\n").split("\n")
                temp = ""
                for sent in sentences:
                    if len(temp) + len(sent) < chunk_size:
                        temp += sent
                    else:
                        if temp:
                            chunks.append(temp)
                        temp = sent
                if temp:
                    current_chunk = temp
                else:
                    current_chunk = ""
            else:
                current_chunk = para

    if current_chunk:
        chunks.append(current_chunk)

    return chunks


# ============================================================
# 2. 向量化 + 存储
# ============================================================

class VectorStore:
    """
    向量数据库封装。

    用 ChromaDB 存储文本块和它们的向量。
    提供"搜索相似文本"的功能。
    """

    def __init__(self, collection_name: str = "rag_docs"):
        # 加载 embedding 模型
        # all-MiniLM-L6-v2: 小巧、快速、效果不错
        # 第一次运行会自动下载
        print("正在加载 embedding 模型...")
        self.model = SentenceTransformer("all-MiniLM-L6-v2")

        # 创建 ChromaDB 客户端（数据存在内存中）
        # 如果想持久化，改成: PersistentClient(path="./chroma_db")
        self.client = chromadb.Client()
        self.collection = self.client.get_or_create_collection(
            name=collection_name,
            metadata={"hnsw:space": "cosine"},  # 用余弦相似度
        )
        print("向量数据库就绪。")

    def add_documents(self, chunks: list[str], source: str = "unknown"):
        """
        把文本块加入向量数据库。

        流程:
            文本块 → embedding 模型 → 向量 → 存入 ChromaDB
        """
        if not chunks:
            return

        # 生成向量
        embeddings = self.model.encode(chunks).tolist()

        # 生成 ID
        ids = [f"{source}_chunk_{i}" for i in range(len(chunks))]

        # 存入 ChromaDB
        self.collection.add(
            documents=chunks,
            embeddings=embeddings,
            ids=ids,
            metadatas=[{"source": source} for _ in chunks],
        )

        print(f"已添加 {len(chunks)} 个文本块（来源: {source}）")

    def search(self, query: str, top_k: int = 3) -> list[dict]:
        """
        搜索和问题最相关的文本块。

        流程:
            问题 → embedding → 和数据库中的向量比较 → 返回最相似的 top_k 个

        返回:
            [{"text": "...", "source": "...", "score": 0.85}, ...]
        """
        # 把问题向量化
        query_embedding = self.model.encode([query]).tolist()

        # 在 ChromaDB 中搜索
        results = self.collection.query(
            query_embeddings=query_embedding,
            n_results=top_k,
        )

        # 整理结果
        found = []
        for i in range(len(results["documents"][0])):
            found.append({
                "text": results["documents"][0][i],
                "source": results["metadatas"][0][i].get("source", "unknown"),
                "distance": results["distances"][0][i] if results["distances"] else 0,
            })

        return found


# ============================================================
# 3. 文件加载
# ============================================================

def load_file(file_path: str) -> str:
    """读取文本文件内容"""
    path = Path(file_path)
    if not path.exists():
        raise FileNotFoundError(f"文件不存在: {file_path}")

    with open(path, "r", encoding="utf-8") as f:
        return f.read()


def load_and_index(file_path: str, store: VectorStore) -> int:
    """
    加载文件 → 切块 → 存入向量数据库。

    返回:
        处理的块数
    """
    content = load_file(file_path)
    chunks = chunk_text(content)
    source = Path(file_path).name
    store.add_documents(chunks, source=source)
    return len(chunks)

main.py — Agent 主程序

[!note]- 展开查看 main.py

"""
Day 3: RAG 文件问答 Agent

核心思想：
  不要让模型"猜"答案，先从文档里找相关内容，再让模型基于真实内容回答。
  这就是 RAG — Retrieval Augmented Generation（检索增强生成）。

流程：
  用户提问 → 向量搜索找到相关段落 → 把段落塞进 prompt → 模型回答

和 Day 2 的区别：
  Day 2: 模型调用工具（天气、计算）— 工具返回精确结果
  Day 3: 模型基于检索到的文档段落回答 — 更适合"知识问答"场景
"""

import os
import json
import requests
from dotenv import load_dotenv
from rag import VectorStore, load_and_index

load_dotenv()

api_key = os.environ.get("OPENROUTER_API_KEY")
if not api_key:
    print("Error: 在 .env 文件中设置 OPENROUTER_API_KEY")
    exit(1)

API_URL = "https://openrouter.ai/api/v1/chat/completions"
MODEL = "openai/gpt-oss-120b:free"


def call_llm(system_prompt: str, user_message: str) -> str:
    """
    调用 LLM（非流式）。

    RAG 场景下，system_prompt 会包含检索到的文档内容。
    """
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }
    payload = {
        "model": MODEL,
        "messages": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_message},
        ],
        "stream": False,
    }
    try:
        r = requests.post(API_URL, headers=headers, json=payload, timeout=(10, 60))
        if r.status_code != 200:
            return f"API 错误 (HTTP {r.status_code})"
        data = r.json()
        return data["choices"][0]["message"]["content"]
    except requests.exceptions.RequestException as e:
        return f"网络请求失败：{e}"


def build_rag_prompt(query: str, store: VectorStore, top_k: int = 3) -> str:
    """
    RAG 的核心：构建包含检索结果的 prompt。

    步骤：
    1. 用问题搜索相关文本块
    2. 把文本块拼成"参考资料"
    3. 告诉模型"基于这些资料回答"
    """
    results = store.search(query, top_k=top_k)

    if not results:
        return (
            "你是一个助手。没有找到相关文档内容。"
            "请根据你的知识回答，如果不确定请说明。"
        )

    context_parts = []
    for i, r in enumerate(results, 1):
        context_parts.append(
            f"--- 参考段落 {i}（来源: {r['source']}）---\n{r['text']}"
        )

    context = "\n\n".join(context_parts)

    system_prompt = f"""你是一个知识问答助手。请根据以下参考资料回答用户的问题。

规则：
1. 优先使用参考资料中的内容回答
2. 如果参考资料中没有相关信息，请明确说明
3. 回答时可以引用来源
4. 保持回答简洁准确

参考资料：
{context}
"""
    return system_prompt


# --- 主程序 ---
print("=" * 60)
print("RAG 文件问答 Agent")
print("=" * 60)

store = VectorStore()

print("\n可用命令:")
print("  load <文件路径>  — 加载文档到知识库")
print("  ask <问题>       — 基于文档回答问题")
print("  search <关键词>  — 直接搜索相关段落（不经过 LLM）")
print("  status           — 查看已加载的文档")
print("  quit             — 退出")
print("-" * 60)

while True:
    user_input = input("\n> ").strip()

    if not user_input:
        continue

    if user_input.lower() in ("quit", "exit", "q"):
        print("Goodbye!")
        break

    if user_input.lower().startswith("load "):
        file_path = user_input[5:].strip().strip('"').strip("'")
        try:
            count = load_and_index(file_path, store)
            print(f"✅ 成功加载，共 {count} 个文本块")
        except FileNotFoundError:
            print(f"❌ 文件不存在: {file_path}")
        except Exception as e:
            print(f"❌ 加载失败: {e}")
        continue

    if user_input.lower().startswith("search "):
        query = user_input[7:].strip()
        results = store.search(query, top_k=3)
        if results:
            for i, r in enumerate(results, 1):
                print(f"\n--- 结果 {i} (来源: {r['source']}, 距离: {r['distance']:.4f}) ---")
                print(r["text"][:300])
        else:
            print("没有找到相关内容。")
        continue

    if user_input.lower() == "status":
        count = store.collection.count()
        print(f"向量数据库中有 {count} 个文本块")
        continue

    # 默认：RAG 问答
    query = user_input
    print("正在检索相关文档...")
    system_prompt = build_rag_prompt(query, store)
    print("正在生成回答...\n")
    answer = call_llm(system_prompt, query)
    print(f"Agent: {answer}")

运行效果

> load data/python-basics.md
已添加 11 个文本块（来源: python-basics.md）
✅ 成功加载，共 11 个文本块

> search 装饰器

--- 结果 1 (来源: python-basics.md, 距离: 0.6526) ---
装饰器常用于日志、权限检查、缓存、计时等场景。

> 什么是装饰器？怎么用
正在检索相关文档...
正在生成回答...

Agent: 装饰器（Decorator）是一种在不修改原函数代码的前提下，
为函数"包装"额外功能的高级技巧...

踩坑记录

Chunking 切块质量直接影响检索效果

搜索”列表和字典区别”时，模型回答”资料中没有字典相关内容”，但实际上文档里有。

原因：chunking 把字典那段归到了另一个块里，搜索时没命中。

这正是 RAG 的核心难点：

切太大 → 不精确，搜到的内容混杂
切太小 → 丢失上下文，每块没有完整信息
overlap 太小 → 重要信息在边界被切断

调优方向：调整 chunk_size、overlap，或改用更智能的切分策略（按语义、按标题）。

HF Hub 警告

Warning: You are sending unauthenticated requests to the HF Hub.

Hugging Face 对匿名用户有速率限制。不影响功能，模型已缓存到本地。想消除就去 https://huggingface.co/settings/tokens 创建免费 token，加到 .env：

HF_TOKEN=hf_xxxxxxxxxxxx

关键流程解析

1. RAG = 先搜后答

传统 LLM:    用户提问 → 模型凭记忆回答（可能编造）
RAG:         用户提问 → 搜索文档 → 相关段落塞进 prompt → 模型基于资料回答

2. Embedding 的本质

把文字变成数字，意思相近的文字数字也相近：

"猫" → [0.8, 0.1, 0.9, ...]    ← 和"狗"很近
"狗" → [0.7, 0.2, 0.85, ...]
"汽车" → [0.1, 0.9, 0.05, ...]  ← 和"猫"很远

384 个数字 = 384 个语义维度。all-MiniLM-L6-v2 是速度和效果的平衡点。

3. 向量搜索 vs 关键词搜索

	关键词搜索	向量搜索
匹配方式	精确匹配文字	语义相似度
”装饰器”搜”decorator”	❌ 找不到	✅ 能找到
”怎么写函数”搜”def”	❌ 找不到	✅ 能找到
实现	简单（grep）	需要 embedding 模型

常见问题

ModuleNotFoundError：uv add sentence-transformers chromadb
首次运行很慢：需要下载 embedding 模型（约 80MB），之后从缓存加载
搜不到相关内容：调整 chunk_size 和 top_k，或检查文档内容是否被正确切块
回答说”资料中没有”：可能是 chunking 导致相关段落没被检索到，用 search 命令验证