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

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

- MySQL
- Milvus
- ElasticSearch

而现在有一个更优的解决方案：TiDB Cloud Starter！TiDB Cloud Starter 兼容 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 Cloud Starter 并不在该服务器上，且需要初始化，总体初始化时间约 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:HTTPS_REVERSE_PROXY_URL}?next=${LAB:"/vscode/?folder=/root/tidb-lab-unify-storage"|urlencode} "VS Code") 并点击代码文件 `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 Cloud Starter。`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 极大的简化了代码的实现。