## Module 5: Use TiDB to Simplify Your AI Application

So far, we have completed a multi-channel RAG application. Its accuracy is quite satisfactory. However, its complexity and maintainability are far from ideal. This is because, in traditional solutions, we have to introduce many databases as the underlying storage. Let's recall which databases we used:

- MySQL
- Milvus
- ElasticSearch

But now there is a better solution: TiDB Serverless! TiDB Serverless is compatible with MySQL syntax and supports vectors as well as full-text search.

In this module, we will rewrite our multi-channel RAG application using the `pytidb` library. Let's see to what extent we can simplify our service.

### Step 1. Run the TiDB AI Application's Streamlit Service

Create a new terminal again:

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

Then run the new TiDB AI application:

```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
```

### Step 2. Use the TiDB AI Application's Demo Web Page

- Open your [TiDB AI Application](http://${LAB:HOST_MONITOR1_PUBLIC_IP}:8504/). Since TiDB Serverless is not on this server and needs to initialize, the total initialization time is about 5 minutes. Please be patient. The page after initialization is as follows:

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

- Open the folder named `All Employees in the Database` to view the current data in TiDB. Since we haven't saved any data yet, it should be empty.

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

- Click `Save Data` to save the data.
- After saving the data, you can see the data in TiDB in the `All Employees in the Database` folder. We have 4 columns: `id`, `name`, `intro`, and `intro_vec`. TiDB has one more column than MySQL, `intro_vec`, because TiDB supports vector types. We can store vectors directly in TiDB.

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

- Click `Retrieval Only`, and enter the question:

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

- Click the `Send` button to start the retrieval.

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

- As you can see, the retrieved information includes some metrics in addition to the information itself: `_distance`, `_score`, and `_match_score`. `pytidb` uses these metrics to calculate text similarity.
- Change the `Select Mode` option to `RAG Q&A`, and then click `Send` to provide an answer.

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

- We have perfectly replicated the multi-channel RAG. Now let's look at the code complexity.

### Step 3. TiDB AI Application Code Reading

Let's go back to [VS Code](${LAB:HTTPS_REVERSE_PROXY_URL}?next=${LAB:"/vscode/?folder=/root/tidb-lab-unify-storage-openai"|urlencode} "VS Code") and click on the code file `tidb_rag.py`.

1. The most significant change is the imported libraries. We no longer need clients for multiple databases like `sqlalchemy`, `pymilvus`, `elasticsearch`. Now, we only need the `pytidb` library, which will serve as a unified interface for interacting with 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. The initialization part is also greatly simplified. We only need to create a `TiDBClient` instance to connect to TiDB Serverless. `EmbeddingFunction` is a helper class from `pytidb` that helps us seamlessly integrate the Embedding model calls with the data table schema.

    ```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. This is the core of the application's simplification. We use `pytidb.schema.TableModel` to define a table named `Employee`.

    - `id`, `name`, `intro` are standard text and numeric fields.
    - `intro_vec`: This is the most critical part. We define it as `text_embedding_function.VectorField` and specify `source_field="intro"`. This tells `pytidb`: when inserting or updating data, automatically fetch the content of the `intro` field, call the predefined Embedding model to convert it into a vector, and then store the result in the `intro_vec` column. The definition of this single field replaces all the logic of manually generating vectors and storing them in 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. Data and index preparation has also become very simple.
    - `db.open_table` and `db.create_table` ensure the table's existence.
    - `table.create_fts_index("intro")`: **This single line of code replaces the entire Elasticsearch setup**. It creates a full-text search index on the `intro` column, enabling us to perform keyword searches.

    ```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. Data operations are simplified to single commands.

    - **Reset button**: `table.truncate()` clears all data with a single command. There is no longer a need to coordinate delete operations across multiple databases.
    - **Save Data button**: `table.bulk_insert(employees)` saves all data with a single command. `pytidb` automatically handles the process in the background: generating a vector for each `employee` record's `intro`, and then writing the complete data, including the vector, into the same table in TiDB. The full-text index is also updated automatically.

    ```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. This is the most amazing simplification. The entire complex logic of multi-channel retrieval is now condensed into a single line of code.
    - `table.search(query, search_type="hybrid").text_column("intro").limit(1)`: This one line of code accomplishes what previously required multiple steps:
        1. It uses the user query `query` for both **vector similarity search** (on the `intro_vec` column) and **full-text search** (on the `intro` column).
        2. The `search_type="hybrid"` parameter instructs `pytidb` to automatically perform both searches and use a built-in fusion algorithm to rank and merge the results from both channels.
        3. It returns a unified, most relevant list of results.
    - We no longer need to manually call the embedding model, query Milvus and Elasticsearch separately, merge the IDs, and finally query MySQL. All of this is elegantly handled by `pytidb` in a single step.

    ```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")
    ```

### Step 4. Summary of the TiDB AI Application

It is clear to see that, by using TiDB to meet the same requirements, the code implementation has been greatly simplified.