{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 12 向量数据库\n",
    "\n",
    "## 学习目标\n",
    "1. 理解向量数据库与传统数据库的区别\n",
    "2. 掌握 ChromaDB 的基本使用方法\n",
    "3. 学会向量数据库的持久化存储和加载\n",
    "4. 理解元数据过滤在检索中的作用\n",
    "5. 能够选择合适的向量数据库方案"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. 为什么需要向量数据库\n",
    "\n",
    "上一节课我们用 `InMemoryVectorStore` 把向量存在内存中，它适合学习和测试，但有几个明显缺点：\n",
    "\n",
    "- **数据无法持久化**：程序关闭后向量就消失了\n",
    "- **无法增量更新**：不能方便地添加、删除文档\n",
    "- **不支持复杂查询**：无法按元数据过滤\n",
    "- **不适合大规模数据**：内存容量有限\n",
    "\n",
    "**向量数据库** 就是为解决这些问题而生的专用数据库。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. 向量数据库 vs 传统数据库\n",
    "\n",
    "| 特性 | 传统数据库（MySQL） | 向量数据库（ChromaDB） |\n",
    "| --- | --- | --- |\n",
    "| 存储内容 | 结构化数据 | 向量 + 原始文本 + 元数据 |\n",
    "| 查询方式 | 精确匹配、范围查询 | 相似度搜索 |\n",
    "| 索引类型 | B+ 树 | HNSW、IVF 等近似最近邻索引 |\n",
    "| 典型应用 | 订单、用户管理 | 语义搜索、RAG、推荐 |\n",
    "\n",
    "向量数据库的核心能力：**根据语义相似度快速找到最相关的向量**。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. 核心概念\n",
    "\n",
    "### 3.1 向量（Vector）\n",
    "文本经过 Embedding 模型编码后变成的高维数组。例如一个 2560 维的浮点数数组。\n",
    "\n",
    "### 3.2 相似度（Similarity）\n",
    "两个向量之间的距离。常用度量方式：\n",
    "- **余弦相似度（Cosine Similarity）**：衡量方向是否一致，范围 -1 到 1\n",
    "- **欧氏距离（Euclidean Distance）**：衡量向量空间中的直线距离\n",
    "\n",
    "### 3.3 索引（Index）\n",
    "为了加速海量向量的相似度搜索，向量数据库会构建专门的索引结构，如 HNSW。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. 常见向量数据库\n",
    "\n",
    "| 数据库 | 特点 | 适用场景 |\n",
    "| --- | --- | --- |\n",
    "| **ChromaDB** | 轻量、易用、支持持久化 | 本地开发、中小规模 |\n",
    "| **FAISS** | Meta 开源、高性能 | 大规模向量检索 |\n",
    "| **Milvus** | 企业级、分布式 | 生产环境、海量数据 |\n",
    "| **Pinecone** | 云端托管、无需运维 | 快速上线、Serverless |\n",
    "| **Weaviate** | 支持多模态、GraphQL | 复杂查询、多模态 |\n",
    "\n",
    "本节课以 **ChromaDB** 为例，因为它安装简单、接口友好，非常适合学习。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. 安装 ChromaDB\n",
    "\n",
    "在虚拟环境中执行：\n",
    "\n",
    "```powershell\n",
    "pip install chromadb -i https://pypi.tuna.tsinghua.edu.cn/simple\n",
    "```\n",
    "\n",
    "安装完成后需要**重启 Jupyter 内核**。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. 使用 ChromaDB 创建向量库\n",
    "\n",
    "下面演示如何把文档存入 ChromaDB 并进行检索。"
   ]
  },
  {
   "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.vectorstores import Chroma\n",
    "from langchain_core.documents import Document\n",
    "from langchain_openai import OpenAIEmbeddings\n",
    "import os\n",
    "\n",
    "# 创建 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",
    "# 准备文档\n",
    "documents = [\n",
    "    Document(page_content='人工智能是计算机科学的一个分支。', metadata={'source': 'ai_book', 'category': 'AI'}),\n",
    "    Document(page_content='深度学习使用多层神经网络。', metadata={'source': 'ai_book', 'category': 'AI'}),\n",
    "    Document(page_content='机器学习让计算机从数据中学习。', metadata={'source': 'ai_book', 'category': 'AI'}),\n",
    "    Document(page_content='香蕉是一种黄色的水果。', metadata={'source': 'food_book', 'category': 'food'}),\n",
    "    Document(page_content='苹果通常是红色或绿色的。', metadata={'source': 'food_book', 'category': 'food'})\n",
    "]\n",
    "\n",
    "# 从文档创建 ChromaDB 向量库\n",
    "# persist_directory 指定数据保存目录\n",
    "vectorstore = Chroma.from_documents(\n",
    "    documents=documents,\n",
    "    embedding=embeddings,\n",
    "    persist_directory='./chroma_db_demo'\n",
    ")\n",
    "\n",
    "print('向量库创建成功！')\n",
    "print(f'文档数量：{vectorstore._collection.count()}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 代码解释\n",
    "\n",
    "- `Chroma.from_documents()`：从 Document 列表直接创建向量库\n",
    "- `embedding=embeddings`：指定向量化模型\n",
    "- `persist_directory`：数据持久化目录，程序关闭后数据仍存在"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. 相似度检索\n",
    "\n",
    "创建检索器后，就可以根据用户问题进行语义搜索。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 创建检索器\n",
    "retriever = vectorstore.as_retriever(search_kwargs={'k': 3})\n",
    "\n",
    "# 检索与问题最相关的文档\n",
    "query = '什么是人工智能？'\n",
    "results = retriever.invoke(query)\n",
    "\n",
    "print(f'问题：{query}\\n')\n",
    "for i, doc in enumerate(results):\n",
    "    print(f'--- 结果 {i+1} ---')\n",
    "    print(f'内容：{doc.page_content}')\n",
    "    print(f'元数据：{doc.metadata}\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. 持久化加载\n",
    "\n",
    "ChromaDB 的数据已经保存在 `./chroma_db_demo` 目录中。下次启动程序时，可以直接加载，不需要重新向量化。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 从已有目录加载向量库\n",
    "loaded_vectorstore = Chroma(\n",
    "    persist_directory='./chroma_db_demo',\n",
    "    embedding_function=embeddings\n",
    ")\n",
    "\n",
    "loaded_retriever = loaded_vectorstore.as_retriever(search_kwargs={'k': 2})\n",
    "results = loaded_retriever.invoke('深度学习和神经网络的关系')\n",
    "\n",
    "print('从持久化数据库加载后检索：')\n",
    "for doc in results:\n",
    "    print(f'- {doc.page_content}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9. 元数据过滤\n",
    "\n",
    "元数据过滤可以让我们在检索时只搜索特定类别的文档。例如只搜索 `category='AI'` 的文档。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 只检索 category 为 AI 的文档\n",
    "filtered_results = loaded_vectorstore.similarity_search(\n",
    "    '水果的颜色',\n",
    "    filter={'category': 'AI'},\n",
    "    k=3\n",
    ")\n",
    "\n",
    "print('过滤后只搜索 AI 类文档：')\n",
    "for doc in filtered_results:\n",
    "    print(f'- [{doc.metadata[\"category\"]}] {doc.page_content}')\n",
    "\n",
    "# 只检索 category 为 food 的文档\n",
    "food_results = loaded_vectorstore.similarity_search(\n",
    "    '人工智能',\n",
    "    filter={'category': 'food'},\n",
    "    k=3\n",
    ")\n",
    "\n",
    "print('\\n过滤后只搜索 food 类文档：')\n",
    "for doc in food_results:\n",
    "    print(f'- [{doc.metadata[\"category\"]}] {doc.page_content}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 代码解释\n",
    "\n",
    "- `filter={'category': 'AI'}`：只返回 category 字段等于 AI 的文档\n",
    "- 即使查询词是「水果的颜色」，由于过滤条件限制，也只会返回 AI 类文档\n",
    "- 元数据过滤常用于多租户、多知识库场景"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 10. 增量更新：添加和删除文档\n",
    "\n",
    "向量数据库支持动态增删文档，不需要每次重新构建整个库。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 添加新文档\n",
    "new_docs = [\n",
    "    Document(page_content='自然语言处理让计算机理解人类语言。', metadata={'source': 'ai_book', 'category': 'AI'})\n",
    "]\n",
    "vectorstore.add_documents(new_docs)\n",
    "print(f'添加后文档数量：{vectorstore._collection.count()}')\n",
    "\n",
    "# 查询刚添加的文档\n",
    "results = vectorstore.similarity_search('NLP 是什么', k=2)\n",
    "for doc in results:\n",
    "    print(f'- {doc.page_content}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 11. 完整示例：基于 ChromaDB 的本地知识库问答\n",
    "\n",
    "把第 11 节课的 RAG 链和 ChromaDB 结合起来，构建一个可持久化的本地知识库。"
   ]
  },
  {
   "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",
    "    ('user', '''上下文：\\n{context}\\n\\n问题：{question}''')\n",
    "])\n",
    "\n",
    "def format_docs(docs):\n",
    "    return '\\n\\n'.join([doc.page_content for doc in docs])\n",
    "\n",
    "# 使用 ChromaDB 检索器\n",
    "chroma_retriever = loaded_vectorstore.as_retriever(search_kwargs={'k': 2})\n",
    "\n",
    "qa_chain = (\n",
    "    {'context': chroma_retriever | format_docs, 'question': RunnablePassthrough()}\n",
    "    | prompt\n",
    "    | llm\n",
    "    | StrOutputParser()\n",
    ")\n",
    "\n",
    "question = '深度学习是什么？'\n",
    "answer = qa_chain.invoke(question)\n",
    "print(f'Q: {question}')\n",
    "print(f'A: {answer}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 12. 向量数据库选型建议\n",
    "\n",
    "| 场景 | 推荐方案 |\n",
    "| --- | --- |\n",
    "| 本地学习、原型开发 | ChromaDB / InMemoryVectorStore |\n",
    "| 大规模生产环境 | Milvus / FAISS |\n",
    "| 不想自己运维 | Pinecone / Weaviate Cloud |\n",
    "| 需要多模态检索 | Weaviate |\n",
    "| 已有 Elasticsearch 集群 | Elasticsearch 向量检索 |\n",
    "\n",
    "选择向量数据库时主要考虑：数据规模、查询性能、运维成本、是否支持元数据过滤。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 13. 本节课练习\n",
    "\n",
    "1. 创建 5 条以上不同类别的文档，存入 ChromaDB，并指定 source 和 category 元数据\n",
    "2. 分别用无过滤和有过滤两种方式检索，对比结果差异\n",
    "3. 关闭 Jupyter 后重新打开，加载已持久化的 ChromaDB，验证数据没有丢失\n",
    "4. 尝试添加一条新文档，然后检索这条新文档相关的内容\n",
    "5. 用 ChromaDB 替换第 11 节课中的 InMemoryVectorStore，重新构建 RAG 链"
   ]
  }
 ],
 "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.14.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}