## Module 3: Run, Experience, and Study the Text2SQL Demo Application

Next demo, we will dive into another kind of AI application. It uses TiDB Cloud Starter Cluster as a SQL verifier. It will:

1. Generate SQL from your queries
2. Run it on TiDB Cloud Starter cluster
3. Get the feedback or error message from TiDB Cloud Starter cluster
4. Correct itself until it finds a solid result

### Step 1. Run the Streamlit Service

Create a new terminal:

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

Then use `streamlit` command line tool to run this script:

```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
```

### Step 2. Use Text2SQL Demo Webpage

- Open your [Text2SQL demo webpage](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)

- Then, we can ask a question about this dataset, e.g.:

  ```text
  How many tables in the database?
  ```

- Click the <kbd>Enter</kbd> to generate SQL and run it on the TiDB Cloud Starter cluster.

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

- Or, you can try another question like:

  ```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)

- It is because we ran the Module 2, and we created `chunks` table at there.
- You can try some questions that you want to ask about this dataset.

### Step 3. Text2SQL Demo Code Reading

Again, let's head back to [VS Code](http://${LAB:HOST_MONITOR1_PUBLIC_IP}:3000/?folder=/root/tidb-lab-ai) and click the code file `openai/text2sql.py`.

1. First, we load environment variables from the `.env` file, which contains sensitive information like API keys and database credentials.

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

2. We define a Pydantic model `QuestionSQLResponse`. This model specifies the structure we expect the OpenAI API to return: the original question, the generated SQL query, and a preliminary markdown representation. Using Pydantic ensures the API response conforms to our needs.

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

3. Set the title for the Streamlit application page.

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

4. Establish a connection to the TiDB Cloud Starter database using credentials loaded from the environment variables. We also initialize the OpenAI client with the API key.

    ```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. Initialize Streamlit's session state. `st.session_state` is a dictionary-like object that persists across user interactions within a single session. We create two empty lists, `generated` and `past`, to store the assistant's responses and the user's inputs, respectively, forming the chat history.

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

6. Fetch the schema definitions (`SHOW CREATE TABLE`) for all tables in the current database. This schema information is crucial context for the LLM to generate correct SQL queries. The database name is also retrieved.

    ```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. Define the `on_submit` function, which is triggered when the user enters a question.

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

8. Inside `on_submit`, if there's user input:

    - Call the OpenAI API (`oai.beta.chat.completions.parse`) using `gpt-4o-mini`.
    - Provide a system prompt instructing the model to act as a senior DBA, write MySQL SQL, use backticks for identifiers, and provide the fetched table definitions and database name as context.
    - Include the user's question.
    - Specify the `response_format` to match our `QuestionSQLResponse` Pydantic model.
    - Parse the first choice's message to get the `QuestionSQLResponse` object.

    ```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"""
                            + "".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. Perform a safety check to prevent potentially harmful SQL operations. If the generated SQL contains `INSERT` or `UPDATE`, display an error and stop execution.

    ```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. Execute the generated SQL query against the TiDB database.

    - Use a `try...except` block to handle potential errors during query execution.
    - If successful, fetch the results (`rows`), format them into a string (`sql_result`).
    - Call the OpenAI API again (`oai.chat.completions.create` with `gpt-4o-mini`). This time, the prompt asks the LLM to format the original question, the generated SQL, the initial markdown, and the actual database result (`sql_result`) into a user-friendly Markdown format, including tables for the data.
    - Store the final formatted answer in `st.session_state.generated`.
    - If an error occurs during SQL execution, store the error message in `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. Set up the user interface elements for the chat.

    - Add a subheader "User Query".
    - Use `st.container` to group the input field and chat history.
    - Create a chat input widget (`st.chat_input`) with a placeholder text. The `key="user_input"` links it to the session state, and `on_submit=on_submit` registers the callback function.
    - Create an empty placeholder (`st.empty()`) where the chat history will be dynamically rendered.

    ```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. Render the chat history within the `chat_placeholder`.

    - Iterate through the `generated` (assistant) and `past` (user) messages in reverse order (newest first).
    - Use `st.chat_message("user")` and `st.chat_message("assistant")` to display messages with appropriate styling.

    ```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])
    ```