## 模块 3: 通过 RAG 改进你的 AI 应用程序
我们在模块 2 中完成了初步的 AI 应用，但是遇到了 LLM 在回答训练数据之外的数据的时候，产生了严重的幻觉问题。那怎么解决呢？答案也很简单。
我们可以预先将信息帮助 LLM 获取回来，放置在调用 LLM 的输入中，使用 LLM 帮我们总结内容，而非依赖它的训练集的知识储备。这样一来，我们可以得到以下好处：
1. 降低幻觉：Transformer 架构的 LLM，使用多头注意力模型在 context window 内进行基于顺序的预测。如果我们把相关知识放入 context window 内，它会生成和相关知识更接近的回答。从表象上来看，它降低了幻觉的产生频率或程度。
2. 提供外部知识：再一次，由于 Transformer 架构的 LLM 拥有一个 context window（在 API 上的体现为 token size 的限制），我们无法将所有的信息全部放入 prompt。因此，我们需要事先挑选出有关的内容。
3. 降低成本：即使现在 LLM 的 context window 越来越大，我们仍需从成本考量，不应为每个问题浪费太多的 token。因为绝大多数 LLM 都是以 token 计费的（你可以把 token 理解为和文本长度正相关的一个变量，你的文本越长，token 就越多）。
那么，现在摆在我们面前的问题就只有一个了：
**怎么实现帮助 LLM 获取相关信息？**
实现它的方案不止一种。我们在模块 3 中会使用其中的一种：文本 Embedding 向量方式。
> [!NOTE]
> **选读部分**
>
> Embedding 向量，是 Embedding 模型将文本转换成的向量。而 Embedding 模型，是一类专门训练来分离语义的的模型。简单的来讲，在对比学习训练时，它会得到非常多组的样本，而每一组样本分为：
>
> - 一个锚点样本：一段文字
> - 一个正样本：一段与锚点样本语义相似的文字
> - 一组负样本：很多段与锚点样本语义不相似的文字
>
> Embedding 模型的任务就是：
>
> - 最小化锚点样本和正样本的输出向量的距离
> - 最大化锚点样本和负样本的输出向量的距离
>
> 更正确的，此处使用了 InfoNCE loss 这个损失函数训练模型，进行了随机梯度下降算法，拟合了样本。
如果你看不懂选读部分，这完全没有问题，它只是原理的解释。你只需要知道：两个文本的 Embedding 向量的余弦距离越远，那么他们的语义就越不相似。而 Embedding 向量的余弦距离越近，那么他们的语义就越相似。
因此，我们可以使用这个特性来完成这一次的改进。在模块 3 中，我们将使用 Mulvis 作为向量的搜索引擎，MySQL 作为关系型数据库存储原始文本。
> [!NOTE]
> **为什么不使用向量数据库的 metadata 来存储原始文本？**
>
> 这在技术上是可行的，但是由于向量数据库的能力缺失，比如事务等。无法保证的数据的完全一致性。在实际的工程中，如果依旧使用关系型数据库和向量数据库的分离架构，建议使用关系型数据库作为数据的 source of truth。同时将原始文本的 ID 与 Embedding 向量这组数据异步的更新至向量数据库。
### 步骤 1. 运行 RAG AI 应用程序的 Streamlit 服务
创建一个新的终端：
然后运行新的 RAG AI 应用程序：
```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
```
### 步骤 2. 使用 RAG AI 应用程序的演示网页
- 打开您的 [RAG AI 应用程序](http://${LAB:HOST_MONITOR1_PUBLIC_IP}:8502/)。
- 打开名为 `All Employees in the MySQL` 和 `All Employees ID and Embedding Mapping in the Milvus` 的文件夹，以查看当前 MySQL 和 Milvus 中的数据。由于我们尚未保存任何数据，所以它们应该是空的。
- 点击 `Save Data` 保存我们为本次演示准备的模拟员工信息数据。您可以在下一步的代码阅读中看到原始数据，它是一个员工对象列表。
- 保存数据后，您可以在 `All Employees in the MySQL` 文件夹中看到 MySQL 中的数据。我们有 3 列，分别是 `id`、`name` 和 `intro`。而在 `All Employees ID and Embedding Mapping in the Milvus` 内，你可以看到我们保存了 `id` 和 `embedding` 的关系。此时 Milvus 内的 `id` 和 MySQL 中的 `id` 是一一映射的。
- 如果您只想检索上下文，可以点击 `Retrieval Only` 选项。然后，输入您的问题。
- 我们来试试 RAG AI 应用程序是否能通过 Embedding 选取出相关的信息，输入问题：
```text
Who is Cheese Wang at PingCAP?
```
- 点击 `Send` 按钮开始检索。
- 如您所见，我们得到了 `Cheese Wang` 的相关信息。
- 接下来，the moment of truth, 我们来检测现在 RAG AI 应用程序是否能回答简易 AI 应用程序无法回答的问题。将 `Select Mode` 选项更改为 `RAG Q&A`，然后点击 `Send`。
- 这次，RAG AI 应用程序将给我们提供答案，而不是信息列表。
- 可以看到，因为我们提供了 LLM 在训练阶段并未获取到的信息。因此 LLM 不再认为我是 PingCAP 的 CEO，而是真实的回答了我的信息。
### 步骤 3. RAG AI 应用程序代码阅读
让我们回到 [VS Code](${LAB:HTTPS_REVERSE_PROXY_URL}?next=${LAB:"/vscode/?folder=/root/tidb-lab-unify-storage"|urlencode} "VS Code") 并点击代码文件 `rag.py`。
1. 脚本首先导入了构建 RAG 应用所需的库：
- `os`, `dotenv`: 用于加载和读取环境变量，例如数据库密码和主机地址。
- `pandas`: 用于数据处理和在 Streamlit 中展示表格。
- `sqlalchemy`, `sqlmodel`: 用于连接关系型数据库（在本例中为 MySQL）并定义数据模型。`SQLModel` 提供了结合 SQLAlchemy 和 Pydantic 的便利性。
- `pymilvus`: Milvus 的 Python 客户端，用于连接和操作 Milvus 向量数据库。
- `litellm`: 一个用于与各种大语言模型（LLM）和 Embedding 模型进行标准化交互的库。
- `streamlit`: 用于创建和运行 Web 用户界面。
- `common`: 从公共模块导入预定义的 RAG 提示模板和 CSS 样式。
```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. 进行初始化和配置：
- `create_engine`: 使用从环境变量中读取的配置（主机、端口、密码等）创建一个到 MySQL 数据库的连接引擎。
- `embedding_model`, `llm_model`: 定义了将要使用的 Embedding 模型和 LLM 的名称。
- `MilvusClient`: 初始化 Milvus 客户端，以连接到在指定 URI 上运行的 Milvus 服务。
- `milvus_collection_name`: 为将在 Milvus 中创建的集合（Collection）命名。
```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. 定义了一个辅助函数 `get_embedding`，它接收一个文本字符串，调用 `litellm.embedding` API 将文本转换为一个 1024 维的向量，并返回该向量。
```python
def get_embedding(text: str) -> List[float]:
response = embedding(
model=embedding_model,
input=[text],
dimensions=embedding_dimensions
)
return response.data[0]['embedding']
```
4. 使用 `SQLModel` 定义了 `Employee` 表的结构，该表将在 MySQL 中创建，用于存储员工的原始信息（ID、姓名和简介）。`SQLModel.metadata.create_all(engine)` 会确保在数据库中创建这张表。
```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. 创建了一个包含三名员工信息的列表 `employees` 作为演示用的样本数据。
```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. 为 Milvus 定义了数据模式（Schema）和索引。
- `MilvusClient.create_schema`: 定义了集合的结构，包含一个主键 `id`（用于与 MySQL 中的 `Employee` 记录关联）和一个 `embedding` 向量字段。
- `milvus_client.prepare_index_params`: 为 `embedding` 字段准备索引参数，使用 `AUTOINDEX` 让 Milvus 自动选择最佳索引类型。
- `milvus_client.create_collection`: 使用定义的模式和索引在 Milvus 中创建一个新的集合。这使得我们能够高效地进行向量相似度搜索。
```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. 设置 Streamlit 界面，提供数据库操作按钮。
- **Reset 按钮**: 点击后，会清空 MySQL 中的 `Employee` 表和 Milvus 集合中的所有相关数据，用于重置演示环境。
- **Save Data 按钮**: 点击后，会检查 MySQL 中是否已有数据。如果没有，它会：
a. 遍历 `employees` 列表，为每个员工的 `intro` 生成 Embedding 向量。
b. 将员工的完整信息（id, name, intro）存入 MySQL。
c. 将员工的 ID 和其对应的 Embedding 向量存入 Milvus。
这个过程展示了如何将结构化元数据和向量数据分别写入 RDB 和向量数据库。
```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. 创建了两个可折叠的部分（`st.expander`），分别用于显示当前存储在 MySQL 和 Milvus 中的数据。这有助于用户直观地看到数据操作的结果。
```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. 这部分是 RAG 应用的核心交互逻辑。当用户输入问题并点击 "Send" 时：
- **向量化查询**: 首先使用 `get_embedding` 函数将用户的问题文本转换为查询向量。
- **向量检索 (Retrieval)**: 使用 `milvus_client.search` 在 Milvus 集合中进行向量相似度搜索（使用余弦距离），找出与查询向量最相似的一条记录，并获取其 `id`。
- **数据获取 (Fetch)**: 使用从 Milvus 中获取的 `id`，到 MySQL 的 `Employee` 表中查询并获取对应的完整员工信息（姓名、简介等）。
- **生成回答 (Generation)**:
- 如果模式是 `Retrieval Only`，则直接展示从 Milvus 和 MySQL 中检索到的原始数据。
- 如果模式是 `RAG Q&A`，则将从 MySQL 获取的员工简介（`intro`）作为上下文（`context`），连同用户原始问题，填入 `RAG_PROMPT_TEMPLATE` 提示模板中。
- 最后，调用 `litellm.completion` 将构建好的提示发送给 LLM，生成最终的回答，并将其以美化的样式显示在界面上。同时，提供一个可折叠部分来展示用于生成该回答的"知识来源"。
- 如果向量搜索没有找到相关信息，则提示用户。
```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")
```
### 步骤 4. RAG AI 应用程序的总结
我们可以很清楚的比对出简易 AI 应用程序和 RAG AI 应用程序的方案优劣势。简易 AI 应用程序简单，但是其背景知识均构建在训练集内，若问题在训练集外，则大概率会产生幻觉。而 RAG AI 应用程序很好的解决了这一点，而且带来了我们在模块 3 开头提出的一系列优势，但是它仅仅在最简易的 demo 应用程序中，也带来了相当高的复杂度。
接下来，我们需要优化 RAG 应用程序，让其准确率进一步提升。