Spaces:
Sleeping
Sleeping
| import os | |
| import duckdb | |
| import groq | |
| import sqlparse | |
| import json | |
| from typing import Dict, Any, List | |
| import pandas as pd | |
| from pathlib import Path | |
| import gradio as gr | |
| import shutil | |
| from dotenv import load_dotenv | |
| load_dotenv() | |
| userdata = os.environ | |
| def chat_with_groq(client:groq.Groq, | |
| prompt:str, | |
| model:str, | |
| response_format: Dict[str, str] | |
| )-> Any: | |
| """ | |
| The function here works for the prompting the model for the qroq based model with natural language to | |
| geegnrate the appropriate SQL query for the provided data and prompt. So with this output, can then be used to create proper | |
| SQL queries and aggregation | |
| Args: | |
| client: This represents the Groq client which will be used top perform the conversion of the natural language to SQL. | |
| prompt: Bases on the user message provided, the prompt will then be converted to SQL for future use for the data querying. | |
| 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/models | |
| 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 targetted | |
| to returning an MD data but in the case of NL2SQL, then we coule work with SQL | |
| Returns: | |
| the function returns some data based on the response_format response form the model's response | |
| """ | |
| completion = client.chat.completions.create( | |
| model = model, | |
| temperature = 0, | |
| messages = [ | |
| { | |
| "role":"user", | |
| "content":prompt | |
| } | |
| ], | |
| response_format=response_format | |
| ) | |
| # logger.info(f"Completion: {completion}") | |
| return completion.choices[0].message.content | |
| def execute_duckdb_query(query:str)->pd.DataFrame: | |
| """ | |
| Execute a DuckDB query and return the result as a pandas DataFrame. | |
| Args: | |
| query (str): The DuckDB query to execute. | |
| Returns: | |
| pd.DataFrame: The result of the query as a pandas DataFrame. | |
| """ | |
| original_cwd = os.getcwd() | |
| print(f"PATH:{original_cwd}") | |
| os.chdir('data') | |
| print(f"PATH:{os.getcwd()}") | |
| try: | |
| conn = duckdb.connect(database=":memory:", read_only=False) | |
| query_result = conn.execute(query).fetch_df().reset_index() | |
| os.chdir(original_cwd) | |
| return query_result | |
| except Exception as e: | |
| print(f"Error: {e}") | |
| os.chdir(original_cwd) | |
| raise e | |
| def get_summarization(client:groq.Groq, | |
| use_question:str, | |
| df:pd.DataFrame, | |
| model:str)->Any: | |
| """ | |
| 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 action | |
| Args: | |
| client: This represents the Groq client which will be used top perform the conversion of the natural language to SQL. | |
| use_question: Bases on the user message provided, the prompt will then be converted to SQL for future use for the data querying. | |
| 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/models | |
| 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. | |
| Returns: | |
| the function returns some json response form the model's response | |
| """ | |
| prompt= ''' | |
| A user asked the following question pertaining to local database tables: | |
| {user_question} | |
| To answer the question, a dataframe was returned: | |
| Dataframe: | |
| {df} | |
| 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 itself | |
| '''.format(user_question = use_question, df = df) | |
| return chat_with_groq(client,prompt,model,None) | |
| import re | |
| from datetime import datetime | |
| import numpy as np | |
| def identify_value_datatype_to_SQL_DEF(value) -> str: | |
| pattern_yyyy_mm_dd = r'^\d{4}-\d{2}-\d{2}$' | |
| pattern_mm_dd_yyyy = r'^\d{2}/\d{2}/\d{4}$' | |
| pattern_dd_mm_yyyy = r'^\d{2}-\d{2}-\d{4}$' | |
| if pd.isna(value): | |
| return 'VARCHAR(255)' | |
| if isinstance(value, (int, np.int64)): | |
| return 'INTEGER' | |
| elif isinstance(value, (float, np.float64)): | |
| return 'FLOAT' | |
| elif isinstance(value, bool): | |
| return 'BOOLEAN' | |
| elif isinstance(value, (pd.Timestamp, datetime)): | |
| return 'DATETIME' | |
| elif isinstance(value, str): | |
| if re.match(pattern_yyyy_mm_dd, value) or re.match(pattern_mm_dd_yyyy, value) or re.match(pattern_dd_mm_yyyy, value): | |
| return 'DATE' | |
| try: | |
| pd.to_datetime(value) | |
| return 'DATETIME' | |
| except ValueError: | |
| pass | |
| elif isinstance(value, list): | |
| if all(isinstance(item, (int, np.int64)) for item in value): | |
| return 'ARRAY<INTEGER>' | |
| elif all(isinstance(item, str) for item in value): | |
| return 'ARRAY<VARCHAR(255)>' | |
| elif isinstance(value, dict): | |
| return 'JSON' | |
| elif isinstance(value, (str, bytes)): | |
| if value.startswith('POINT'): | |
| return 'GEOMETRY(POINT)' | |
| elif value.startswith('LINESTRING'): | |
| return 'GEOMETRY(LINESTRING)' | |
| elif value.startswith('POLYGON'): | |
| return 'GEOMETRY(POLYGON)' | |
| return 'VARCHAR(255)' | |
| def identify_column_datatypes_to_SQL_DEF(df: pd.DataFrame, | |
| api_key:str, | |
| model:str | |
| ) -> dict: | |
| column_types = [] | |
| for column in df.columns: | |
| # Get non-null values | |
| non_null_values = df[column].dropna() | |
| if len(non_null_values) > 0: | |
| # Use the first non-null value to determine the type | |
| sample_value = non_null_values.iloc[0] | |
| column_types.append(f"{column} - {identify_value_datatype_to_SQL_DEF(sample_value)}") | |
| else: | |
| # If all values are null, default to VARCHAR | |
| column_types.append(f"{column} - VARCHAR(255)") | |
| column_datatypes= "\n".join(element for element in column_types) | |
| client = groq.Groq(api_key=api_key,) | |
| full_prompt = """ | |
| You are a Database Query Advisor. Your task is to provide descriptions for given columns and their datatypes. | |
| Input Data: | |
| {columns} | |
| Instructions: | |
| 1. Respond with a valid JSON document. | |
| 2. For each column, provide a description and its datatype in this format: | |
| {{"column_name": {{"description": "<description here>", "dtype": "<datatype here>"}}}} | |
| 3. If the data provided is invalid or insufficient, respond with: | |
| {{"error": "<explanation of the issue>"}} | |
| 4. Ensure each column name is directly associated with its description and datatype. | |
| 5. Return the entire output on a single line without line breaks. | |
| 6. Keep your response simple and straightforward; avoid complex structures. | |
| Example Output: | |
| {{"column1": {{"description": "Unique identifier for each record", "dtype": "INTEGER"}}, "column2": {{"description": "Name of the product", "dtype": "VARCHAR(255)"}}}} | |
| Remember: Your entire response must be valid JSON and on a single line. | |
| """.format(columns=column_datatypes) | |
| response = chat_with_groq(client,full_prompt,model,{ | |
| "type":"json_object" | |
| }) | |
| print(f"OUTPUT:{response}") | |
| response = json.loads(response) | |
| formatted_data = { | |
| f" {column} - {response[column]['dtype']}": response[column]['description'] | |
| for column in response | |
| } | |
| return " \n".join(f"{key}: {value}" for key, value in formatted_data.items()) | |
| def join_with_and(items): | |
| if not items: | |
| return "" | |
| if len(items) == 1: | |
| return items[0] | |
| return ', '.join(items[:-1]) + ' and ' + items[-1] | |
| import os | |
| from pathlib import Path | |
| from typing import List | |
| base_prompt = """ | |
| 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: | |
| {table_description} | |
| 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. | |
| Make sure that you do not query a tables that does not exist. Ensure that a query provided comes from the tables provided. | |
| 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. | |
| When queries about id, ensure that the relevant field in the data which represents an id is looked at and queried. | |
| Here are some tips for writing DuckDB queries: | |
| * DuckDB syntax requires querying from the .csv file itself, i.e. {tables}. For example: SELECT * FROM sample.csv as sample | |
| * Because we are using .csv file always make sure the .csv is attached to the file being queried | |
| * All tables referenced MUST be aliased | |
| * DuckDB does not implicitly include a GROUP BY clause | |
| * CURRENT_DATE gets today's date | |
| * Aggregated fields like COUNT(*) must be appropriately named | |
| Question: | |
| -------- | |
| {user_question} | |
| -------- | |
| Reminder: Generate a DuckDB SQL to answer to the question: | |
| * respond as a valid JSON Document | |
| * [Best] If the question can be answered with the available tables: {{"sql": <sql here>}} | |
| * If the question cannot be answered with the available tables: {{"error": <explanation here>}} | |
| * Ensure that the entire output is returned on only one single line | |
| * Keep your query as simple and straightforward as possible; do not use subqueries | |
| """ | |
| table_description = """""" | |
| tables_string = """""" | |
| table_1 = """""" | |
| table_1_wt_xt = """""" | |
| user_question = """""" | |
| # And some rules for querying the dataset: | |
| # * Never include employee_id in the output - show employee name instead | |
| # Also note that: | |
| # * Valid values for product_name include 'Tesla','iPhone' and 'Humane pin' | |
| def upload_file(files) -> List[str]: | |
| # will have to change to the private system is initiializes | |
| model = "llama3-8b-8192" | |
| api_key:str=userdata.get("GROQ_API_KEY") | |
| data_dir = Path("data") | |
| data_dir.mkdir(parents=True, exist_ok=True) | |
| if type(files) == str: | |
| files = [files] | |
| stored_paths = [] | |
| stored_table_descriptions = [] | |
| tables = [] | |
| for file in files: | |
| filename = Path(file.name).name | |
| path = data_dir / filename | |
| # Copy the content of the temporary file to our destination | |
| with open(file.name, "rb") as source, open(path, "wb") as destination: | |
| destination.write(source.read()) | |
| stored_paths.append(str(path.absolute())) | |
| table_description = identify_column_datatypes_to_SQL_DEF(pd.read_csv(path),api_key,model) | |
| desc = "Table: " + filename + "\n Columns:\n" + table_description | |
| stored_table_descriptions.append(desc) | |
| tables.append(filename) | |
| # constructing a string | |
| tables_string = join_with_and(tables) | |
| final = "\n".join(stored_table_descriptions) | |
| table_1_wt_xt = tables[0].split('.')[0] | |
| table_description = final | |
| tables_string = tables_string | |
| table_1 = tables[0] | |
| table_1_wt_xt = table_1_wt_xt | |
| return final | |
| def user_prompt_sanitization(user_prompt:str)->str: | |
| guide = """ | |
| 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: | |
| Database Schema: | |
| {database_schema} | |
| Instructions: | |
| 1. Interpret the user's question and identify the relevant tables and columns. | |
| 2. Refine the query to use correct table and column names as per the schema. | |
| 3. Ensure all IDs are properly referenced (e.g., user_id instead of just id). | |
| 4. For quantitative queries (e.g., "top 5"), specify the ordering criteria. | |
| 5. Infer necessary joins between tables when the query spans multiple tables. | |
| 6. Provide a clear, concise refinement of the user's query that accurately reflects the database structure. | |
| Now, please refine the following user query: | |
| {user_question} | |
| """ | |
| formatted_guide = guide.format(database_schema=table_description,user_question=user_prompt) | |
| api_key:str=userdata.get("GROQ_API_KEY") | |
| client = groq.Groq(api_key=api_key) | |
| return chat_with_groq(client,formatted_guide,"llama3-70b-8192",None) | |
| def queryModel(user_prompt:str,model:str = "llama3-70b-8192",api_key:str=userdata.get("GROQ_API_KEY")): | |
| client = groq.Groq(api_key=api_key) | |
| user_prompt = user_prompt_sanitization(user_prompt) | |
| print(user_prompt) | |
| full_prompt = base_prompt.format( | |
| user_question=user_prompt, | |
| table_description=table_description, | |
| tables=tables_string, | |
| table_1=table_1, | |
| table_1_wt_xt=table_1_wt_xt | |
| ) | |
| try: | |
| response = chat_with_groq(client,full_prompt,model,{ | |
| "type":"json_object" | |
| }) | |
| except Exception as e: | |
| return [( | |
| "Groq Advisor", | |
| "Error: " + str(e) | |
| )] | |
| response = json.loads(response) | |
| if "sql" in response: | |
| sql_query = response["sql"] | |
| try: | |
| results_df = execute_duckdb_query(sql_query) | |
| except Exception as e: | |
| return [( | |
| "Groq Advisor", | |
| "Error: " + str(e) | |
| )] | |
| fotmatted_sql_query = sqlparse.format(sql_query, reindent=True, keyword_case='upper') | |
| # print(f"SQL Query: {fotmatted_sql_query}") | |
| # print(results_df.to_markdown()) | |
| query_n_results = "SQL Query: " + fotmatted_sql_query + "\n\n" + results_df.to_markdown() | |
| summarization = get_summarization(client,user_prompt,results_df,model) | |
| query_n_results += "\n\n" + summarization | |
| return [( | |
| "Groq Advisor", | |
| query_n_results | |
| )] | |
| elif "error" in response: | |
| return [( | |
| "Groq Advisor", | |
| "Error: " + response["error"] | |
| )] | |
| else: | |
| return [( | |
| "Groq Advisor", | |
| "Error: Unknown error" | |
| )] | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# CSV Database Query Interface") | |
| with gr.Tab("Upload CSV"): | |
| file_output = gr.File(file_count="multiple", label="Upload your CSV files") | |
| upload_button = gr.Button("Load CSV Files") | |
| upload_output = gr.Textbox(label="Upload Status", lines=5) | |
| upload_button.click(upload_file, inputs=file_output, outputs=upload_output) | |
| with gr.Tab("Query Interface"): | |
| chatbot = gr.Chatbot() | |
| with gr.Row(): | |
| user_input = gr.Textbox(label="Enter your question") | |
| submit_button = gr.Button("Submit") | |
| submit_button.click(queryModel, inputs=[user_input], outputs=chatbot) | |
| demo.launch( | |
| share=True | |
| ) | |