## モジュール2: RAGデモを実行

準備の後、TiDB Cloud Starter クラスタに接続し、埋め込みを保存し、関連するデータを取得し、回答を生成するRAGデモを実行することができます。

### ステップ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
  ```

  ![90%](https://public-pingcap-innovation.s3.us-west-2.amazonaws.com/demo-409/409-rag-streamlit-start-0427.jpg)

### ステップ2. RAGデモウェブページを使用

- RAGデモウェブページを開きます [RAG demo webpage](http://${LAB:HOST_MONITOR1_PUBLIC_IP}:8501/)。

  ![90%](https://public-pingcap-innovation.s3.us-west-2.amazonaws.com/demo-409/409-rag-frontend.jpg)

- `All Chunks in the Database` フォルダーを開いて、現在 TiDB Cloud Starter クラスターにあるデータを確認してください。まだデータは保存していないため、空になっているはずです。

  ![90%](https://public-pingcap-innovation.s3.us-west-2.amazonaws.com/demo-409/409-rag-start-data.jpg)

- `Save Data`をクリックして、このデモのために用意されたテキストデータを保存します。次のステップで元のデータを見ることができ、文字列のリストです。
- データを保存した後、`All Chunks in the Database`フォルダで TiDB Cloud Starter クラスタ内のデータが表示されているのを確認できます。そして、3つのカラムがあり、それぞれ`id`、`text`、および`text_vec`です。

  ![90%](https://public-pingcap-innovation.s3.us-west-2.amazonaws.com/demo-409/409-rag-saved-data.jpg)

- 文脈のみを取得したい場合は、`Retrieval Only`オプションをクリックします。そして、質問を入力してください。
- 例えば、次の質問をしたいと思います:

  ```text
  What is TiDB?
  ```

- `送信`ボタンをクリックして取得を開始します。

  ![90%](https://public-pingcap-innovation.s3.us-west-2.amazonaws.com/demo-409/409-rag-retrieve.jpg)

- ご覧のとおり、`TiDBはHTAP機能を持つ分散型SQLデータベースです。`が最も高い類似度スコアを持っています。

> **ヒント:**
>
> ボタンを押してカラムのヘッドをドラッグするとテーブルの順序を変更できます。
> データをCSVテーブルとしてダウンロードすることもできます。
> ![70%](https://public-pingcap-innovation.s3.us-west-2.amazonaws.com/demo-409/409-rag-tips-download-0427.jpg)

- 最後に、RAGに回答を生成させてみましょう。
- 左上の`Retrieval Limit`を`1`に変更すると、入力した質問に対して最も類似したテキストのみが取得されます。

  ![90%](https://public-pingcap-innovation.s3.us-west-2.amazonaws.com/demo-409/409-rag-retrival-limit-1.jpg)

- オプションを`RAG Q&A`に変更し、`送信`をクリックします。
- 今回のデモでは、マテリアルリストではなく回答が与えられます。

  ![90%](https://public-pingcap-innovation.s3.us-west-2.amazonaws.com/demo-409/409-rag-answer-0427.jpg)

### ステップ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プロンプトテンプレート
    RAG_PROMPT_TEMPLATE = """以下の参考情報に基づいて質問に答えてください。

    参考情報:
    {context}

    質問: {question}

    回答をお願いします:"""
    ```

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. 埋め込みモデルとLLMの名前を定義しました。次に、この埋め込みモデルを使って`Chunk`テーブルを定義します。これはベクトルインデックスの長さが固定されているため、テーブルを作成する前に使用するモデルを決定する必要があるからです。

    ```python
    text_embed = EmbeddingFunction("openai/text-embedding-3-small")
    llm_model = "openai/gpt-4o"

    # Chunkテーブルを定義します
    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. サンプルチャンクを文字列のリストとして書きます。実世界では、全体のドキュメントを分割する戦略を立てる必要があるかもしれません。このプロセスはRAGのシナリオで「チャンク化」と呼ばれます。

    ```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`に設定し、「Database Operations」というサブヘッダーを追加します。

    ```python
    st.title("🔍 RAG Demo")

    st.subheader("Database Operations")
    ```

7. データベース操作ボタンのための2つのカラムを作成します。「Reset」ボタンがクリックされると、テーブルはトランケート（全データを削除）され、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 <span style='color: red;'>will delete all data</span></h3>", 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`がラベルされたエクスパンダーセクションが作成されます。それはテーブルからすべての行をクエリし、そのデータを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`）をユーザーに選択させます。サイドバーのスライダーでは、取得するチャンクの数を調整できます（`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`に基づいて、アプリケーションはRAW抽出結果を表示するか、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(
                        """
                    <style>
                    .llm-response {
                        background: rgba(255, 255, 255, 0.05);
                        padding: 25px;
                        border-radius: 15px;
                        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.2);
                        border: 1px solid rgba(255, 255, 255, 0.1);
                        margin: 15px 0;
                        font-size: 1.1em;
                        line-height: 1.6;
                        color: #e1e4e8;
                    }
                    .llm-response:hover {
                        background: rgba(255, 255, 255, 0.08);
                        box-shadow: 0 6px 12px rgba(0, 0, 0, 0.3);
                        transition: all 0.3s ease;
                    }
                    </style>
                    """,
                        unsafe_allow_html=True,
                    )

                    # show the response
                    st.markdown(
                        f'<div class="llm-response">{response.choices[0].message.content}</div>',
                        unsafe_allow_html=True,
                    )

                    with st.expander("📚 Retrieved Knowledge", expanded=False):
                        st.dataframe(res.to_pandas())
            else:
                st.info("No relevant information found")
    ```
