## 模块 3：运行 Text2SQL 示例

接下来的示例，我们将探索另一种 AI 应用。它使用 TiDB Cloud Starter 集群作为 SQL 校验器。流程如下：

1. 从你的查询生成 SQL
2. 在 TiDB Cloud Starter 集群上执行
3. 获取 TiDB Cloud Starter 集群的反馈或错误信息
4. 自动修正，直到得到正确结果

### 第一步：运行 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
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.17.0.2:8502
External URL: http://${LAB:HOST_MONITOR1_PUBLIC_IP}:8502
```

### 第二步：使用 Text2SQL 示例网页

- 打开你的 [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
  How many tables in the database?
  ```

- 按下 <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
  What is the table name in the database?
  ```

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

- 这是因为在模块 2 中，我们创建了 `chunks` 表。
- 你可以根据这个数据集自由尝试想问的问题。

### 第三步：阅读 Text2SQL 示例代码

接下来回到 [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 以及初步的 Markdown 表达。使用 Pydantic 确保 API 响应符合要求。

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

3. 设置 Streamlit 应用页面的标题。

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

4. 使用加载的环境变量连接 TiDB Cloud Starter 数据库，同时初始化 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`，分别保存助手的回答和用户的输入，形成聊天历史。

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

6. 获取当前数据库中所有表的建表语句（`SHOW CREATE TABLE`）。这些表定义信息是生成正确 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 调用与处理）...
    ```

8. 在 `on_submit` 中，如果用户有输入：

    - 使用 `gpt-4o-mini` 调用 OpenAI API（`oai.beta.chat.completions.parse`）；
    - 提供系统提示词，指示模型扮演资深 DBA，撰写符合 MySQL 规范的 SQL；
    - 传递表定义和数据库名作为上下文；
    - 指定返回格式为 `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 查询。

    - 使用 `try...except` 捕获查询异常；
    - 成功后，提取结果并格式化成字符串；
    - 再次调用 OpenAI API，将问题、SQL、初始 Markdown 和实际结果整合为最终的 Markdown 格式回答；
    - 将最终回答存入 `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. 设置聊天界面元素。

    - 添加小标题 "User Query"；
    - 使用 `st.container` 分组输入框和聊天记录；
    - 创建聊天输入控件 (`st.chat_input`)，绑定到 `user_input`；
    - 创建空占位元素 (`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])
    ```