Building a Flask API with Azure OpenAI and SQLAlchemy for Dynamic SQL Query Generation

Admin
Python
Oct 30, 2024

Table of Content

Hiral Patel

VP of Technology at Prioxis | Cloud Evangelist

Today, companies need fast, flexible ways to access data, and APIs are at the heart of making that happen. In fact, 85% of businesses now consider APIs essential for growth, and API usage is climbing over 30% every year. But here’s the big question: what if your API could understand questions in natural language, like “What were last month’s top sales?” without needing any SQL skills?

That’s exactly what happens when you combine Flask, Azure OpenAI, and SQLAlchemy. Together, they create an API that can take everyday language, turn it into a smart SQL query, and get you the data you need right away. It’s not just for developers, it’s a setup that makes finding data easier for anyone, technical or not.

In this post, we will walk you through how to build a Flask API that dynamically generates and executes SQL queries using Azure OpenAI, SQLAlchemy, and LangChain. This solution is particularly useful for building dynamic reporting tools or applications that require users to query a database through natural language inputs.

Table of Contents:

Understanding Azure OpenAI
Setting Up the Development Environment
Connecting to PostgreSQL with SQLAlchemy
Integrating Azure OpenAI for SQL Querying
Using LangChain for SQL Generation
Natural Language to SQL Pipeline
Handling Responses and Token Usage
Conclusion

What is Azure OpenAI

Azure OpenAI Service lets you use OpenAI's advanced language models through the Azure platform. These models can understand and create text that sounds like it was written by a person, making them great for tasks like analyzing language, generating text, and even converting plain language into SQL queries.

Why This Approach?

This project leverages the power of OpenAI’s language models to convert natural language queries into SQL, making it easier for non-technical users to interact with databases.

Talk to Your Data – No SQL Required

With Azure OpenAI, just type a question like “Show me last month’s orders,” and the API converts it into the right SQL query. Perfect for users who need data fast, no SQL skills needed.

SQLAlchemy Keeps Data Simple

SQLAlchemy makes database work cleaner and easier in Python. Forget complex SQL code—this tool lets developers manage data smoothly, speeding up the process.

Ask Anything, Get Data Instantly

Dynamic query generation means you’re not stuck with predefined questions. The API adapts, fetching whatever data you need, on demand.

Azure Means Big-League Scaling and Security

Hosting on Azure keeps the API fast, reliable, and secure, no matter the traffic. As your data needs grow, Azure has you covered.

Less Coding, More Results

With Azure OpenAI and SQLAlchemy, developers don’t need to build separate endpoints for every data request. This setup saves time and keeps the API versatile.

Smart Data Access Without the Risks

AI-driven queries and SQLAlchemy reduce security risks, making data access both smarter and safer. Azure adds a robust layer of protection.

It is a smart approach for apps that need flexible and easy access to data. Here’s why:

The API is built using:

Flask to provide RESTful endpoints,
SQLAlchemy for interacting with a PostgreSQL database
LangChain for structuring AI-based prompts,
Azure OpenAI to generate SQL queries from user inputs.

Prerequisites

Basic knowledge of coding and Python.
Familiarity with full stack development.
Intermediate understanding of Bash and command-line usage.

Let’s dive into how you can create this API step by step.

1. Setting Up the Development Environment

Before we begin, we need to set up our environment. Below are the key libraries required for this project:

pip install Flask SQLAlchemy pandas python-dotenv langchain_openai psycopg2

In addition, make sure you have an Azure OpenAI account set up and PostgreSQL running. You'll need access to environment variables to securely load API keys and database credentials.

Environment Setup

Create a .env file to store sensitive information like your database credentials and API keys:

DB_USER=
DB_PASSWORD=
DB_HOST=
DB_PORT=
DB_NAME=
AZURE_OPENAI_ENDPOINT=
AZURE_OPENAI_API_KEY=

2. Connecting to PostgreSQL with SQLAlchemy

SQLAlchemy is a powerful library for working with databases in Python. In our project, we’ll use it to connect to PostgreSQL.

from sqlalchemy import create_engine
from sqlalchemy.engine.url import URL
import os
from dotenv import load_dotenv
load_dotenv()

# Load environment variables
username = os.getenv("DB_USER")
password = os.getenv("DB_PASSWORD")
server_name = os.getenv("DB_HOST")
port = os.getenv("DB_PORT")
database_name = os.getenv("DB_NAME")

# Create PostgreSQL connection URL
postgres_url = URL.create(
    drivername='postgresql+psycopg2',
    username=username,
    password=password,
    host=server_name,
    port=port,
    database=database_name
)
# Create SQLAlchemy engine
engine = create_engine(postgres_url)

This code sets up a connection to the PostgreSQL database using credentials stored in environment variables. The connection is encapsulated within SQLAlchemy’s create_engine() method, which will allow us to run SQL queries later.

Loading Table Descriptions

We assume that a CSV file contains metadata about the database’s tables, which we use to help the AI determine which tables are relevant to the user’s query. This data is fetched using Pandas:

import pandas as pd
def get_table_details() -> str:
    file_path = os.path.join('Data', 'database_table_descriptions.csv')
    print(f"Reading file from: {file_path}")
    table_description = pd.read_csv(file_path)
    table_names = ', '.join(table_description['Table'].tolist())
    return table_names

3. Integrating Azure OpenAI for SQL Querying

The core of our solution involves using Azure OpenAI to convert user questions into SQL queries. We set up the Azure OpenAI client using the langchain_openai library.

from langchain_openai import AzureChatOpenAI
# Initialize Azure OpenAI client
azure_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
api_key = os.getenv("AZURE_OPENAI_API_KEY")
llm = AzureChatOpenAI(
    deployment_name="gpt-35-turbo-0613",
    temperature=0,
    max_tokens=500,
    api_version="2024-05-01-preview",
    azure_endpoint=azure_endpoint,
    api_key=api_key,
)

This configuration sets up the OpenAI model to process SQL queries. The model is responsible for understanding the user's question and generating a syntactically correct SQL query for our database.

    # Get table details
    table_details = get_table_details()
    # Define table extraction prompt

    table_details_prompt = f"""Return the names of ALL the SQL tables that MIGHT be relevant to the user question.
    The tables are: {table_details}
    Remember to include ALL POTENTIALLY RELEVANT tables, even if you're not sure that they're needed.
    Return the table names as a comma-separated list."""
    table_extraction_prompt = ChatPromptTemplate.from_messages([
        ("system", table_details_prompt),
        ("human", "{input}")
    ])
    list_parser = CommaSeparatedListOutputParser()

    # Invoke table extraction and capture the response
    tables_response = table_extraction_prompt | llm | list_parser
    tables = tables_response.invoke({"input": user_question})

In this example, we're using a large language model (LLM) to extract relevant SQL table names based on user queries. The code first retrieves potential tables, builds a system prompt, and passes it to the LLM along with the user's question. The model then returns a list of suggested tables. The result is parsed using a CommaSeparatedListOutputParser to ensure we get a structured response. This approach helps streamline the process of identifying relevant tables for complex queries, reducing manual effort and improving accuracy.

4. Using LangChain for SQL Generation

We use the LangChain library to build prompt templates that guide the language model to generate SQL queries. It provides a structured way to give the AI model context and examples of how it should respond.

Creating a Few-Shot Prompt

Few-shot learning helps the model understand the context by providing a few examples of how to convert questions into SQL queries.

examples = [
        {"input": "List all employees",
         "query": """select * from employee"""}
    ]
    example_prompt = ChatPromptTemplate.from_messages(
        [("human", "{input}\nSQLQuery:"), ("ai", "{query}")]
    )
    few_shot_prompt = FewShotChatMessagePromptTemplate(
        example_prompt=example_prompt,
        examples=examples,
        input_variables=["input"]
    )

Here, we define some example questions and their corresponding SQL queries. This gives the model an idea of how it should generate queries when a new question is presented.

5. Natural Language to SQL Pipeline

The central part of the application is the pipeline that takes the user’s question and generates an SQL query:

final_prompt = ChatPromptTemplate.from_messages(
        [
            ("system", "You are a PostgreSQL expert. Given an input question, create a syntactically correct PostgreSQL query to run. "
                       "Here is the relevant table info: {table_info}.\n"
                       "Consider the top {top_k} results. Below are a number of examples of questions and their corresponding SQL queries."),
            few_shot_prompt,
            ("human", "{input}"),
        ]
    )
    # Create the SQL query chain
    generate_query = create_sql_query_chain(llm, db, final_prompt)
    # Execute the query using QuerySQLDataBaseTool
    execute_query_tool = QuerySQLDataBaseTool(db=db)
    answer_prompt = PromptTemplate.from_template(
        """
        Given the following user question, corresponding SQL query, and SQL result, answer the user question.
        Question: {question}
       SQL Query: {query}
        SQL Result: {result}
        Answer:
        """
    )
    # Prepare the answer prompt processing
    rephrase_answer = answer_prompt | llm | StrOutputParser()
    # Prepare the chain using RunnableMap
    chain = (
        RunnablePassthrough.assign(table_info=lambda _: table_details) |
        RunnablePassthrough.assign(query=generate_query).assign(
            result=itemgetter("query") | execute_query_tool
        ) |
        # Store the raw response before parsing
        RunnablePassthrough.assign(raw_response=lambda x: rephrase_answer.invoke({
            "question": x["question"],
            "query": x["query"],
            "result": x["result"]
        }))  # Invoke rephrase_answer here to get raw LLM response
    )
    # Invoke the final chain with the input_data
    final_response = chain.invoke({
        "question": user_question,
        "top_k": top_k_default,
    })

This code sets up a chain to generate and execute SQL queries based on user input. It utilizes a language model to create the SQL query, processes it through a query execution tool, and then returns the result in a structured format. The RunnableMap ensures that table details and query execution are passed through the pipeline efficiently. Finally, the response is rephrased and returned to the user in a clear, understandable format, making SQL query generation seamless and user-friendly.

6. Handling Responses and Token Usage

Finally, we ensure that the results are formatted correctly before sending them back to the user. We also calculate token usage to help monitor API costs when interacting with Azure OpenAI.

# Calculate token usage
    input_tokens = len(user_question.split())  # Estimate input tokens based on word count
    output_tokens = len(raw_response.split())  # Estimate output tokens based on raw response word count
    # Return the response along with estimated token usage
    return jsonify({
        "result": output_values,
        "token_usage": {
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            "total_tokens": input_tokens + output_tokens
        }
    })

This approach not only provides the user with the data they’re seeking but also gives insight into the computational resources being used.

Conclusion

By combining Flask, SQLAlchemy, Azure OpenAI, and LangChain, we can dynamically convert user queries into SQL and retrieve data from a PostgreSQL database. This project is a great example of how AI can bridge the gap between complex systems and non-technical users, allowing them to interact with databases through natural language.

Key Takeaways:

Flask makes it easy to build RESTful APIs.
SQLAlchemy abstracts away much of the complexity of working with SQL databases.
LangChain helps create structured prompts to guide AI models.
Azure OpenAI allows us to harness the power of large language models for complex tasks like SQL generation.

This API could be extended to support more advanced SQL queries, error handling, or even other types of databases. The potential applications are vast, and this is just the beginning!

Get in touch