{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 13 RAG 构建\n",
    "\n",
    "## 学习目标\n",
    "1. 掌握从零构建完整 RAG 应用的整体流程\n",
    "2. 学会将文档加载、分割、向量化、存储、检索、生成串联成系统\n",
    "3. 理解 RAG 系统的评估方法和常见优化策略\n",
    "4. 能够将 RAG 代码模块化封装，便于复用\n",
    "5. 完成一个可运行的多文档问答系统"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. RAG 应用整体架构\n",
    "\n",
    "一个可落地的 RAG 应用通常包含以下模块：\n",
    "\n",
    "```\n",
    "数据准备阶段：\n",
    "  原始文档 -> 文档加载器 -> 文本分割器 -> 文本块\n",
    "\n",
    "索引阶段：\n",
    "  文本块 -> Embedding 模型 -> 向量 -> 向量数据库\n",
    "\n",
    "查询阶段：\n",
    "  用户问题 -> Embedding 模型 -> 向量检索 -> 相关文本块\n",
    "  相关文本块 + 用户问题 -> Prompt -> LLM -> 最终回答\n",
    "```\n",
    "\n",
    "本节课将以「公司产品手册问答系统」为例，完整演示上述流程。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. 准备知识库文档\n",
    "\n",
    "先创建两份示例文档，模拟真实的产品手册。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# 创建产品手册目录\n",
    "os.makedirs('product_docs', exist_ok=True)\n",
    "\n",
    "# 文档 1：智能手表使用手册\n",
    "watch_manual = \"\"\"\n",
    "智能手表使用手册\n",
    "\n",
    "1. 开机与配对\n",
    "长按右侧电源键 3 秒开机。首次使用需在手机上下载 HealthApp，打开蓝牙后搜索设备并完成配对。\n",
    "\n",
    "2. 心率监测\n",
    "手表背面配备 PPG 心率传感器。在表盘界面下滑进入功能菜单，点击「心率」即可开始测量。测量时请保持手腕静止。\n",
    "\n",
    "3. 睡眠监测\n",
    "佩戴手表入睡后，系统会自动记录睡眠数据。次日可在 HealthApp 中查看深睡、浅睡和 REM 睡眠时长。\n",
    "\n",
    "4. 防水说明\n",
    "本手表支持 5ATM 防水，可佩戴游泳，但不适用于潜水、热水淋浴或桑拿。\n",
    "\"\"\"\n",
    "\n",
    "# 文档 2：无线耳机使用手册\n",
    "earphone_manual = \"\"\"\n",
    "无线耳机使用手册\n",
    "\n",
    "1. 开机与配对\n",
    "打开充电盒，耳机自动开机并进入配对模式。在手机蓝牙设置中选择 SoundPod 即可完成配对。\n",
    "\n",
    "2. 触控操作\n",
    "单击左耳或右耳可播放/暂停音乐。双击右耳切换下一首，双击左耳切换上一首。长按 2 秒唤醒语音助手。\n",
    "\n",
    "3. 充电说明\n",
    "耳机单次续航 6 小时，充电盒可提供额外 24 小时续航。使用 Type-C 线缆为充电盒充电，约 1.5 小时充满。\n",
    "\n",
    "4. 降噪功能\n",
    "耳机支持主动降噪和环境音模式。在连接状态下，长按右耳 3 秒可在两种模式间切换。\n",
    "\"\"\"\n",
    "\n",
    "with open('product_docs/watch_manual.txt', 'w', encoding='utf-8') as f:\n",
    "    f.write(watch_manual.strip())\n",
    "\n",
    "with open('product_docs/earphone_manual.txt', 'w', encoding='utf-8') as f:\n",
    "    f.write(earphone_manual.strip())\n",
    "\n",
    "print('产品手册已创建：')\n",
    "for f in os.listdir('product_docs'):\n",
    "    print(f' - {f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. 完整 RAG 构建流程\n",
    "\n",
    "下面按照「加载 -> 分割 -> 向量化 -> 存储 -> 检索 -> 生成」的顺序，一步步构建系统。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "warnings.filterwarnings('ignore', category=DeprecationWarning)\n",
    "\n",
    "from langchain_community.document_loaders import DirectoryLoader, TextLoader\n",
    "from langchain_text_splitters import RecursiveCharacterTextSplitter\n",
    "from langchain_openai import OpenAIEmbeddings\n",
    "from langchain_community.vectorstores import Chroma\n",
    "\n",
    "# 1. 加载目录下所有 txt 文件\n",
    "loader = DirectoryLoader(\n",
    "    'product_docs',\n",
    "    glob='*.txt',\n",
    "    loader_cls=TextLoader,\n",
    "    loader_kwargs={'encoding': 'utf-8'}\n",
    ")\n",
    "documents = loader.load()\n",
    "print(f'加载文档数：{len(documents)}')\n",
    "\n",
    "# 2. 分割文档\n",
    "splitter = RecursiveCharacterTextSplitter(\n",
    "    chunk_size=150,\n",
    "    chunk_overlap=30,\n",
    "    separators=['\\n\\n', '\\n', '。', ' ']\n",
    ")\n",
    "chunks = splitter.split_documents(documents)\n",
    "print(f'分割后文本块数：{len(chunks)}')\n",
    "\n",
    "# 3. 创建 embedding 模型\n",
    "embeddings = OpenAIEmbeddings(\n",
    "    model='qwen3-embedding',\n",
    "    openai_api_base=os.environ.get('OPENAI_BASE_URL'),\n",
    "    openai_api_key=os.environ.get('OPENAI_API_KEY')\n",
    ")\n",
    "\n",
    "# 4. 存入 ChromaDB\n",
    "vectorstore = Chroma.from_documents(\n",
    "    documents=chunks,\n",
    "    embedding=embeddings,\n",
    "    persist_directory='./product_knowledge_db'\n",
    ")\n",
    "\n",
    "print('向量数据库创建成功！')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 代码解释\n",
    "\n",
    "- `DirectoryLoader`：批量加载目录下所有匹配文件\n",
    "- `glob='*.txt'`：只加载 txt 文件\n",
    "- `loader_cls=TextLoader`：指定使用 TextLoader 加载每个文件\n",
    "- `RecursiveCharacterTextSplitter`：按语义边界切分文档\n",
    "- `Chroma.from_documents`：一次性完成向量化和存储"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. 构建问答链\n",
    "\n",
    "使用 LCEL 把检索器和 LLM 组合起来。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_openai import ChatOpenAI\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.runnables import RunnablePassthrough\n",
    "\n",
    "llm = ChatOpenAI(model='qwen3.6-35b-A3b', temperature=0.1)\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages([\n",
    "    ('system', '''你是公司产品客服助手。请严格根据提供的产品手册内容回答用户问题。\n",
    "如果手册中没有相关信息，请明确告知「根据现有手册，我无法回答这个问题」，不要编造。'''),\n",
    "    ('user', '''相关手册内容：\\n{context}\\n\\n用户问题：{question}''')\n",
    "])\n",
    "\n",
    "def format_docs(docs):\n",
    "    return '\\n\\n'.join([f'[来自 {doc.metadata[\"source\"]}] {doc.page_content}' for doc in docs])\n",
    "\n",
    "retriever = vectorstore.as_retriever(search_kwargs={'k': 3})\n",
    "\n",
    "qa_chain = (\n",
    "    {'context': retriever | format_docs, 'question': RunnablePassthrough()}\n",
    "    | prompt\n",
    "    | llm\n",
    "    | StrOutputParser()\n",
    ")\n",
    "\n",
    "# 测试\n",
    "questions = [\n",
    "    '智能手表怎么测心率？',\n",
    "    '耳机的降噪功能怎么用？',\n",
    "    '这款手表支持潜水吗？',\n",
    "    '手机怎么连接打印机？'  # 不在手册中\n",
    "]\n",
    "\n",
    "for q in questions:\n",
    "    print(f'Q: {q}')\n",
    "    print(f'A: {qa_chain.invoke(q)}\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. 模块化封装\n",
    "\n",
    "把 RAG 系统封装成一个类，便于在项目中复用。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ProductQASystem:\n",
    "    \"\"\"基于 ChromaDB 的产品手册问答系统。\"\"\"\n",
    "\n",
    "    def __init__(self, persist_dir, llm_model='qwen3.6-35b-A3b', embedding_model='qwen3-embedding'):\n",
    "        self.embeddings = OpenAIEmbeddings(\n",
    "            model=embedding_model,\n",
    "            openai_api_base=os.environ.get('OPENAI_BASE_URL'),\n",
    "            openai_api_key=os.environ.get('OPENAI_API_KEY')\n",
    "        )\n",
    "        self.vectorstore = Chroma(\n",
    "            persist_directory=persist_dir,\n",
    "            embedding_function=self.embeddings\n",
    "        )\n",
    "        self.llm = ChatOpenAI(model=llm_model, temperature=0.1)\n",
    "        self.chain = self._build_chain()\n",
    "\n",
    "    def _format_docs(self, docs):\n",
    "        return '\\n\\n'.join([doc.page_content for doc in docs])\n",
    "\n",
    "    def _build_chain(self):\n",
    "        prompt = ChatPromptTemplate.from_messages([\n",
    "            ('system', '你是产品客服助手，请根据手册内容回答问题。'),\n",
    "            ('user', '''手册内容：\\n{context}\\n\\n问题：{question}''')\n",
    "        ])\n",
    "        retriever = self.vectorstore.as_retriever(search_kwargs={'k': 3})\n",
    "        return (\n",
    "            {'context': retriever | self._format_docs, 'question': RunnablePassthrough()}\n",
    "            | prompt\n",
    "            | self.llm\n",
    "            | StrOutputParser()\n",
    "        )\n",
    "\n",
    "    def ask(self, question):\n",
    "        \"\"\"回答用户问题。\"\"\"\n",
    "        return self.chain.invoke(question)\n",
    "\n",
    "    def add_documents(self, file_paths):\n",
    "        \"\"\"增量添加新文档。\"\"\"\n",
    "        loaders = [TextLoader(path, encoding='utf-8') for path in file_paths]\n",
    "        docs = []\n",
    "        for loader in loaders:\n",
    "            docs.extend(loader.load())\n",
    "        splitter = RecursiveCharacterTextSplitter(chunk_size=150, chunk_overlap=30)\n",
    "        chunks = splitter.split_documents(docs)\n",
    "        self.vectorstore.add_documents(chunks)\n",
    "        return f'已添加 {len(chunks)} 个文本块'\n",
    "\n",
    "# 使用封装类\n",
    "qa_system = ProductQASystem('./product_knowledge_db')\n",
    "print(qa_system.ask('耳机续航多久？'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. 评估检索效果\n",
    "\n",
    "RAG 系统的效果取决于两个环节：检索是否找到相关文档、生成是否基于文档准确回答。下面通过打印检索结果来评估检索环节。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 评估检索效果\n",
    "test_queries = [\n",
    "    '智能手表怎么测心率？',\n",
    "    '耳机怎么切换下一首歌？',\n",
    "    '手表防水吗？',\n",
    "    '耳机充电要多久？'\n",
    "]\n",
    "\n",
    "for q in test_queries:\n",
    "    print(f'\\n问题：{q}')\n",
    "    docs = vectorstore.as_retriever(search_kwargs={'k': 2}).invoke(q)\n",
    "    for i, doc in enumerate(docs):\n",
    "        print(f'  {i+1}. [{doc.metadata[\"source\"]}] {doc.page_content[:60]}...')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. 常见问题与优化策略\n",
    "\n",
    "### 7.1 检索不到相关内容\n",
    "- 检查 embedding 模型是否正常工作\n",
    "- 调整 `chunk_size` 和 `chunk_overlap`\n",
    "- 尝试增加 `k` 值，返回更多候选\n",
    "\n",
    "### 7.2 回答包含错误信息\n",
    "- 加强 Prompt 中的约束，如「不知道就说不知道」\n",
    "- 检查检索片段是否包含答案\n",
    "- 使用更大的语言模型\n",
    "\n",
    "### 7.3 回答过长或偏离问题\n",
    "- 在 Prompt 中明确要求简洁回答\n",
    "- 限制 `max_tokens`\n",
    "\n",
    "### 7.4 多文档冲突\n",
    "- 在 `format_docs` 中标注每个片段的来源\n",
    "- 使用元数据过滤区分不同产品或版本"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. 完整项目：多文档知识库问答系统\n",
    "\n",
    "下面把前面的代码整合成一个可复用的完整项目结构。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 完整流程整合\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore', category=DeprecationWarning)\n",
    "\n",
    "from langchain_community.document_loaders import DirectoryLoader, TextLoader\n",
    "from langchain_text_splitters import RecursiveCharacterTextSplitter\n",
    "from langchain_openai import ChatOpenAI, OpenAIEmbeddings\n",
    "from langchain_community.vectorstores import Chroma\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.runnables import RunnablePassthrough\n",
    "\n",
    "def build_knowledge_base(docs_dir, db_dir):\n",
    "    \"\"\"构建知识库：加载、分割、向量化、持久化。\"\"\"\n",
    "    loader = DirectoryLoader(docs_dir, glob='*.txt', loader_cls=TextLoader, loader_kwargs={'encoding': 'utf-8'})\n",
    "    documents = loader.load()\n",
    "    splitter = RecursiveCharacterTextSplitter(chunk_size=150, chunk_overlap=30)\n",
    "    chunks = splitter.split_documents(documents)\n",
    "    embeddings = OpenAIEmbeddings(\n",
    "        model='qwen3-embedding',\n",
    "        openai_api_base=os.environ.get('OPENAI_BASE_URL'),\n",
    "        openai_api_key=os.environ.get('OPENAI_API_KEY')\n",
    "    )\n",
    "    vectorstore = Chroma.from_documents(chunks, embeddings, persist_directory=db_dir)\n",
    "    return vectorstore\n",
    "\n",
    "def build_qa_chain(vectorstore, k=3):\n",
    "    \"\"\"构建 RAG 问答链。\"\"\"\n",
    "    llm = ChatOpenAI(model='qwen3.6-35b-A3b', temperature=0.1)\n",
    "    prompt = ChatPromptTemplate.from_messages([\n",
    "        ('system', '你是产品客服助手，请根据手册内容准确回答，不知道就说不知道。'),\n",
    "        ('user', '''参考内容：\\n{context}\\n\\n问题：{question}''')\n",
    "    ])\n",
    "    retriever = vectorstore.as_retriever(search_kwargs={'k': k})\n",
    "    def format_docs(docs):\n",
    "        return '\\n\\n'.join([doc.page_content for doc in docs])\n",
    "    return (\n",
    "        {'context': retriever | format_docs, 'question': RunnablePassthrough()}\n",
    "        | prompt\n",
    "        | llm\n",
    "        | StrOutputParser()\n",
    "    )\n",
    "\n",
    "# 构建或加载知识库\n",
    "if not os.path.exists('product_knowledge_db'):\n",
    "    vectorstore = build_knowledge_base('product_docs', 'product_knowledge_db')\n",
    "else:\n",
    "    embeddings = OpenAIEmbeddings(\n",
    "        model='qwen3-embedding',\n",
    "        openai_api_base=os.environ.get('OPENAI_BASE_URL'),\n",
    "        openai_api_key=os.environ.get('OPENAI_API_KEY')\n",
    "    )\n",
    "    vectorstore = Chroma(persist_directory='product_knowledge_db', embedding_function=embeddings)\n",
    "\n",
    "# 构建问答链并提问\n",
    "chain = build_qa_chain(vectorstore)\n",
    "print(chain.invoke('无线耳机怎么配对？'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9. 本节课练习\n",
    "\n",
    "1. 在 `product_docs` 目录下添加第三份产品手册（如「智能音箱使用手册」），重新构建知识库\n",
    "2. 修改 `format_docs` 函数，在 Prompt 中显示每个片段来自哪个文件\n",
    "3. 分别测试 `k=1`、`k=3`、`k=5` 时的回答质量差异\n",
    "4. 设计 5 个产品手册中的问题和 2 个手册外的问题，评估系统表现\n",
    "5. 把 `ProductQASystem` 类封装保存到单独的 `qa_system.py` 文件中，实现命令行问答"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}