## 模块 2：运行、体验和学习 Text2SQL 演示应用程序

接下来的演示，我们将深入探索另一种 GenAI 应用程序。它使用 TiDB Cloud Starter 集群作为 SQL 验证器。它将：

1. 从您的查询生成 SQL。
2. 在 TiDB Cloud Starter 集群上运行它。
3. 从 TiDB Cloud Starter 集群获取反馈或错误消息。
4. 自我纠正直到找到可靠的结果。

### 步骤 1. 运行 Text2SQL Streamlit 服务

创建一个新的终端：

![90%](https://pingcap-edu-us-west-2.s3.us-west-2.amazonaws.com/quick-demo/demo-421/421-text2sql-new-terminal-0514.jpg)

然后使用 `streamlit` 命令行工具运行此脚本：

   ```bash
   streamlit run bedrock/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
   ```

### 步骤 2. 使用 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
  数据库中有多少个表？
  ```

- 点击 <kbd>Enter</kbd> 键生成 SQL 并在 TiDB Cloud Starter 集群上运行它。

  ![90%](https://pingcap-edu-us-west-2.s3.us-west-2.amazonaws.com/quick-demo/demo-421/421-text2sql-query-0514.jpg)

- 或者，您可以尝试另一个问题，如：

  ```text
  数据库中的表名是什么？
  ```

  ![text2SQL-result.png 85%](https://lab-static.pingcap.com/images/2026/1/14/3a4d7b5d1a68de898647b0841a8ca98cc11a67b5.png)

- 这是因为我们在本实验开始前在数据库中创建了 `TEXT2SQL` 表。
- 您可以尝试一些您想要询问有关此数据集的问题。

### 步骤 3. Text2SQL 演示代码阅读

再次，让我们回到 [VS Code](${LAB:HTTPS_REVERSE_PROXY_URL}?next=${LAB:"/vscode/?folder=/root/tidb-lab-ai"|urlencode} "VS Code") 并点击代码文件 `bedrock/text2sql.py`。

1. 脚本导入必要的库：`os` 用于环境变量，`boto3` 用于与 AWS 服务交互（特别是 Bedrock），`streamlit` 用于 UI，以及来自 `pytidb` 的 `TiDBClient` 用于数据库连接。

    ```python
    #!/usr/bin/env python
    # -*- coding: utf-8 -*-
    import os

    import boto3
    import streamlit as st
    from pytidb import TiDBClient
    ```

2. 为 Streamlit 应用程序页面设置标题 `## 📖 Text2SQL`。

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

3. 建立了一个到 TiDB Cloud Starter 集群的 `TiDBClient` 连接。连接详情（主机、端口、用户名、密码、数据库）从环境变量中读取。连接启用了 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. 初始化了一个用于 Amazon Bedrock 运行时的 `boto3` 客户端，特别是针对 `us-west-2` 区域。这个客户端将用于调用 Bedrock 模型。

    ```python
    client = boto3.client(
        service_name="bedrock-runtime",
        region_name="us-west-2"
    )
    ```

5. 初始化 Streamlit 的会话状态以存储聊天历史。如果它们尚不存在，则创建两个列表，`generated`（用于助手的消息）和 `past`（用于用户的消息）。

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

6. 从连接的 TiDB 数据库中获取表定义。脚本遍历所有表名，对每个表执行 `SHOW CREATE TABLE`，并将 DDL 语句存储在 `table_definitions` 列表中。同时还检索了 `current_database` 名称。这些信息为 LLM 提供了架构上下文。

    ```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. 定义了一个辅助函数 `format_answer`。该函数通过移除前导/尾随代码块标记（如 \`\`\`sql 或 \`\`\`）来清理 LLM 的响应，确保 SQL 或 markdown 是干净的。

    ```python
    def format_answer(answer: str) -> str:
        # Handle ```.* specifically
        if answer.startswith("```"):
            # Find the first newline after the opening ```
            first_newline = answer.find("\n")
            if first_newline != -1:
                # Remove everything from start to the newline
                answer = answer[first_newline:].strip()

        # Remove closing ``` if present
        if answer.endswith("```"):
            answer = answer[:-len("```")].strip()

        return answer
    ```

8. 定义了 `on_submit` 函数。当用户提交问题时会调用此函数。
    - 它从会话状态中检索 `user_input`。
    - 它使用 `client.converse` 调用 Amazon Bedrock 模型 `us.amazon.nova-pro-v1:0` 来生成 SQL 查询。系统提示指导模型扮演高级 DBA 角色，使用提供的表定义和数据库名称，并直接以 SQL 形式响应，不包含额外文本。
    - 从响应中提取生成的 SQL 并使用 `format_answer` 进行清理。
    - 用户的输入被添加到 `st.session_state.past` 列表中。

    ```python
    def on_submit():
        user_input = st.session_state.user_input
        if user_input:
            generated_sql = (
                client.converse(
                    modelId="us.amazon.nova-pro-v1:0",
                    messages=[
                        {"role": "user", "content": [{"text": f"Question: {user_input}\n Respond in SQL directly, without any other text and format."}]},
                    ],
                    system=[
                        {"text": 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)}
                    ],
                )["output"]["message"]["content"][0]["text"]
            )

            generated_sql = format_answer(generated_sql)
            st.session_state.past.append(user_input)
    ```

9. 对 `generated_sql` 进行安全检查。如果它包含 `INSERT` 或 `UPDATE`（不区分大小写），将显示错误消息，并且脚本停止执行，防止执行潜在有害的 DML 语句。

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

10. 脚本尝试在 TiDB 数据库上执行 `generated_sql`：
    - 使用 `try-except` 块处理查询执行过程中可能出现的错误。
    - 如果成功，获取数据库查询结果（`rows`）并将其格式化为字符串 `sql_result`。
    - 再次调用 Bedrock 模型 `us.amazon.nova-pro-v1:0`。这次，LLM 被提示充当 markdown 格式化器。它接收原始问题、生成的 SQL 和 `sql_result` 作为输入，并被要求将它们格式化为可读的 markdown 响应，包括数据的 markdown 表格。它还被指示避免使用额外的格式化，如代码块。
    - LLM 的格式化答案使用 `format_answer` 进行清理，并添加到 `st.session_state.generated` 中。
    - 如果在此过程中发生任何异常，错误消息将被添加到 `st.session_state.generated` 中。

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

                answer = (
                    client.converse(
                        modelId="us.amazon.nova-pro-v1:0",
                        messages=[
                            {"role": "user", "content": [{"text": f"""
                            Question: {user_input}\n\n
                            SQL: {generated_sql}\n\n
                            Result: {sql_result}"""}]},
                        ],
                        system=[
                            {"text": """You are a markdown formatter,
                                format the Question and SQL to markdown text,
                                format the data row into markdown tables.
                                NOTE: Without any other format like ``` etc.
                                """}
                        ],
                    )["output"]["message"]["content"][0]["text"]
                )
                answer = format_answer(answer)
                st.session_state.generated.append(answer)
            except Exception as e:
                st.session_state.generated.append(f"Error: {e}")
    ```

11. 设置 Streamlit UI 用于用户交互：
    - 显示一个子标题 `##### User Query`。
    - 使用 `st.container` 将聊天输入和聊天显示区域分组。
    - `st.chat_input` 创建一个文本输入字段用于用户提问。`key="user_input"` 将其绑定到会话状态，`on_submit=on_submit` 将其链接到回调函数。
    - `st.empty()` 创建一个占位符 (`chat_placeholder`)，用于渲染聊天历史。

    ```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` 容器内渲染：

    - 它以倒序（最新的在前）遍历 `st.session_state["generated"]`（助手回复）和 `st.session_state["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])
    ```

### 步骤 4. 前往 TiDB Cloud 控制台开始您自己的旅程

- 创建您自己的 AWS 账户并配置 Amazon Bedrock 超出了本实验的范围。有关详细信息，请参阅 [Amazon Bedrock 文档](https://docs.aws.amazon.com/bedrock/latest/userguide/what-is-bedrock.html)。

- TiDB 及其配套的 GenAI SDK `pytidb` 通过为 GenAI 应用程序提供统一的数据存储来简化开发。您现在可以在[这里](https://tidbcloud.com/clusters)开始您的 TiDB Cloud 之旅。如果您在本实验中使用了自己的免费层集群，您会在同一页面上找到它。下个实验见！