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)