blesspearl/NL2SQL
1
1import os2import duckdb3import groq4import sqlparse5import json6from typing import Dict, Any, List7import pandas as pd8from pathlib import Path9import gradio as gr10import shutil11from dotenv import load_dotenv12load_dotenv()13userdata = os.environ14 15 16def chat_with_groq(client:groq.Groq,17 prompt:str,18 model:str,19 response_format: Dict[str, str]20 )-> Any:21 """22 The function here works for the prompting the model for the qroq based model with natural language to23 geegnrate the appropriate SQL query for the provided data and prompt. So with this output, can then be used to create proper24 SQL queries and aggregation25 26 Args:27 client: This represents the Groq client which will be used top perform the conversion of the natural language to SQL.28 prompt: Bases on the user message provided, the prompt will then be converted to SQL for future use for the data querying.29 model: The Model here represents the based open source model that is used to make the query on. examples include llama3-70b-8192 or llama3-8b-8192 find out more here https://console.groq.com/docs/models30 response_format: As model has the possibility of supporting JSON and MD outputs, this is a dictionary which specifies the the required return response format. by default this for the model is usually targetted31 to returning an MD data but in the case of NL2SQL, then we coule work with SQL32 33 Returns:34 the function returns some data based on the response_format response form the model's response35 """36 completion = client.chat.completions.create(37 model = model,38 temperature = 0,39 messages = [40 {41 "role":"user",42 "content":prompt43 }44 ],45 response_format=response_format46 )47 # logger.info(f"Completion: {completion}")48 return completion.choices[0].message.content49 50def execute_duckdb_query(query:str)->pd.DataFrame:51 """52 Execute a DuckDB query and return the result as a pandas DataFrame.53 54 Args:55 query (str): The DuckDB query to execute.56 57 Returns:58 pd.DataFrame: The result of the query as a pandas DataFrame.59 """60 original_cwd = os.getcwd()61 print(f"PATH:{original_cwd}")62 os.chdir('data')63 print(f"PATH:{os.getcwd()}")64 65 try:66 conn = duckdb.connect(database=":memory:", read_only=False)67 query_result = conn.execute(query).fetch_df().reset_index()68 os.chdir(original_cwd)69 return query_result70 except Exception as e:71 print(f"Error: {e}")72 os.chdir(original_cwd)73 raise e74def get_summarization(client:groq.Groq,75 use_question:str,76 df:pd.DataFrame,77 model:str)->Any:78 """79 For this query, the user input is better summarized around the provided Dataframe. This give a better contextual structure for the user to perfom the action80 81 Args:82 client: This represents the Groq client which will be used top perform the conversion of the natural language to SQL.83 use_question: Bases on the user message provided, the prompt will then be converted to SQL for future use for the data querying.84 model: The Model here represents the based open source model that is used to make the query on. examples include llama3-70b-8192 or llama3-8b-8192 find out more here https://console.groq.com/docs/models85 df: this is a pandas dataframe which containe the database daat which will then be used to summarize the base query around it, making the prompt more realistic.86 87 Returns:88 the function returns some json response form the model's response89 """90 prompt= '''91 A user asked the following question pertaining to local database tables:92 93 {user_question}94 95 To answer the question, a dataframe was returned:96 97 Dataframe:98 {df}99 100 In a few sentences, summarize the data in the table as it pertains to the original user question. Avoid qualifiers like "based on the data" and do not comment on the structure or metadata of the table itself101 '''.format(user_question = use_question, df = df)102 return chat_with_groq(client,prompt,model,None)103 104 105import re106from datetime import datetime107import numpy as np108 109def identify_value_datatype_to_SQL_DEF(value) -> str:110 111 pattern_yyyy_mm_dd = r'^\d{4}-\d{2}-\d{2}$'112 pattern_mm_dd_yyyy = r'^\d{2}/\d{2}/\d{4}$'113 pattern_dd_mm_yyyy = r'^\d{2}-\d{2}-\d{4}$'114 115 if pd.isna(value):116 return 'VARCHAR(255)'117 118 if isinstance(value, (int, np.int64)):119 return 'INTEGER'120 121 122 elif isinstance(value, (float, np.float64)):123 return 'FLOAT'124 125 126 elif isinstance(value, bool):127 return 'BOOLEAN'128 129 130 elif isinstance(value, (pd.Timestamp, datetime)):131 return 'DATETIME'132 133 134 elif isinstance(value, str):135 if re.match(pattern_yyyy_mm_dd, value) or re.match(pattern_mm_dd_yyyy, value) or re.match(pattern_dd_mm_yyyy, value):136 return 'DATE'137 138 139 try:140 pd.to_datetime(value)141 return 'DATETIME'142 except ValueError:143 pass144 145 146 elif isinstance(value, list):147 if all(isinstance(item, (int, np.int64)) for item in value):148 return 'ARRAY<INTEGER>'149 elif all(isinstance(item, str) for item in value):150 return 'ARRAY<VARCHAR(255)>'151 152 153 elif isinstance(value, dict):154 return 'JSON'155 156 elif isinstance(value, (str, bytes)):157 if value.startswith('POINT'):158 return 'GEOMETRY(POINT)'159 elif value.startswith('LINESTRING'):160 return 'GEOMETRY(LINESTRING)'161 elif value.startswith('POLYGON'):162 return 'GEOMETRY(POLYGON)'163 164 return 'VARCHAR(255)'165 166def identify_column_datatypes_to_SQL_DEF(df: pd.DataFrame,167 api_key:str,168 model:str169 ) -> dict:170 column_types = []171 172 for column in df.columns:173 # Get non-null values174 non_null_values = df[column].dropna()175 176 if len(non_null_values) > 0:177 # Use the first non-null value to determine the type178 sample_value = non_null_values.iloc[0]179 column_types.append(f"{column} - {identify_value_datatype_to_SQL_DEF(sample_value)}")180 else:181 # If all values are null, default to VARCHAR182 column_types.append(f"{column} - VARCHAR(255)")183 184 column_datatypes= "\n".join(element for element in column_types)185 client = groq.Groq(api_key=api_key,)186 full_prompt = """187 You are a Database Query Advisor. Your task is to provide descriptions for given columns and their datatypes.188 189 Input Data:190 {columns}191 192 Instructions:193 1. Respond with a valid JSON document.194 2. For each column, provide a description and its datatype in this format:195 {{"column_name": {{"description": "<description here>", "dtype": "<datatype here>"}}}}196 3. If the data provided is invalid or insufficient, respond with:197 {{"error": "<explanation of the issue>"}}198 4. Ensure each column name is directly associated with its description and datatype.199 5. Return the entire output on a single line without line breaks.200 6. Keep your response simple and straightforward; avoid complex structures.201 202 Example Output:203 {{"column1": {{"description": "Unique identifier for each record", "dtype": "INTEGER"}}, "column2": {{"description": "Name of the product", "dtype": "VARCHAR(255)"}}}}204 205 Remember: Your entire response must be valid JSON and on a single line.206 """.format(columns=column_datatypes)207 208 response = chat_with_groq(client,full_prompt,model,{209 "type":"json_object"210 })211 print(f"OUTPUT:{response}")212 response = json.loads(response)213 214 formatted_data = {215 f" {column} - {response[column]['dtype']}": response[column]['description']216 for column in response217 }218 return " \n".join(f"{key}: {value}" for key, value in formatted_data.items())219 220def join_with_and(items):221 if not items:222 return ""223 if len(items) == 1:224 return items[0]225 return ', '.join(items[:-1]) + ' and ' + items[-1]226 227 228import os229from pathlib import Path230from typing import List231 232base_prompt = """233 You are Groq Advisor, and you are tasked with generating SQL queries for DuckDB based on user questions about data stored in two tables derived from CSV files:234 235 {table_description}236 237 Given a user's question about this data, write a valid DuckDB SQL query that accurately extracts or calculates the requested information from these tables and adheres to SQL best practices for DuckDB, optimizing for readability and performance where applicable.238 Make sure that you do not query a tables that does not exist. Ensure that a query provided comes from the tables provided.239 When asked about quantitative values, ensure that you look at the numerical data types and not the string data types such as when asked about top purchase, the amount can be a field to be looked at.240 When queries about id, ensure that the relevant field in the data which represents an id is looked at and queried.241 242 Here are some tips for writing DuckDB queries:243 * DuckDB syntax requires querying from the .csv file itself, i.e. {tables}. For example: SELECT * FROM sample.csv as sample244 * Because we are using .csv file always make sure the .csv is attached to the file being queried245 * All tables referenced MUST be aliased246 * DuckDB does not implicitly include a GROUP BY clause247 * CURRENT_DATE gets today's date248 * Aggregated fields like COUNT(*) must be appropriately named249 250 251 Question:252 --------253 {user_question}254 --------255 Reminder: Generate a DuckDB SQL to answer to the question:256 * respond as a valid JSON Document257 * [Best] If the question can be answered with the available tables: {{"sql": <sql here>}}258 * If the question cannot be answered with the available tables: {{"error": <explanation here>}}259 * Ensure that the entire output is returned on only one single line260 * Keep your query as simple and straightforward as possible; do not use subqueries261 """262table_description = """"""263tables_string = """"""264table_1 = """"""265table_1_wt_xt = """"""266user_question = """"""267 268# And some rules for querying the dataset:269# * Never include employee_id in the output - show employee name instead270 271# Also note that:272# * Valid values for product_name include 'Tesla','iPhone' and 'Humane pin'273 274 275def upload_file(files) -> List[str]:276 # will have to change to the private system is initiializes277 model = "llama3-8b-8192"278 api_key:str=userdata.get("GROQ_API_KEY")279 data_dir = Path("data")280 data_dir.mkdir(parents=True, exist_ok=True)281 if type(files) == str:282 files = [files]283 stored_paths = []284 stored_table_descriptions = []285 tables = []286 for file in files:287 filename = Path(file.name).name288 path = data_dir / filename289 290 # Copy the content of the temporary file to our destination291 with open(file.name, "rb") as source, open(path, "wb") as destination:292 destination.write(source.read())293 294 stored_paths.append(str(path.absolute()))295 table_description = identify_column_datatypes_to_SQL_DEF(pd.read_csv(path),api_key,model)296 desc = "Table: " + filename + "\n Columns:\n" + table_description297 stored_table_descriptions.append(desc)298 tables.append(filename)299 # constructing a string300 tables_string = join_with_and(tables)301 final = "\n".join(stored_table_descriptions)302 table_1_wt_xt = tables[0].split('.')[0]303 table_description = final304 tables_string = tables_string305 table_1 = tables[0]306 table_1_wt_xt = table_1_wt_xt307 return final308 309def user_prompt_sanitization(user_prompt:str)->str:310 guide = """311 You are an AI assistant specializing in database queries. Your task is to interpret user questions about a database and refine them based on the available table structures. Use the following table descriptions to guide your responses:312 313 Database Schema:314 {database_schema}315 316 Instructions:317 1. Interpret the user's question and identify the relevant tables and columns.318 2. Refine the query to use correct table and column names as per the schema.319 3. Ensure all IDs are properly referenced (e.g., user_id instead of just id).320 4. For quantitative queries (e.g., "top 5"), specify the ordering criteria.321 5. Infer necessary joins between tables when the query spans multiple tables.322 6. Provide a clear, concise refinement of the user's query that accurately reflects the database structure.323 324 Now, please refine the following user query:325 {user_question}326 327 """328 329 formatted_guide = guide.format(database_schema=table_description,user_question=user_prompt)330 api_key:str=userdata.get("GROQ_API_KEY")331 client = groq.Groq(api_key=api_key)332 return chat_with_groq(client,formatted_guide,"llama3-70b-8192",None)333 334def queryModel(user_prompt:str,model:str = "llama3-70b-8192",api_key:str=userdata.get("GROQ_API_KEY")):335 client = groq.Groq(api_key=api_key)336 user_prompt = user_prompt_sanitization(user_prompt)337 print(user_prompt)338 full_prompt = base_prompt.format(339 user_question=user_prompt,340 table_description=table_description,341 tables=tables_string,342 table_1=table_1,343 table_1_wt_xt=table_1_wt_xt344 )345 try:346 response = chat_with_groq(client,full_prompt,model,{347 "type":"json_object"348 })349 except Exception as e:350 return [(351 "Groq Advisor",352 "Error: " + str(e)353 )]354 response = json.loads(response)355 if "sql" in response:356 sql_query = response["sql"]357 try:358 results_df = execute_duckdb_query(sql_query)359 except Exception as e:360 return [(361 "Groq Advisor",362 "Error: " + str(e)363 )]364 365 fotmatted_sql_query = sqlparse.format(sql_query, reindent=True, keyword_case='upper')366 # print(f"SQL Query: {fotmatted_sql_query}")367 # print(results_df.to_markdown())368 query_n_results = "SQL Query: " + fotmatted_sql_query + "\n\n" + results_df.to_markdown()369 summarization = get_summarization(client,user_prompt,results_df,model)370 query_n_results += "\n\n" + summarization371 372 return [(373 "Groq Advisor",374 query_n_results375 )]376 elif "error" in response:377 return [(378 "Groq Advisor",379 "Error: " + response["error"]380 )]381 else:382 return [(383 "Groq Advisor",384 "Error: Unknown error"385 )]386 387with gr.Blocks() as demo:388 gr.Markdown("# CSV Database Query Interface")389 390 with gr.Tab("Upload CSV"):391 file_output = gr.File(file_count="multiple", label="Upload your CSV files")392 upload_button = gr.Button("Load CSV Files")393 upload_output = gr.Textbox(label="Upload Status", lines=5)394 395 upload_button.click(upload_file, inputs=file_output, outputs=upload_output)396 with gr.Tab("Query Interface"):397 chatbot = gr.Chatbot()398 with gr.Row():399 user_input = gr.Textbox(label="Enter your question")400 submit_button = gr.Button("Submit")401 submit_button.click(queryModel, inputs=[user_input], outputs=chatbot)402 403 404 405demo.launch(406 share=True407)408