app.py

import os
import duckdb
import gradio as gr
import matplotlib.pyplot as plt
from transformers import HfEngine, ReactCodeAgent
from transformers.agents import Tool
from langsmith import traceable
# 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')

print('Connecting to DB...')
# Connect to DB
conn = duckdb.connect(f"md:my_db?motherduck_token={md_token}", read_only=True)

models = ["Qwen/Qwen2.5-72B-Instruct","meta-llama/Meta-Llama-3-70B-Instruct",
          "meta-llama/Llama-3.1-70B-Instruct"]

model_loaded = False 
for model in models:
  try:
      llm_engine = HfEngine(model=model)
      info = llm_engine.client.get_endpoint_info()
      model_loaded = True
      break
  except Exception as e:
      print(f"Error for model {model}: {e}")
      continue

if not model_loaded:
    gr.Warning(f"❌ None of the model form {models} are available. {e}")

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, full_path


def get_visualization(question, tool, schema, table_name):
    agent = ReactCodeAgent(tools=[tool], llm_engine=llm_engine, add_base_tools=True,
                           additional_authorized_imports=['matplotlib.pyplot',
                                                 'pandas', 'plotly.express',
                                                 'seaborn'], max_iterations=10)
    results = agent.run(

    task= f'''

    Here are the steps you should follow while writing code for Visualization:
    1. You have access to the database with the `sql_engine` tool, which allows you to run DuckDB SQL queries and return results as a df.
    2. Query the database using `sql_engine`, print the first 5 rows to inspect the data.
    3. Select the most appropriate chart type for the data:
       - Use bar charts for categorical comparisons, line charts for trends over time, scatter plots for relationships between variables, pie charts for proportions, histograms for distribution, and box plots for data spread and outliers.
    4. Analyze the data and choose the best visualization type to answer the query.
    5. Always include a plot in your answer.
    6. Use Seaborn for the plots.
    7. In the end, return a dictionary containing the final figure (`fig` key), the generated SQL (`sql` key), and the data as a dataframe (`data` key) using the `final_answer` tool, e.g. `final_answer(answer={{"fig": 'fig.png', "sql": sql, "data": data}})`.

    Example:

    ```python
    # Input query
    query_description = 'Average tip amount based on the ride time length in minutes.'

    # SQL Query to get ride time length and average tip amount
    query = """
        SELECT 
            EXTRACT(EPOCH FROM (tpep_dropoff_datetime - tpep_pickup_datetime)) / 60 AS ride_time_length,
            AVG(tip_amount) AS avg_tip_amount
        FROM 
            sample_data.nyc.taxi
        GROUP BY 
            EXTRACT(EPOCH FROM (tpep_dropoff_datetime - tpep_pickup_datetime)) / 60
    """

    # Execute the query using the sql_engine tool
    df = sql_engine(query=query)

    # Print the result to observe the data
    print(df)

    # Create a line plot using seaborn
    import seaborn as sns
    import matplotlib.pyplot as plt

    plt.figure(figsize=(10,6))
    sns.lineplot(x="ride_time_length", y="avg_tip_amount", data=df)

    # Set the title and labels
    plt.title("Average Tip Amount vs Ride Time Length")
    plt.xlabel("Ride Time Length (minutes)")
    plt.ylabel("Average Tip Amount")

    # Print the plot to observe the results
    print("Plot created")

    # Since we are required to return a fig, sql, and data, let's store the plot in a variable
    fig = plt.gcf()

    # Store the query in a variable
    sql = query

    # Store the dataframe in a variable
    data = df

    # Return the final answer
    final_answer(answer={{"fig": fig, "sql": sql, "data": data}})
    ```


    Here is the query you should generate a plot for: '{question}'.
    Here is the schema: '{schema}' and here is the table name: '{table_name}

      
    '''
    )

    return results


@traceable()
def query_response(input_prompt, generated_sql):
  return generated_sql


class SQLExecutorTool(Tool):
    name = "sql_engine"
    inputs = {
        "query": {
            "type": "text",
            "description": f"The query to perform. This should be correct DuckDB SQL.",
        }
    }
    description = """Allows you to perform SQL queries on the table. Returns a pandas dataframe representation of the result."""
    output_type = "pandas.core.frame.DataFrame"

    def forward(self, query: str) -> str:
        output_df = conn.sql(query).df()
        return output_df
    
tool = SQLExecutorTool()

def main(table, text_query):
    # Empty Fig
    fig, ax = plt.subplots()
    ax.set_axis_off()  
    
    schema, table_name = get_table_schema(table)
    
    
    try:
        output = get_visualization(question=text_query, tool=tool, schema=schema, table_name=table_name)
        fig = output.get('fig', None)
        generated_sql = output.get('sql', None)
        data = output.get('data', None)
        input_prompt = text_query + '\n' + table_name + '\n' + schema
        _ = query_response(input_prompt, generated_sql)
    except Exception as e:
        gr.Warning(f"❌ Unable to generate the visualization. {e}")
        
    return fig, generated_sql, data
    
    

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

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;'>DataViz Agent</strong>
    <br>
    <span style='font-size: 20px;'>Visualize SQL queries based on a given text for the dataset.</span>
    </div>
    """)

    with gr.Row():

        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=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("Plot"):
            result_plot = gr.Plot()
        with gr.Tab("SQL"):    
            generated_sql = gr.Textbox(lines=TAB_LINES, label="Generated SQL", value="", interactive=False,
                                              autoscroll=False)
        with gr.Tab("Data"):    
            data = gr.Dataframe(label="Data", interactive=False)

        schema_dropdown.change(update_tables, inputs=schema_dropdown, outputs=tables_dropdown)
        generate_query_button.click(main, inputs=[tables_dropdown, query_input], outputs=[result_plot, generated_sql, data])

if __name__ == "__main__":
    demo.launch(debug=True)