## 模块 2：运行 RAG 演示
准备好之后，我们可以运行 RAG 演示来连接到 TiDB Cloud Starter 集群，保存您的 Embedding，检索相关数据，并生成答案。
### 步骤 1. 运行 Streamlit 服务
然后使用 `streamlit` 命令行工具运行此脚本：
```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
```
### 步骤 2. 使用 RAG 演示网页
- 打开您的 [RAG 示例网页](http://${LAB:HOST_MONITOR1_PUBLIC_IP}:8501/)。
- 打开名为 `All Chunks in the Database` 的文件夹，以查看当前 TiDB Cloud Starter 集群中的数据。由于我们尚未保存任何数据，所以它应该是空的。
- 点击 `Save Data` 以保存我们为此演示准备的文本数据。您可以在下一步中看到原始数据，它是一个字符串列表。
- 保存数据后，您会看到 TiDB Cloud Starter 集群中的数据显示在 `All Chunks in the Database` 文件夹中。我们有 3 列，分别是 `id`，`text` 和 `text_vec`。
- 如果您只想检索上下文，可以点击 `Retrieval Only` 选项。然后输入您的问题。
- 假设我想问一个问题：
```text
What is TiDB?
```
- 点击 `Send` 按钮开始检索。
- 如您所见，`TiDB is a distributed SQL database with HTAP capabilities.` 的相似度得分最高。
> **提示：**
>
> 您可以按下按钮并拖动列的顶部以更改表格的顺序。
> 您还可以将数据下载为 CSV 表格。
> 
- 最后，让我们尝试让 RAG 给我们答案。
- 在左上角将 `Retrieval Limit` 改为 `1`，这意味着我们只会检索与您输入的问题最相似的文本。
- 将选项更改为 `RAG Q&A`，然后点击 `Send`。
- 这次，演示将给我们答案而不是材料列表。
### 步骤 3. RAG 演示代码阅读
让我们回到 [VS Code](http://${LAB:HOST_MONITOR1_PUBLIC_IP}:3000/?folder=/root/tidb-lab-ai) 并点击代码文件 `openai/rag.py`。
1. 使用 `dotenv.load_dotenv()` 从 `.env` 加载参数后，下一步是我们的提示模板。可以看到我们将使用这个模板来插入 `{context}` 和 `{question}`，然后让 LLM 根据上下文给出答案。
```python
# RAG prompt template
RAG_PROMPT_TEMPLATE = """Answer the question based on the following reference information.
Reference Information:
{context}
Question: {question}
Please answer:"""
```
2. 我们将连接到 TiDB Cloud Starter 集群。实验室 VM 启动时我们已经设置了这些环境。如果您需要在 TiDB 实验室之外使用这段代码，您可能需要在 [TiDB Cloud](https://tidbcloud.com/clusters) 中找到这些参数，并手动设置。
```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. 我们定义了 embedding 模型和 LLM 的名称。然后用这个 embedding 模型定义一个 `Chunk` 表。这是因为向量索引的长度是固定的，我们必须在创建表之前决定要使用的模型。
```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. 我们将示例 chunks 写为字符串列表。在实际应用中，您可能需要制定策略来分隔整个文档，这个过程在 RAG 场景中称为“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.",
"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. 我们尝试打开名为 `chunks` 的表。如果数据库中尚不存在该表（这将是首次运行时的情况），我们会根据之前定义的 `Chunk` 模式创建它。这确保了表已准备好用于存储和查询数据。
```python
table = db.open_table(table_name)
if table is None:
table = db.create_table(schema=Chunk)
```
6. 下一部分使用 Streamlit 设置用户界面。我们将网页的标题设置为 `🔍 RAG Demo` 并添加子标题“数据库操作”。
```python
st.title("🔍 RAG Demo")
st.subheader("Database Operations")
```
7. 我们为数据库操作按钮创建两个列。“重置”按钮单击后，会截断表（删除所有数据）并重新运行 Streamlit 应用以刷新视图。警告信息显示在其下方。“Save Data” 按钮会在当前表为空时插入 `sample_chunks`，然后重新运行应用。保存 `text` 到表中时，`pytidb` 会调用嵌入模型向量化您的 `text` 数据。这是由表字段定义 `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. 创建一个标记为 `📁 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.")
```
9. 我们添加指导性文本并设置用户输入元素。单选按钮让用户选择操作模式（`Retrieval Only` 或 `RAG Q&A`）。侧边栏的滑块让用户调整检索的 chunk 数量（`Retrieval Limit`）。最后，一个文本输入框提示用户输入他们的问题。
```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. 这一块为处理用户点击 "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)
```
11. 根据选择的 `mode`，应用程序要么直接展示检索结果，要么继续执行 RAG（检索增强生成）流程。
- 如果选择了 `Retrieval Only`（仅检索）：会直接将检索到的块（转换成 Pandas DataFrame）展示出来。
- 如果选择了 `RAG Q&A`：
- 从检索结果中提取 `text` 内容；
- 通过拼接提取的文本，构建 `context`；
- 使用生成的 `context` 和用户的 `question` 填充 `RAG_PROMPT_TEMPLATE`；
- 调用 LLM（使用 `litellm.completion`）发送构建好的提示词以获取答案；
- 使用自定义 CSS 样式展示 LLM 的回复；
- 并包含一个可展开的区域（"📚 Retrieved Knowledge"），显示用于生成答案的检索块。
如果初步搜索 (`res`) 没有返回任何结果，会显示一条信息提示。
```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")
```