Initial Commit

README.md

@@ -1,10 +1,10 @@
 ---
 title: DataViz Agent
-emoji: 💬
+emoji: 📊
 colorFrom: yellow
 colorTo: purple
 sdk: gradio
-sdk_version: 4.36.1
+sdk_version: 4.42.0
 app_file: app.py
 pinned: false
 license: mit

app.py

@@ -1,63 +1,323 @@
+import os
+import json
+import torch
+import duckdb
+import spaces
+import pandas as pd
 import gradio as gr
-from huggingface_hub import InferenceClient
+import matplotlib.pyplot as plt
+from visualization_prompt import graph_instructions
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_huggingface.llms import HuggingFacePipeline
+from plot_utils import plot_bar_chart, plot_horizontal_bar_chart, plot_line_graph, plot_scatter, plot_pie
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, pipeline
+from dotenv import load_dotenv
 
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+load_dotenv()
 
 
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
+# Height of the Tabs Text Area
+TAB_LINES = 8
+# Load Token
+md_token = os.getenv('MD_TOKEN')
+os.environ['HF_TOKEN'] = os.getenv('HF_TOKEN')
 
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
+print('Connecting to DB...')
+# Connect to DB
+conn = duckdb.connect(f"md:my_db?motherduck_token={md_token}", read_only=True)
 
-    messages.append({"role": "user", "content": message})
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+    print(f"Using GPU: {torch.cuda.get_device_name(device)}")
+else:
+    device = torch.device("cpu")
+    print("Using CPU")
 
-    response = ""
+print('Loading Model...')
 
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
+tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")
 
-        response += token
-        yield response
+quantization_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_compute_dtype=torch.bfloat16,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type= "nf4")
 
+model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3", quantization_config=quantization_config,
+                                             device_map="auto", torch_dtype=torch.bfloat16)
+pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, return_full_text=False, max_new_tokens=512)
+llm = HuggingFacePipeline(pipeline=pipe)
+
+
+print('Model Loaded...')
+print(f'Model Device: {model.device}')
+
+
+def generate_output(schema, prompt, generated_sql_query, result_output):
+    return {
+        table_schema: schema,
+        input_prompt: prompt,
+        generated_query: generated_sql_query,
+        result_output: result_output,
+    }
+
+# Get Databases
+def get_schemas():
+    schemas = conn.execute("""
+    SELECT DISTINCT schema_name
+    FROM information_schema.schemata
+    WHERE schema_name NOT IN ('information_schema', 'pg_catalog')
+    """).fetchall()
+    return [item[0] for item in schemas]
+
+# Get Tables
+def get_tables(schema_name):
+    tables = conn.execute(f"SELECT table_name FROM information_schema.tables WHERE table_schema = '{schema_name}'").fetchall()
+    return [table[0] for table in tables]
+
+# Update Tables
+def update_tables(schema_name):
+    tables = get_tables(schema_name)
+    return gr.update(choices=tables)
+
+# Get Schema
+def get_table_schema(table):
+    result = conn.sql(f"SELECT sql, database_name, schema_name FROM duckdb_tables() where table_name ='{table}';").df()
+    ddl_create = result.iloc[0,0]
+    parent_database = result.iloc[0,1]
+    schema_name = result.iloc[0,2]
+    full_path = f"{parent_database}.{schema_name}.{table}"
+    if schema_name != "main":
+        old_path = f"{schema_name}.{table}"
+    else:
+        old_path = table
+    ddl_create = ddl_create.replace(old_path, full_path)
+    return ddl_create
+
+# Get Prompt
+def get_prompt(schema, query_input):
+    text = f"""
+    ### Instruction:
+    Your task is to generate valid duckdb SQL query to answer the following question. Only Respond with SQL query do not inlcude anything else.
+    ### Input:
+    Here is the database schema that the SQL query will run on:
+    {schema}
+
+    ### Question:
+    {query_input}
+    ### Response (use duckdb shorthand if possible):
+    """
+    return text
+
+@spaces.GPU(duration=60)
+def run_llm(prompt):
+    result = llm.invoke(prompt)
+    return result
+
+
+def get_visualization_type(text_query, sql_query, sql_result):
+    system_message = '''
+You are an AI assistant that recommends appropriate data visualizations. Based on the user's question, SQL query, and query results, suggest the most suitable type of graph or chart to visualize the data. If no visualization is appropriate, indicate that.
+
+Available chart types and their use cases:
+- Bar Graphs: Best for comparing categorical data or showing changes over time when categories are discrete and the number of categories is more than 2. Use for questions like "What are the sales figures for each product?" or "How does the population of cities compare? or "What percentage of each city is male?"
+- Horizontal Bar Graphs: Best for comparing categorical data or showing changes over time when the number of categories is small or the disparity between categories is large. Use for questions like "Show the revenue of A and B?" or "How does the population of 2 cities compare?" or "How many men and women got promoted?" or "What percentage of men and what percentage of women got promoted?" when the disparity between categories is large.
+- Scatter Plots: Useful for identifying relationships or correlations between two numerical variables or plotting distributions of data. Best used when both x axis and y axis are continuous. Use for questions like "Plot a distribution of the fares (where the x axis is the fare and the y axis is the count of people who paid that fare)" or "Is there a relationship between advertising spend and sales?" or "How do height and weight correlate in the dataset? Do not use it for questions that do not have a continuous x axis."
+- Pie Charts: Ideal for showing proportions or percentages within a whole. Use for questions like "What is the market share distribution among different companies?" or "What percentage of the total revenue comes from each product?"
+- Line Graphs: Best for showing trends and distributionsover time. Best used when both x axis and y axis are continuous. Used for questions like "How have website visits changed over the year?" or "What is the trend in temperature over the past decade?". Do not use it for questions that do not have a continuous x axis or a time based x axis.
+
+Consider these types of questions when recommending a visualization:
+1. Aggregations and Summarizations (e.g., "What is the average revenue by month?" - Line Graph)
+2. Comparisons (e.g., "Compare the sales figures of Product A and Product B over the last year." - Line or Column Graph)
+3. Plotting Distributions (e.g., "Plot a distribution of the age of users" - Scatter Plot)
+4. Trends Over Time (e.g., "What is the trend in the number of active users over the past year?" - Line Graph)
+5. Proportions (e.g., "What is the market share of the products?" - Pie Chart)
+6. Correlations (e.g., "Is there a correlation between marketing spend and revenue?" - Scatter Plot)
+
+Provide your response in the following format:
+Recommended Visualization: [Chart type or "None"]. ONLY use the following names: bar, horizontal_bar, line, pie, scatter, none
+Reason: [Brief explanation for your recommendation]
+'''
+    human_message = '''
+User question: {question}
+SQL query: {sql_query}
+Query results: {results}
+
+Recommend a visualization:
+'''
+
+    prompt = ChatPromptTemplate.from_messages([
+                    ("system", system_message),
+                    ("human", human_message),
+                ])
+
+    final_prompt = prompt.format_prompt(question=text_query,
+                                        sql_query=sql_query, results=sql_result)
+    print(final_prompt)
+    response = run_llm(final_prompt)
+
+    print(f"Visualization Type: {response}")
+
+    lines = response.strip().split('\n')
+    visualization = lines[0].split(': ')[1]
+    reason = lines[1].split(': ')[1]
+
+
+    print(visualization, reason)
+    return visualization, reason
+
+def format_data(text_query, sql_query, sql_result, visualization_type):
+    instruction = graph_instructions[visualization_type]
+
+    template = ChatPromptTemplate.from_messages([
+    ("system", "You are a Data expert who formats data according to the required needs. You are given the question asked by the user, it's sql query, the result of the query and the format you need to format it in."),
+    ("human", "For the given question: {question}\n\nSQL query: {sql_query}\n\Result: {results}\n\nUse the following example to structure the data: {instructions}. If there is None in Result please change it to '0'. Just give the json string. Do not format it or add any label or text."),
+    ])
+
+
+    prompt = template.format_prompt(question=text_query, sql_query=sql_query,
+                                    results=sql_result, instructions=instruction)
+    print(f'New Template for Data fromat {prompt}')
+
+    formatted_data = run_llm(prompt)
+    print(f'Formatted Data {formatted_data}')
+
+    return json.loads(formatted_data.replace('.', '').strip())
+
+def visualize_result(text_query, visualization_type, reason, sql_query,
+                      sql_result):
+
+    if visualization_type == 'bar':
+        data = format_data(text_query=text_query, sql_query=sql_query,
+                           sql_result=sql_result, visualization_type=visualization_type)
+        print(data)
+        return plot_bar_chart(data)
+
+    elif visualization_type == 'horizontal_bar':
+        data = format_data(text_query=text_query, sql_query=sql_query,
+                           sql_result=sql_result, visualization_type=visualization_type)
+        return plot_horizontal_bar_chart(data)
+
+    elif visualization_type == 'line':
+        data = format_data(text_query=text_query, sql_query=sql_query,
+                            sql_result=sql_result, visualization_type=visualization_type)
+        return plot_line_graph(data)
+
+    elif visualization_type == 'pie':
+        data = format_data(text_query=text_query, sql_query=sql_query,
+                           sql_result=sql_result, visualization_type=visualization_type)
+        return plot_pie(data)
+
+    elif visualization_type == 'scatter':
+        data = format_data(text_query=text_query, sql_query=sql_query,
+                           sql_result=sql_result, visualization_type=visualization_type)
+        return plot_scatter(data)
+
+    elif visualization_type == 'none':
+        fig, ax = plt.subplots()
+        ax.set_visible(False) 
+        return fig
+
+
+def main(table, text_query):
+    if table is None:
+        return ["", "", "", pd.DataFrame([{"error": "❌ Table is None."}])]
+    fig, ax = plt.subplots()
+    ax.set_visible(False)
+
+    schema = get_table_schema(table)
+    prompt = get_prompt(schema, text_query)
+    try:
+        generated_sql_query = run_llm(prompt)
+        print(f'Generated SQL Query: {generated_sql_query}')
+    except Exception as e:
+        return generate_output(schema, prompt, '', fig)
+
+    try:
+        sql_query_result = conn.sql(generated_sql_query).fetchall()
+        print(f"SQL Query Result: {sql_query_result}")
+    except Exception as e:
+        return generate_output(schema, prompt, generated_sql_query, fig)
+
+    try:
+      print('Generating Visualization Type')
+      visualization_type, reason = get_visualization_type(text_query=text_query,
+                                                  sql_query=generated_sql_query, sql_result=sql_query_result)
+
+      print(f"Visualization Type: {visualization_type}")
+
+    except Exception as e:
+        return generate_output(schema, prompt, generated_sql_query,
+                               fig)
+    try:
+
+        plot = visualize_result(text_query=text_query, sql_query=generated_sql_query,
+                      sql_result=sql_query_result, visualization_type=visualization_type, reason=reason)
+        print(f"Plot: {plot}")
+    except Exception as e:
+        return generate_output(schema, prompt, generated_sql_query, fig)
+
+    return generate_output(schema, prompt, generated_sql_query, plot)
+
+def generate_output(schema, prompt, generated_sql_query, result_output):
+    # Returning values as a list in the order expected by Gradio
+    return [schema, prompt, generated_sql_query, result_output]
+
+custom_css = """
+.gradio-container {
+    background-color: #f0f4f8;
+}
+.logo {
+    max-width: 200px;
+    margin: 20px auto;
+    display: block;
+}
+.gr-button {
+    background-color: #4a90e2 !important;
+}
+.gr-button:hover {
+    background-color: #3a7bc8 !important;
+}
 """
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
-)
 
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="purple", secondary_hue="indigo"), css=custom_css) as demo:
+    gr.Image("logo.png", label=None, show_label=False, container=False, height=100)
+
+    gr.Markdown("""
+    <div style='text-align: center;'>
+    <strong style='font-size: 36px;'>Datajoi SQL Agent</strong>
+    <br>
+    <span style='font-size: 20px;'>Generate and Run SQL queries based on a given text for the dataset.</span>
+    </div>
+    """)
+
+    with gr.Row():
+
+        with gr.Column(scale=1, variant='panel'):
+            schema_dropdown = gr.Dropdown(choices=get_schemas(), label="Select Schema", interactive=True)
+            tables_dropdown = gr.Dropdown(choices=[], label="Available Tables", value=None)
+
+        with gr.Column(scale=2):
+            query_input = gr.Textbox(lines=5, label="Text Query", placeholder="Enter your text query here...")
+            with gr.Row():
+                with gr.Column(scale=7):
+                    pass
+                with gr.Column(scale=1):
+                    generate_query_button = gr.Button("Run Query", variant="primary")
+
+            with gr.Tabs():
+                with gr.Tab("Plot"):
+                    result_output = gr.Plot()
+                    # result_output = gr.Textbox(lines=TAB_LINES, label="Generated SQL Query", value="", interactive=False)
+                with gr.Tab("SQL Query"):
+                    generated_query = gr.Textbox(lines=TAB_LINES, label="Generated SQL Query", value="", interactive=False)
+                with gr.Tab("Prompt"):
+                    input_prompt = gr.Textbox(lines=TAB_LINES, label="Input Prompt", value="", interactive=False)
+                with gr.Tab("Schema"):
+                    table_schema = gr.Textbox(lines=TAB_LINES, label="Table Schema", value="", interactive=False)
+
+        schema_dropdown.change(update_tables, inputs=schema_dropdown, outputs=tables_dropdown)
+        generate_query_button.click(main, inputs=[tables_dropdown, query_input], outputs=[table_schema, input_prompt, generated_query, result_output])
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(debug=True)

plot_utils.py

@@ -0,0 +1,78 @@
+import matplotlib.pyplot as plt
+
+def plot_bar_chart(data):
+
+    fig, ax = plt.subplots()
+    labels = data['labels']
+    values = data['values']
+    
+    for value in values:
+        ax.bar(labels, value['data'], label=value['label'])
+    
+    ax.set_title('Bar Chart')
+    ax.set_xlabel('Labels')
+    ax.set_ylabel('Values')
+    ax.legend()
+    return fig
+    
+def plot_horizontal_bar_chart(data):
+    fig, ax = plt.subplots()
+    labels = data['labels']
+    values = data['values']
+    
+    for value in values:
+        ax.barh(labels, value['data'], label=value['label'])
+    
+    ax.set_title('Horizontal Bar Chart')
+    ax.set_xlabel('Values')
+    ax.set_ylabel('Labels')
+    ax.legend()
+    return fig
+    
+def plot_line_graph(data):
+    fig, ax = plt.subplots()
+    x_values = data['xValues']
+    y_values_list = data['yValues']
+    
+    for y_values in y_values_list:
+        label = y_values.get('label', None)
+        ax.plot(x_values, y_values['data'], label=label)
+    
+    ax.set_title('Line Graph')
+    ax.set_xlabel('X Values')
+    ax.set_ylabel('Y Values')
+    
+    if any('label' in y_values for y_values in y_values_list):
+        ax.legend()
+    
+    return fig
+
+def plot_scatter(data):
+    fig, ax = plt.subplots()
+    for series in data['series']:
+        x_values = [point['x'] for point in series['data']]
+        y_values = [point['y'] for point in series['data']]
+        label = series.get('label', None)
+        ax.scatter(x_values, y_values, label=label)
+
+    ax.set_title('Scatter Plot')
+    ax.set_xlabel('X Values')
+    ax.set_ylabel('Y Values')
+
+    if any('label' in series for series in data['series']):
+        ax.legend()
+
+    return fig
+    
+def plot_pie(data):
+    values = [item['value'] for item in data]
+    labels = [item['label'] for item in data]
+
+    fig, ax = plt.subplots()
+    wedges, texts, autotexts = ax.pie(values, labels=labels, autopct='%1.1f%%', startangle=140)
+    ax.set_title('Pie Chart')
+
+    for text in texts + autotexts:
+        text.set_fontsize(10)
+
+    return fig

visualization_prompt.py

@@ -0,0 +1,151 @@
+barGraphIntstruction = '''
+
+  Where data is: {
+    labels: string[]
+    values: {\data: number[], label: string}[]
+  }
+
+// Examples of usage:
+Each label represents a column on the x axis.
+Each array in values represents a different entity. 
+
+Here we are looking at average income for each month.
+{
+  labels: ['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun'],
+  values: [{data:[21.5, 25.0, 47.5, 64.8, 105.5, 133.2], label: 'Income'}],
+}
+
+Here we are looking at the performance of american and european players for each series. Since there are two entities, we have two arrays in values.
+{
+  labels: ['series A', 'series B', 'series C'],
+  values: [{data:[10, 15, 20], label: 'American'}, {data:[20, 25, 30], label: 'European'}],
+}
+'''
+
+horizontalBarGraphIntstruction = '''
+
+  Where data is: {
+    labels: string[]
+    values: {\data: number[], label: string}[]
+  }
+
+// Examples of usage:
+Each label represents a column on the x axis.
+Each array in values represents a different entity. 
+
+Here we are looking at average income for each month.
+{
+  labels: ['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun'],
+  values: [{data:[21.5, 25.0, 47.5, 64.8, 105.5, 133.2], label: 'Income'}],
+}
+
+Here we are looking at the performance of american and european players for each series. Since there are two entities, we have two arrays in values.
+{
+  labels: ['series A', 'series B', 'series C'],
+  values: [{data:[10, 15, 20], label: 'American'}, {data:[20, 25, 30], label: 'European'}],
+}
+
+'''
+
+
+lineGraphIntstruction = '''
+
+  Where data is: {
+  xValues: number[] | string[]
+  yValues: { data: number[]; label: string }[]
+}
+
+// Examples of usage:
+
+Here we are looking at the momentum of a body as a function of mass.
+{
+  xValues: ['2020', '2021', '2022', '2023', '2024'],
+  yValues: [
+    { data: [2, 5.5, 2, 8.5, 1.5]},
+  ],
+}
+
+Here we are looking at the performance of american and european players for each year. Since there are two entities, we have two arrays in yValues.
+{
+  xValues: ['2020', '2021', '2022', '2023', '2024'],
+  yValues: [
+    { data: [2, 5.5, 2, 8.5, 1.5], label: 'American' },
+    { data: [2, 5.5, 2, 8.5, 1.5], label: 'European' },
+  ],
+}
+'''
+
+pieChartIntstruction = '''
+
+  Where data is: {
+    labels: string
+    values: number
+  }[]
+
+// Example usage:
+[
+        { id: 0, value: 10, label: 'series A' },
+        { id: 1, value: 15, label: 'series B' },
+        { id: 2, value: 20, label: 'series C' },
+      ],
+'''
+
+scatterPlotIntstruction = '''
+Where data is: {
+  series: {
+    data: { x: number; y: number; id: number }[]
+    label: string
+  }[]
+}
+
+// Examples of usage:
+1. Here each data array represents the points for a different entity. 
+We are looking for correlation between amount spent and quantity bought for men and women.
+{
+  series: [
+    {
+      data: [
+        { x: 100, y: 200, id: 1 },
+        { x: 120, y: 100, id: 2 },
+        { x: 170, y: 300, id: 3 },
+      ],
+      label: 'Men',
+    },
+    {
+      data: [
+        { x: 300, y: 300, id: 1 },
+        { x: 400, y: 500, id: 2 },
+        { x: 200, y: 700, id: 3 },
+      ],
+      label: 'Women',
+    }
+  ],
+}
+
+2. Here we are looking for correlation between the height and weight of players.
+{
+  series: [
+    {
+      data: [
+        { x: 180, y: 80, id: 1 },
+        { x: 170, y: 70, id: 2 },
+        { x: 160, y: 60, id: 3 },
+      ],
+      label: 'Players',
+    },
+  ],
+}
+
+// Note: Each object in the 'data' array represents a point on the scatter plot.
+// The 'x' and 'y' values determine the position of the point, and 'id' is a unique identifier.
+// Multiple series can be represented, each as an object in the outer array.
+'''
+
+
+graph_instructions = {
+    "bar": barGraphIntstruction,
+    "horizontal_bar": horizontalBarGraphIntstruction,
+    "line": lineGraphIntstruction,
+    "pie": pieChartIntstruction,
+    "scatter": scatterPlotIntstruction
+}