## Module 3: Improve Your AI Application via RAG
In Module 2, we completed a preliminary AI application, but we encountered a serious hallucination problem when the LLM answered questions outside its training data. So how do we solve this? The answer is also simple.
We can pre-fetch the information for the LLM, place it in the input when calling the LLM, and use the LLM to summarize the content, rather than relying on the knowledge from its training set. By doing this, we can get the following benefits:
1. Reduce Hallucinations: Transformer-based LLMs use a multi-head attention model to make sequence-based predictions within a context window. If we put relevant knowledge into the context window, it will generate answers that are closer to the relevant knowledge. On the surface, it reduces the frequency or severity of hallucinations.
2. Provide External Knowledge: Again, because a Transformer-based LLM has a context window (reflected as a token size limit in the API), we cannot put all the information into the prompt. Therefore, we need to select the relevant content beforehand.
3. Reduce Costs: Even though LLMs' context windows are getting larger, we still need to consider costs and should not waste too many tokens on each question. Because most LLMs charge by the token (you can think of a token as a variable positively correlated with the length of the text; the longer your text, the more tokens).
So, the only problem we face now is:
**How to help the LLM retrieve relevant information?**
There is more than one way to achieve this. In Module 3, we will use one of them: the text Embedding vector method.
> [!NOTE]
> **Optional Reading**
>
> An Embedding vector is a vector that an Embedding model converts text into. An Embedding model is a type of model specially trained to distinguish semantics. Simply put, during contrastive learning training, it gets many sets of samples, and each set is divided into:
>
> - An anchor sample: a piece of text
> - A positive sample: a piece of text semantically similar to the anchor sample
> - A set of negative samples: many pieces of text semantically dissimilar to the anchor sample
>
> The task of the Embedding model is to:
>
> - Minimize the distance between the output vectors of the anchor sample and the positive sample
> - Maximize the distance between the output vectors of the anchor sample and the negative samples
>
> More correctly, the InfoNCE loss function is used here to train the model, performing the stochastic gradient descent algorithm to fit the samples.
If you don't understand the optional reading part, that's completely fine; it's just an explanation of the principle. You just need to know: the greater the cosine distance between the Embedding vectors of two texts, the less similar their semantics are. And the smaller the cosine distance between the Embedding vectors, the more similar their semantics are.
Therefore, we can use this feature to make this improvement. In Module 3, we will use Milvus as the vector search engine and MySQL as the relational database to store the original text.
> [!NOTE]
> **Why not use the metadata of a vector database to store the original text?**
>
> This is technically feasible, but due to the lack of capabilities in vector databases, such as transactions, complete data consistency cannot be guaranteed. In actual engineering, if you still use a separate architecture of a relational database and a vector database, it is recommended to use the relational database as the source of truth for the data. At the same time, the data pair of the original text ID and the Embedding vector should be updated to the vector database asynchronously.
### Step 1. Run the RAG AI Application's Streamlit Service
Create a new terminal:
Then run the new RAG AI application:
```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
```
### Step 2. Use the RAG AI Application's Demo Web Page
- Open your [RAG AI Application](http://${LAB:HOST_MONITOR1_PUBLIC_IP}:8502/).
- Open the folders named `All Employees in the MySQL` and `All Employees ID and Embedding Mapping in the Milvus` to view the current data in MySQL and Milvus. Since we haven't saved any data yet, they should be empty.
- Click `Save Data` to save the mock employee information data we prepared for this demo. You can see the original data in the next code reading step, which is a list of employee objects.
- After saving the data, you can see the data in MySQL in the `All Employees in the MySQL` folder. We have 3 columns: `id`, `name`, and `intro`. And in `All Employees ID and Embedding Mapping in the Milvus`, you can see that we have saved the relationship between `id` and `embedding`. At this time, the `id` in Milvus and the `id` in MySQL are one-to-one mapped.
- If you only want to retrieve the context, you can click the `Retrieval Only` option. Then, enter your question.
- Let's try if the RAG AI application can select relevant information through Embedding. Enter the question:
```text
Who is Cheese Wang at PingCAP?
```
- Click the `Send` button to start the retrieval.
- As you can see, we got the relevant information about `Cheese Wang`.
- Next, the moment of truth, let's test if the RAG AI application can now answer the question that the simple AI application couldn't. Change the `Select Mode` option to `RAG Q&A`, and then click `Send`.
- This time, the RAG AI application will give us an answer, not a list of information.
- As you can see, because we provided the information that the LLM did not get during its training phase, the LLM no longer thinks I am the CEO of PingCAP, but answers my information truthfully.
### Step 3. 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"|urlencode} "VS Code") and click on the code file `rag.py`.
1. The script first imports the many libraries needed to build the RAG application:
- `os`, `dotenv`: Used to load and read environment variables, such as database passwords and host addresses.
- `pandas`: Used for data processing and displaying tables in Streamlit.
- `sqlalchemy`, `sqlmodel`: Used to connect to a relational database (in this case, MySQL) and define the data model. `SQLModel` provides the convenience of combining SQLAlchemy and Pydantic.
- `pymilvus`: The Python client for Milvus, used to connect and operate the Milvus vector database.
- `litellm`: A library for standardized interaction with various Large Language Models (LLMs) and Embedding models.
- `streamlit`: Used to create and run the web user interface.
- `common`: Imports predefined RAG prompt templates and CSS styles from a common module.
```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. Initialization and configuration:
- `create_engine`: Creates a connection engine to the MySQL database using configurations read from environment variables (host, port, password, etc.).
- `embedding_model`, `llm_model`: Defines the names of the Embedding model and LLM to be used.
- `MilvusClient`: Initializes the Milvus client to connect to the Milvus service running at the specified URI.
- `milvus_collection_name`: Names the collection to be created in Milvus.
```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. Defines a helper function `get_embedding`, which takes a text string, calls the `litellm.embedding` API to convert the text into a 1024-dimension vector, and returns the vector.
```python
def get_embedding(text: str) -> List[float]:
response = embedding(
model=embedding_model,
input=[text],
dimensions=embedding_dimensions
)
return response.data[0]['embedding']
```
4. Uses `SQLModel` to define the structure of the `Employee` table, which will be created in MySQL to store the original employee information (ID, name, and intro). `SQLModel.metadata.create_all(engine)` ensures that this table is created in the database.
```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. Creates a list `employees` containing information for three employees as sample data for the demo.
```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. Defines the data schema and index for Milvus.
- `MilvusClient.create_schema`: Defines the structure of the collection, containing a primary key `id` (to associate with the `Employee` record in MySQL) and an `embedding` vector field.
- `milvus_client.prepare_index_params`: Prepares the index parameters for the `embedding` field, using `AUTOINDEX` to let Milvus automatically select the best index type.
- `milvus_client.create_collection`: Creates a new collection in Milvus using the defined schema and index. This allows us to perform efficient vector similarity searches.
```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. Sets up the Streamlit interface, providing database operation buttons.
- **Reset button**: When clicked, it will clear all related data from the `Employee` table in MySQL and the Milvus collection, used to reset the demo environment.
- **Save Data button**: When clicked, it checks if data already exists in MySQL. If not, it will:
a. Iterate through the `employees` list to generate an Embedding vector for each employee's `intro`.
b. Save the complete employee information (id, name, intro) into MySQL.
c. Save the employee's ID and its corresponding Embedding vector into Milvus.
This process demonstrates how to write structured metadata and vector data into an RDB and a vector database, respectively.
```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 will delete all data", 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. Creates two expandable sections (`st.expander`) to display the data currently stored in MySQL and Milvus, respectively. This helps users intuitively see the results of data operations.
```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. This part is the core interaction logic of the RAG application. When the user enters a question and clicks "Send":
- **Vectorize Query**: First, use the `get_embedding` function to convert the user's question text into a query vector.
- **Vector Retrieval**: Use `milvus_client.search` to perform a vector similarity search (using cosine distance) in the Milvus collection to find the most similar record and get its `id`.
- **Data Fetch**: Use the `id` obtained from Milvus to query the `Employee` table in MySQL and fetch the corresponding complete employee information (name, intro, etc.).
- **Generation**:
- If the mode is `Retrieval Only`, directly display the raw data retrieved from Milvus and MySQL.
- If the mode is `RAG Q&A`, use the employee intro (`intro`) fetched from MySQL as the context, along with the user's original question, to fill in the `RAG_PROMPT_TEMPLATE`.
- Finally, call `litellm.completion` to send the constructed prompt to the LLM, generate the final answer, and display it on the interface in a beautified style. At the same time, provide an expandable section to show the "knowledge source" used to generate the answer.
- If the vector search finds no relevant information, inform the user.
```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'{response.choices[0].message.content}
',
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")
```
### Step 4. Summary of the RAG AI Application
We can clearly compare the pros and cons of the simple AI application and the RAG AI application. The simple AI application is straightforward, but its background knowledge is all built into the training set. If a question is outside the training set, it is likely to produce hallucinations. The RAG AI application solves this well and brings a series of advantages that we proposed at the beginning of Module 3, but even in the simplest demo application, it also brings considerable complexity.
Next, we need to optimize the RAG application to further improve its accuracy.