## Module 4: Improve Your RAG AI Application with Multi-channel Retrieval

At the beginning of Module 3, we mentioned that RAG can use more than just Embedding vectors as our retrieval channel. You can use methods like:

- BoW
- TF-IDF
- Word2Vec

to generate text vectors, and then still calculate semantic similarity.

At the same time, you can introduce a graph database, extract the text to a graph. Then searching the graph. i.e. GraphRAG.

Of course, you can also use the method described below, using Full-text Search as a second retrieval channel. This will give your RAG another angle for retrieval, enriching the retrieved content.

Full-text search is a type of technology that allows you to quickly search for text chunks in a large amount of text using keywords. It is often used in businesses like search engines and advertising. A common implementation of full-text search is using an inverted index. Its core process is:

1. When storing text, use a tokenizer to split it into individual words.
2. Create an index with the words as keys and the text IDs as values.
3. When searching for keywords, use this index to quickly retrieve a list of text IDs.

Fortunately, we don't have to implement this process ourselves. In the traditional approach, we can introduce a search engine. ElasticSearch is a commonly used search engine, and we can use it as the underlying storage for our multi-channel retrieval.

### Step 1. Run the Multi-channel RAG 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 multi-channel RAG AI application:

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

### Step 2. Use the Multi-channel RAG AI Application's Demo Web Page

- Open your [RAG AI Application](http://${LAB:HOST_MONITOR1_PUBLIC_IP}:8503/).

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

- Open the folders named `All Employees in the MySQL`, `All Employees ID and Embedding Mapping in the Milvus`, and `All Employees in the Elasticsearch` to view the current data in MySQL, Milvus, and ElasticSearch. Since we saved data in MySQL and Milvus in Module 2, there should be no data in ElasticSearch currently, while data exists in MySQL and Milvus.

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

- Click `Reset` to delete the data.
- You can then see that the folders `All Employees in the MySQL` and `All Employees ID and Embedding Mapping in the Milvus` are empty.

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

- Click `Save Data` to save the data. This will save data to MySQL, Milvus, and ElasticSearch.

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

- Click `Retrieval Only` for context retrieval.
- Enter the question:

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

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

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

- For this retrieval, we used two different retrieval channels, so we have both the retrieved ID from Milvus's Embedding Vector and the score from ElasticSearch. ElasticSearch uses the BM25 algorithm based on TF-IDF by default to calculate similarity scores. You can consider that the higher the score, the more similar the text.
- Next, change the `Select Mode` option to `RAG Q&A`, and then click `Send`.

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

> [!NOTE]
> In this Demo, the information from multi-channel retrieval is the same as from single-channel retrieval, which is due to the lack of complexity in this example. Due to space constraints, a proof will not be given here, but you can refer to the following:
>
> - 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) (This article suggests that hybrid search improves answer accuracy by up to 35% compared to traditional vector-only retrieval)

### Step 3. Multi-channel RAG 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 `improved_rag.py`.

1. Compared to the previous RAG application, the main change in the script's import section is the addition of the `elasticsearch` library for interacting with the Elasticsearch service.

    ```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. In the configuration and initialization section, in addition to the existing clients for MySQL (RDB) and Milvus (vector database), the initialization of the Elasticsearch client is added. Also, the index name `ES_INDEX_NAME` and mapping for Elasticsearch are defined, specifying the structure of the documents to be stored, including fields like `id`, `name`, `intro` and their types.

    ```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. This code block ensures that the required index exists in Elasticsearch. If the index does not exist, it creates a new one using the predefined `es_mapping`.

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

4. The database operation logic has also been extended.
    - **Reset button**: In addition to clearing data from MySQL and Milvus, it now also deletes and recreates the index in Elasticsearch, ensuring a complete reset of the environment.
    - **Save Data button**: After saving data to MySQL and Milvus, it iterates through the `employees` list and indexes each employee's information as a document in Elasticsearch. This makes the employee's intro content available for full-text search.

    ```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. To demonstrate the effect of multi-channel retrieval, a new expandable section is added to the UI to display all employee data stored in 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. This is the core logic of the multi-channel RAG application, which combines two different retrieval methods:
    - **Vector Retrieval (Milvus)**: As before, the user's question is converted into an embedding, and then a semantic similarity search is performed in Milvus to get the first batch of candidate IDs (`milvus_ids`).
    - **Full-text Retrieval (Elasticsearch)**: The user's original question is used as a keyword to perform a `match` query on the `intro` field in Elasticsearch, conducting a full-text search to get the second batch of candidate IDs (`es_ids`) and relevance scores.
    - **Result Merging**: The ID lists from both Milvus and Elasticsearch are merged, and `set` is used for deduplication to get a more comprehensive set of candidate document IDs.
    - **Data Fetching and Generation**: The merged `ids` are used to fetch the complete employee information from MySQL, then a prompt is constructed, and the LLM is called to generate an answer. This process is the same as the previous RAG application, but now it is based on a richer context.

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

### Step 4. Summary of the Multi-channel RAG AI Application

We have improved the capability of the RAG application and further enhanced its accuracy! However, the shortcomings are also obvious. We now have three data stores, and the code has become extremely complex. Every data update requires writing code in three places, and data synchronization and error handling have become incredibly complicated. Yes, we made the AI application better, but at what cost? We sacrificed the application's maintainability in exchange for improved service capability.

So, is there a way to improve service capability without sacrificing application maintainability? Yes! In the next step, we will use TiDB as our storage to simplify our service.