## モジュール 3: RAG デモアプリケーションの実行、体験、および学習
TiDB Cloud Starter に接続する RAG デモアプリケーションを実行し、埋め込みを保存し、関連データを取得し、クエリに対する回答を生成します。
### 手順 1. RAG 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`をクリックして、このデモ用に準備したサンプルテキストデータを保存します。次の手順でサンプルデータを確認できます。これは文字列のリストです。
- データを保存した後、TiDB Cloud Starter クラスタ内のデータが `All Chunks in the Database` フォルダに表示されます。そして、`id`、`text`、`text_vec` の3つの列があります。
- コンテキストのみを取得したい場合は、`Retrieval Only`（検索のみ）オプションをクリックします。その後、質問を入力します。
- 例えば、次のような質問をしたいとします：
```text
What is TiDB?
```
- `Send`ボタンをクリックして検索を開始します。
- ご覧のように、`TiDB is a distributed SQL database with HTAP capabilities.`が最も高い類似度スコアを持っています。
> [!NOTE]
> 列の見出しのボタンを押してドラッグすることで、テーブルの順序を変更できます。
> また、データを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) に戻り、コードファイル `bedrock/rag.py` をクリックしましょう。
1. スクリプトは必要なライブラリをインポートすることから始まります：環境変数用の`os`、データ操作用の`pandas`、LLMとの対話用の`litellm`、Webインターフェース用の`streamlit`、そしてTiDBを統一された永続データストアとして活用するための`pytidb`からのコンポーネントです。
```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_embed`関数によって生成された`text`フィールドの埋め込みを自動的に格納するベクトルフィールド。
```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`（文字列）のリストが提供されています。実際のアプリケーションでは、これらのチャンクは通常、「チャンキング」と呼ばれるプロセスを通じて、より大きな文書から派生します。
```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`をテーブルに挿入します。`pytidb`は、`text_vec`フィールドで定義されているように、挿入時に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`。
- サイドバーのスライダー（`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`の場合：
- 取得されたチャンクから`text`コンテンツが抽出されます。
- これらのテキストを結合して`context`が構築されます。
- `RAG_PROMPT_TEMPLATE`がこの`context`とユーザーの`question`でフォーマットされます。
- Bedrock LLM（`llm_model`で指定）が`litellm.completion`を通じて構築されたプロンプトで呼び出され、回答を生成します。
- LLMの応答はカスタムCSSでスタイリングされて表示されます。
- 展開可能なセクション（`📚 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")
```