While I’m sure Snowflake and Databricks are frantically adding this functionality into their products as we speak, in the meantime, there aren’t many public examples of how to actually do this with existing LLM infrastructure.
In this tutorial, we will show an end-to-end example of natural language to database query using OpenAI’s new function calls and a demo Snowflake instance. By the end, you should be able to ask a question like “How many orders did Acme, Inc. place last month?”, and receive an answer.
1. Understand how OpenAI’s function calls work
OpenAI has a pretty good writeup on how function calls work so I won’t go into too much detail here, but the gist is that they allow you to define functions with structured parameters.
These functions are then passed in alongside your query as context. If the function is appropriate for the query, the LLM will automatically parse out the structured parameters from the query. You can then run the function with those parameters, and pass the response back to the LLM.
This is relevant for our use case of converting natural language to SQL because 1) We need to hook up the LLM to Snowflake to actually execute the query. 2) More importantly, we need the “pure” SQL query as an input parameter. If we just use the regular OpenAI API, oftentimes, a response is preceded with a bunch of expository or explanatory text, something like, “Here is your SQL query: <sql-query>.”
2. Set up our data source (Snowflake)
For the purposes of this tutorial, I set up a demo Snowflake instance and uploaded the following sample data, which mimics customer orders for an e-commerce company.
3. Gather metadata from your data warehouse
Vanilla ChatGPT already knows the basics of SQL syntax and how to convert natural language to SQL, but as it admits in the conversation below, if you don’t tell it anything about the table structure, it’s forced to make a bunch of assumptions about the database schema, and probably won’t be very accurate.
So before we start, we need to gather metadata from our warehouse; I just hit describe table <table-name> for each of the tables, which gave me some basic information about column names and data types. This information will be included as context in the LLM query. I copied and pasted it over since it was only a small amount of demo data, but there are probably more automated ways of exporting this information out of a data warehouse.
4. Install openai and snowflake-python-connector packages
We create a global string variable to store the table_metadata collected above.
import openai
import json
import snowflake.connector
table_metadata = """
Table name: agents
name type kind null? default primary key unique key check expression comment policy name
AGENT_CODE VARCHAR(6) COLUMN N Y N
AGENT_NAME VARCHAR(40) COLUMN Y N N
WORKING_AREA VARCHAR(35) COLUMN Y N N
COMMISSION NUMBER(10,2) COLUMN Y N N
PHONE_NO VARCHAR(15) COLUMN Y N N
COUNTRY VARCHAR(25) COLUMN Y N N
Table name: customer
name type kind null? default primary key unique key check expression comment policy name
CUST_CODE VARCHAR(6) COLUMN N Y N
CUST_NAME VARCHAR(40) COLUMN N N N
CUST_CITY VARCHAR(35) COLUMN Y N N
WORKING_AREA VARCHAR(35) COLUMN N N N
CUST_COUNTRY VARCHAR(20) COLUMN N N N
GRADE NUMBER(38,0) COLUMN Y N N
OPENING_AMT NUMBER(12,2) COLUMN N N N
RECEIVE_AMT NUMBER(12,2) COLUMN N N N
PAYMENT_AMT NUMBER(12,2) COLUMN N N N
OUTSTANDING_AMT NUMBER(12,2) COLUMN N N N
PHONE_NO VARCHAR(17) COLUMN N N N
AGENT_CODE VARCHAR(6) COLUMN N N N
Table name: orders
name type kind null? default primary key unique key check expression comment policy name
ORD_NUM NUMBER(6,0) COLUMN N Y N
ORD_AMOUNT NUMBER(12,2) COLUMN N N N
ADVANCE_AMOUNT NUMBER(12,2) COLUMN N N N
ORD_DATE DATE COLUMN N N N
CUST_CODE VARCHAR(6) COLUMN N N N
AGENT_CODE VARCHAR(6) COLUMN N N N
ORD_DESCRIPTION VARCHAR(60) COLUMN N N N
"""
6. Write our function
Now, diving into the code. The first thing we want to define is a function called sql_query that takes in a string parameter target_query (the SQL query). This function will open a connection to Snowflake, and execute the target_query.
def sql_query(target_query):
"""Execute the target_query in Snowflake"""
# Create a connection object
conn = snowflake.connector.connect(
user='<user>',
password='<password>',
account=<account>,
warehouse='COMPUTE_WH',
database='SAMPLE_DATA',
schema='TPCH_SF1',
role='ACCOUNTADMIN'
)
# Create a cursor object from the connection
cur = conn.cursor()
# Execute a query
cur.execute(target_query)
answer = []
# Fetch the result
for row in cur:
print(row)
answer.append(row)
# Close the cursor and the connection
cur.close()
conn.close()
# convert answer to string
answer = json.dumps(answer)
return answer
7. Call the LLM to determine whether the function should be called
Next, we construct our prompt, enriching it with the data warehouse metadata that we retrieved above. We also pass along the function, sql_query, that we constructed above, and receive a response.
def run_conversation(query):
# Step 1: send the conversation and available functions to GPT
prompt = f"""
Here is the context for how the tables are structured:
{table_metadata}
Now please convert the query below into working SQL and execute it:
{query}
"""
messages = [{"role": "user", "content": prompt}]
functions = [
{
"name": "sql_query",
"description": "Execute the given SQL query and return the results",
"parameters": {
"type": "object",
"properties": {
"target_query": {
"type": "string",
"description": "The SQL query to execute",
} },
"required": ["target_query"],
},
}
]
response = openai.ChatCompletion.create(
model="gpt-3.5-turbo-0613",
messages=messages,
functions=functions,
function_call="auto", # auto is default, but we'll be explicit
)
response_message = response["choices"][0]["message"]
8. Make a function call
Upon receiving the response_message back from the OpenAI call, we first check whether the sql_query function has been triggered.
When we print out the response_message, as you can see, the sql_query was correctly triggered and the target_query argument was correctly identified.
>>> run_conversation("How many orders did customers in Mumbai place in June of 2008?")
{
"role": "assistant",
"content": null,
"function_call": {
"name": "sql_query",
"arguments": "{\n \"target_query\": \"SELECT COUNT(*) as order_count FROM orders o JOIN customer c ON o.CUST_CODE = c.CUST_CODE WHERE c.CUST_CITY = 'Mumbai' AND EXTRACT(MONTH FROM o.ORD_DATE) = 6 AND EXTRACT(YEAR FROM o.ORD_DATE) = 2008\"\n}"
}
}
If it has been triggered, then we actually call sql_query, passing along the target_query parameter that was parsed out.
# Step 2: check if GPT wanted to call a function
if response_message.get("function_call"):
# Step 3: call the function
# Note: the JSON response may not always be valid; be sure to handle errors
available_functions = {
"sql_query": sql_query,
} # only one function in this example, but you can have multiple
function_name = response_message["function_call"]["name"]
fuction_to_call = available_functions[function_name]
function_args = json.loads(response_message["function_call"]["arguments"])
function_response = fuction_to_call(
target_query=function_args.get("target_query"),
)
When we run the code and print out the function_response, you see that the function has responded with (3,), which is the result of the SQL function that was run in the warehouse.
9. Feed function response back into LLM conversation
We can now pass the function_response back into the next message of the LLM.
# Step 4: send the info on the function call and function response to GPT
messages.append(response_message) # extend conversation with assistant's reply
messages.append(
{
"role": "function",
"name": function_name,
"content": function_response,
}
) # extend conversation with function response
second_response = openai.ChatCompletion.create(
model="gpt-3.5-turbo-0613",
messages=messages,
) # get a new response from GPT where it can see the function response
return second_response["choices"][0]["message"]["content"]
When we run the code and print out the second_response, you now have a natural language answer for how many orders customers based in Mumbai placed in June of 2008.
>>> run_conversation("How many orders did customers in Mumbai place in June of 2008?")
(3,)
[[3]]
<OpenAIObject at 0x143332090> JSON: {
"role": "assistant",
"content": "The customers in Mumbai placed 3 orders in June of 2008."
}
Natural language to SQL, or in fact natural language to any kind of DSL, can be incredibly powerful.
The benefits aren’t just in developer productivity. Making it easy for people to get the information they need from data warehouses reduces unnecessarily duplicated tables, leads to more accurate and consistent reporting, and cuts costs.
Hope this helps so that you can implement it in your own product!One of the most intriguing LLM use cases is querying databases in natural language. Imagine that you’re a non-technical person, say an ops manager at Citibike, and you want some quick answers to questions like, “Where is there congestion in our bike network at 3 p.m. on a Tuesday?” In the past, you would have to spend two hours finding the right tables in the data warehouse and another half an hour painstakingly constructing a complex SQL query. But in the near future, with LLMs, you’ll be able to type in that question in your data warehouse and get a correct, natural language answer back.
Yiren Lu
Software engineer and writer currently running Frindle, a technical writing agency.
