front/tabs/tab_drias.py

import gradio as gr
from typing import TypedDict, List, Optional
import os
import pandas as pd

from climateqa.engine.talk_to_data.main import ask_drias
from climateqa.engine.talk_to_data.drias.config import DRIAS_MODELS, DRIAS_UI_TEXT

class DriasUIElements(TypedDict):
    tab: gr.Tab
    details_accordion: gr.Accordion
    examples_hidden: gr.Textbox
    examples: gr.Examples
    image_examples: gr.Row
    drias_direct_question: gr.Textbox
    result_text: gr.Textbox
    table_names_display: gr.Radio
    query_accordion: gr.Accordion
    drias_sql_query: gr.Textbox
    chart_accordion: gr.Accordion
    plot_information: gr.Markdown
    model_selection: gr.Dropdown
    drias_display: gr.Plot
    table_accordion: gr.Accordion
    drias_table: gr.DataFrame


async def ask_drias_query(query: str, index_state: int, user_id: str):
    result = await ask_drias(query, index_state, user_id)
    return result


def show_results(sql_queries_state, dataframes_state, plots_state, table_names):
    if not sql_queries_state or not dataframes_state or not plots_state:
        # If all results are empty, show "No result"
        return (
            gr.update(visible=True),
            gr.update(visible=False),
            gr.update(visible=False),
            gr.update(visible=False),
            gr.update(visible=False),
        )
    else:
        # Show the appropriate components with their data
        return (
            gr.update(visible=False),
            gr.update(visible=True),
            gr.update(visible=True),
            gr.update(visible=True),
            gr.update(choices=table_names, value=table_names[0], visible=True),
        )


def filter_by_model(dataframes, figures, index_state, model_selection):
    df = dataframes[index_state]
    if df.empty:
        return df, None
    if "model" not in df.columns:
        return df, figures[index_state](df)
    if model_selection != "ALL":
        df = df[df["model"] == model_selection]
        if df.empty:
            return df, None
    figure = figures[index_state](df)
    return df, figure


def on_table_click(selected_label, table_names, sql_queries, dataframes, plot_informations, plots):
    index = table_names.index(selected_label)
    figure = plots[index](dataframes[index])
    return (
        sql_queries[index],
        dataframes[index],
        figure,
        plot_informations[index],
        index,
    )


def create_drias_ui() -> DriasUIElements:
    """Create and return all UI elements for the DRIAS tab."""
    with gr.Tab("France - Talk to DRIAS", elem_id="tab-vanna", id=6) as tab:
        with gr.Accordion(label="❓ How to use?", elem_id="details") as details_accordion:
            gr.Markdown(DRIAS_UI_TEXT)
            
        # Add examples for common questions
        examples_hidden = gr.Textbox(visible=False, elem_id="drias-examples-hidden")
        examples = gr.Examples(
            examples=[
                ["What will the temperature be like in Paris?"],
                ["What will be the total rainfall in France in 2030?"],
                ["How frequent will extreme events be in Lyon?"],
                ["Comment va évoluer la température en France entre 2030 et 2050 ?"]
            ],
            label="Example Questions",
            inputs=[examples_hidden],
            outputs=[examples_hidden],
        )
        
        with gr.Row():
            drias_direct_question = gr.Textbox(
                label="Direct Question",
                placeholder="You can write direct question here",
                elem_id="direct-question",
                interactive=True,
            )
        
        
        with gr.Row(visible=True, elem_id="example-img-container") as image_examples:
            gr.Markdown("### Examples of possible visualizations")

            with gr.Row():
                gr.Image("./front/assets/talk_to_drias_winter_temp_paris_example.png", label="Evolution of Mean Winter Temperature in Paris", elem_classes=["example-img"])
                gr.Image("./front/assets/talk_to_drias_annual_temperature_france_example.png", label="Mean Annual Temperature in 2030 in France", elem_classes=["example-img"])
                gr.Image("./front/assets/talk_to_drias_frequency_remarkable_precipitation_lyon_example.png", label="Frequency of Remarkable Daily Precipitation in Lyon", elem_classes=["example-img"])

        result_text = gr.Textbox(
            label="", elem_id="no-result-label", interactive=False, visible=True
        )
        
        with gr.Row():
            table_names_display = gr.Radio(
                choices=[], 
                label="Relevant figures created",
                interactive=True,
                elem_id="table-names",
                visible=False
            )

            with gr.Accordion(label="SQL Query Used", visible=False) as query_accordion:
                drias_sql_query = gr.Textbox(
                    label="", elem_id="sql-query", interactive=False
                )

    
        with gr.Accordion(label="Chart", visible=False) as chart_accordion:
            with gr.Row():
                model_selection = gr.Dropdown(
                    label="Model", choices=DRIAS_MODELS, value="ALL", interactive=True
                )
                with gr.Accordion(label="Informations about the plot", open=False):
                    plot_information = gr.Markdown(value = "")
                
            drias_display = gr.Plot(elem_id="vanna-plot")

        with gr.Accordion(
            label="Data used", open=False, visible=False
        ) as table_accordion:
            drias_table = gr.DataFrame([], elem_id="vanna-table")

        return DriasUIElements(
            tab=tab,
            details_accordion=details_accordion,
            examples_hidden=examples_hidden,
            examples=examples,
            image_examples=image_examples,
            drias_direct_question=drias_direct_question,
            result_text=result_text,
            table_names_display=table_names_display,
            query_accordion=query_accordion,
            drias_sql_query=drias_sql_query,
            chart_accordion=chart_accordion,
            plot_information=plot_information,
            model_selection=model_selection,
            drias_display=drias_display,
            table_accordion=table_accordion,
            drias_table=drias_table,
        )



def setup_drias_events(ui_elements: DriasUIElements, share_client=None, user_id=None) -> None:
    """Set up all event handlers for the DRIAS tab."""
    # Create state variables
    sql_queries_state = gr.State([])
    dataframes_state = gr.State([])
    plots_state = gr.State([])
    plot_informations_state = gr.State([])
    index_state = gr.State(0)
    table_names_list = gr.State([])
    user_id = gr.State(user_id)

    # Handle direct question submission - trigger the same workflow by setting examples_hidden
    ui_elements["drias_direct_question"].submit(
        lambda x: gr.update(value=x),
        inputs=[ui_elements["drias_direct_question"]],
        outputs=[ui_elements["examples_hidden"]],
    )

    # Handle example selection
    ui_elements["examples_hidden"].change(
        lambda x: (gr.Accordion(open=False), gr.Textbox(value=x)),
        inputs=[ui_elements["examples_hidden"]],
        outputs=[ui_elements["details_accordion"], ui_elements["drias_direct_question"]]
    ).then(
        lambda : gr.update(visible=False),
        inputs=None,
        outputs=ui_elements["image_examples"]
    ).then(
        ask_drias_query,
        inputs=[ui_elements["examples_hidden"], index_state, user_id],
        outputs=[
            ui_elements["drias_sql_query"],
            ui_elements["drias_table"],
            ui_elements["drias_display"],
            ui_elements["plot_information"],
            sql_queries_state,
            dataframes_state,
            plots_state,
            plot_informations_state,
            index_state,
            table_names_list,
            ui_elements["result_text"],
        ],
    ).then(
        show_results,
        inputs=[sql_queries_state, dataframes_state, plots_state, table_names_list],
        outputs=[
            ui_elements["result_text"],
            ui_elements["query_accordion"],
            ui_elements["table_accordion"],
            ui_elements["chart_accordion"],
            ui_elements["table_names_display"],
        ],
    )

   
    # Handle model selection change
    ui_elements["model_selection"].change(
        filter_by_model,
        inputs=[dataframes_state, plots_state, index_state, ui_elements["model_selection"]],
        outputs=[ui_elements["drias_table"], ui_elements["drias_display"]],
    )


    # Handle table selection
    ui_elements["table_names_display"].change(
        fn=on_table_click,
        inputs=[ui_elements["table_names_display"], table_names_list, sql_queries_state, dataframes_state, plot_informations_state, plots_state],
        outputs=[ui_elements["drias_sql_query"], ui_elements["drias_table"], ui_elements["drias_display"], ui_elements["plot_information"], index_state],
    )

def create_drias_tab(share_client=None, user_id=None):
    """Create the DRIAS tab with all its components and event handlers."""
    ui_elements = create_drias_ui()
    setup_drias_events(ui_elements, share_client=share_client, user_id=user_id)