src/callbacks/chartgpt/graph.py

# %pip install chartgpt
from dotenv import load_dotenv

load_dotenv()

import os

import chartgpt as cg
import dash_mantine_components as dmc
import pandas as pd
from dash import ALL, MATCH, Input, Output, State, callback, ctx, dcc, html, no_update
from dash_iconify import DashIconify
from plotly.graph_objects import Figure

DATA_FOLDER = os.getenv("DATA_FOLDER", "data")
df_full = pd.read_pickle(os.path.join(DATA_FOLDER, "demo.pkl"))
chart = cg.Chart(df_full, model_name="gpt-3.5-turbo")


@callback(
    Output("history-list", "children"),
    Input("history", "data"),
)
def update_history(history):
    if not history:
        return []
    messages_list = []
    for i, message in enumerate(history):
        messages_list.append(
            dmc.AccordionItem(
                [
                    dmc.AccordionControl(
                        message.get("message"),
                    ),
                ],
                value=str(i),
            )
        )
    return messages_list


@callback(
    Output("button", "disabled"),
    Output("button", "variant"),
    Output("button", "color"),
    Input("prompt", "value"),
)
def update_button(prompt):
    if prompt:
        return False, "light", "primary"
    return True, "transparent", "gray"


def get_message(prompt, fig, code):
    message = {
        "message": prompt,
        "fig": fig,
        "code": code,
    }
    return message


@callback(
    Output("container", "children"),
    Output("history", "data"),
    Input("button", "n_clicks"),
    Input("history-list", "value"),
    State("prompt", "value"),
    State("history", "data"),
    prevent_initial_call=True,
)
def update_chart(submit_clicked, history_clicked, prompt, history):
    triggered_id = ctx.triggered_id
    if triggered_id == "button" and prompt:
        fig = chart.plot(prompt, return_fig=True, debug=False)
        code = chart.last_run_code
        message = get_message(prompt, fig, code)
        out = children(fig, code)
        return out, [message] + history
    elif triggered_id == "history-list" and history_clicked is not None:
        fig = history[int(history_clicked)].get("fig", None)
        if fig is None:
            return no_update, history
        fig = Figure(fig)
        code = history[int(history_clicked)].get("code")
        out = children(fig, code)
        return out, history
    elif triggered_id == "history-list" and history_clicked is None:
        return [], history
    else:
        raise Exception("This should never happen")


def children(fig, code):
    return dmc.Stack(
        [
            dcc.Graph(figure=fig),
            dmc.Accordion(
                variant="separated",
                chevronPosition="right",
                radius="md",
                children=[
                    dmc.AccordionItem(
                        [
                            dmc.AccordionControl(
                                "Show code", icon=DashIconify(icon="solar:code-bold")
                            ),
                            dmc.AccordionPanel(
                                dmc.Prism(
                                    code,
                                    language="python",
                                    withLineNumbers=True,
                                ),
                            ),
                        ],
                        value="customization",
                    )
                ],
            ),
        ]
    )