## モジュール 5: TiDBによるAIアプリケーションの簡素化

これまでに、マルチチャネルRAGアプリケーションを完成させました。その精度はかなり満足のいくものです。しかし、その複雑さと保守性は理想とはほど遠いものです。これは、従来のソリューションでは、基盤となるストレージとして多くのデータベースを導入する必要があるためです。使用したデータベースを思い出してみましょう：

- MySQL
- Milvus
- ElasticSearch

しかし、今ではより良いソリューションがあります：TiDB Serverless！TiDB ServerlessはMySQL構文と互換性があり、ベクトルと全文検索をサポートしています。

このモジュールでは、`pytidb`ライブラリを使用してマルチチャネルRAGアプリケーションを書き直します。サービスをどの程度簡素化できるか見てみましょう。

### 手順 1. TiDB AIアプリケーションのStreamlitサービスの実行

もう一度新しいターミナルを作成します：

![90%](https://lab-static.pingcap.com/images/2025/7/24/51b77b5a782acbdfb63d2624ab21160896be7451.png)

次に、新しい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アプリケーションのデモWebページの使用

- [TiDB AIアプリケーション](http://${LAB:HOST_MONITOR1_PUBLIC_IP}:8504/)を開きます。TiDB Serverlessはこのサーバー上にはなく、初期化が必要なため、全体の初期化には約5分かかります。しばらくお待ちください。初期化後のページは次のようになります：

    ![90%](https://lab-static.pingcap.com/images/2025/7/24/0cd42e0af28ea7ee8399e477bee7f6eff1ddbf8e.png)

- `All Employees in the Database`という名前のフォルダを開き、TiDBの現在のデータを確認します。まだデータを保存していないので、空のはずです。

    ![90%](https://lab-static.pingcap.com/images/2025/7/24/895d46dd4b5c5c2eb80b8b721ea9b3ea2ab4fe56.png)

- `Save Data`をクリックしてデータを保存します。
- データを保存した後、`All Employees in the Database`フォルダでTiDBのデータを見ることができます。`id`、`name`、`intro`、`intro_vec`の4つの列があります。TiDBはMySQLよりも1つ多い列`intro_vec`を持っています。これは、TiDBがベクトル型をサポートしているためです。ベクトルを直接TiDBに保存できます。

    ![90%](https://lab-static.pingcap.com/images/2025/7/24/acdadeafd99195d48d4bcc4e738682af226be454.png)

- `Retrieval Only`をクリックし、質問を入力します：

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

- `Send`ボタンをクリックして検索を開始します。

    ![90%](https://lab-static.pingcap.com/images/2025/7/24/5ff98d1438d0eae3abd7e2fca36841e5d9c75f00.png)

- ご覧の通り、取得された情報には情報自体に加えて、`_distance`、`_score`、`_match_score`などのメトリクスが含まれています。`pytidb`はこれらのメトリクスを使用してテキストの類似性を計算します。
- `Select Mode`オプションを`RAG Q&A`に変更し、`Send`をクリックして回答を提供します。

    ![90%](https://lab-static.pingcap.com/images/2025/7/24/792b0b07af6ffb5fdbecb3519944827546b5499a.png)

- マルチチャネルRAGを完璧に再現しました。では、コードの複雑さを見てみましょう。

### 手順 3. TiDB AIアプリケーションのコード読解

[VS Code](${LAB:HTTPS_REVERSE_PROXY_URL}?next=${LAB:"/vscode/?folder=/root/tidb-lab-unify-storage-openai"|urlencode} "VS Code")に戻り、`tidb_rag.py`コードファイルをクリックしましょう。

1. 最も重要な変更点は、インポートされるライブラリです。`sqlalchemy`、`pymilvus`、`elasticsearch`のような複数のデータベース用のクライアントはもう必要ありません。今や、TiDBとの対話のための統一インターフェースとして機能する`pytidb`ライブラリだけが必要です。

    ```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. 初期化部分も大幅に簡素化されています。TiDB Serverlessに接続するために`TiDBClient`インスタンスを作成するだけです。`EmbeddingFunction`は`pytidb`のヘルパークラスで、埋め込みモデルの呼び出しをデータテーブルスキーマとシームレスに統合するのに役立ちます。

    ```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 = "openai/text-embedding-3-small"
    llm_model = "openai/gpt-4o"

    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`フィールドの内容を自動的に取得し、事前定義された埋め込みモデルを呼び出してベクトルに変換し、その結果を`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. データ操作は単一のコマンドに簡素化されました。

    - **リセットボタン**: `table.truncate()` は1つのコマンドで全てのデータをクリアします。複数のデータベースにまたがる削除操作を調整する必要はもうありません。
    - **データ保存ボタン**: `table.bulk_insert(employees)` は1つのコマンドで全てのデータを保存します。`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. 統一された、最も関連性の高い結果のリストを返します。
    - 埋め込みモデルを手動で呼び出し、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を使用することで、コードの実装が大幅に簡素化されたことが明らかです。