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

# 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 = ["meta-llama/Meta-Llama-3-70B-Instruct",
          "meta-llama/Llama-3.1-70B-Instruct"]

for model in models:
  llm_engine = HfEngine(model=model)
  model_status = llm_engine.client.get_model_status()

  if model_status.loaded:
    break

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):
    agent = ReactCodeAgent(tools=[tool], llm_engine=llm_engine, add_base_tools=True,
                           additional_authorized_imports=['matplotlib.pyplot',
                                                 'pandas', 'plotly.express',
                                                 'seaborn'], max_iterations=10)
    fig = agent.run(
        
        task=f'''
    THINK STEP BY STEP
    Here are the steps you should follow while writing code for Visualization:
    1. Select the most appropriate chart type for 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 analysis, and box plots for visualizing data spread and outliers.
    2. Ensure clear and appropriate labels, colors, and design elements, keeping visual elements legible and uncluttered.
    3. Follow best practices, avoiding unnecessary visual distractions (chartjunk).
    4. Ensure the code is error-free, with correct fields, transformations, and aesthetics.
    5. Use descriptive and accurate x and y axis labels that reflect the data.
    6. Ensure units of measurement are clearly indicated on axes (e.g., %, $, cm).
    7. Ensure that categorical data is plotted on one axis and numerical data on the other, with appropriate labels that clearly represent the data being visualized.
    8. When plotting categorical data, arrange categories in a meaningful order (e.g., by size, time, or frequency) rather than randomly.
    9. Ensure that the categorical data are plotted on the x-axis, and the frequencies (numerical data) are plotted on the y-axis.
    10. Use seaborn
    11. In the end you have to return a dict which contain final fig as fig key, Generated SQL as sql key, Data as a dataframe with data key using the `final_answer` tool.
    Here is the task:
    task: {question}
    ''',
    )

    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"

    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, _ = 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:
        output = get_visualization(question=text_query, tool=tool)
        fig = output.get('fig', None)
        generated_sql = output.get('sql', None)
        data = output.get('data', None)
    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)