# 使用 TiDB 作为 AI 应用的统一存储

## 概述

在本实验的前四个模块中，您将使用多个异构数据存储（如 MySQL、Milvus 和 ElasticSearch）来满足 GenAI 应用程序的不同需求。这将需要多个 Python 包，让您亲身体验开发和维护的挑战。然而，在模块 5 中，您将只使用一个 Python 包 - pytidb - 和一个数据存储 - TiDB - 作为 GenAI 的统一数据存储。在这里，您将亲眼看到使用 TiDB 为 GenAI 应用程序提供支持是多么简单和高效。

> [!IMPORTANT]
> 本实验将授予您临时访问 Amazon Bedrock 的权限，以调用必要的模型。请严格遵守实验材料中概述的使用指南，并避免在任何其他情境中使用此访问权限。对于任何导致财务损失的 Amazon Bedrock 权限滥用行为，我们保留追究用户法律责任的权利。
>
> 目前，TiDB Full-Text Search with SQL 功能仅在 TiDB Serverless 法兰克福 (eu-central-1) 和新加坡 (ap-southeast-1) 区域可用。如果您在本实验中使用自己的免费 TiDB Cloud 数据库，请确保您的 TiDB 集群部署在法兰克福或新加坡。

![70%](https://lab-static.pingcap.com/quick-demo/demo-422/422-with-bg.png)

生成式 AI 目前是一个前沿领域，吸引了许多渴望探索的开发者。然而，几个障碍常常阻碍他们进入这个领域，例如：

- 缺乏 AI 基础知识，使他们不确定从何处开始。
- 缺少全面、循序渐进且互动的构建生成式 AI 应用程序指南。
- 学习所需的基础设施有限，如数据库和大型模型 API。

为了解决这些挑战，我们开发了一系列实验，包括本实验，通过这些实验，您将学习：

- 如何使用传统方法编写一个 AI 应用程序，基于以下存储，构建一个完整的应用程序:

    - 关系型数据库: 例如 `MySQL`
    - 向量数据库: 例如 `Milvus`
    - 全文搜索引擎: 例如 `ElasticSearch`

- 如何通过 `pytidb`，使用 **TiDB Cloud Serverless** 作为统一存储层，来简化你的 AI 应用程序。

> [!NOTE]
> 在本实验中，我们将使用 Streamlit 作为快速应用程序开发框架，以简化和加速数据应用程序的构建。
> 无论您选择自己的 TiDB Cloud Serverless 集群还是沙盒 TiDB Cloud Serverless 选项来开始实验，都将为此动手实验创建一个专用数据库。

### 任务

1. 构建你的第一个简易 AI 应用程序：调用 LLM 并给出问题的回答。
2. 使用 RAG 来改进你的 AI 应用程序，让你的应用程序拥有记忆力：

    - 使用 MySQL 作为文字块的存储
    - 使用 Milvus 作为向量存储及搜索引擎

3. 提升你的 AI 应用的表现，是时候引入多一个召回通道了。全文搜索是一个好的选择。这将会给你的 AI 应用程序以不同的搜索角度：

    - 使用 ElasticSearch 作为全文搜索引擎

4. **使用 TiDB 来重写你的应用程序，简化你的整体架构。**

### 实验时长

- 完成此练习大约需要 90 分钟。

> [!TIP]
> 点击 **开始** 启动实验。实验环境准备就绪后。点击底部栏上的 **前进** 进入 **模块 1**。

## 模块 1: 准备运行 Python 代码

### 步骤 1. 检查实验环境

- 点击左侧的 `实验资源` 按钮查看为您的实践设置的实验环境。您将看到：

    - 一个名为 monitor1 的服务器，运行着 VS Code 服务器，它既是开发服务器又是应用服务器。
    - TiDB Cloud Serverless 集群的连接详情 — 根据您的选择，可能来自您自己的 TiDB Cloud 账户或实验的沙盒环境。
    - 此外，此动手实验的数据库名称为 `${LAB:SERVERLESS_CLUSTER_DATABASE_NAME}`，它将在您选择的 TiDB Cloud Serverless 集群中代您创建。

- 关闭 `实验资源` 弹出窗口。

### 步骤 2. 在 IDE 中打开终端

- 在浏览器中打开 [VS Code](${LAB:SECURITY_ACCESS_PROXY_URL}?next=${LAB:"/vscode/?folder=/root/tidb-lab-unify-storage"|urlencode})。

- 在弹出的窗口中点击 `Yes, I trust the authors`（是的，我信任作者）。

    ![70%](https://lab-static.pingcap.com/quick-demo/vscode-trust.jpg)

- 点击左上角的 `≡`，选择 `Terminal`，然后点击 `New Terminal` 打开一个新的终端窗口。

    ![90%](https://public-pingcap-innovation.s3.us-west-2.amazonaws.com/demo-405/405-open-terminal-0307.jpg)

### 步骤 3. 安装 Python 依赖

在 VS Code 终端中运行以下命令来安装所需的依赖。

```bash
poetry install
```

```output
Creating virtualenv tidb-lab-unify-storage-wZek_onr-py3.11 in /root/.cache/pypoetry/virtualenvs
Installing dependencies from lock file

Package operations: 93 installs, 0 updates, 0 removals

- Installing attrs (25.3.0)
- Installing rpds-py (0.25.1)
- Installing typing-extensions (4.13.2)
- Installing certifi (2025.4.26)
...
```

## 模块 2: 构建你的第一个简易 AI 应用程序

让我们从最简单的 AI 应用程序开始。在这一步里，我们将调用 AWS 的 Nova 模型，这是一种大语言模型（LLM）。它在被训练的时候，就已经了解了很多知识。关于如何将原始知识转换为大语言模型（LLM），请参考以下材料。您可以稍后保存这些材料以供学习，让我们直接从 步骤 1 开始。

- Transformer 奠基论文 [Attention Is All You Need](https://arxiv.org/abs/1706.03762)
- Youtuber 3Blue1Brown 将 Transformer 可视化：[Transformers (how LLMs work) explained visually](https://www.youtube.com/watch?v=wjZofJX0v4M)

### 步骤 1. 运行简易 AI 应用程序的 Streamlit 服务

- 在 VS Code 终端中，使用 `streamlit` 命令行工具运行以下脚本来启动 Streamlit 服务。你可能注意到了，我们在 `streamlit` 前，加入了 `poetry run`。这让我们可以使用 poetry 来管理依赖，而非默认的 pip。

    ```bash
    poetry run streamlit run vanilla_ai_app.py --server.port 8501
    ```

    ```output
    Collecting usage statistics. To deactivate, set browser.gatherUsageStats to false.


    You can now view your Streamlit app in your browser.

    Local URL: http://localhost:8501
    Network URL: http://172.18.0.7:8501
    External URL: http://${LAB:HOST_MONITOR1_PUBLIC_IP}:8501
    ```

    ![90%](https://lab-static.pingcap.com/quick-demo/demo-422/422-vanilla-ai-app-streamlit-start-0604.jpg)

### 步骤 2. 使用简易 AI 应用程序的演示网页

- 点击访问您的 [RAG 演示网页](http://${LAB:HOST_MONITOR1_PUBLIC_IP}:8501/)。

    ![90%](https://lab-static.pingcap.com/quick-demo/demo-422/422-vanilla-ai-app-frontend-0604.jpg)

- 现在，你就可以在 `Enter your question:` 下的输入框内输入你的问题了，比如：

    ```txt
    What is TiDB?
    ```

- 随后，点击 `Send` 按钮。稍做等候，你就可以得到回答了！

    ![90%](https://lab-static.pingcap.com/quick-demo/demo-422/422-vanilla-ai-app-answer-0604.jpg)

- 恭喜🎉！你已经完成了你的第一个 AI 应用程序。你可以继续问 AI 一些自己想知道的问题。甚至你可以尝试一下更困难的问题，来看看我们初步的 AI 应用程序的能力边界在哪里。
- 接下来，让我们看看代码内到底发生了什么？

### 步骤 3. 简易 AI 应用程序代码阅读

让我们回到 [VS Code](${LAB:SECURITY_ACCESS_PROXY_URL}?next=${LAB:"/vscode/?folder=/root/tidb-lab-unify-storage"|urlencode}) 并点击代码文件 `vanilla_ai_app.py`。

1. 脚本首先导入必要的库：`litellm` 用于与 LLM 交互，`streamlit` 用于构建网页界面。此外，它还从 `common` 模块导入 `LLM_RESPONSE_STYLE`，这是一个预定义的 CSS 样式，用于美化 LLM 的响应输出。

    ```python
    from litellm import completion
    import streamlit as st

    from common import LLM_RESPONSE_STYLE
    ```

2. 定义了要使用的 LLM 模型。我们将模型设置为 `bedrock/us.amazon.nova-lite-v1:0`，`litellm` 使用的模型名称为两段式的，`/` 前被为 provider，是模型的提供者，在这个例子中，为 Amazon Bedrock，而后方是模型的名称，由 Amazon 内部定义。

    ```python
    llm_model = "bedrock/us.amazon.nova-lite-v1:0"
    ```

3. 设置 Streamlit 页面的标题为 "🔍 Vanilla AI Q&A Demo"。

    ```python
    st.title("🔍 Vanilla AI Q&A Demo")
    ```

4. 在页面上显示一段指导性文字，告知用户此应用程序的功能：用户输入问题，系统将基于其内置知识进行回答。

    ```python
    st.write(
        "Enter your question, and the system will answer it based on the built-in knowledge"
    )
    ```

5. 使用 `st.text_input` 创建一个文本输入框，允许用户输入他们的问题。输入的问题存储在 `query` 变量中。

    ```python
    query = st.text_input("Enter your question:", "")
    ```

6. 此代码块处理用户点击 "Send" 按钮并且 `query` (用户输入的问题) 不为空时的逻辑：
    - 使用 `with st.spinner("Processing...")` 显示一个加载动画，告知用户系统正在处理请求。
    - 构建一个简单的提示，其内容是 "Answer the question: " 加上用户提出的问题 `query`。

    ```python
    if st.button("Send") and query:
        with st.spinner("Processing..."):
            prompt = f"Answer the question: {query}"
    ```

7. 调用 `litellm.completion` 函数与 LLM 进行交互以生成答案。
    - `model` 参数指定了要使用的 LLM 模型 (即 `llm_model`)。
    - `messages` 参数是一个列表，其中包含一个字典，定义了消息内容 (`prompt`) 和角色 (`user`)。

    ```python
    # Call LLM to generate answer
    response = completion(
        model=llm_model,
        messages=[{"content": prompt, "role": "user"}],
    )
    ```

8. 在页面上显示 LLM 的模型名称和一个机器人图标。

    ```python
    st.markdown(f"### 🤖 {llm_model}")
    ```

9. 应用预定义的 CSS 样式 `LLM_RESPONSE_STYLE` 来美化 LLM 响应的显示效果。`unsafe_allow_html=True` 参数允许在 Markdown 中嵌入 HTML 和 CSS。

    ```python
    st.markdown(LLM_RESPONSE_STYLE, unsafe_allow_html=True)
    ```

10. 显示 LLM 生成的答案。答案从 `response.choices[0].message.content` 中提取，并使用 HTML `div` 标签和 `llm-response` 类进行包装，以便应用样式。

    ```python
    # show the response
    st.markdown(
        f'<div class="llm-response">{response.choices[0].message.content}</div>',
        unsafe_allow_html=True,
    )
    ```

### 步骤 4. 简易 AI 应用程序的总结

我们可以看到，简易 AI 应用程序非常的简单。总的来说，就是单纯的调用了一下 Amazon Bedrock 的 API，获取了回答。而一切的知识储备，都是 LLM 在训练的时候就输入的。比如说上方的示例问题：

```text
What is TiDB?
```

之所以 LLM 可以回答这个问题，是因为 Amazon 在训练这个模型的时候就有相关的文档或介绍被给到了 LLM。那我们不禁会想到一个问题：如果 LLM 在训练的时候没有这个数据怎么办？举个例子：

```text
Who is Cheese Wang at PingCAP?
```

![90%](https://lab-static.pingcap.com/quick-demo/demo-422/422-vanilla-ai-app-hallucination-0604.jpg)

可以看到，LLM "编造"了一些信息出来。这是由于在 LLM 训练的时候，根本没有类似的信息输入。因此它无法正确的得知该员工和 PingCAP 之间的关系。只能"猜测"被询问的人为 CEO，因为 CEO 的名字和公司的联系会更加紧密，被询问的概率更大。

而且，这个答案是错误的，因为本实验的作者是 ***Cheese Wang***，他并不是 PingCAP 的 CEO。

这种 LLM 的行为，被称为 **幻觉 (Hallucination)**，这是我们在使用 LLM 的时候需要格外注意的一个现象。我们将在下一个模块，改进我们的 AI 应用程序，降低它的幻觉出现频率，同时给予它额外的知识储备，让他可以回答训练数据之外的数据，而无需更改模型本身。

## 模块 3: 通过 RAG 改进你的 AI 应用程序

我们在模块 1 中完成了初步的 AI 应用，但是遇到了 LLM 在回答训练数据之外的数据的时候，产生了严重的幻觉问题。那怎么解决呢？答案也很简单。

我们可以预先将信息帮助 LLM 获取回来，放置在调用 LLM 的输入中，使用 LLM 帮我们总结内容，而非依赖它的训练集的知识储备。这样一来，我们可以得到以下好处：

1. 降低幻觉：Transformer 架构的 LLM，使用多头注意力模型在 context window 内进行基于顺序的预测。如果我们把相关知识放入 context window 内，它会生成和相关知识更接近的回答。从表象上来看，它降低了幻觉的产生频率或程度。
2. 提供外部知识：再一次，由于 Transformer 架构的 LLM 拥有一个 context window（在 API 上的体现为 token size 的限制），我们无法将所有的信息全部放入 prompt。因此，我们需要事先挑选出有关的内容。
3. 降低成本：即使现在 LLM 的 context window 越来越大，我们仍需从成本考量，不应为每个问题浪费太多的 token。因为绝大多数 LLM 都是以 token 计费的（你可以把 token 理解为和文本长度正相关的一个变量，你的文本越长，token 就越多）。

那么，现在摆在我们面前的问题就只有一个了：

**怎么实现帮助 LLM 获取相关信息？**

实现它的方案不止一种。我们在模块 2 中会使用其中的一种：文本 Embedding 向量方式。

> [!NOTE]
> **选读部分**
>
> Embedding 向量，是 Embedding 模型将文本转换成的向量。而 Embedding 模型，是一类专门训练来分离语义的的模型。简单的来讲，在对比学习训练时，它会得到非常多组的样本，而每一组样本分为：
>
> - 一个锚点样本：一段文字
> - 一个正样本：一段与锚点样本语义相似的文字
> - 一组负样本：很多段与锚点样本语义不相似的文字
>
> Embedding 模型的任务就是：
>
> - 最小化锚点样本和正样本的输出向量的距离
> - 最大化锚点样本和负样本的输出向量的距离
>
> 更正确的，此处使用了 InfoNCE loss 这个损失函数训练模型，进行了随机梯度下降算法，拟合了样本。

如果你看不懂选读部分，这完全没有问题，它只是原理的解释。你只需要知道：两个文本的 Embedding 向量的余弦距离越远，那么他们的语义就越不相似。而 Embedding 向量的余弦距离越近，那么他们的语义就越相似。

因此，我们可以使用这个特性来完成这一次的改进。在模块 2 中，我们将使用 Mulvis 作为向量的搜索引擎，MySQL 作为关系型数据库存储原始文本。

> [!NOTE]
> **为什么不使用向量数据库的 metadata 来存储原始文本？**
>
> 这在技术上是可行的，但是由于向量数据库的能力缺失，比如事务等。无法保证的数据的完全一致性。在实际的工程中，如果依旧使用关系型数据库和向量数据库的分离架构，建议使用关系型数据库作为数据的 source of truth。同时将原始文本的 ID 与 Embedding 向量这组数据异步的更新至向量数据库。

### 步骤 1. 运行 RAG AI 应用程序的 Streamlit 服务

创建一个新的终端：

![90%](https://lab-static.pingcap.com/quick-demo/demo-422/422-rag-streamlit-start-0606.jpg)

然后运行新的 RAG AI 应用程序：

```bash
poetry run streamlit run rag.py --server.port 8502
```

```output
Collecting usage statistics. To deactivate, set browser.gatherUsageStats to false.


You can now view your Streamlit app in your browser.

Local URL: http://localhost:8502
Network URL: http://172.18.0.7:8502
External URL: http://${LAB:HOST_MONITOR1_PUBLIC_IP}:8502
```

### 步骤 2. 使用 RAG AI 应用程序的演示网页

- 打开您的 [RAG AI 应用程序](http://${LAB:HOST_MONITOR1_PUBLIC_IP}:8502/)。

    ![90%](https://lab-static.pingcap.com/quick-demo/demo-422/422-rag-streamlit-ui-0606.jpg)

- 打开名为 `All Employees in the MySQL` 和 `All Employees ID and Embedding Mapping in the Milvus` 的文件夹，以查看当前 MySQL 和 Milvus 中的数据。由于我们尚未保存任何数据，所以它们应该是空的。

    ![90%](https://lab-static.pingcap.com/quick-demo/demo-422/422-rag-streamlit-start-data-0606.jpg)

- 点击 `Save Data` 保存我们为本次演示准备的模拟员工信息数据。您可以在下一步的代码阅读中看到原始数据，它是一个员工对象列表。
- 保存数据后，您可以在 `All Employees in the MySQL` 文件夹中看到 MySQL 中的数据。我们有 3 列，分别是 `id`、`name` 和 `intro`。而在 `All Employees ID and Embedding Mapping in the Milvus` 内，你可以看到我们保存了 `id` 和 `embedding` 的关系。此时 Milvus 内的 `id` 和 MySQL 中的 `id` 是一一映射的。

    ![90%](https://lab-static.pingcap.com/quick-demo/demo-422/422-rag-streamlit-saved-data-0613.jpg)

- 如果您只想检索上下文，可以点击 `Retrieval Only` 选项。然后，输入您的问题。
- 我们来试试 RAG AI 应用程序是否能通过 Embedding 选取出相关的信息，输入问题：

    ```text
    Who is Cheese Wang at PingCAP?
    ```

- 点击 `Send` 按钮开始检索。

    ![90%](https://lab-static.pingcap.com/quick-demo/demo-422/422-rag-streamlit-retrieve-0606.jpg)

- 如您所见，我们得到了 `Cheese Wang` 的相关信息。
- 接下来，the moment of truth, 我们来检测现在 RAG AI 应用程序是否能回答简易 AI 应用程序无法回答的问题。将 `Select Mode` 选项更改为 `RAG Q&A`，然后点击 `Send`。
- 这次，RAG AI 应用程序将给我们提供答案，而不是信息列表。

    ![90%](https://lab-static.pingcap.com/quick-demo/demo-422/422-rag-streamlit-qna-0606.jpg)

- 可以看到，因为我们提供了 LLM 在训练阶段并为获取到的信息。因此 LLM 不再认为我是 PingCAP 的 CEO，而是真实的回答了我的信息。

### 步骤 3. RAG AI 应用程序代码阅读

让我们回到 [VS Code](${LAB:SECURITY_ACCESS_PROXY_URL}?next=${LAB:"/vscode/?folder=/root/tidb-lab-unify-storage"|urlencode}) 并点击代码文件 `rag.py`。

1. 脚本首先导入了构建 RAG 应用所需的库：
    - `os`, `dotenv`: 用于加载和读取环境变量，例如数据库密码和主机地址。
    - `pandas`: 用于数据处理和在 Streamlit 中展示表格。
    - `sqlalchemy`, `sqlmodel`: 用于连接关系型数据库（在本例中为 MySQL）并定义数据模型。`SQLModel` 提供了结合 SQLAlchemy 和 Pydantic 的便利性。
    - `pymilvus`: Milvus 的 Python 客户端，用于连接和操作 Milvus 向量数据库。
    - `litellm`: 一个用于与各种大语言模型（LLM）和 Embedding 模型进行标准化交互的库。
    - `streamlit`: 用于创建和运行 Web 用户界面。
    - `common`: 从公共模块导入预定义的 RAG 提示模板和 CSS 样式。

    ```python
    import os
    import pandas as pd
    from sqlalchemy import func
    from sqlmodel import Field, SQLModel, Session, create_engine, select, delete
    from pymilvus import MilvusClient, DataType

    from litellm import completion, embedding
    import streamlit as st

    from typing import List, Optional, Any
    from pydantic_core import MultiHostUrl
    from common import LLM_RESPONSE_STYLE, RAG_PROMPT_TEMPLATE
    from dotenv import load_dotenv

    load_dotenv()
    ```

2. 进行初始化和配置：
    - `create_engine`: 使用从环境变量中读取的配置（主机、端口、密码等）创建一个到 MySQL 数据库的连接引擎。
    - `embedding_model`, `llm_model`: 定义了将要使用的 Embedding 模型和 LLM 的名称。
    - `MilvusClient`: 初始化 Milvus 客户端，以连接到在指定 URI 上运行的 Milvus 服务。
    - `milvus_collection_name`: 为将在 Milvus 中创建的集合（Collection）命名。

    ```python
    engine = create_engine(
        str(
            MultiHostUrl.build(
                scheme="mysql+pymysql",
                username="root",
                password=os.getenv("MYSQL_ROOT_PASSWORD"),
                host=os.getenv("MYSQL_HOST"),
                port=int(os.getenv("MYSQL_PORT", "4000")),
                path=os.getenv("MYSQL_DATABASE")
            )
        )
    )
    embedding_model = "bedrock/amazon.titan-embed-text-v2:0"
    embedding_dimensions = 1024

    llm_model = "bedrock/us.amazon.nova-lite-v1:0"
    milvus_client = MilvusClient(
        uri=f"http://{os.getenv('MILVUS_HOST')}:{os.getenv('MILVUS_GRPC_PORT')}",
        token="root:Milvus"
    )
    milvus_collection_name = "employee_id_mapping"
    ```

3. 定义了一个辅助函数 `get_embedding`，它接收一个文本字符串，调用 `litellm.embedding` API 将文本转换为一个 1024 维的向量，并返回该向量。

    ```python
    def get_embedding(text: str) -> List[float]:
        response = embedding(
            model=embedding_model,
            input=[text],
            dimensions=embedding_dimensions
        )
        return response.data[0]['embedding']
    ```

4. 使用 `SQLModel` 定义了 `Employee` 表的结构，该表将在 MySQL 中创建，用于存储员工的原始信息（ID、姓名和简介）。`SQLModel.metadata.create_all(engine)` 会确保在数据库中创建这张表。

    ```python
    class Employee(SQLModel, table=True):
        __table_args__ = {"extend_existing": True}
        id: int | None = Field(default=None, primary_key=True)
        name: str
        intro: str

    SQLModel.metadata.create_all(engine)
    ```

5. 创建了一个包含三名员工信息的列表 `employees` 作为演示用的样本数据。

    ```python
    employees = [
        Employee(id=1, name="Cheese Wang", intro="Cheese is a Senior Software Engineer of AI Application at PingCAP, the company behind TiDB"),
        Employee(id=2, name="Hao Huo", intro="Hao Huo is an engineer of PingCAP, the company behind TiDB | Build pingcap/autoflow / OSSInsight.io"),
        Employee(id=3, name="Humble BGL", intro="Humble BGL is the Head of Engagement Innovation at PingCAP, the company behind TiDB")
    ]
    ```

6. 为 Milvus 定义了数据模式（Schema）和索引。
    - `MilvusClient.create_schema`: 定义了集合的结构，包含一个主键 `id`（用于与 MySQL 中的 `Employee` 记录关联）和一个 `embedding` 向量字段。
    - `milvus_client.prepare_index_params`: 为 `embedding` 字段准备索引参数，使用 `AUTOINDEX` 让 Milvus 自动选择最佳索引类型。
    - `milvus_client.create_collection`: 使用定义的模式和索引在 Milvus 中创建一个新的集合。这使得我们能够高效地进行向量相似度搜索。

    ```python
    schema = MilvusClient.create_schema(auto_id=False, enable_dynamic_field=True)
    schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True)
    schema.add_field(field_name="embedding", datatype=DataType.FLOAT_VECTOR, dim=embedding_dimensions)

    index_params = milvus_client.prepare_index_params()
    index_params.add_index(field_name="embedding", index_type="AUTOINDEX")

    milvus_client.create_collection(
        collection_name=milvus_collection_name,
        schema=schema,
        index_params=index_params
    )
    ```

7. 设置 Streamlit 界面，提供数据库操作按钮。
    - **Reset 按钮**: 点击后，会清空 MySQL 中的 `Employee` 表和 Milvus 集合中的所有相关数据，用于重置演示环境。
    - **Save Data 按钮**: 点击后，会检查 MySQL 中是否已有数据。如果没有，它会：
        a. 遍历 `employees` 列表，为每个员工的 `intro` 生成 Embedding 向量。
        b. 将员工的完整信息（id, name, intro）存入 MySQL。
        c. 将员工的 ID 和其对应的 Embedding 向量存入 Milvus。
        这个过程展示了如何将结构化元数据和向量数据分别写入 RDB 和向量数据库。

    ```python
    st.title("🔍 RAG Demo")

    st.subheader("Database Operations")

    left, right = st.columns(2)

    if left.button("Reset", type="primary"):
        left.write("Resetting...")
        with Session(engine) as session:
            session.exec(delete(Employee))
            session.commit()

        # Delete the employee id mapping data from Milvus
        milvus_client.delete(
            collection_name=milvus_collection_name,
            ids=[employee.id for employee in employees]
        )

        st.rerun()

    left.markdown("This option <span style='color: red;'>will delete all data</span></h3>", unsafe_allow_html=True)

    if right.button("Save Data"):
        right.write("Saving...")

        # Prepare RDB data
        with Session(engine) as session:
            count = session.exec(select(func.count()).select_from(Employee)).one()

            if count == 0:
                employee_intro_embedding = [
                    {
                        "id": employee.id,
                        "embedding": get_embedding(employee.intro)
                    } for employee in employees
                ]
                session.add_all(employees)
                session.commit()

                # Insert the employee id mapping data into Milvus
                milvus_client.insert(
                    collection_name=milvus_collection_name,
                    data=employee_intro_embedding
                )

                st.rerun()
            else:
                st.error("Data already exists")
    ```

8. 创建了两个可折叠的部分（`st.expander`），分别用于显示当前存储在 MySQL 和 Milvus 中的数据。这有助于用户直观地看到数据操作的结果。

    ```python
    with st.expander("📁 All Employees in the MySQL", expanded=False):
        with Session(engine) as session:
            employees = session.exec(select(Employee)).all()
            if employees:
                data = [{'id': employee.id, 'name': employee.name, 'intro': employee.intro} for employee in employees]
                df = pd.DataFrame(data)
                st.dataframe(df)
            else:
                st.info("No data found in MySQL.")

    with st.expander("📁 All Employees ID and Embedding Mapping in the Milvus", expanded=False):
        employees = milvus_client.get(
            collection_name=milvus_collection_name,
            ids=[employee.id for employee in employees],
            output_fields=["id", "embedding"]
        )

        if employees:
            data = [{'id': employee['id'], 'embedding': employee['embedding']} for employee in employees]
            df = pd.DataFrame(data)
            st.dataframe(df)
        else:
            st.info("No data found in Milvus.")
    ```

9. 这部分是 RAG 应用的核心交互逻辑。当用户输入问题并点击 "Send" 时：
    - **向量化查询**: 首先使用 `get_embedding` 函数将用户的问题文本转换为查询向量。
    - **向量检索 (Retrieval)**: 使用 `milvus_client.search` 在 Milvus 集合中进行向量相似度搜索（使用余弦距离），找出与查询向量最相似的一条记录，并获取其 `id`。
    - **数据获取 (Fetch)**: 使用从 Milvus 中获取的 `id`，到 MySQL 的 `Employee` 表中查询并获取对应的完整员工信息（姓名、简介等）。
    - **生成回答 (Generation)**:
        - 如果模式是 `Retrieval Only`，则直接展示从 Milvus 和 MySQL 中检索到的原始数据。
        - 如果模式是 `RAG Q&A`，则将从 MySQL 获取的员工简介（`intro`）作为上下文（`context`），连同用户原始问题，填入 `RAG_PROMPT_TEMPLATE` 提示模板中。
        - 最后，调用 `litellm.completion` 将构建好的提示发送给 LLM，生成最终的回答，并将其以美化的样式显示在界面上。同时，提供一个可折叠部分来展示用于生成该回答的"知识来源"。
    - 如果向量搜索没有找到相关信息，则提示用户。

    ```python
    st.write(
        "Enter your question, and the system will retrieve relevant knowledge and generate an answer"
    )
    mode = st.radio("Select Mode:", ["Retrieval Only", "RAG Q&A"])

    query = st.text_input("Enter your question:", "")

    if st.button("Send") and query:
        with st.spinner("Processing..."):
            # Retrieve relevant chunks
            query_embedding = get_embedding(query)
            res = milvus_client.search(
                collection_name=milvus_collection_name,
                anns_field="embedding",
                data=[query_embedding],
                limit=1,
                search_params={"metric_type": "COSINE"},
                output_fields=["id", "embedding"]
            )

            retrieved_milvus = [{"id": item.entity.id, "embedding": item.entity.embedding} for item in res[0]]
            ids = [item.entity.id for item in res[0]]

            retrieved_mysql = []
            with Session(engine) as session:
                employees = session.exec(select(Employee).where(Employee.id.in_(ids))).all()
                retrieved_mysql = [{'id': employee.id, 'name': employee.name, 'intro': employee.intro} for employee in employees]

            if retrieved_milvus and retrieved_mysql:
                if mode == "Retrieval Only":
                    st.write("### Retrieved Embedding and ID in Milvus:")
                    st.dataframe(pd.DataFrame(retrieved_milvus))

                    st.write("### Retrieved Data in MySQL:")
                    st.dataframe(pd.DataFrame(retrieved_mysql))
                else:
                    intro = [employee['intro'] for employee in retrieved_mysql]

                    # Build RAG prompt
                    context = "\n".join(intro)
                    prompt = RAG_PROMPT_TEMPLATE.format(context=context, question=query)

                    # Call LLM to generate answer
                    response = completion(
                        model=llm_model,
                        messages=[{"content": prompt, "role": "user"}],
                    )

                    st.markdown(f"### 🤖 {llm_model}")
                    st.markdown(LLM_RESPONSE_STYLE, unsafe_allow_html=True)

                    # show the response
                    st.markdown(
                        f'<div class="llm-response">{response.choices[0].message.content}</div>',
                        unsafe_allow_html=True,
                    )

                    with st.expander("📚 Retrieved Knowledge", expanded=False):
                        st.write("### Retrieved Embedding and ID in Milvus:")
                        st.dataframe(pd.DataFrame(retrieved_milvus))

                        st.write("### Retrieved Data in MySQL:")
                        st.dataframe(pd.DataFrame(retrieved_mysql))

            else:
                st.info("No relevant information found")
    ```

### 步骤 4. RAG AI 应用程序的总结

我们可以很清楚的比对出简易 AI 应用程序和 RAG AI 应用程序的方案优劣势。简易 AI 应用程序简单，但是其背景知识均构建在训练集内，若问题在训练集外，则大概率会产生幻觉。而 RAG AI 应用程序很好的解决了这一点，而且带来了我们在模块 2 开头提出的一系列优势，但是它仅仅在最简易的 demo 应用程序中，也带来了相当高的复杂度。

接下来，我们需要优化 RAG 应用程序，让其准确率进一步提升。

## 模块 4: 使用多路召回改进你的 RAG AI 应用程序

在模块 2 的开头，我们提到了 RAG 不仅仅可以使用 Embedding 向量作为我们的召回渠道。你可以使用如：

- BoW
- TF-IDF
- Word2Vec

等方法来生成文本向量，随后依旧计算语义相似度。
同时，你也可以引入图数据库，将文本抽取为图后，进行图搜索。换个词来说，GraphRAG。
当然，你也可以用下面我们叙述的方法，使用全文检索（Full-text Search）的方法进行第二路召回。这将给你的 RAG 以另一个角度进行召回，丰富召回的内容。

全文检索，是一类可以让你在大量文本中，快速通过关键词搜索到文本块的技术。常被用于搜索引擎、广告等业务中。一种常见全文检索的实现方式是使用倒排索引。它的核心流程是：

1. 在文本存储的时候，使用分词器将其分开成一个个的词。
2. 以词为 key，文本的 ID 为 value，建立索引
3. 在搜索关键词时，利用此索引，快速检索到文本的 ID 列表

幸运的是，我们不用自己实现这一套流程，在传统方法中，我们可以引入一个搜索引擎。ElasticSearch 就是一个常用的搜索引擎，我们可以使用它作为我们多路召回的底层存储。

### 步骤 1. 运行多路召回 RAG AI 应用程序的 Streamlit 服务

再次创建一个新的终端：

![90%](https://lab-static.pingcap.com/quick-demo/demo-422/422-rag-streamlit-start-0606.jpg)

然后运行新的多路召回 RAG AI 应用程序：

```bash
poetry run streamlit run improved_rag.py --server.port 8503
```

```output
Collecting usage statistics. To deactivate, set browser.gatherUsageStats to false.


You can now view your Streamlit app in your browser.

Local URL: http://localhost:8503
Network URL: http://172.18.0.7:8503
External URL: http://${LAB:HOST_MONITOR1_PUBLIC_IP}:8503
```

### 步骤 2. 使用多路召回 RAG AI 应用程序的演示网页

- 打开您的 [RAG AI 应用程序](http://${LAB:HOST_MONITOR1_PUBLIC_IP}:8503/)。

    ![90%](https://lab-static.pingcap.com/quick-demo/demo-422/422-improved-rag-streamlit-ui-0606.jpg)

- 打开名为 `All Employees in the MySQL`, `All Employees ID and Embedding Mapping in the Milvus` 和 `All Employees in the Elasticsearch` 的文件夹，以查看当前 MySQL, Milvus 和 ElasticSearch 中的数据。由于我们在模块 2 中，在 MySQL 和 Milvus 中保存过数据，当前应在 ElasticSearch 中没有数据，而 MySQL 和 Milvus 存在数据。

    ![90%](https://lab-static.pingcap.com/quick-demo/demo-422/422-improved-rag-streamlit-data-start-0613.jpg)

- 点击 `Reset` 删除数据。
- 随后可以看到 `All Employees in the MySQL` 和 `All Employees ID and Embedding Mapping in the Milvus` 的文件夹内为空。

    ![90%](https://lab-static.pingcap.com/quick-demo/demo-422/422-improved-rag-streamlit-data-reset-0606.jpg)

- 点击 `Save Data` 保存数据。此时会保存数据到 MySQL，Mulvus 和 ElasticSearch。

    ![90%](https://lab-static.pingcap.com/quick-demo/demo-422/422-improved-rag-streamlit-data-saved-0613.jpg)

- 点击 `Retrieval Only` 进行上下文检索。
- 输入问题：

    ```text
    Who is Cheese Wang at PingCAP?
    ```

- 点击 `Send` 按钮开始检索。

    ![90%](https://lab-static.pingcap.com/quick-demo/demo-422/422-improved-rag-streamlit-retrieve-0606.jpg)

- 这次的检索，我们得使用了两个不同的召回途径，因此既有 Milvus 的 Embedding Vector 召回 ID，又有 ElasticSearch 的 score。ElasticSearch 默认使用基于 TF-IDF 的 BM25 算法计算相似度得分。你可以认为 score 越高，文本越相似。
- 接下来，将 `Select Mode` 选项更改为 `RAG Q&A`，然后点击 `Send`。

    ![90%](https://lab-static.pingcap.com/quick-demo/demo-422/422-improved-rag-streamlit-qna-0606.jpg)

> [!NOTE]
> 在这个 Demo 中，多路召回与单路召回的信息相同，这是由于这个示例不够复杂所导致的。此处由于篇幅限制，不会给出证明，但你可以参考以下内容：
>
> - MicroSoft: [Doing RAG? Vector search is \*not\* enough](https://techcommunity.microsoft.com/blog/azuredevcommunityblog/doing-rag-vector-search-is-not-enough/4161073)
> - DataStax: [How to Implement Hybrid Search in RAG Pipelines for LLMs](https://www.datastax.com/guides/hybrid-search-rag-pipelines) （此篇文章提出，与传统的仅向量检索相比，混合搜索的答案精度提高了最多 35%）

### 步骤 3. 多路召回 RAG AI 应用程序代码阅读

让我们回到 [VS Code](${LAB:SECURITY_ACCESS_PROXY_URL}?next=${LAB:"/vscode/?folder=/root/tidb-lab-unify-storage"|urlencode}) 并点击代码文件 `improved_rag.py`。

1. 与上一个 RAG 应用相比，脚本的导入部分最主要的变化是增加了 `elasticsearch` 库，用于与 Elasticsearch 服务进行交互。

    ```python
    import os
    import pandas as pd
    from pydantic_core import MultiHostUrl
    from sqlalchemy import func
    from sqlmodel import Field, SQLModel, Session, create_engine, select, delete
    from pymilvus import MilvusClient, DataType
    from elasticsearch import Elasticsearch

    from litellm import completion, embedding
    import streamlit as st

    from typing import List

    from common import LLM_RESPONSE_STYLE, RAG_PROMPT_TEMPLATE
    ```

2. 在配置和初始化部分，除了原有的 MySQL (RDB) 和 Milvus (向量数据库) 的客户端外，新增了 Elasticsearch 客户端的初始化。同时，定义了 Elasticsearch 的索引名称 `ES_INDEX_NAME` 和映射（mapping），该映射指定了存入 Elasticsearch 的文档结构，包括 `id`, `name`, `intro` 等字段及其类型。

    ```python
    # Elasticsearch
    es = Elasticsearch(f"http://{os.getenv('ELASTICSEARCH_HOST')}:{os.getenv('ELASTICSEARCH_PORT')}")
    ES_INDEX_NAME = "employee_id_mapping"
    es_mapping = {
        "mappings": {
            "properties": {
                "id": {"type": "integer"},
                "name": {"type": "text"},
                "intro": {"type": "text"}
            }
        }
    }
    ```

3. 此代码块确保在 Elasticsearch 中存在所需的索引。如果索引不存在，它会使用预定义的 `es_mapping` 创建一个新的索引。

    ```python
    if not es.indices.exists(index=ES_INDEX_NAME):
        es.indices.create(index=ES_INDEX_NAME, body=es_mapping)
    ```

4. 数据库操作的逻辑也进行了扩展。
    - **Reset 按钮**：除了清空 MySQL 和 Milvus 的数据外，现在还会删除并重建 Elasticsearch 中的索引，确保环境的完全重置。
    - **Save Data 按钮**：在将数据存入 MySQL 和 Milvus 之后，会遍历 `employees` 列表，并将每条员工信息作为一个文档索引到 Elasticsearch 中。这使得员工的简介（`intro`）内容可以被全文检索。

    ```python
    if left.button("Reset", type="primary"):
        left.write("Resetting...")
        with Session(engine) as session:
            session.exec(delete(Employee))
            session.commit()
    
        # Delete the employee id mapping data from Milvus
        milvus_client.delete(
            collection_name=milvus_collection_name,
            ids=[employee.id for employee in employees]
        )
    
        if es.indices.exists(index=ES_INDEX_NAME):
            es.indices.delete(index=ES_INDEX_NAME)
            es.indices.create(index=ES_INDEX_NAME, body=es_mapping)
    
        st.rerun()
    
    if right.button("Save Data"):
        right.write("Saving...")
    
        # Prepare RDB data
        with Session(engine) as session:
            count = session.exec(select(func.count()).select_from(Employee)).one()
    
            if count == 0:
                employee_intro_embedding = [
                    {
                        "id": employee.id,
                        "embedding": get_embedding(employee.intro)
                    } for employee in employees
                ]
                session.add_all(employees)
                session.commit()
    
                # Insert the employee id mapping data into Milvus
                milvus_client.insert(
                    collection_name=milvus_collection_name,
                    data=employee_intro_embedding
                )
    
                for employee in employees:
                    es.index(
                        index=ES_INDEX_NAME,
                        id=employee.id,
                        document={"id": employee.id, "name": employee.name, "intro": employee.intro}
                    )
    
                st.rerun()
            else:
                st.error("Data already exists")
    ```

5. 为了展示多路召回的效果，UI 上增加了一个新的可折叠部分，用于显示 Elasticsearch 中存储的所有员工数据。

    ```python
    with st.expander("📁 All Employees in the Elasticsearch", expanded=False):
        response = es.search(index=ES_INDEX_NAME, body={"query": {"match_all": {}}})
        employees = [{"id": item["_id"], "name": item["_source"]["name"], "intro": item["_source"]["intro"]} for item in response["hits"]["hits"]]
        st.dataframe(pd.DataFrame(employees))
    ```

6. 这是多路召回 RAG 应用的核心逻辑，它组合了两种不同的检索方式：
    - **向量检索 (Milvus)**: 与之前一样，将用户问题转换为 embedding，然后在 Milvus 中进行语义相似度搜索，获得第一批候选 ID (`milvus_ids`)。
    - **全文检索 (Elasticsearch)**: 将用户的原始问题作为关键词，在 Elasticsearch 的 `intro` 字段上执行 `match` 查询，进行全文搜索，获得第二批候选 ID (`es_ids`) 和相关性得分。
    - **结果合并**: 将从 Milvus 和 Elasticsearch 两路召回的 ID 列表进行合并，并使用 `set` 去重，得到一个更全面的候选文档 ID 集合。
    - **数据获取与生成**: 使用合并后的 `ids` 从 MySQL 中获取完整的员工信息，然后构建提示并调用 LLM 生成答案。这个过程与之前的 RAG 应用相同，但现在它基于一个更丰富的上下文信息。

    ```python
    if st.button("Send") and query:
        with st.spinner("Processing..."):
            # Retrieve relevant chunks in Milvus
            query_embedding = get_embedding(query)
            res = milvus_client.search(
                collection_name=milvus_collection_name,
                anns_field="embedding",
                data=[query_embedding],
                limit=1,
                search_params={"metric_type": "COSINE"},
                output_fields=["id", "embedding"]
            )
    
            retrieved_milvus = [{"id": item.entity.id, "embedding": item.entity.embedding} for item in res[0]]
            milvus_ids = [item.entity.id for item in res[0]]
    
            # Retrieve relevant chunks in Elasticsearch
            es_query = {
                "query": {
                    "match": {
                        "intro": query
                    }
                },
                "size": 1
            }
            res = es.search(index=ES_INDEX_NAME, body=es_query)
            retrieved_es = [{"id": item["_id"], "score": item["_score"]} for item in res["hits"]["hits"]]
            es_ids = [int(item["_id"]) for item in res["hits"]["hits"]]
    
            # Merge the ids from Milvus and Elasticsearch
            ids = list(set(milvus_ids + es_ids))
    
            # Retrieve relevant chunks in RDB
            retrieved_mysql = []
            with Session(engine) as session:
                employees = session.exec(select(Employee).where(Employee.id.in_(ids))).all()
                retrieved_mysql = [{'id': employee.id, 'name': employee.name, 'intro': employee.intro} for employee in employees]
    
            if retrieved_milvus and retrieved_mysql:
                if mode == "Retrieval Only":
                    st.write("### Retrieved Embedding and ID in Milvus:")
                    st.dataframe(pd.DataFrame(retrieved_milvus))
    
                    st.write("### Retrieved ID and Score in Elasticsearch:")
                    st.dataframe(pd.DataFrame(retrieved_es))
                    
                    st.write("### Retrieved Data in MySQL:")
                    st.dataframe(pd.DataFrame(retrieved_mysql))
                else:
                    intro = [employee['intro'] for employee in retrieved_mysql]
    
                    # Build RAG prompt
                    context = "\n".join(intro)
                    prompt = RAG_PROMPT_TEMPLATE.format(context=context, question=query)
    
                    # Call LLM to generate answer
                    response = completion(
                        model=llm_model,
                        messages=[{"content": prompt, "role": "user"}],
                    )
    
                    st.markdown(f"### 🤖 {llm_model}")
                    st.markdown(LLM_RESPONSE_STYLE, unsafe_allow_html=True)
    
                    # show the response
                    st.markdown(
                        f'<div class="llm-response">{response.choices[0].message.content}</div>',
                        unsafe_allow_html=True,
                    )
    
                    with st.expander("📚 Retrieved Knowledge", expanded=False):
                        st.write("### Retrieved Embedding and ID in Milvus:")
                        st.dataframe(pd.DataFrame(retrieved_milvus))
    
                        st.write("### Retrieved ID and Score in Elasticsearch:")
                        st.dataframe(pd.DataFrame(retrieved_es))
    
                        st.write("### Retrieved Data in MySQL:")
                        st.dataframe(pd.DataFrame(retrieved_mysql))
    
            else:
                st.info("No relevant information found")
    ```

### 步骤 4. 多路召回 RAG AI 应用程序的总结

我们改进了 RAG 应用的能力，让准确率进一步提升了！但是缺点也很明显，我们现在有三个数据存储，代码也变得极其复杂，每次数据的更新都需要同时编写三处，数据同步和 Error Handling 也变得复杂无比。是的，我们让 AI 应用程序变得更优秀了，but at what cost? 我们牺牲了应用程序的可维护性，换来了服务的能力提升。

那我们有没有一个办法，在不牺牲应用程序的可维护性的前提下，提升服务能力呢？有的！下一步，我们将使用 TiDB 作为存储，简化我们的服务。

## 模块 5: Use TiDB to simplify Your AI Application

到此为止，我们已经完成了一个多路召回的 RAG 应用。它的准确率已经让我们颇为满意了。但是，它的复杂性和可维护性却差强人意。这是由于我们在传统的解决方案里，不得不引入非常多的数据库作为存储底层。回想一下，我们有用到了哪些数据库？

- MySQL
- Milvus
- ElasticSearch

而现在有一个更优的解决方案：TiDB Serverless！TiDB Serverless 兼容 MySQL 语法，而且支持向量，以及全文检索。

我们在这个模块中，将使用 `pytidb` 库重写我们的多路召回应用。观察我们可以将服务简化到什么程度。

### 步骤 1. 运行 TiDB AI 应用程序的 Streamlit 服务

再次创建一个新的终端：

![90%](https://lab-static.pingcap.com/quick-demo/demo-422/422-rag-streamlit-start-0606.jpg)

然后运行新的 TiDB AI 应用程序：

```bash
poetry run streamlit run tidb_rag.py --server.port 8504
```

```output
Collecting usage statistics. To deactivate, set browser.gatherUsageStats to false.


You can now view your Streamlit app in your browser.

Local URL: http://localhost:8504
Network URL: http://172.18.0.7:8504
External URL: http://${LAB:HOST_MONITOR1_PUBLIC_IP}:8504
```

### 步骤 2. 使用 TiDB AI 应用程序的演示网页

- 打开您的 [TiDB AI 应用程序](http://${LAB:HOST_MONITOR1_PUBLIC_IP}:8504/)，因为 TiDB Serverless 并不在该服务器上，且需要初始化，总体初始化时间约 5 分钟，请耐心等待。初始化结束后的页面如下：

    ![90%](https://lab-static.pingcap.com/quick-demo/demo-422/422-tidb-rag-streamlit-ui-0606.jpg)

- 打开名为 `All Employees in the Database` 的文件夹，以查看当前 TiDB 中的数据。由于我们尚未保存任何数据，所以它应该是空的。

    ![90%](https://lab-static.pingcap.com/quick-demo/demo-422/422-tidb-rag-streamlit-data-start-0606.jpg)

- 点击 `Save Data` 保存数据。
- 保存数据后，您可以在 `All Employees in the Database` 文件夹中看到 TiDB 中的数据。我们有 4 列，分别是 `id`, `name`, `intro` 和 `intro_vec`。TiDB 比 MySQL 多了一列 `intro_vec`，这是因为 TiDB 支持向量类型。我们可以直接存储向量在 TiDB 中。

    ![90%](https://lab-static.pingcap.com/quick-demo/demo-422/422-tidb-rag-streamlit-data-saved-0606.jpg)

- 点击 `Retrieval Only`，输入问题：

    ```text
    Who is Cheese Wang at PingCAP?
    ```

- 点击 `Send` 按钮开始检索。

    ![90%](https://lab-static.pingcap.com/quick-demo/demo-422/422-tidb-rag-streamlit-retrieve-0606.jpg)

- 可以看到，我们召回的信息中，除了信息本身，也存在一些指标：`_distance`, `_score` 和 `_match_score`。 `pytidb` 就是使用这些指标计算是否文本是否相近的。
- 将 `Select Mode` 选项更改为 `RAG Q&A`，然后点击 `Send` 来提供答案。

    ![90%](https://lab-static.pingcap.com/quick-demo/demo-422/422-tidb-rag-streamlit-qna-0606.jpg)

- 我们完美复刻了多路召回的 RAG，下面让我们看看代码的复杂度吧。

### 步骤 3. TiDB AI 应用程序代码阅读

让我们回到 [VS Code](${LAB:SECURITY_ACCESS_PROXY_URL}?next=${LAB:"/vscode/?folder=/root/tidb-lab-unify-storage"|urlencode}) 并点击代码文件 `tidb_rag.py`。

1. 最显著的变化是导入的库。我们不再需要 `sqlalchemy`, `pymilvus`, `elasticsearch` 等多个数据库的客户端。现在，我们只需要 `pytidb` 库，它将作为与 TiDB 交互的统一接口。

    ```python
    import os
    import pandas as pd
    
    from litellm import completion
    import streamlit as st
    
    from typing import Optional, Any
    from pytidb import TiDBClient
    from pytidb.schema import TableModel, Field
    from pytidb.embeddings import EmbeddingFunction
    
    from common import LLM_RESPONSE_STYLE, RAG_PROMPT_TEMPLATE
    ```

2. 初始化部分也得到了极大的简化。我们只需要创建一个 `TiDBClient` 实例来连接到 TiDB Serverless。`EmbeddingFunction` 是一个来自 `pytidb` 的辅助类，它帮助我们将 Embedding 模型的调用与数据表模式无缝集成。

    ```python
    db = TiDBClient.connect(
        host=os.getenv("SERVERLESS_CLUSTER_HOST"),
        port=int(os.getenv("SERVERLESS_CLUSTER_PORT")),
        username=os.getenv("SERVERLESS_CLUSTER_USERNAME"),
        password=os.getenv("SERVERLESS_CLUSTER_PASSWORD"),
        database=os.getenv("SERVERLESS_CLUSTER_DATABASE_NAME"),
        enable_ssl=True,
    )
    
    embedding_model = "bedrock/amazon.titan-embed-text-v2:0"
    llm_model = "bedrock/us.amazon.nova-lite-v1:0"
    
    text_embedding_function = EmbeddingFunction(
        embedding_model,
        timeout=60
    )
    ```

3. 这是整个应用简化的核心。我们使用 `pytidb.schema.TableModel` 定义了一个名为 `Employee` 的表。

    - `id`, `name`, `intro` 是标准的文本和数字字段。
    - `intro_vec`: 这是最关键的部分。我们将其定义为 `text_embedding_function.VectorField`，并指定 `source_field="intro"`。这告诉 `pytidb`：当插入或更新数据时，自动获取 `intro` 字段的内容，调用预定义的 Embedding 模型将其转换为向量，然后将结果存储在 `intro_vec` 列中。这一个字段的定义就替代了之前手动生成向量并将其存入 Milvus 的所有逻辑。

    ```python
    table_name = "employees"
    class Employee(TableModel, table=True):
        __tablename__ = table_name
        __table_args__ = {"extend_existing": True}
        id: int | None = Field(default=None, primary_key=True)
        name: str
        intro: str
        intro_vec: Optional[Any] = text_embedding_function.VectorField(
            source_field="intro",
        )
    ```

4. 数据和索引的准备也变得非常简单。
    - `db.open_table` 和 `db.create_table` 确保表的存在。
    - `table.create_fts_index("intro")`: **这一行代码就替代了整个 Elasticsearch 的设置**。它在 `intro` 列上创建了一个全文检索引擎，使我们能够进行关键词搜索。

    ```python
    table = db.open_table(table_name)
    if table is None:
        table = db.create_table(schema=Employee)
    
    if not table.has_fts_index("intro"):
        table.create_fts_index("intro")
    ```

5. 数据操作被简化为单个命令。

    - **Reset 按钮**: `table.truncate()` 一条命令即可清空所有数据。不再需要在多个数据库之间协调删除操作。
    - **Save Data 按钮**: `table.bulk_insert(employees)` 一条命令即可完成所有数据的保存。`pytidb` 在后台会自动处理：为每条 `employee` 记录的 `intro` 生成向量，然后将包括向量在内的完整数据写入 TiDB 的同一张表中。全文索引也会自动更新。

    ```python
    if left.button("Reset", type="primary"):
        left.write("Resetting...")
        table.truncate()
        st.rerun()

    if right.button("Save Data"):
        right.write("Saving...")
        if table.rows() == 0:
            table.bulk_insert(employees)
            st.rerun()
    ```

6. 这是最惊人的简化。整个多路召回的复杂逻辑现在被浓缩成了一行代码。
    - `table.search(query, search_type="hybrid").text_column("intro").limit(1)`: 这一行代码完成了之前需要多个步骤才能实现的功能：
        1. 它将用户查询 `query` 同时用于**向量相似度搜索**（在 `intro_vec` 列上）和**全文检索**（在 `intro` 列上）。
        2. `search_type="hybrid"` 参数指示 `pytidb` 自动执行这两种搜索，并使用内置的融合算法对两路结果进行排序和合并。
        3. 返回一个统一的、最相关的结果列表。
    - 我们不再需要手动调用 embedding 模型、分别查询 Milvus 和 Elasticsearch、合并 ID、最后再查询 MySQL。所有这些都由 `pytidb` 在一个步骤中优雅地完成了。

    ```python
    if st.button("Send") and query:
        with st.spinner("Processing..."):
            # Retrieve relevant employees
            results = table.search(query, search_type="hybrid").text_column("intro").limit(1)
    
            if results:
                if mode == "Retrieval Only":
                    st.write("### Hybrid Search Results:")
                    st.dataframe(results.to_pandas())
                else:
                    intro = [employee.intro for employee in results.to_rows()]
    
                    # Build RAG prompt
                    context = "\n".join(intro)
                    prompt = RAG_PROMPT_TEMPLATE.format(context=context, question=query)
    
                    # Call LLM to generate answer
                    response = completion(
                        model=llm_model,
                        messages=[{"content": prompt, "role": "user"}],
                    )
    
                    st.markdown(f"### 🤖 {llm_model}")
                    st.markdown(LLM_RESPONSE_STYLE, unsafe_allow_html=True)
    
                    # show the response
                    st.markdown(
                        f'<div class="llm-response">{response.choices[0].message.content}</div>',
                        unsafe_allow_html=True,
                    )
    
                    with st.expander("📚 Hybrid Search Retrieved Knowledge", expanded=False):
                        st.dataframe(results.to_pandas())
            else:
                st.info("No relevant information found")
    ```

### 步骤 4. TiDB AI 应用程序的总结

可以清楚的看到，完成了同样的需求的情况下，使用 TiDB 极大的简化了代码的实现。

## 结论

在 AI 应用程序的开发中，关系型数据库、向量计算、全文索引变成了经常使用的一些特性。而由于市面上解决方案的不统一，我们需要引入关系型数据库、向量数据库、搜索引擎。这造成了我们的解决方案异常复杂。而这种复杂性实际上是由数据库的能力缺失导致的。

![70%](https://lab-static.pingcap.com/quick-demo/demo-422/422-with-bg.png)

TiDB 解决了此问题，让代码变得更简单、可维护。开发者不必使用数个数据存储，并在存储之间转移数据。而是可以一站式的完成数据的存储和读取。从而在 AI 的领域中，极大的简化了开发应用程序的门槛。
