## モジュール 3: RAGによるAIアプリケーションの改善
モジュール2では、予備的なAIアプリケーションを完成させましたが、LLMがトレーニングデータ外の質問に答える際に深刻なハルシネーション問題に遭遇しました。では、これをどう解決すればよいでしょうか？答えは簡単です。
LLMを呼び出す際に入力に情報を事前に取得して配置し、LLMに内容を要約させることで、トレーニングセットの知識に頼るのではなく、LLMを活用することができます。これにより、以下の利点が得られます：
1. ハルシネーションの削減：TransformerベースのLLMは、コンテキストウィンドウ内でシーケンスベースの予測を行うためにマルチヘッドアテンションモデルを使用します。関連知識をコンテキストウィンドウに含めると、その知識により近い回答が生成されます。表面的には、ハルシネーションの頻度や深刻度が減少します。
2. 外部知識の提供：繰り返しになりますが、TransformerベースのLLMにはコンテキストウィンドウ（APIではトークンサイズの制限として反映される）があるため、すべての情報をプロンプトに含めることはできません。したがって、関連するコンテンツを事前に選択する必要があります。
3. コスト削減：LLMのコンテキストウィンドウが大きくなっているにもかかわらず、コストを考慮し、各質問にあまりにも多くのトークンを浪費すべきではありません。なぜなら、ほとんどのLLMはトークン単位で課金されるためです（トークンはテキストの長さに正比例する変数と考えることができます。テキストが長ければ長いほど、トークンも多くなります）。
したがって、今我々が直面している唯一の問題は：
**LLMが関連情報を検索するのをどのように助けるか？**
これを達成する方法は一つではありません。モジュール3では、そのうちの一つであるテキスト埋め込みベクトル方式を使用します。
> [!NOTE]
> **補足資料**
>
> 埋め込みベクトルとは、埋め込みモデルがテキストを変換して生成するベクトルです。埋め込みモデルとは、意味を区別するために特別に訓練されたモデルの一種です。簡単に言うと、対照学習の訓練中に、多数のサンプルセットを取得し、各セットは以下のように分けられます：
>
> - アンカーサンプル：テキストの一部
> - ポジティブサンプル：アンカーサンプルと意味的に類似したテキストの一部
> - ネガティブサンプルのセット：アンカーサンプルと意味的に類似しない多数のテキストの一部
>
> 埋め込みモデルのタスクは：
>
> - アンカーサンプルとポジティブサンプルの出力ベクトル間の距離を最小化すること
> - アンカーサンプルとネガティブサンプルの出力ベクトル間の距離を最大化すること
>
> より正確に言えば、ここではInfoNCE損失関数を使用してモデルを訓練し、確率的勾配降下法を実行してサンプルに適合させます。
任意読解部分が理解できなくても全く問題ありません。それは原理の説明に過ぎません。ただ知っておくべきことは：2つのテキストの埋め込みベクトル間のコサイン距離が大きいほど、それらの意味は類似しておらず、コサイン距離が小さいほど、それらの意味は類似しているということです。
したがって、この特性を利用して今回の改善を行うことができます。モジュール3では、Milvusをベクトル検索エンジンとして、MySQLをリレーショナルデータベースとして元のテキストを保存します。
> [!NOTE]
> **なぜベクトルデータベースのメタデータを使って元のテキストを保存しないのか？**
>
> これは技術的には可能ですが、トランザクションなど、ベクトルデータベースの能力が不足しているため、完全なデータ一貫性を保証できません。実際のエンジニアリングでは、リレーショナルデータベースとベクトルデータベースの分離アーキテクチャを依然として使用する場合、リレーショナルデータベースをデータの信頼できる情報源として使用することをお勧めします。同時に、元のテキストIDと埋め込みベクトルのデータペアを非同期でベクトルデータベースに更新する必要があります。
### 手順 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アプリケーションのデモWebページの使用
- [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のデータを確認できます。`id`、`name`、`intro` の3つの列があります。また、`All Employees ID and Embedding Mapping in the Milvus` では、`id` と `embedding` の関係が保存されていることがわかります。この時点で、Milvusの `id` とMySQLの `id` は1対1でマッピングされています。
- コンテキストの検索のみを行いたい場合は、`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）や埋め込みモデルと標準化された対話を行うためのライブラリです。
- `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`: 使用する埋め込みモデルとLLMの名前を定義します。
- `MilvusClient`: 指定されたURIで実行されているMilvusサービスに接続するためにMilvusクライアントを初期化します。
- `milvus_collection_name`: Milvusで作成されるコレクションに名前を付けます。
```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. テキスト文字列を受け取り、`litellm.embedding` APIを呼び出してテキストを1024次元のベクトルに変換し、そのベクトルを返すヘルパー関数 `get_embedding` を定義します。
```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. デモ用のサンプルデータとして、3人の従業員の情報を含むリスト `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のデータスキーマとインデックスを定義します。
- `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インターフェースを設定し、データベース操作ボタンを提供します。
- **リセットボタン**: クリックすると、MySQLの `Employee` テーブルとMilvusコレクションから関連するすべてのデータがクリアされ、デモ環境をリセットするために使用されます。
- **データ保存ボタン**: クリックすると、MySQLにデータが既に存在するかどうかを確認します。存在しない場合、次の処理を行います：
a. `employees` リストを反復処理し、各従業員の `intro` の埋め込みベクトルを生成します。
b. 完全な従業員情報（ID、名前、紹介文）をMySQLに保存します。
c. 従業員のIDとそれに対応する埋め込みベクトルを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. 2つの展開可能なセクション（`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アプリケーションのコアな対話ロジックです。ユーザーが質問を入力して「送信」をクリックすると：
- **クエリのベクトル化**: まず、`get_embedding` 関数を使用してユーザーの質問テキストをクエリベクトルに変換します。
- **ベクトル検索**: `milvus_client.search` を使用してMilvusコレクションでベクトル類似性検索（コサイン距離を使用）を行い、最も類似したレコードを見つけてその `id` を取得します。
- **データ取得**: Milvusから取得した `id` を使用してMySQLの `Employee` テーブルをクエリし、対応する完全な従業員情報（名前、紹介文など）を取得します。
- **生成**:
- モードが `Retrieval Only` の場合、MilvusとMySQLから取得した生のデータを直接表示します。
- モードが `RAG Q&A` の場合、MySQLから取得した従業員の紹介文（`intro`）をコンテキストとして使用し、ユーザーの元の質問とともに `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の冒頭で提案した一連の利点をもたらしますが、最も単純なデモアプリケーションでさえ、かなりの複雑さを伴います。
次に、RAGアプリケーションを最適化して、その精度をさらに向上させる必要があります。