## Module 2: Run, Experience, and Study the Text2SQL Demo Application Next demo, we will dive into another kind of GenAI 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 Text2SQL Streamlit service Create a new terminal: ![90%](https://pingcap-edu-us-west-2.s3.us-west-2.amazonaws.com/quick-demo/demo-421/421-text2sql-new-terminal-0514.jpg) Then use `streamlit` command line tool to run this script: ```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 ``` ### Step 2. Use Text2SQL demo web page - 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 Enter to generate SQL and run it on the TiDB Cloud Starter cluster. ![90%](https://pingcap-edu-us-west-2.s3.us-west-2.amazonaws.com/quick-demo/demo-421/421-text2sql-query-0514.jpg) - Or, you can try another question like: ```text What is the table name in the database? ``` ![text2SQL-result.png 85%](https://lab-static.pingcap.com/images/2026/1/14/3a4d7b5d1a68de898647b0841a8ca98cc11a67b5.png) - This is because we created the table `TEXT2SQL` in the database before we started this demo. - 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 `bedrock/text2sql.py`. 1. The script imports necessary libraries: `os` for environment variables, `boto3` to interact with AWS services (specifically Bedrock), `streamlit` for the UI, and `TiDBClient` from `pytidb` for database connection. ```python #!/usr/bin/env python # -*- coding: utf-8 -*- import os import boto3 import streamlit as st from pytidb import TiDBClient ``` 2. A title `## 📖 Text2SQL` is set for the Streamlit application page. ```python st.markdown("## 📖 Text2SQL") ``` 3. A `TiDBClient` connection is established to the TiDB Cloud Starter cluster. Connection details (host, port, user, password, database) are read from environment variables. SSL is enabled for the connection. ```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. A `boto3` client for the Amazon Bedrock runtime is initialized, specifically for the `us-west-2` region. This client will be used to invoke Bedrock models. ```python client = boto3.client( service_name="bedrock-runtime", region_name="us-west-2" ) ``` 5. Streamlit's session state is initialized to store chat history. Two lists, `generated` (for assistant's messages) and `past` (for user's messages), are created if they don't already exist. ```python for item in ["generated", "past"]: if item not in st.session_state: st.session_state[item] = [] ``` 6. Table definitions are fetched from the connected TiDB database. The script iterates through all table names, executes `SHOW CREATE TABLE` for each, and stores the DDL statements in the `table_definitions` list. The `current_database` name is also retrieved. This information provides schema context to the 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. A helper function `format_answer` is defined. This function cleans up the LLM's response by removing leading/trailing code block markers (like \`\`\`sql or \`\`\`) to ensure the SQL or markdown is clean. ```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. The `on_submit` function is defined. This function is called when the user submits their question. - It retrieves the `user_input` from the session state. - It calls the Amazon Bedrock model `us.amazon.nova-pro-v1:0` using `client.converse` to generate an SQL query. The system prompt instructs the model to act as a senior DBA, use the provided table definitions and database name, and respond directly with SQL without extra text. - The generated SQL is extracted from the response and cleaned using `format_answer`. - The user's input is added to the `st.session_state.past` list. ```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. A safety check is performed on the `generated_sql`. If it contains `INSERT` or `UPDATE` (case-insensitive), an error message is displayed, and the script stops, preventing execution of potentially harmful DML statements. ```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. The script attempts to execute the `generated_sql` against the TiDB database: - A `try-except` block handles potential errors during query execution. - If successful, the database query results (`rows`) are fetched and formatted into a string `sql_result`. - Another call to the Bedrock model `us.amazon.nova-pro-v1:0` is made. This time, the LLM is prompted to act as a markdown formatter. It takes the original question, the generated SQL, and the `sql_result` as input and is asked to format them into a readable markdown response, including a markdown table for the data. It is also instructed to avoid extra formatting like code blocks. - The LLM's formatted answer is cleaned using `format_answer` and appended to `st.session_state.generated`. - If any exception occurs during this process, the error message is appended to `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. The Streamlit UI for user interaction is set up: - A subheader `##### User Query` is displayed. - A `st.container` groups the chat input and the chat display area. - `st.chat_input` creates a text input field for the user's question. The `key="user_input"` binds it to the session state, and `on_submit=on_submit` links it to the callback function. - `st.empty()` creates a placeholder (`chat_placeholder`) where the chat history will be 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. The chat history is rendered within the `chat_placeholder` container: - It iterates through the `st.session_state["generated"]` (assistant responses) and `st.session_state["past"]` (user inputs) lists in reverse chronological order (newest first). - `st.chat_message("user")` and `st.chat_message("assistant")` are used to display each message 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]) ``` ### Step 4. Go to the TiDB Cloud Console to start your own journey + Creating your own AWS account and configuring Amazon Bedrock is out of the scope of this lab. For details, refer to the [Amazon Bedrock Documentation](https://docs.aws.amazon.com/bedrock/latest/userguide/what-is-bedrock.html). + TiDB and its companion GenAI SDK, `pytidb`, simplify development by providing a unified data store for GenAI applications. You can now start your TiDB Cloud journey [here](https://tidbcloud.com/clusters). If you had been using your own free-tier cluster for this lab, you'll find it listed on the same page, the cluster name is ${LAB:SERVERLESS_CLUSTER_NAME}.