## Module 2: 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 Streamlit service
- In the VS Code terminal, use `streamlit` command line tool to run the following script to start Streamlit service.
```bash
streamlit run openai/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 text data that we prepared for this demo. You can see the original 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.
> **Tips:**
>
> 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 `openai/rag.py`.
1. After using `dotenv.load_dotenv()` to load parameters from `.env`, the next step is our prompt template. As you can see, we will use this template to insert the `{context}` and `{question}` later. Then let LLM to give us an answer based on the context.
```python
# RAG prompt template
RAG_PROMPT_TEMPLATE = """Answer the question based on the following reference information.
Reference Information:
{context}
Question: {question}
Please answer:"""
```
2. We will connect to TiDB Cloud Starter cluster, we already set these environment when the lab VM was starting. If you need to use this code snippet outside of TiDB Labs, you may need to find the parameters in your [TiDB Cloud](https://tidbcloud.com/clusters), and set it manually.
```python
db = TiDBClient.connect(
host=os.getenv("TIDB_HOST"),
port=int(os.getenv("TIDB_PORT")),
username=os.getenv("TIDB_USERNAME"),
password=os.getenv("TIDB_PASSWORD"),
database=os.getenv("TIDB_DATABASE"),
enable_ssl=True,
)
```
3. We defined the names of embedding model and LLM. Then define a `Chunk` table by using this embedding model. This is because the length of the vector index is fixed, we have to decide the a model you want to use before creating the table.
```python
text_embed = EmbeddingFunction("openai/text-embedding-3-small")
llm_model = "openai/gpt-4o"
# 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",
)
```
4. We write the sample chunks as a list of string. In the real world, you might need to make a strategy to seperate a whole document, and this process is called "chunking" in RAG scenario.
```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.",
"TiDB is a distributed SQL database with HTAP capabilities.",
"Einstein's theory of relativity revolutionized modern physics.",
"The Great Wall of China stretches over 13,000 miles.",
"Ollama enables local deployment of large language models.",
"HTTP/3 uses QUIC protocol for improved web performance.",
"Kubernetes orchestrates containerized applications across clusters.",
"Blockchain technology enables decentralized transaction systems.",
"GPT-4 demonstrates remarkable few-shot learning capabilities.",
"Machine learning algorithms improve with more training data.",
"Quantum computing uses qubits instead of traditional bits.",
"Neural networks are inspired by the human brain's structure.",
"Docker containers package applications with their dependencies.",
"Cloud computing provides on-demand computing resources.",
"Artificial intelligence aims to mimic human cognitive functions.",
"Cybersecurity protects systems from digital attacks.",
"Big data analytics extracts insights from large datasets.",
"Internet of Things connects everyday objects to the internet.",
"Augmented reality overlays digital content on the real world.",
]
```
5. We attempt to open the table named `chunks`. If the table doesn't exist in the database yet (which would be the case on the first run), we create it based on the `Chunk` schema we defined earlier. This ensures the table is ready for storing and querying data.
```python
table = db.open_table(table_name)
if table is None:
table = db.create_table(schema=Chunk)
```
6. The next part sets up the user interface using Streamlit. We set the title of the webpage to `🔍 RAG Demo` and add a subheader "Database Operations".
```python
st.title("🔍 RAG Demo")
st.subheader("Database Operations")
```
7. We create two columns for database operation buttons. The "Reset" button, when clicked, truncates the table (deletes all data) and reruns the Streamlit app to refresh the view. A warning message is displayed below it. The "Save Data" button inserts the `sample_chunks` into the table if it's currently empty and then reruns the app. When you save the `text` into the table, `pytidb` will call the embedding model to vectorize your `text` data. It is defined by the field defination of table `text_vec: Optional[Any] = text_embed.VectorField(source_field="text")`.
```python
left, right = st.columns(2)
if left.button("Reset", type="primary"):
left.write("Resetting...")
table.truncate()
st.rerun()
left.markdown("This option will delete all data", unsafe_allow_html=True)
if right.button("Save Data"):
right.write("Saving...")
# insert sample chunks
if table.rows() == 0:
chunks = [Chunk(text=text) for text in sample_chunks]
table.bulk_insert(chunks)
st.rerun()
```
8. An expander section labeled `📁 All Chunks in the Database` is created. It queries all rows from the `table`, converts the data into a Pandas DataFrame, and displays it within the expander. If the table is empty, it shows an informational message.
```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.")
```
9. We add instructional text and set up the user input elements. A radio button allows the user to select the operation mode (`Retrieval Only` or `RAG Q&A`). A slider in the sidebar lets the user adjust the number of chunks to retrieve (`Retrieval Limit`). Finally, a text input box 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)
query = st.text_input("Enter your question:", "")
```
10. This block handles the logic when the user clicks the "Send" button and has entered a question. It shows a spinner while processing. The core operation is `table.search(query).limit(query_limit)`, which 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)
```
11. Based on the selected `mode`, the application either displays the raw retrieval results or proceeds with the RAG process.
- If `Retrieval Only`: It displays the retrieved chunks (converted to a Pandas DataFrame) directly.
- If `RAG Q&A`:
- It extracts the `text` content from the retrieved chunks.
- It constructs the `context` by joining the extracted texts.
- It formats the `RAG_PROMPT_TEMPLATE` with the generated `context` and the user's `question`.
- It calls the LLM (using `litellm.completion`) with the constructed prompt to get an answer.
- It displays the LLM's response using custom CSS for styling.
- It includes an expander section ("📚 Retrieved Knowledge") showing the chunks used to generate the answer.
If the initial search (`res`) returns no results, an informational message 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")
```