## Module 3: Run, Experience, and Study the RAG Demo Application
Run a RAG demo application that connects to TiDB Cloud Starter, save your embeddings, retrieve related data, and generate the answer to your query.
### Step 1. Run the RAG Streamlit service
- Create a new terminal:
+ In the VS Code terminal, use `streamlit` command line tool to run the following script to start Streamlit service.
```bash
streamlit run bedrock/rag.py --server.port 8501
```
```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:8501
Network URL: http://172.17.0.2:8501
External URL: http://${LAB:HOST_MONITOR1_PUBLIC_IP}:8501
```
### Step 2. Use RAG demo web page
- Click to access your [RAG demo web page](http://${LAB:HOST_MONITOR1_PUBLIC_IP}:8501/). It will take around **5 minutes** to complete the initialization based on the lab environment setup.
- Open the folder of `All Chunks in the Database` to see the data currently in the TiDB Cloud Starter cluster. We did not save any data yet, so it should be empty.
- Click `Save Data` to save the sample text data that we prepared for this demo. You can see the sample data in the next step, which is a list of string.
- After saving the data, you can see the data in the TiDB Cloud Starter cluster is showing in the `All Chunks in the Database` folder. And we have 3 columns, which are `id`, `text`, and `text_vec`.
- If you want to retrieve the context only, you can click the `Retrieval Only` option. Then, enter your question.
- Say I want to ask a question:
```text
What is TiDB?
```
- Click the `Send` button to start to retrieve.
- As you can see, the `TiDB is a distributed SQL database with HTAP capabilities.` has the highest similarity score.
> [!NOTE]
> You can press the button and drag the head of column to change the order of the table.
> You can also download the data as a CSV table.
- At the end, let's try to make the RAG to give us the answer.
- Change the `Retrieval Limit` at the left-top corner to `1`, which means we will only retrieve the most similar text towords the question that you inputed.
- Change the option to `RAG Q&A`, then click `Send`.
- This time, the demo will give us the answer rather than a material list.
### Step 3. RAG demo code reading
Let's head back to [VS Code](http://${LAB:HOST_MONITOR1_PUBLIC_IP}:3000/?folder=/root/tidb-lab-ai) and click the code file `bedrock/rag.py`.
1. The script begins by importing necessary libraries: `os` for environment variables, `pandas` for data manipulation, `litellm` for interacting with LLMs, `streamlit` for the web interface, and components from `pytidb` for leveraging TiDB as a unified persistent data store.
```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
```
2. A prompt template named `RAG_PROMPT_TEMPLATE` is defined. This template will be used to structure the input for the LLM, providing it with context and the user's question to generate an answer.
```python
# RAG prompt template
RAG_PROMPT_TEMPLATE = """Answer the question based on the following reference information.
Reference Information:
{context}
Question: {question}
Please answer:"""
```
3. A connection to the TiDB Cloud Starter cluster is established using `TiDBClient.connect()`. Connection parameters such as host, port, username, password, and database name are retrieved from environment variables (e.g., `SERVERLESS_CLUSTER_HOST`). SSL is enabled for a secure connection.
```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,
)
```
4. The embedding model and the LLM are defined. For this lab, we use Amazon Bedrock models:
- `text_embed`: An `EmbeddingFunction` is initialized with `bedrock/amazon.titan-embed-text-v2:0` for generating text embeddings. A timeout of 60 seconds is set for the embedding process.
- `llm_model`: The variable `llm_model` is set to `bedrock/us.amazon.nova-lite-v1:0`, which will be used for generating answers.
```python
text_embed = EmbeddingFunction(
"bedrock/amazon.titan-embed-text-v2:0",
timeout=60
)
llm_model = "bedrock/us.amazon.nova-lite-v1:0"
```
5. A `Chunk` table schema is defined using `pytidb.schema.TableModel`. This table, named `chunks`, will store our data.
- `id`: An integer field serving as the primary key.
- `text`: A string field to store the raw text chunks.
- `text_vec`: A vector field that will automatically store the embeddings of the `text` field, generated by the `text_embed` function defined earlier.
```python
# Define the Chunk table
table_name = "chunks"
class Chunk(TableModel, table=True):
__tablename__ = table_name
__table_args__ = {"extend_existing": True}
id: int = Field(primary_key=True)
text: str = Field()
text_vec: Optional[Any] = text_embed.VectorField(
source_field="text",
)
```
6. A list of `sample_chunks` (strings) is provided. In a real-world application, these chunks would typically be derived from a larger document through a process called "chunking".
```python
sample_chunks = [
"Llamas are camelids known for their soft fur and use as pack animals.",
"Python's GIL ensures only one thread executes bytecode at a time.",
# ... (other sample chunks)
"Augmented reality overlays digital content on the real world.",
]
```
7. The script attempts to open the `chunks` table. If the table does not exist in the database (e.g., on the first run), it creates the table based on the `Chunk` schema defined previously. This ensures the table is ready for data operations.
```python
table = db.open_table(table_name)
if table is None:
table = db.create_table(schema=Chunk)
```
8. The Streamlit user interface is set up. The page title is set to `🔍 RAG Demo`, and a subheader `Database Operations` is added.
```python
st.title("🔍 RAG Demo")
st.subheader("Database Operations")
```
9. A database operation button is created:
- `Save Data`: When clicked, this button inserts the `sample_chunks` into the table if the table is currently empty. `pytidb` automatically calls the Bedrock embedding model to vectorize the `text` data upon insertion, as defined in the `text_vec` field. The app then reruns.
```python
if st.button("Save Data", type="primary"):
st.write("Saving...")
# insert sample chunks
if table.rows() == 0:
chunks = [Chunk(text=text) for text in sample_chunks]
table.bulk_insert(chunks)
st.rerun()
```
10. An expander section labeled `📁 All Chunks in the Database` is created. When expanded, it queries all rows from the `table`, converts the data into a Pandas DataFrame, and displays it. If the table is empty, an informational message is shown.
```python
with st.expander("📁 All Chunks in the Database", expanded=False):
chunks = table.query()
if chunks:
data = [{'id': chunk.id, 'text': chunk.text, 'text_vec': chunk.text_vec} for chunk in chunks]
df = pd.DataFrame(data)
st.dataframe(df)
else:
st.info("No data found in database.")
```
11. Instructional text and user input elements are added:
- A radio button (`mode`) allows the user to select the operation mode: `Retrieval Only` or `RAG Q&A`.
- A slider (`query_limit`) in the sidebar lets the user adjust the number of chunks to retrieve (default is 5, min 1, max 20).
- A `Reset` button in the sidebar that truncates (clears all data from) the `chunks` table and reruns the Streamlit app to refresh the view. A warning message is displayed above it.
- A text input box (`query`) prompts the user to enter their question.
```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_limit = st.sidebar.slider("Retrieval Limit", min_value=1, max_value=20, value=5)
st.sidebar.markdown("---")
st.sidebar.markdown("This option will delete all data", unsafe_allow_html=True)
if st.sidebar.button("Reset"):
st.sidebar.write("Resetting...")
table.truncate()
st.rerun()
query = st.text_input("Enter your question:", "")
```
12. This block handles the logic when the user clicks the `Send` button and has entered a query:
- A spinner is shown while processing.
- The core operation `table.search(query).limit(query_limit)` performs a vector similarity search on the `text_vec` column using the user's query and retrieves the top `query_limit` most similar chunks.
```python
if st.button("Send") and query:
with st.spinner("Processing..."):
# Retrieve relevant chunks
res = table.search(query).limit(query_limit)
```
13. Based on the selected `mode` and if results (`res`) are found:
- If `Retrieval Only`: It displays the retrieved chunks (converted to a Pandas DataFrame) directly.
- If `RAG Q&A`:
- The `text` content is extracted from the retrieved chunks.
- The `context` is constructed by joining these texts.
- The `RAG_PROMPT_TEMPLATE` is formatted with this `context` and the user's `question`.
- The Bedrock LLM (specified by `llm_model`) is called via `litellm.completion` with the constructed prompt to generate an answer.
- The LLM's response is displayed using custom CSS for styling.
- An expander section (`📚 Retrieved Knowledge`) shows the chunks that were used to generate the answer.
- If the initial search returns no results, an informational message `No relevant information found` is displayed.
```python
if res:
if mode == "Retrieval Only":
st.write("### Retrieval Results:")
st.dataframe(res.to_pandas())
else:
text = [chunk.text for chunk in res.to_rows()]
# Build RAG prompt
context = "\n".join(text)
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(
"""
""",
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.dataframe(res.to_pandas())
else:
st.info("No relevant information found")
```