## 模块 3: 运行、体验和研究 RAG 演示应用程序
运行一个 RAG 演示应用程序，该应用程序连接到 TiDB Cloud Starter，保存您的嵌入向量，检索相关数据，并生成对您查询的回答。
### 步骤 1. 运行 RAG Streamlit 服务
- 创建一个新的终端：
- 在 VS Code 终端中，使用 `streamlit` 命令行工具运行以下脚本来启动 Streamlit 服务。
```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
```
### 步骤 2. 使用 RAG 演示网页
- 点击访问您的 [RAG 演示网页](http://${LAB:HOST_MONITOR1_PUBLIC_IP}:8501/)。根据实验环境设置，初始化完成大约需要 **5 分钟**。
- 打开名为 `All Chunks in the Database` 的文件夹，以查看当前 TiDB Cloud Starter 集群中的数据。由于我们尚未保存任何数据，所以它应该是空的。
- 点击 `Save Data` 保存我们为此演示准备的示例文本数据。您可以在下一步中看到示例数据，这是一个字符串列表。
- 保存数据后，您可以在 `All Chunks in the Database` 文件夹中看到 TiDB Cloud Starter 集群中的数据。我们有 3 列，分别是 `id`、`text` 和 `text_vec`。
- 如果您只想检索上下文，可以点击 `Retrieval Only` 选项。然后，输入您的问题。
- 假设我想问一个问题：
```text
What is TiDB?
```
- 点击 `Send` 按钮开始检索。
- 如您所见，`TiDB is a distributed SQL database with HTAP capabilities.`（TiDB 是一个具有 HTAP 能力的分布式 SQL 数据库）具有最高的相似度分数。
> [!NOTE]
> 您可以按下按钮并拖动列头来更改表格的排序。
> 您还可以将数据下载为 CSV 表格。
- 最后，让我们尝试让 RAG 给我们提供答案。
- 在左上角将 `Retrieval Limit`（检索限制）更改为 `1`，这意味着我们将只检索与您输入的问题最相似的文本。
- 将选项更改为 `RAG Q&A`，然后点击 `Send`。
- 这次，演示将给我们提供答案，而不是材料列表。
### 步骤 3. RAG 演示代码阅读
让我们回到 [VS Code](${LAB:HTTPS_REVERSE_PROXY_URL}?next=${LAB:"/vscode/?folder=/root/tidb-lab-ai"|urlencode} "VS Code") 并点击代码文件 `bedrock/rag.py`。
1. 脚本首先导入必要的库：用于环境变量的 `os`，用于数据操作的 `pandas`，用于与 LLM 交互的 `litellm`，用于 Web 界面的 `streamlit`，以及来自 `pytidb` 的组件，用于将 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
```
2. 定义了一个名为 `RAG_PROMPT_TEMPLATE` 的提示模板。这个模板将用于构建输入给 LLM 的结构，为其提供上下文和用户的问题以生成答案。
```python
# RAG prompt template
RAG_PROMPT_TEMPLATE = """Answer the question based on the following reference information.
Reference Information:
{context}
Question: {question}
Please answer:"""
```
3. 使用 `TiDBClient.connect()` 建立与 TiDB Cloud Starter 集群的连接。连接参数如主机、端口、用户名、密码和数据库名称都是从环境变量（例如 `SERVERLESS_CLUSTER_HOST`）中获取的。启用了 SSL 以确保连接安全。
```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. 定义了嵌入模型和 LLM。对于本实验，我们使用 Amazon Bedrock 模型：
- `text_embed`：使用 `bedrock/amazon.titan-embed-text-v2:0` 初始化了一个 `EmbeddingFunction`，用于生成文本嵌入。为嵌入过程设置了 60 秒的超时时间。
- `llm_model`：变量 `llm_model` 设置为 `bedrock/us.amazon.nova-lite-v1:0`，将用于生成答案。
```python
text_embed = EmbeddingFunction(
"bedrock/amazon.titan-embed-text-v2:0",
timeout=60
)
llm_model = "bedrock/us.amazon.nova-lite-v1:0"
```
5. 使用 `pytidb.schema.TableModel` 定义了一个 `Chunk` 表模式。这个名为 `chunks` 的表将用于存储我们的数据。
- `id`：一个整数字段，作为主键。
- `text`：一个字符串字段，用于存储原始文本块。
- `text_vec`：一个向量字段，将自动存储 `text` 字段的嵌入向量，这些嵌入向量由前面定义的 `text_embed` 函数生成。
```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. 提供了一个 `sample_chunks`（字符串）列表。在实际应用中，这些文本块通常是通过一个称为"分块"（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. 脚本尝试打开 `chunks` 表。如果该表在数据库中不存在（例如，在首次运行时），它会根据之前定义的 `Chunk` 模式创建表。这确保表已准备好进行数据操作。
```python
table = db.open_table(table_name)
if table is None:
table = db.create_table(schema=Chunk)
```
8. 设置了 Streamlit 用户界面。页面标题设置为 `🔍 RAG Demo`，并添加了一个子标题 `Database Operations`。
```python
st.title("🔍 RAG Demo")
st.subheader("Database Operations")
```
9. 创建了一个数据库操作按钮：
- `Save Data`：点击时，如果表当前为空，此按钮会将 `sample_chunks` 插入表中。根据 `text_vec` 字段的定义，`pytidb` 在插入时会自动调用 Bedrock 嵌入模型来向量化 `text` 数据。然后应用程序会重新运行。
```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. 创建了一个标记为 `📁 All Chunks in the Database`（数据库中的所有文本块）的可展开部分。当展开时，它会查询 `table` 中的所有行，将数据转换为 Pandas DataFrame，并显示出来。如果表为空，则显示一条信息提示。
```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. 添加了指导文本和用户输入元素：
- 单选按钮（`mode`）允许用户选择操作模式：`Retrieval Only`（仅检索）或 `RAG Q&A`（RAG 问答）。
- 侧边栏中的滑块（`query_limit`）让用户可以调整要检索的文本块数量（默认为 5，最小 1，最大 20）。
- 侧边栏中的 `Reset` 按钮会截断（清除所有数据）`chunks` 表，并重新运行 Streamlit 应用程序以刷新视图。在它上方显示了一条警告消息。
- 文本输入框（`query`）提示用户输入他们的问题。
```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. 此代码块处理用户点击 `Send` 按钮并输入查询时的逻辑：
- 处理过程中显示一个加载动画。
- 核心操作 `table.search(query).limit(query_limit)` 使用用户的查询在 `text_vec` 列上执行向量相似度搜索，并检索相似度最高的前 `query_limit` 个文本块。
```python
if st.button("Send") and query:
with st.spinner("Processing..."):
# Retrieve relevant chunks
res = table.search(query).limit(query_limit)
```
13. 根据所选的 `mode`（模式）和是否找到结果（`res`）：
- 如果是 `Retrieval Only`（仅检索）：直接显示检索到的文本块（转换为 Pandas DataFrame）。
- 如果是 `RAG Q&A`（RAG 问答）：
- 从检索到的文本块中提取 `text` 内容。
- 通过连接这些文本构建 `context`（上下文）。
- 使用这个 `context` 和用户的 `question`（问题）格式化 `RAG_PROMPT_TEMPLATE`。
- 通过 `litellm.completion` 调用 Bedrock LLM（由 `llm_model` 指定）并传入构建的提示来生成答案。
- 使用自定义 CSS 样式显示 LLM 的响应。
- 一个可展开的部分（`📚 Retrieved Knowledge`）显示用于生成答案的文本块。
- 如果初始搜索没有返回结果，则显示信息提示 `No relevant information found`（未找到相关信息）。
```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")
```