Agent Execution

app.py

@@ -1,16 +1,8 @@
 import os
-import json
-import torch
 import duckdb
-import spaces
-import pandas as pd
 import gradio as gr
-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 transformers import HfEngine, ReactCodeAgent
+from transformers.agents import Tool
 
 # Height of the Tabs Text Area
 TAB_LINES = 8
@@ -22,32 +14,7 @@ print('Connecting to DB...')
 # Connect to DB
 conn = duckdb.connect(f"md:my_db?motherduck_token={md_token}", read_only=True)
 
-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")
-
-print('Loading Model...')
-
-
-tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-9b-it")
-
-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("google/gemma-2-9b-it", 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}')
+llm_engine = HfEngine(model="meta-llama/Meta-Llama-3-70B-Instruct")
 
 def get_schemas():
     schemas = conn.execute("""
@@ -79,225 +46,54 @@ def get_table_schema(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 a valid DuckDB SQL query to answer the following question. Select only the columns that are necessary for visualization or that will make visualization easier (e.g., numeric or categorical data). Respond only with the SQL query and do not include anything else.
-    
-    ### Input:
-    Here is the database schema that the SQL query will run on:
-    {schema}
+    return ddl_create, full_path
+
+def get_visualization(question, tool):
+    agent = ReactCodeAgent(tools=[tool], llm_engine=llm_engine, add_base_tools=True,
+                           additional_authorized_imports=['matplotlib.pyplot',
+                                                 'pandas', 'plotly.express',
+                                                 'seaborn'], max_iterations=20)
+    fig = agent.run(
+        task=f'''
+    Use seaborn. Always 
+    Question: {question}
+    Always use the right colors.
+    If the question is about showing n number of rows return empty figure.
+    In the end you have to return a final fig using the `final_answer` tool
+    ''',
+    )
     
-    ### 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)
-
-Generate a JSON object. The JSON object should have the following structure:
-
-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)
-    response = run_llm(final_prompt)
-    response = response.replace('```', '')
-    response = response.replace('json', '')
-    json_data = json.loads(response)
-    visualization = json_data['Recommended Visualization']
-    reason = json_data['Reason']
-    print(visualization, reason)
-    return visualization, reason
-
-def format_bar_data(results, question, sql_query):
-    if isinstance(results, str):
-        results = eval(results)
-
-    if len(results[0]) == 2:
-        labels = [str(row[0]) for row in results]
-        data = [float(row[1]) if row[1] is not None else 0 for row in results]
-
-        prompt = ChatPromptTemplate.from_messages([
-            ("system", "You are a data labeling expert. Given a question, SQL Query used, and some data, provide a concise and relevant label for the data series."),
-            ("human", "Question: {question}\n SQL Query: {sql_query} \nData (first few rows): {data}\n\n Provide a concise label name for this y axis. Just give me the name."),
-        ])
-        prompt = prompt.format_prompt(question=question, data=str(results[:2]),
-                                     sql_query=sql_query)
-        label = run_llm(prompt)
-        label = label.replace('**Answer:**', '')
-        values = [{"data": data, "label": label}]
-    elif len(results[0]) == 3:
-        categories = set(row[1] for row in results)
-        labels = list(categories)
-        entities = set(row[0] for row in results)
-        values = []
-        for entity in entities:
-            entity_data = [float(row[2]) for row in results if row[0] == entity]
-            values.append({"data": entity_data, "label": str(entity)})
-    else:
-        raise ValueError("Unexpected data format in results")
-
-    formatted_data = {
-        "labels": labels,
-        "values": values
+    return fig
+
+class SQLExecutorTool(Tool):
+    name = "sql_engine"
+    inputs = {
+        "query": {
+            "type": "text",
+            "description": f"The query to perform. This should be correct DuckDB SQL.",
+        }
     }
+    output_type = "pandas.core.frame.DataFrame"
 
-    return formatted_data
-
-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,Given the question, SQL query, and results, your job is to Understand the question and SQL query. Ensure the result matches the query and can be easily visualized."),
-    ("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. Do not give backticks."),
-    ])
-
-
-    prompt = template.format_prompt(question=text_query, sql_query=sql_query,
-                                    results=sql_result, instructions=instruction)
-    formatted_data = run_llm(prompt)
-    formatted_data = formatted_data.replace('```', '')
-    formatted_data = formatted_data.replace('json', '')
-    print(f'Formatted Data {formatted_data}')
-    return json.loads(formatted_data.replace('.', '').strip())
-
-def visualize_result(text_query, visualization_type, sql_query,
-                      sql_result):
-
-    if visualization_type == 'bar':
-        try:
-            data =  format_bar_data(sql_result, text_query, sql_query)
-        except Exception as e:
-            data =  format_data(text_query=text_query, sql_query=sql_query,
-                            sql_result=sql_result, visualization_type=visualization_type)
-        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 forward(self, query: str) -> str:
+        with duckdb.connect(f"md:my_db?motherduck_token={md_token}", read_only=True) as con:
+            output_df = conn.sql(query).df()
+        return output_df
+    
+tool = SQLExecutorTool()
 
 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, pd.DataFrame([{"error": f"❌ Unable to generate the SQL query. {e}"}]))
-
-    try:
-        sql_query_result_raw = conn.sql(generated_sql_query)
-        sql_query_result  = sql_query_result_raw.fetchall()
-        sql_query_df = sql_query_result_raw.df()
-        
-        print(f"SQL Query Result: {sql_query_result}")
-    except Exception as e:
-        return generate_output(schema, prompt, generated_sql_query, fig, pd.DataFrame([{"error": f"❌ Unable to execute the SQL query. {e}"}]))
-
-    if len(sql_query_result) >= 100:
-        gr.Warning(f"⚠️ Data is too large for visualization. Please refine your query.")
-        return generate_output(schema, prompt, generated_sql_query,
-                               fig, sql_query_df)
+    schema, _ = get_table_schema(table)
+    tool.description = f"""Allows you to perform SQL queries on the table. Returns a pandas dataframe representation of the result.
+    The table schema is as follows: \n{schema}"""
     
     try:
-      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}")
-      print(f"Reason: {reason}")
-
+        fig = get_visualization(question=text_query, tool=tool)
     except Exception as e:
-        return generate_output(schema, prompt, generated_sql_query,
-                               fig, sql_query_df)
+        gr.Warning(f"❌ Unable to generate the visualization. {e}")
+        
+    return fig
     
-    if visualization_type != 'none':
-        try:
-            plot = visualize_result(text_query=text_query, sql_query=generated_sql_query,
-                        sql_result=sql_query_result, visualization_type=visualization_type)
-
-        except Exception as e:
-            gr.Warning(f"⚠️ {e}")
-            return generate_output(schema, prompt, generated_sql_query, fig, sql_query_df)
-    else:
-      gr.Warning(f"⚠️ {reason}")
-      return generate_output(schema, prompt, generated_sql_query, fig, sql_query_df)
-
-    return generate_output(schema, prompt, generated_sql_query, plot, sql_query_df)
-
-def generate_output(schema, prompt, generated_sql_query, result_plot, sql_query_df):
-    return [schema, prompt, generated_sql_query, result_plot, sql_query_df]
 
 custom_css = """
 .gradio-container {
@@ -312,6 +108,7 @@ custom_css = """
     background-color: #4a90e2 !important;
 }
 .gr-button:hover {
+    
     background-color: #3a7bc8 !important;
 }
 """
@@ -329,32 +126,24 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="purple", secondary_hue="indigo"
 
     with gr.Row():
 
-        with gr.Column(scale=1, variant='panel'):
+        with gr.Column(scale=1):
             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...")
+            query_input = gr.Textbox(lines=3, 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("Result"):
-                    query_result_output = gr.DataFrame(label="Query Results", value=[], interactive=False)
-                with gr.Tab("Plot"):
-                    result_plot = gr.Plot()
-                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)
+    with gr.Tabs():
+        with gr.Tab("Plot"):
+            result_plot = gr.Plot()
 
         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_plot, query_result_output])
+        generate_query_button.click(main, inputs=[tables_dropdown, query_input], outputs=[result_plot])
 
 if __name__ == "__main__":
     demo.launch(debug=True)

plot_utils.py

@@ -1,85 +0,0 @@
-import matplotlib.pyplot as plt
-import numpy as np
-
-def plot_bar_chart(data):
-    fig, ax = plt.subplots()
-    labels = data['labels']
-    values = data['values']
-    
-    colors = ['darkviolet', 'indigo', 'blueviolet']
-    width = 0.2  
-    x = np.arange(len(labels)) 
-    
-    for i, value in enumerate(values):
-        ax.bar(x + i * width, value['data'], width, label=value['label'], color=colors[i])
-    
-    ax.set_xticks(x + width / 2 * (len(values) - 1))  
-    ax.set_xticklabels(labels)
-    
-    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'], color=['darkviolet', 'indigo', 'blueviolet'])
-    
-    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

requirements.txt

@@ -1,8 +1,3 @@
 huggingface_hub
-accelerate
-bitsandbytes
 transformers
-duckdb
-langchain-huggingface
-langchain-core
-sentencepiece
+duckdb

visualization_prompt.py

@@ -1,154 +0,0 @@
-barGraphIntstruction = '''
-
-  Where data is: {
-    labels: string[]
-    values: {\data: number[], label: string}[]
-  }
-    
-The output must follow this format strictly, even if the input data differs from the examples below.   
-// 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'}],
-}
-
-The output format must be consistent with this structure, regardless of the specific input data.
-'''
-
-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
-}