NLP to database query

gistfile1.txt

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	
"""


# Example dummy function hard coded to return the same weather
# In production, this could be your backend API or an external API
def sql_query(target_query):
    """Execute the target_query in Snowflake"""

    # Create a connection object
    conn = snowflake.connector.connect(
        user='<user>',
        password='<password>',
        account='ddhefdg-ep68604',
        warehouse='COMPUTE_WH',
        database='REN_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


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"]

    # 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"),
        )

        # 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"]