Creating an SQL Agent Using OpenAI and Python

In today’s data-driven world, accessing databases and retrieving information efficiently is crucial. However, not everyone is proficient in SQL. That’s where SQL agents come in. These tools bridge the gap between natural language and SQL, making data access more intuitive for everyone.

In this blog, we’ll build a SQL agent that takes natural language queries, converts them into SQL using OpenAI’s GPT model, and executes them against a database using SQLAlchemy.

Why Build a SQL AI Agent?

• Simplify Database Queries: Users don’t need to know SQL to interact with databases.
• Increase Productivity: Reduces the time spent writing SQL queries.
• Expand Access: Enables non-technical users to access and manipulate data.

How It Works
The SQL AI agent works in three key steps:

Interpret the User’s Query: OpenAI’s language model translates the user’s natural language input into SQL.
Execute the SQL Query: The generated SQL is run against an MS SQL Server database using SQLAlchemy.
Display Results: The results of the query are returned to the user in a human-readable format.

Code Walkthrough

import openai
from sqlalchemy import create_engine, text
import re

# Set OpenAI API Key
openai.api_key = "your_openai_api_key"

# Database connection string
DATABASE_URI = "your_db_connection string"

engine = create_engine(DATABASE_URI)

def interpret_prompt_with_ai(prompt):
    try:
        response = openai.chat.completions.create(
            model="gpt-4o",
            messages=[{
                "role": "user",
                "content": f"Convert this user query into an SQL statement for MS SQL Server: {prompt}"
            }],
            temperature=0.5,
            max_tokens=150,
        )
        response_text = response.choices[0].message.content
        sql_query = re.search(r"```

sql\n(.*?)\n

```", response_text, re.DOTALL)
        if sql_query:
            return sql_query.group(1).strip()
        else:
            return response_text
    except Exception as e:
        print("Error interpreting prompt:", e)
        return None

def execute_query(sql_query):
    try:
        with engine.connect() as connection:
            result = connection.execute(text(sql_query))
            return result.fetchall()
    except Exception as e:
        print("Error executing query:", e)
        return None

def main():
    print("Welcome! Please specify your query in plain English:")
    user_prompt = input("> ")
    sql_query = interpret_prompt_with_ai(user_prompt)

    if not sql_query:
        print("Sorry, I couldn't generate a query from your input.")
        return

    print("\nGenerated SQL Query:")
    print(sql_query)

    results = execute_query(sql_query)
    if results is None:
        print("Sorry, the query could not be executed.")
        return

    print("\nQuery Results:")
    for row in results:
        print(row)

if __name__ == "__main__":
    main()

Key Features

1. Natural Language Input: Users can input plain English queries.
2. Dynamic SQL Generation: AI generates accurate SQL queries tailored to user requests.
3. Database Execution: SQLAlchemy ensures safe and efficient query execution.
4. Error Handling: The application gracefully handles errors in AI interpretation or SQL execution.

With Python, OpenAI, and SQLAlchemy, you can create a tool that’s both powerful and user-friendly.
Try it out, and let me know how you plan to extend its functionality! 🚀