app.py

import time
from dataclasses import dataclass, field
from typing import Any, Callable, List, Optional

import pandas as pd
import streamlit as st

import utils
from prompts import (
    TOOLS,
    DataDescription,
    get_data_description,
    get_modified_code,
    get_task_code,
    get_task_summary,
)

MAX_COLUMNS = 15

utils.page_setup()


@dataclass
class ChartingRequest:
    library: str
    instruction: Optional[str] = None
    generated_code: Optional[str] = None
    method_name: Optional[str] = None
    chart_obj: Optional[Any] = None
    exception: Optional[str] = None


@dataclass
class DatasetState:
    dataset: Optional[pd.DataFrame] = None
    data_description: Optional[DataDescription] = None
    column_descriptions: Optional[List[str]] = None
    history: List[ChartingRequest] = field(default_factory=list)


if "dataset" not in st.session_state:
    st.session_state["dataset"] = DatasetState()
charting_state: DatasetState = st.session_state["dataset"]


def clear_history() -> None:
    st.session_state["dataset"].history = []


def clear_dataset() -> DatasetState:
    st.session_state["dataset"] = DatasetState()
    return st.session_state["dataset"]


with st.sidebar:
    library_mapping = {library.name: library for library in TOOLS}

    selected_library = st.selectbox(
        label="Select a Chart Library",
        key="library_selection",
        options=list(library_mapping.keys()),
        on_change=clear_history,
    )
    assert selected_library is not None
    selected_tool = library_mapping[selected_library]

    show_system_messages = st.checkbox("Show System Messages", value=False)


st.title("📈 ChartGPT")

with st.chat_message("assistant"):
    st.markdown(
        "Hi 👋 I'm ChartGPT. I can help you create charts from your data. To get started, please load your dataset below:"
    )

    dataset_type = st.selectbox(
        label="Data Source",
        options=["Toy Datasets", "CSV File"],
        on_change=clear_dataset,  # type: ignore
    )

    with st.form(key="load_data"):
        load_data: Optional[Callable] = None
        if dataset_type == "Toy Datasets":
            toy_dataset = st.selectbox(
                label="Select a Toy Dataset", options=utils.get_toy_datasets(), index=2
            )
            if toy_dataset is not None:
                from vega_datasets import data as vega_data

                dataset_name = toy_dataset.replace("-", "_")
                dataset = getattr(vega_data, dataset_name)
                load_data = dataset
        elif dataset_type == "CSV File":
            csv_file = st.file_uploader("Upload CSV File", type=["csv"])
            if csv_file is not None:
                load_data = lambda: pd.read_csv(csv_file)

        if st.form_submit_button("📂 Load Data"):
            charting_state = clear_dataset()
            with st.spinner("📥 Loading data..."):
                charting_state.dataset = load_data() if load_data else None
                time.sleep(1)

if charting_state.dataset is None:
    st.stop()

with st.chat_message("assistant"):
    st.markdown("Thanks, I was able to load your data 🎉 Here is an overview:")

    col1, col2 = st.columns(2)

    with col1:
        st.metric("Number of Rows", value=len(charting_state.dataset))

    with col2:
        st.metric("Number of Columns", value=len(charting_state.dataset.columns))

    st.dataframe(
        charting_state.dataset.head(min(len(charting_state.dataset), 5000)),
        use_container_width=True,
    )

with st.chat_message("assistant"):
    with st.spinner("🔬 Analyzing your data..."):
        if charting_state.data_description is None:
            charting_state.data_description = get_data_description(
                charting_state.dataset,
                openai_model=st.session_state["openai_model"],
                show_system_messages=show_system_messages,
            )
    st.markdown(charting_state.data_description["data_description"])
    st.markdown("Here is my interpretation of every column:")
    charting_state.column_descriptions = st.data_editor(
        pd.DataFrame(charting_state.data_description["columns"]),
        column_config={
            "name": "Column",
            "description": st.column_config.TextColumn(
                "Description",
                help="Description of the column generated by ChartGPT. "
                "You can edit these descriptions to help the visualization task.",
            ),
        },
        use_container_width=True,
        disabled=["name"],
        hide_index=True,
    )
    markdown = "I also found some interesting **observations** about the data:\n\n"
    for observation in charting_state.data_description["observations"]:
        markdown += f"- {observation}\n"
    st.markdown(markdown)
    st.markdown("And these are my top **4 charting ideas** based on the provided data:")
    preset_prompt = None
    col1, col2 = st.columns(2)
    for i, chart_idea in enumerate(charting_state.data_description["chart_ideas"]):
        chart_idea = chart_idea.strip()
        if i % 2 == 0:
            col = col1
        else:
            col = col2

        if col.button("💡 " + chart_idea, help=chart_idea):
            clear_history()
            preset_prompt = chart_idea

    st.markdown("What would you like to visualize?")


def render_charting_response(chart_request: ChartingRequest) -> None:
    st.expander("Generated code").code(chart_request.generated_code, language="python")
    if chart_request.chart_obj:
        selected_tool.handle_output(chart_request.chart_obj)
    if chart_request.exception:
        st.error(chart_request.exception)


def load_code(new_chart: ChartingRequest) -> None:
    assert new_chart.method_name is not None
    assert new_chart.generated_code is not None
    try:
        loaded_module = utils.load_code_as_module(new_chart.generated_code)
        task_method = getattr(loaded_module, new_chart.method_name)
        new_chart.chart_obj = task_method(data=charting_state.dataset)
    except Exception as ex:
        exception_text = str(ex)
        new_chart.exception = exception_text
    charting_state.history.append(new_chart)
    render_charting_response(new_chart)


for chart_request in charting_state.history:
    st.chat_message("user").write(chart_request.instruction)
    with st.chat_message("assistant"):
        render_charting_response(chart_request)

chat_input_placeholder = st.empty()

if not charting_state.history:
    prompt = (
        chat_input_placeholder.chat_input("What do you like to visualize?")
        or preset_prompt
    )
    if prompt:
        st.chat_message("user").write(prompt)

        new_chart = ChartingRequest(library=selected_library)
        new_chart.instruction = prompt

        with st.chat_message("assistant"):
            with st.spinner("Analyzing instruction..."):
                task_summary = get_task_summary(
                    charting_state.dataset,
                    prompt,
                    openai_model=st.session_state["openai_model"],
                    column_desc_df=charting_state.column_descriptions,
                    show_system_messages=show_system_messages,
                )
                name = task_summary["emoji"] + " " + task_summary["name"]
                new_chart.method_name = utils.name_to_variable(
                    task_summary["method_name"]
                )

                task_code = selected_tool.get_task_prompt(
                    task_name=new_chart.method_name,
                    description=task_summary["task_description"],
                )
                with st.spinner(f"Generating code for: {name}"):
                    new_chart.generated_code = get_task_code(
                        charting_state.dataset,
                        task_instruction=prompt,
                        task_code_template=task_code,
                        task_library=new_chart.library,
                        openai_model=st.session_state["openai_model"],
                        column_desc_df=charting_state.column_descriptions,
                        show_system_messages=show_system_messages,
                    )
                with st.spinner(f"Creating chart for: {name}"):
                    load_code(new_chart)

if charting_state.history:
    latest_chart_request = charting_state.history[-1]
    assert latest_chart_request is not None
    assert latest_chart_request.generated_code is not None
    assert latest_chart_request.method_name is not None

    prompt = chat_input_placeholder.chat_input("What do you like to modify?")
    if prompt:
        new_chart = ChartingRequest(library=latest_chart_request.library)
        new_chart.instruction = prompt
        new_chart.method_name = latest_chart_request.method_name

        st.chat_message("user").write(prompt)

        with st.chat_message("assistant"):
            with st.spinner(
                "🔧 Fixing code..."
                if latest_chart_request.exception
                else "🔨 Modifying code..."
            ):
                modified_code = get_modified_code(
                    task_code=latest_chart_request.generated_code,
                    instruction=prompt,
                    dataset_df=charting_state.dataset,
                    exception_message=latest_chart_request.exception,
                    openai_model=st.session_state["openai_model"],
                    column_desc_df=charting_state.column_descriptions,
                    show_system_messages=show_system_messages,
                )

            if modified_code:
                new_chart.generated_code = modified_code
                with st.spinner(f"🔮 Generating code..."):
                    load_code(new_chart)
            else:
                st.markdown(
                    "Sorry, I couldn't find a way to modify the code. Please try again."
                )