climateqa/engine/talk_to_data/objects/states.py

from typing import Any, Callable, Optional, TypedDict
from plotly.graph_objects import Figure
import pandas as pd
from climateqa.engine.talk_to_data.objects.plot import Plot

class TTDOutput(TypedDict):
    status: str
    plot: Plot
    table: str
    sql_query: Optional[str]
    dataframe: Optional[pd.DataFrame]
    figure: Optional[Callable[..., Figure]]
    plot_information: Optional[str]
class State(TypedDict):
    user_input: str
    plots: list[str]
    outputs: dict[str, TTDOutput]
    error: Optional[str]