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

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

1. あなたのクエリからSQLを生成します。
2. TiDB Cloud Starter クラスタで実行します。
3. TiDB Cloud Starter クラスタからフィードバックまたはエラーメッセージを取得します。
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 Starter クラスタで実行します。

  ![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 Cloud Starterデータベースに接続を確立します。また、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\n"""
                            + "\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])
    ```
