split front element and event listening

front/tabs/tab_drias.py

@@ -1,9 +1,30 @@
 import gradio as gr
+from typing import TypedDict, List, Optional
 
 from climateqa.engine.talk_to_data.main import ask_drias
 from climateqa.engine.talk_to_data.config import DRIAS_MODELS, DRIAS_UI_TEXT
 
 
+class DriasUIElements(TypedDict):
+    tab: gr.Tab
+    details_accordion: gr.Accordion
+    examples_hidden: gr.Textbox
+    examples: gr.Examples
+    drias_direct_question: gr.Textbox
+    result_text: gr.Textbox
+    table_names_display: gr.DataFrame
+    query_accordion: gr.Accordion
+    drias_sql_query: gr.Textbox
+    chart_accordion: gr.Accordion
+    model_selection: gr.Dropdown
+    drias_display: gr.Plot
+    table_accordion: gr.Accordion
+    drias_table: gr.DataFrame
+    pagination_display: gr.Markdown
+    prev_button: gr.Button
+    next_button: gr.Button
+
+
 async def ask_drias_query(query: str, index_state: int):
     return await ask_drias(query, index_state)
 
@@ -91,18 +112,13 @@ def on_table_click(evt: gr.SelectData, table_names, sql_queries, dataframes, plo
     )
 
 
-def create_drias_tab():
-    sql_queries_state = gr.State([])
-    dataframes_state = gr.State([])
-    plots_state = gr.State([])
-    index_state = gr.State(0)
-    table_names_list = gr.State([])
-
-    with gr.Tab("Beta - Talk to DRIAS", elem_id="tab-vanna", id=6):
+def create_drias_ui() -> DriasUIElements:
+    """Create and return all UI elements for the DRIAS tab."""
+    with gr.Tab("Beta - Talk to DRIAS", elem_id="tab-vanna", id=6) as tab:
         with gr.Accordion(label="Details") as details_accordion:
             gr.Markdown(DRIAS_UI_TEXT)
             
-                # Add examples for common questions
+        # Add examples for common questions
         examples_hidden = gr.Textbox(visible=False, elem_id="drias-examples-hidden")
         examples = gr.Examples(
             examples=[
@@ -116,7 +132,6 @@ def create_drias_tab():
             outputs=[examples_hidden],
         )
         
-    
         with gr.Row():
             drias_direct_question = gr.Textbox(
                 label="Direct Question",
@@ -125,7 +140,6 @@ def create_drias_tab():
                 interactive=True,
             )
 
-
         result_text = gr.Textbox(
             label="", elem_id="no-result-label", interactive=False, visible=True
         )
@@ -158,124 +172,161 @@ def create_drias_tab():
             prev_button = gr.Button("Previous", visible=False)
             next_button = gr.Button("Next", visible=False)
 
-        # Handle example selection
-        examples_hidden.change(
-            lambda x: (gr.Accordion(open=False), gr.Textbox(value=x)),
-            inputs=[examples_hidden],
-            outputs=[details_accordion, drias_direct_question]
-        ).then(
-            ask_drias_query,
-            inputs=[examples_hidden, index_state],
-            outputs=[
-                drias_sql_query,
-                drias_table,
-                drias_display,
-                sql_queries_state,
-                dataframes_state,
-                plots_state,
-                index_state,
-                table_names_list,
-                result_text,
-            ],
-        ).then(
-            show_results,
-            inputs=[sql_queries_state, dataframes_state, plots_state],
-            outputs=[
-                result_text,
-                query_accordion,
-                table_accordion,
-                chart_accordion,
-                prev_button,
-                next_button,
-                pagination_display,
-                table_names_display,
-            ],
-        ).then(
-            update_pagination,
-            inputs=[index_state, sql_queries_state],
-            outputs=[pagination_display],
-        ).then(
-            display_table_names,
-            inputs=[table_names_list],
-            outputs=[table_names_display],
+        return DriasUIElements(
+            tab=tab,
+            details_accordion=details_accordion,
+            examples_hidden=examples_hidden,
+            examples=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,
+            model_selection=model_selection,
+            drias_display=drias_display,
+            table_accordion=table_accordion,
+            drias_table=drias_table,
+            pagination_display=pagination_display,
+            prev_button=prev_button,
+            next_button=next_button
         )
 
-        drias_direct_question.submit(
-            lambda: gr.Accordion(open=False),
-            inputs=None,
-            outputs=[details_accordion]
-        ).then(
-            ask_drias_query,
-            inputs=[drias_direct_question, index_state],
-            outputs=[
-                drias_sql_query,
-                drias_table,
-                drias_display,
-                sql_queries_state,
-                dataframes_state,
-                plots_state,
-                index_state,
-                table_names_list,
-                result_text,
-            ],
-        ).then(
-            show_results,
-            inputs=[sql_queries_state, dataframes_state, plots_state],
-            outputs=[
-                result_text,
-                query_accordion,
-                table_accordion,
-                chart_accordion,
-                prev_button,
-                next_button,
-                pagination_display,
-                table_names_display,
-            ],
-        ).then(
-            update_pagination,
-            inputs=[index_state, sql_queries_state],
-            outputs=[pagination_display],
-        ).then(
-            display_table_names,
-            inputs=[table_names_list],
-            outputs=[table_names_display],
-        )
+def setup_drias_events(ui_elements: DriasUIElements) -> 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([])
+    index_state = gr.State(0)
+    table_names_list = gr.State([])
 
-        model_selection.change(
-            filter_by_model,
-            inputs=[dataframes_state, plots_state, index_state, model_selection],
-            outputs=[drias_table, drias_display],
-        )
+    # 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(
+        ask_drias_query,
+        inputs=[ui_elements["examples_hidden"], index_state],
+        outputs=[
+            ui_elements["drias_sql_query"],
+            ui_elements["drias_table"],
+            ui_elements["drias_display"],
+            sql_queries_state,
+            dataframes_state,
+            plots_state,
+            index_state,
+            table_names_list,
+            ui_elements["result_text"],
+        ],
+    ).then(
+        show_results,
+        inputs=[sql_queries_state, dataframes_state, plots_state],
+        outputs=[
+            ui_elements["result_text"],
+            ui_elements["query_accordion"],
+            ui_elements["table_accordion"],
+            ui_elements["chart_accordion"],
+            ui_elements["prev_button"],
+            ui_elements["next_button"],
+            ui_elements["pagination_display"],
+            ui_elements["table_names_display"],
+        ],
+    ).then(
+        update_pagination,
+        inputs=[index_state, sql_queries_state],
+        outputs=[ui_elements["pagination_display"]],
+    ).then(
+        display_table_names,
+        inputs=[table_names_list],
+        outputs=[ui_elements["table_names_display"]],
+    )
 
-        prev_button.click(
-            show_previous,
-            inputs=[index_state, sql_queries_state, dataframes_state, plots_state],
-            outputs=[drias_sql_query, drias_table, drias_display, index_state],
-        ).then(
-            update_pagination,
-            inputs=[index_state, sql_queries_state],
-            outputs=[pagination_display],
-        )
+    # Handle direct question submission
+    ui_elements["drias_direct_question"].submit(
+        lambda: gr.Accordion(open=False),
+        inputs=None,
+        outputs=[ui_elements["details_accordion"]]
+    ).then(
+        ask_drias_query,
+        inputs=[ui_elements["drias_direct_question"], index_state],
+        outputs=[
+            ui_elements["drias_sql_query"],
+            ui_elements["drias_table"],
+            ui_elements["drias_display"],
+            sql_queries_state,
+            dataframes_state,
+            plots_state,
+            index_state,
+            table_names_list,
+            ui_elements["result_text"],
+        ],
+    ).then(
+        show_results,
+        inputs=[sql_queries_state, dataframes_state, plots_state],
+        outputs=[
+            ui_elements["result_text"],
+            ui_elements["query_accordion"],
+            ui_elements["table_accordion"],
+            ui_elements["chart_accordion"],
+            ui_elements["prev_button"],
+            ui_elements["next_button"],
+            ui_elements["pagination_display"],
+            ui_elements["table_names_display"],
+        ],
+    ).then(
+        update_pagination,
+        inputs=[index_state, sql_queries_state],
+        outputs=[ui_elements["pagination_display"]],
+    ).then(
+        display_table_names,
+        inputs=[table_names_list],
+        outputs=[ui_elements["table_names_display"]],
+    )
 
-        next_button.click(
-            show_next,
-            inputs=[index_state, sql_queries_state, dataframes_state, plots_state],
-            outputs=[drias_sql_query, drias_table, drias_display, index_state],
-        ).then(
-            update_pagination,
-            inputs=[index_state, sql_queries_state],
-            outputs=[pagination_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"]],
+    )
 
-        table_names_display.select(
-            fn=on_table_click,
-            inputs=[table_names_list, sql_queries_state, dataframes_state, plots_state],
-            outputs=[drias_sql_query, drias_table, drias_display, index_state],
+    # Handle pagination buttons
+    ui_elements["prev_button"].click(
+        show_previous,
+        inputs=[index_state, sql_queries_state, dataframes_state, plots_state],
+        outputs=[ui_elements["drias_sql_query"], ui_elements["drias_table"], ui_elements["drias_display"], index_state],
+    ).then(
+        update_pagination,
+        inputs=[index_state, sql_queries_state],
+        outputs=[ui_elements["pagination_display"]],
+    )
 
-        ).then(
-            update_pagination,
-            inputs=[index_state, sql_queries_state],
-            outputs=[pagination_display],
-        )
-        
-       
+    ui_elements["next_button"].click(
+        show_next,
+        inputs=[index_state, sql_queries_state, dataframes_state, plots_state],
+        outputs=[ui_elements["drias_sql_query"], ui_elements["drias_table"], ui_elements["drias_display"], index_state],
+    ).then(
+        update_pagination,
+        inputs=[index_state, sql_queries_state],
+        outputs=[ui_elements["pagination_display"]],
+    )
+
+    # Handle table selection
+    ui_elements["table_names_display"].select(
+        fn=on_table_click,
+        inputs=[table_names_list, sql_queries_state, dataframes_state, plots_state],
+        outputs=[ui_elements["drias_sql_query"], ui_elements["drias_table"], ui_elements["drias_display"], index_state],
+    ).then(
+        update_pagination,
+        inputs=[index_state, sql_queries_state],
+        outputs=[ui_elements["pagination_display"]],
+    )
+
+def create_drias_tab():
+    """Main function to create the DRIAS tab with UI and event handling."""
+    ui_elements = create_drias_ui()
+    setup_drias_events(ui_elements)
+    
+