# OpenAI を使用した RAG および Text2SQL AI アプリケーションの構築

## 概要

AIは現在最先端の分野であり、多くの開発者がその学習に意欲的です。しかし、以下のようないくつかの障害が学習プロセスを妨げることがあります。

- 基礎的なAI知識の不足
- 過度に複雑な基盤アーキテクチャ
- AIアプリケーションに対する包括的で段階的なドキュメントの欠如

これらの問題に対処するため、私たちはこの演習を開発しました。この演習では次のことを学びます。

- **TiDB Serverless**をAIアプリケーションのデータストレージレイヤーとして使用し、AI開発プロセスを高速化する方法。
- 埋め込みとその機能について。
- 埋め込みベクトルを取得後にどのように活用するか。
- 大量のデータを扱う際に、データ回収のためのRAGをどのように使用するか。
- データセットを探索するための有用なText2SQL AIアプリをどのように作成するか。

> **注意:**
>
> この演習に入る前に、課金可能な OpenAI アカウントを登録し、`OPENAI_API_KEY`を取得する必要があります。この[OpenAIのドキュメント](https://help.openai.com/en/articles/4936850-where-do-i-find-my-openai-api-key)から API キーの取得方法を確認できます。

### タスク

- **RAG**

  - 埋め込みベクトルと他のデータをTiDB Serverless クラスタに保存します。
  - RAGアプリケーションのパフォーマンスをテストするためにstreamlitウェブページを開始します。
  - RAGデモのコードを読みます。

- **Text2SQL**

  - TiDB ServerlessクラスタをSQLランナーとして使用してLLMを訂正します。
  - Text2SQLアプリケーションを使用するためにstreamlitウェブページを開始します。
  - Text2SQLのコードを読みます。

### 所要時間

- この演習を完了するためには約1時間かかります。

> [!TIP]
> 演習環境がデプロイされたら、下部のバーで**次へ**をクリックして**モジュール1**を開始します。

## モジュール1: Python コードを実行する準備

### ステップ1. ターミナルを開く

- ブラウザで自分の [VS Code](${LAB:SECURITY_ACCESS_PROXY_URL}?next=${LAB:"/vscode/?folder=/root/tidb-lab-ai"|urlencode}) を開きます。

- ポップアップウィンドウが表示されたら、「Yes, I trust the authors（はい、著者を信頼します）」をクリックしてください。

  ![90%](https://lab-static.pingcap.com/quick-demo/vscode-trust.jpg)

- 左上の `≡` をクリックし、次に `Terminal`、そして `New Terminal` をクリックして新しいターミナルを開いてください。

  ![90%](https://public-pingcap-innovation.s3.us-west-2.amazonaws.com/demo-405/405-open-terminal-0307.jpg)

### ステップ2. Python 依存関係のインストール

Python の依存パッケージをインストールします。

```bash
pip install -r requirements.txt
```

```output
Collecting aiohappyeyeballs==2.6.1 (from -r requirements.txt (line 1))
  Downloading aiohappyeyeballs-2.6.1-py3-none-any.whl.metadata (5.9 kB)
Collecting aiohttp==3.11.16 (from -r requirements.txt (line 2))
  Downloading aiohttp-3.11.16-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (7.7 kB)
Collecting aiosignal==1.3.2 (from -r requirements.txt (line 3))
  Downloading aiosignal-1.3.2-py2.py3-none-any.whl.metadata (3.8 kB)
...
[notice] A new release of pip is available: 24.0 -> 25.1
[notice] To update, run: pip install --upgrade pip
```

### ステップ3. .envファイルを作成

- `openai`フォルダに`.env`を作成します。

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

- `.env`ファイルに以下の内容を追加し、保存します。`<Your OpenAI API Key>`のプレースホルダーをあなたのReal OpenAI APIキーに置き換えてください。

  ```properties
  OPENAI_API_KEY=<Your OpenAI API Key>

  TIDB_HOST=${LAB:SERVERLESS_CLUSTER_HOST}
  TIDB_PORT=${LAB:SERVERLESS_CLUSTER_PORT}
  TIDB_USERNAME='${LAB:SERVERLESS_CLUSTER_USERNAME}'
  TIDB_PASSWORD='${LAB:SERVERLESS_CLUSTER_PASSWORD}'
  TIDB_DATABASE=${LAB:SERVERLESS_CLUSTER_DATABASE_NAME}
  ```

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

準備の後、TiDB Cloud Serverless クラスタに接続し、埋め込みを保存し、関連するデータを取得し、回答を生成する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 Serverless クラスターにあるデータを確認してください。まだデータは保存していないため、空になっているはずです。

  ![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 Serverless クラスタ内のデータが表示されているのを確認できます。そして、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 Serverless クラスタに接続します。これこれらの環境設定は演習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")
  ```

## Module 3: テキストから SQL へのデモの実行

次のデモでは、もう一つの種類のAIアプリケーションについて学びます。これはTiDB Cloud Serverless ClusterをSQLの確認ツールとして使用します。主な手順は以下の通りです。

1. あなたのクエリからSQLを生成します。
2. TiDB Serverless クラスタで実行します。
3. TiDB Serverless クラスタからフィードバックまたはエラーメッセージを取得します。
4. 固い結論が出るまで自分を修正します。

### ステップ 1. Streamlit サービスの実行

新しいターミナルを作成します。

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

その後、`streamlit`コマンドラインツールを使用してこのスクリプトを実行します。

```bash
streamlit run openai/text2sql.py --server.port 8502
```

```output
使用状況統計を収集しています。無効にするには、browser.gatherUsageStats を false に設定してください。


ブラウザで Streamlit アプリを表示できます。

ローカル URL: http://localhost:8502
ネットワーク URL: http://172.17.0.2:8502
外部 URL: http://${LAB:HOST_MONITOR1_PUBLIC_IP}:8502
```

### ステップ 2. テキストから SQL へのデモウェブページの使用

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

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

- その後、このデータセットに関する質問をしてみましょう、例:

  ```text
  データベースにはいくつのテーブルがありますか?
  ```

- <kbd>Enter</kbd> をクリックしてSQLを生成し、TiDB Cloud Serverless クラスタで実行します。

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

- または、別の質問を試してみましょう:

  ```text
  データベースのテーブル名は何ですか?
  ```

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

- Module 2 を実行し、そこで `chunks` テーブルを作成しました。そのため、これをお試しいただけます。
- このデータセットについて知りたい質問をいくつか試してみてください。

### ステップ 3. テキストから SQL へのデモコード読み取り

再び[VS Code](http://${LAB:HOST_MONITOR1_PUBLIC_IP}:3000/?folder=/root/tidb-lab-ai)に戻り、コードファイル `openai/text2sql.py` をクリックします。

1. 最初に、環境変数は `.env` ファイルから読み込まれます。これにはAPIキーやデータベース資格情報のような機密情報が含まれます。

  ```python
  dotenv.load_dotenv()
  ```

2. Pydantic モデル `QuestionSQLResponse` を定義します。このモデルは、OpenAI APIが返すことを期待する構造を指定します：元の質問、生成されたSQLクエリ、および初期マークダウン表現。Pydanticを使用することでAPIレスポンスがニーズに合うようになります。

  ```python
  class QuestionSQLResponse(BaseModel):
      question: str
      sql: str
      markdown: str
  ```

3. Streamlit アプリケーションページのタイトルを設定します。

  ```python
  st.markdown("## 📖 Text2SQL")
  ```

4. 環境変数からロードされた資格情報を使ってTiDB Serverlessデータベースに接続を確立します。また、APIキーを使用してOpenAIクライアントを初期化します。

  ```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,
  )
  oai = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
  ```

5. Streamlit のセッションステートを初期化します。`st.session_state` は、単一のセッション内でユーザーのインタラクションを超えて持続する辞書のようなオブジェクトです。アシスタントの応答とユーザーの入力を順次保存するために、`generated` と `past` の2つの空のリストを作成します。

  ```python
  for item in ["generated", "past"]:
      if item not in st.session_state:
          st.session_state[item] = []
  ```

6. 現在のデータベースのすべてのテーブルのスキーマ定義（`SHOW CREATE TABLE`）を取得します。LLMが正しいSQLクエリを生成するために、このスキーマ情報は重要です。データベース名も取得します。

  ```python
  table_definitions = []
  current_database = db._db_engine.url.database
  for table_name in db.table_names():
      table_definitions.append(db.query(f"SHOW CREATE TABLE `{table_name}`").to_rows()[0])
  ```

7. ユーザーが質問を入力したときにトリガーされる `on_submit` 関数を定義します。

  ```python
  def on_submit():
      user_input = st.session_state.user_input
      if user_input:
          # ... (LLM call and processing) ...
  ```

8. `on_submit` 内で、ユーザー入力がある場合：

  - `oai.beta.chat.completions.parse` を使用して OpenAI API を呼び出し、`gpt-4o-mini` を利用します。
  - モデルにシステムプロンプトを提供し、モデルが上級DBAとして動作し、MySQLのSQLを書き、識別子に対してバッククォートを使用し、提供されたテーブル定義とデータベース名を文脈として含めます。
  - ユーザーの質問を含めます。
  - `response_format` で `QuestionSQLResponse` Pydantic モデルに一致させます。
  - 最初の選択肢のメッセージを解析し `QuestionSQLResponse` オブジェクトを取得します。

  ```python
          response = (
              oai.beta.chat.completions.parse(
                  model="gpt-4o-mini",
                  messages=[
                      {
                          "role": "system",
                          "content": f"""
                          You are a very senior database administrator who can write SQL very well,
                          please write MySQL SQL to answer user question,
                          Use backticks to quote table names and column names,
                          here are some table definitions in database,
                          the database name is {current_database}

"""
                          + "\n".join("|".join(t) for t in table_definitions),
                      },
                      {"role": "user", "content": f"Question: {user_input}\n"},
                  ],
                  response_format=QuestionSQLResponse,
              )
              .choices[0]
              .message.parsed
          )
          st.session_state.past.append(user_input)
  ```

9. 潜在的に危険なSQL操作を防ぐための安全チェックを実行します。生成されたSQLに `INSERT` または `UPDATE` が含まれている場合は、エラーを表示して実行を停止します。

  ```python
          if "insert" in response.sql.lower() or "update" in response.sql.lower():
              st.error(
                  "The generated SQL is not a SELECT statement, please check it carefully before running it."
              )
              st.stop()
  ```

10. 生成されたSQLクエリをTiDBデータベースに対して実行します。

  - クエリ実行中に発生する可能性のあるエラーを処理するために、`try...except` ブロックを使用します。
  - 成功した場合、結果（`rows`）を取得し、文字列（`sql_result`）にフォーマットします。
  - OpenAI API を再び呼び出します（`oai.chat.completions.create` with `gpt-4o-mini`）。このとき、プロンプトはLLMに元の質問、生成されたSQL、初期マークダウン、および実際のデータベース結果（`sql_result`）をユーザーフレンドリーなMarkdownフォーマットに変換し、データのテーブルを含むように指示します。
  - 最終的なフォーマットされた回答を `st.session_state.generated` に保存します。
  - SQL実行中にエラーが発生した場合、エラーメッセージを `st.session_state.generated` に保存します。

  ```python
          # Execute the SQL query and set the result
          answer = None
          try:
              rows = db.query(response.sql).to_rows()
              sql_result = "\n".join(str(row) for row in rows)

              answer = (
                  oai.chat.completions.create(
                      model="gpt-4o-mini",
                      messages=[
                          {
                              "role": "system",
                              "content": "You are a markdown formatter, format the user input to markdown, format the data row into markdown tables.",
                          },
                          {
                              "role": "user",
                              "content": f"""
                          Question: {response.question}


                          SQL: {response.sql}


                          Markdown: {response.markdown}


                          Result: {sql_result}""",
                          },
                      ],
                  )
                  .choices[0]
                  .message.content
              )
              st.session_state.generated.append(answer)
          except Exception as e:
              st.session_state.generated.append(f"Error: {e}")
  ```

11. チャット用のユーザーインターフェース要素を設定します。

  - 「ユーザーのクエリ」というサブヘッダーを追加します。
  - `st.container` を使って入力フィールドとチャット履歴をグループ化します。
  - チャット入力ウィジェット（`st.chat_input`）をプレースホルダーテキストと共に作成します。`key="user_input"`でセッションステートにリンクされ、`on_submit=on_submit`でコールバック関数が登録されます。
  - チャット履歴を動的にレンダリングするための空のプレースホルダーを作成します（`st.empty()`）。

  ```python
  st.markdown("##### User Query")
  with st.container():
      st.chat_input(
          "Input your question, e.g. how many tables in the database?",
          key="user_input",
          on_submit=on_submit,
      )

      chat_placeholder = st.empty()
  ```

12. `chat_placeholder` 内でチャット履歴をレンダリングします。

  - `generated`（アシスタント）と`past`（ユーザー）メッセージを逆順（最新のものから）に繰り返し処理します。
  - `st.chat_message("user")` と `st.chat_message("assistant")` を使って適切なスタイルでメッセージを表示します。

  ```python
      with chat_placeholder.container():
          for i in range(len(st.session_state["generated"]) - 1, -1, -1):
              with st.chat_message("user"):
                  st.write(st.session_state["past"][i])
              with st.chat_message("assistant"):
                  st.write(st.session_state["generated"][i])
  ```

## 結論

この演習では、Python、TiDB Serverless、およびOpenAIを使用して2種類のAIアプリケーションを構築する方法を学びました。

- **リトリーバル強化生成（RAG）:** `pytidb`を使ってTiDB Serverlessに接続し、OpenAIの`text-embedding-3-small`モデルが生成したベクトル埋め込みとともにテキストデータを格納し、ベクトル類似性検索を実行して関連コンテキストを取得し、`gpt-4o`のような大規模言語モデル（LLM）を使って、取り出されたコンテキストに基づく回答を生成する方法を学びました。また、Streamlitインターフェースを構築してRAGシステムと対話しました。
- **テキストからSQLへ:** 自然言語の質問をSQLクエリに翻訳するためにLLM（`gpt-4o-mini`）を活用する方法を探りました。TiDBからテーブルスキーマをLLMに文脈として提供して、正確なSQL生成を行い、生成されたクエリをTiDBに対して実行し、別のLLM呼び出しを使用して結果をフォーマットしました。このアプリケーションもStreamlitインターフェースを介して提供されました。

データベース操作、ベクトル検索、およびLLM機能を統合して実用的なAIツールを作成する実践的な経験を積みました。

次のステップとしては、以下を試してみてください。

- この演習のソースコードをさらに探索してください [tidb-samples/tidb-lab-ai](https://github.com/tidb-samples/tidb-lab-ai)。
- `pytidb`とその機能についてもっと学ぶことができます [Githubリポジトリ](https://github.com/pingcap/pytidb)。
- 自分のTiDB Cloud Serverless クラスタにアクセス [https://tidbcloud.com](https://tidbcloud.com)。
- `tidb-samples`の組織内で他の例を発見することができます [GitHub](https://github.com/tidb-samples)。
