--- title: "Build RAG and Text2SQL AI Applications with OpenAI" version: "v2" status: "GA" description: "In this lab, you will learn how to: Use TiDB Serverless as a storage layer for AI applications; Generate and utilize Embeddings with OpenAI models; Perform Retrieval-Augmented Generation (RAG) to retrieve information from large datasets; Build a Text-to-SQL AI application that translates natural language into SQL queries and executes them on TiDB Serverless; Develop interactive AI applications with Streamlit, integrating database operations, vector search, and LLM capabilities to create practical tools." --- # Build RAG and Text2SQL AI Applications with OpenAI ## Overview > [!IMPORTANT] > Before you start this lab, you need to register a billable OpenAI account and get the `OPENAI_API_KEY`. You can find the guidance of getting your keys from this [document from OpenAI](https://help.openai.com/en/articles/4936850-where-do-i-find-my-openai-api-key). Generative AI is currently a cutting-edge field that has attracted many developers eager to explore it. However, several barriers often hinder their entry into this domain, such as: - A lack of foundational AI knowledge, leaving them unsure where to begin. - The absence of comprehensive, step-by-step, and interactive guides for building generative AI applications. - Limited access to essential infrastructure needed for learning, such as databases and large model APIs. To address these challenges, we developed a series of labs, including this one, through which you will learn: - How to leverage pytidb and **TiDB Cloud Serverless** as a unified data storage layer for AI applications to streamline development. - What Embedding is and its capabilities. - How to utilize embedded vectors after they have been obtained. - How to employ RAG (Retrieval-Augmented Generation) for information retrieval when dealing with large volumes of data, and working with LLM (Large Language Model). - How to make a useful Text2SQL AI application to explore the dataset. > [!NOTE] > In this lab, we'll use Streamlit as a rapid application development framework to streamline and accelerate data application building. > Whether you choose your own TiDB Cloud Serverless cluster or the sandbox TiDB Cloud Serverless option to start the lab, a dedicated database will be created for this hands-on lab. ### Tasks - **Build your first RAG application** - Save your embedding vector and the other data to TiDB Serverless cluster. - Start a Streamlit webpage to test the performance of your RAG application. - Read the code of RAG Demo. - **Build your first Text2SQL application** - Use TiDB Serverless cluster as a SQL runner to correct LLM. - Start a Streamlit webpage to use your Text2SQL application. - Read the code of Text2SQL. ### Duration - This lab should take you approximately 90 minutes to complete. > [!TIP] > Once the lab environment is deployed, click **Next** on the bottom bar to begin **Module 1**. ## Module 1: Prepare to Run the Python Code ### Step 1. Check the lab environment - Click the `Resources` button on the left to view the lab environment set up for your hands-on practice. You will see: - A server named monitor1 running a VS Code server, which serves as both the development and application server. - Connection details for a TiDB Cloud Serverless cluster — either from your own TiDB Cloud account or the lab's sandbox environment, depending on your selection. - Additionally, the database name for this hands-on lab is `${LAB:SERVERLESS_CLUSTER_DATABASE_NAME}`, it will be created on your behalf in the selected TiDB Cloud Serverless cluster. - Close the `Resources` popup window. ### Step 2. Open a terminal in IDE - Enter your own [VS Code](${LAB:SECURITY_ACCESS_PROXY_URL}?next=${LAB:"/vscode/?folder=/root/tidb-lab-ai"|urlencode}) in the browser. - Click `Yes, I trust the authors` in the popup window. ![90%](https://lab-static.pingcap.com/quick-demo/vscode-trust.jpg) - Click on `≡`, `Terminal`, then `New Terminal` to open a new terminal in the VS Code. ![90%](https://public-pingcap-innovation.s3.us-west-2.amazonaws.com/demo-405/405-open-terminal-0307.jpg) ### Step 3. Install Python dependencies - Run the following commands in the VS Code terminal to install the required dependencies. ```bash pip install --upgrade pip pip install -r requirements.txt ``` ```output ... Collecting aiohappyeyeballs==2.6.1 (from -r requirements.txt (line 1)) Downloading aiohappyeyeballs-2.6.1-py3-none-any.whl.metadata (5.9 kB) Collecting aiohttp==3.11.16 (from -r requirements.txt (line 2)) Downloading aiohttp-3.11.16-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (7.7 kB) Collecting aiosignal==1.3.2 (from -r requirements.txt (line 3)) Downloading aiosignal-1.3.2-py2.py3-none-any.whl.metadata (3.8 kB) ... ``` ### Step 4. Create a .env file - Create an empty `.env` file under `openai` folder. ![90%](https://public-pingcap-innovation.s3.us-west-2.amazonaws.com/demo-409/409-new-env-0427.jpg) - Add the following content to the `.env` file and save it. Please change the `` placeholder to your actual OpenAI API Key. ```properties OPENAI_API_KEY= TIDB_HOST=${LAB:SERVERLESS_CLUSTER_HOST} TIDB_PORT=${LAB:SERVERLESS_CLUSTER_PORT} TIDB_USERNAME='${LAB:SERVERLESS_CLUSTER_USERNAME}' TIDB_PASSWORD='${LAB:SERVERLESS_CLUSTER_PASSWORD}' TIDB_DATABASE=${LAB:SERVERLESS_CLUSTER_DATABASE_NAME} ``` ## Module 2: Run, Experience, and Study the RAG Demo Application Run a RAG demo application that connects to TiDB Cloud Serverless, save your embeddings, retrieve related data, and generate the answer to your query. ### Step 1. Run the Streamlit service - In the VS Code terminal, use `streamlit` command line tool to run the following script to start Streamlit service. ```bash streamlit run openai/rag.py --server.port 8501 ``` ```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:8501 Network URL: http://172.17.0.2:8501 External URL: http://${LAB:HOST_MONITOR1_PUBLIC_IP}:8501 ``` ![90%](https://public-pingcap-innovation.s3.us-west-2.amazonaws.com/demo-409/409-rag-streamlit-start-0427.jpg) ### Step 2. Use RAG demo web page - Click to access your [RAG demo web page](http://${LAB:HOST_MONITOR1_PUBLIC_IP}:8501/). It will take around 5 minutes to complete the initialization based on the lab environment setup. ![90%](https://public-pingcap-innovation.s3.us-west-2.amazonaws.com/demo-409/409-rag-frontend.jpg) - Open the folder of `All Chunks in the Database` to see the data currently in the TiDB Serverless cluster. We did not save any data yet, so it should be empty. ![90%](https://public-pingcap-innovation.s3.us-west-2.amazonaws.com/demo-409/409-rag-start-data.jpg) - Click `Save Data` to save the text data that we prepared for this demo. You can see the original data in the next step, which is a list of string. - After saving the data, you can see the data in the TiDB Serverless cluster is showing in the `All Chunks in the Database` folder. And we have 3 columns, which are `id`, `text`, and `text_vec`. ![90%](https://public-pingcap-innovation.s3.us-west-2.amazonaws.com/demo-409/409-rag-saved-data.jpg) - If you want to retrieve the context only, you can click the `Retrieval Only` option. Then, enter your question. - Say I want to ask a question: ```text What is TiDB? ``` - Click the `Send` button to start to retrieve. ![90%](https://public-pingcap-innovation.s3.us-west-2.amazonaws.com/demo-409/409-rag-retrieve.jpg) - As you can see, the `TiDB is a distributed SQL database with HTAP capabilities.` has the highest similarity score. > **Tips:** > > You can press the button and drag the head of column to change the order of the table. > You can also download the data as a CSV table. ![70%](https://public-pingcap-innovation.s3.us-west-2.amazonaws.com/demo-409/409-rag-tips-download-0427.jpg) - At the end, let's try to make the RAG to give us the answer. - Change the `Retrieval Limit` at the left-top corner to `1`, which means we will only retrieve the most similar text towords the question that you inputed. ![90%](https://public-pingcap-innovation.s3.us-west-2.amazonaws.com/demo-409/409-rag-retrival-limit-1.jpg) - Change the option to `RAG Q&A`, then click `Send`. - This time, the demo will give us the answer rather than a material list. ![90%](https://public-pingcap-innovation.s3.us-west-2.amazonaws.com/demo-409/409-rag-answer-0427.jpg) ### Step 3. RAG Demo Code Reading 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/rag.py`. 1. After using `dotenv.load_dotenv()` to load parameters from `.env`, the next step is our prompt template. As you can see, we will use this template to insert the `{context}` and `{question}` later. Then let LLM to give us an answer based on the context. ```python # RAG prompt template RAG_PROMPT_TEMPLATE = """Answer the question based on the following reference information. Reference Information: {context} Question: {question} Please answer:""" ``` 2. We will connect to TiDB Serverless cluster, we already set these environment when the lab VM was starting. If you need to use this code snippet outside of TiDB Labs, you may need to find the parameters in your [TiDB Cloud](https://tidbcloud.com/clusters), and set it manually. ```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, ) ``` 3. We defined the names of embedding model and LLM. Then define a `Chunk` table by using this embedding model. This is because the length of the vector index is fixed, we have to decide the a model you want to use before creating the table. ```python text_embed = EmbeddingFunction("openai/text-embedding-3-small") llm_model = "openai/gpt-4o" # Define the Chunk table table_name = "chunks" class Chunk(TableModel, table=True): __tablename__ = table_name __table_args__ = {"extend_existing": True} id: int = Field(primary_key=True) text: str = Field() text_vec: Optional[Any] = text_embed.VectorField( source_field="text", ) ``` 4. We write the sample chunks as a list of string. In the real world, you might need to make a strategy to seperate a whole document, and this process is called "chunking" in RAG scenario. ```python sample_chunks = [ "Llamas are camelids known for their soft fur and use as pack animals.", "Python's GIL ensures only one thread executes bytecode at a time.", "TiDB is a distributed SQL database with HTAP capabilities.", "Einstein's theory of relativity revolutionized modern physics.", "The Great Wall of China stretches over 13,000 miles.", "Ollama enables local deployment of large language models.", "HTTP/3 uses QUIC protocol for improved web performance.", "Kubernetes orchestrates containerized applications across clusters.", "Blockchain technology enables decentralized transaction systems.", "GPT-4 demonstrates remarkable few-shot learning capabilities.", "Machine learning algorithms improve with more training data.", "Quantum computing uses qubits instead of traditional bits.", "Neural networks are inspired by the human brain's structure.", "Docker containers package applications with their dependencies.", "Cloud computing provides on-demand computing resources.", "Artificial intelligence aims to mimic human cognitive functions.", "Cybersecurity protects systems from digital attacks.", "Big data analytics extracts insights from large datasets.", "Internet of Things connects everyday objects to the internet.", "Augmented reality overlays digital content on the real world.", ] ``` 5. We attempt to open the table named `chunks`. If the table doesn't exist in the database yet (which would be the case on the first run), we create it based on the `Chunk` schema we defined earlier. This ensures the table is ready for storing and querying data. ```python table = db.open_table(table_name) if table is None: table = db.create_table(schema=Chunk) ``` 6. The next part sets up the user interface using Streamlit. We set the title of the webpage to `🔍 RAG Demo` and add a subheader "Database Operations". ```python st.title("🔍 RAG Demo") st.subheader("Database Operations") ``` 7. We create two columns for database operation buttons. The "Reset" button, when clicked, truncates the table (deletes all data) and reruns the Streamlit app to refresh the view. A warning message is displayed below it. The "Save Data" button inserts the `sample_chunks` into the table if it's currently empty and then reruns the app. When you save the `text` into the table, `pytidb` will call the embedding model to vectorize your `text` data. It is defined by the field defination of table `text_vec: Optional[Any] = text_embed.VectorField(source_field="text")`. ```python left, right = st.columns(2) if left.button("Reset", type="primary"): left.write("Resetting...") table.truncate() st.rerun() left.markdown("This option will delete all data", unsafe_allow_html=True) if right.button("Save Data"): right.write("Saving...") # insert sample chunks if table.rows() == 0: chunks = [Chunk(text=text) for text in sample_chunks] table.bulk_insert(chunks) st.rerun() ``` 8. An expander section labeled `📁 All Chunks in the Database` is created. It queries all rows from the `table`, converts the data into a Pandas DataFrame, and displays it within the expander. If the table is empty, it shows an informational message. ```python with st.expander("📁 All Chunks in the Database", expanded=False): chunks = table.query() if chunks: data = [{'id': chunk.id, 'text': chunk.text, 'text_vec': chunk.text_vec} for chunk in chunks] df = pd.DataFrame(data) st.dataframe(df) else: st.info("No data found in database.") ``` 9. We add instructional text and set up the user input elements. A radio button allows the user to select the operation mode (`Retrieval Only` or `RAG Q&A`). A slider in the sidebar lets the user adjust the number of chunks to retrieve (`Retrieval Limit`). Finally, a text input box prompts the user to enter their question. ```python st.write( "Enter your question, and the system will retrieve relevant knowledge and generate an answer" ) mode = st.radio("Select Mode:", ["Retrieval Only", "RAG Q&A"]) query_limit = st.sidebar.slider("Retrieval Limit", min_value=1, max_value=20, value=5) query = st.text_input("Enter your question:", "") ``` 10. This block handles the logic when the user clicks the "Send" button and has entered a question. It shows a spinner while processing. The core operation is `table.search(query).limit(query_limit)`, which performs a vector similarity search on the `text_vec` column using the user's query and retrieves the top `query_limit` most similar chunks. ```python if st.button("Send") and query: with st.spinner("Processing..."): # Retrieve relevant chunks res = table.search(query).limit(query_limit) ``` 11. Based on the selected `mode`, the application either displays the raw retrieval results or proceeds with the RAG process. - If `Retrieval Only`: It displays the retrieved chunks (converted to a Pandas DataFrame) directly. - If `RAG Q&A`: - It extracts the `text` content from the retrieved chunks. - It constructs the `context` by joining the extracted texts. - It formats the `RAG_PROMPT_TEMPLATE` with the generated `context` and the user's `question`. - It calls the LLM (using `litellm.completion`) with the constructed prompt to get an answer. - It displays the LLM's response using custom CSS for styling. - It includes an expander section ("📚 Retrieved Knowledge") showing the chunks used to generate the answer. If the initial search (`res`) returns no results, an informational message is displayed. ```python if res: if mode == "Retrieval Only": st.write("### Retrieval Results:") st.dataframe(res.to_pandas()) else: text = [chunk.text for chunk in res.to_rows()] # Build RAG prompt context = "\n".join(text) prompt = RAG_PROMPT_TEMPLATE.format(context=context, question=query) # Call LLM to generate answer response = completion( model=llm_model, messages=[{"content": prompt, "role": "user"}], ) st.markdown(f"### 🤖 {llm_model}") st.markdown( """ """, unsafe_allow_html=True, ) # show the response st.markdown( f'
{response.choices[0].message.content}
', unsafe_allow_html=True, ) with st.expander("📚 Retrieved Knowledge", expanded=False): st.dataframe(res.to_pandas()) else: st.info("No relevant information found") ``` ## 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 Serverless Cluster as a SQL verifier. It will: 1. Generate SQL from your queries 2. Run it on TiDB Cloud Serverless cluster 3. Get the feedback or error message from TiDB Cloud Serverless 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 Enter to generate SQL and run it on the TiDB Cloud Serverless 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 Serverless 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} """ + " ".join("|".join(t) for t in table_definitions), }, {"role": "user", "content": f"Question: {user_input} "}, ], 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 = " ".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]) ``` ## Conclusion In this lab, you learned how to build two types of AI applications using Python, TiDB Serverless, and OpenAI: - **Retrieval-Augmented Generation (RAG):** You learned how to connect to TiDB Serverless using `pytidb`, store text data along with its vector embeddings generated by OpenAI's `text-embedding-3-small` model, perform vector similarity searches to retrieve relevant context, and use a Large Language Model (LLM) like `gpt-4o` to generate answers based on the retrieved context. You also built a Streamlit interface to interact with the RAG system. - **Text-to-SQL:** You explored how to leverage an LLM (`gpt-4o-mini`) to translate natural language questions into SQL queries. You learned how to provide table schemas from TiDB as context to the LLM for accurate SQL generation, execute the generated queries against TiDB, and format the results using another LLM call. This application was also presented through a Streamlit interface. ```hbsptform { "ui": { "titleLevel": 3 }, "portalId": "4466002", "formId": "89fa8368-eba1-4b84-b91e-eb9eb6613ddf", "region": "na1", "sfdcCampaignId": "701RC00000VqCODYA3" } ```