feat: added list of tables and reduced execution time of TTD

app.py

@@ -12,11 +12,10 @@ from climateqa.engine.reranker import get_reranker
 from climateqa.engine.graph import make_graph_agent, make_graph_agent_poc
 from climateqa.engine.chains.retrieve_papers import find_papers
 from climateqa.chat import start_chat, chat_stream, finish_chat
-from climateqa.engine.talk_to_data.main import ask_drias, DRIAS_MODELS
-from climateqa.engine.talk_to_data.myVanna import MyVanna
 
 from front.tabs import create_config_modal, cqa_tab, create_about_tab
 from front.tabs import MainTabPanel, ConfigPanel
+from front.tabs.tab_drias import create_drias_tab
 from front.utils import process_figures
 from gradio_modal import Modal
 
@@ -111,8 +110,6 @@ agent_poc = make_graph_agent_poc(
 #     return ask_vanna(vn, db_vanna_path, query)
 
 
-def ask_drias_query(query: str, index_state: int):
-    return ask_drias(query, index_state)
 
 
 async def chat(
@@ -201,211 +198,6 @@ def update_sources_number_display(
     )
 
 
-# def create_drias_tab():
-#     with gr.Tab("Beta - Talk to DRIAS", elem_id="tab-vanna", id=6) as tab_vanna:
-#         vanna_direct_question = gr.Textbox(label="Direct Question", placeholder="You can write direct question here",elem_id="direct-question", interactive=True)
-#         with gr.Accordion("Details",elem_id = 'vanna-details', open=False) as vanna_details :
-#             vanna_sql_query = gr.Textbox(label="SQL Query Used", elem_id="sql-query", interactive=False)
-#             show_vanna_table = gr.Button("Show Table", elem_id="show-table")
-#             with Modal(visible=False) as vanna_table_modal:
-#                 vanna_table = gr.DataFrame([], elem_id="vanna-table")
-#                 close_vanna_modal = gr.Button("Close", elem_id="close-vanna-modal")
-#                 close_vanna_modal.click(lambda: Modal(visible=False),None, [vanna_table_modal])
-#             show_vanna_table.click(lambda: Modal(visible=True),None ,[vanna_table_modal])
-
-#         vanna_display = gr.Plot()
-#         vanna_direct_question.submit(ask_drias_query, [vanna_direct_question], [vanna_sql_query ,vanna_table, vanna_display])
-
-
-def show_results(sql_queries_state, dataframes_state, plots_state):
-    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),
-            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(visible=True),
-            gr.update(visible=True),
-            gr.update(visible=True),
-        )
-
-
-def filter_by_model(dataframes, figures, index_state, model_selection):
-    df = dataframes[index_state]
-    if model_selection != "ALL":
-        df = df[df["model"] == model_selection]
-    figure = figures[index_state](df)
-    return df, figure
-
-
-def update_pagination(index, sql_queries):
-    pagination = f"{index + 1}/{len(sql_queries)}" if sql_queries else ""
-    return pagination
-
-
-def create_drias_tab():
-    details_text = """
-Hi, I'm **Talk to Drias**, designed to answer your questions using [**DRIAS - TRACC 2023**](https://www.drias-climat.fr/accompagnement/sections/401) data.  
-I'll answer by displaying a list of SQL queries, graphs and data most relevant to your question.
-
-❓ **How to use?**  
-You can ask me anything about these climate indicators: **temperature**, **precipitation** or **drought**.  
-You can specify **location** and/or **year**.  
-You can choose from a list of climate models. By default, we take the **average of each model**.
-
-For example, you can ask:  
-- 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?  
-
-**Example of indicators in the data**:  
-- Mean temperature (annual, winter, summer)  
-- Total precipitation (annual, winter, summer)  
-- Number of days with remarkable precipitations, with dry ground, with temperature above 30°C  
-
-⚠️ **Limitations**:  
-- You can't ask anything that isn't related to **DRIAS - TRACC 2023** data.  
-- You can only ask about **locations in France**.  
-- If you specify a year, there may be **no data for that year for some models**.  
-- You **cannot compare two models**.  
-
-🛈 **Information**  
-Please note that we **log your questions for meta-analysis purposes**, so avoid sharing any sensitive or personal information.
-"""
-    with gr.Tab("Beta - Talk to DRIAS", elem_id="tab-vanna", id=6):
-
-        with gr.Accordion(label="Details"):
-            gr.Markdown(details_text)
-
-        with gr.Row():
-            drias_direct_question = gr.Textbox(
-                label="Direct Question",
-                placeholder="You can write direct question here",
-                elem_id="direct-question",
-                interactive=True,
-            )
-
-        result_text = gr.Textbox(
-            label="", elem_id="no-result-label", interactive=False, visible=True
-        )
-
-        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:
-            model_selection = gr.Dropdown(
-                label="Model", choices=DRIAS_MODELS, value="ALL", interactive=True
-            )
-            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")
-
-        pagination_display = gr.Markdown(value="", visible=False, elem_id="pagination-display")
-
-        with gr.Row():
-            prev_button = gr.Button("Previous", visible=False)
-            next_button = gr.Button("Next", visible=False)
-
-        sql_queries_state = gr.State([])
-        dataframes_state = gr.State([])
-        plots_state = gr.State([])
-        index_state = gr.State(0)
-
-        drias_direct_question.submit(
-            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,
-                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
-            ],
-        ).then(
-            update_pagination,
-            inputs=[index_state, sql_queries_state],
-            outputs=[pagination_display],
-        )
-
-        model_selection.change(
-            filter_by_model,
-            inputs=[dataframes_state, plots_state, index_state, model_selection],
-            outputs=[drias_table, drias_display],
-        )
-
-        def show_previous(index, sql_queries, dataframes, plots):
-            if index > 0:
-                index -= 1
-            return (
-                sql_queries[index],
-                dataframes[index],
-                plots[index](dataframes[index]),
-                index,
-            )
-
-        def show_next(index, sql_queries, dataframes, plots):
-            if index < len(sql_queries) - 1:
-                index += 1
-            return (
-                sql_queries[index],
-                dataframes[index],
-                plots[index](dataframes[index]),
-                index,
-            )
-
-        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],
-        )
-
-        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],
-        )
-
-
 def config_event_handling(
     main_tabs_components: list[MainTabPanel], config_componenets: ConfigPanel
 ):

climateqa/engine/talk_to_data/main.py

@@ -18,10 +18,13 @@ def ask_drias(query:str, index_state: int = 0):
     sql_queries = []
     result_dataframes = []
     figures = []
+    table_list = []
 
     for plot_state in final_state['plot_states'].values():
         for table_state in plot_state['table_states'].values():
             if table_state['status'] == 'OK':
+                if 'table_name' in table_state: 
+                    table_list.append(' '.join(table_state['table_name'].capitalize().split('_')))
                 if 'sql_query' in table_state and table_state['sql_query'] is not None:
                     sql_queries.append(table_state['sql_query'])
                 
@@ -35,8 +38,9 @@ def ask_drias(query:str, index_state: int = 0):
 
     sql_query = sql_queries[index_state]
     dataframe = result_dataframes[index_state]
-    figure = figures[index_state](dataframe)       
-    return sql_query, dataframe, figure, sql_queries, result_dataframes, figures, index_state, ""
+    figure = figures[index_state](dataframe)
+
+    return sql_query, dataframe, figure, sql_queries, result_dataframes, figures, index_state, table_list, ""
 
 DRIAS_MODELS = [
     'ALL',

climateqa/engine/talk_to_data/workflow.py

@@ -87,19 +87,23 @@ def drias_workflow(user_input: str) -> State:
 
         plot_state['tables'] = relevant_tables
 
+        params = {}
+        for param_name in plot['params']:
+            param = find_param(state, param_name, relevant_tables[0])
+            if param:
+                params.update(param)
+
         for n, table in enumerate(plot_state['tables']):
             if n > 2:
                 break
 
             table_state: TableState = {
                 'table_name': table,
-                'params': {},
+                'params': params,
                 'status': 'OK'
             } 
-            for param_name in plot['params']:
-                param = find_param(state, param_name, table)
-                if param:
-                    table_state['params'].update(param)
+
+            table_state["params"]['indicator_column'] = find_indicator_column(table)
             
             sql_query = plot['sql_query'](table, table_state['params'])
 
@@ -111,7 +115,7 @@ def drias_workflow(user_input: str) -> State:
 
             table_state['sql_query'] = sql_query
             df = execute_sql_query(sql_query)
-            
+
             if len(df) > 0:
                 have_dataframe = True
 
@@ -158,9 +162,9 @@ def find_param(state: State, param_name:str, table: str) -> dict[str, Any] | Non
     if param_name == 'location':
         location = find_location(state['user_input'], table)
         return location
-    if param_name == 'indicator_column':
-        indicator_column = find_indicator_column(table)
-        return {'indicator_column': indicator_column} 
+    # if param_name == 'indicator_column':
+    #     indicator_column = find_indicator_column(table)
+    #     return {'indicator_column': indicator_column} 
     if param_name == 'year':
         year = find_year(state['user_input'])
         return {'year': year}

front/tabs/tab_drias.py

@@ -0,0 +1,242 @@
+import gradio as gr
+
+from climateqa.engine.talk_to_data.main import ask_drias, DRIAS_MODELS
+
+
+def ask_drias_query(query: str, index_state: int):
+    return ask_drias(query, index_state)
+
+
+def show_results(sql_queries_state, dataframes_state, plots_state):
+    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),
+            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(visible=True),
+            gr.update(visible=True),
+            gr.update(visible=True),
+            gr.update(visible=True),
+        )
+
+
+def filter_by_model(dataframes, figures, index_state, model_selection):
+    df = dataframes[index_state]
+    if model_selection != "ALL":
+        df = df[df["model"] == model_selection]
+    figure = figures[index_state](df)
+    return df, figure
+
+
+def update_pagination(index, sql_queries):
+    pagination = f"{index + 1}/{len(sql_queries)}" if sql_queries else ""
+    return pagination
+
+
+def show_previous(index, sql_queries, dataframes, plots):
+    if index > 0:
+        index -= 1
+    return (
+        sql_queries[index],
+        dataframes[index],
+        plots[index](dataframes[index]),
+        index,
+    )
+
+
+def show_next(index, sql_queries, dataframes, plots):
+    if index < len(sql_queries) - 1:
+        index += 1
+    return (
+        sql_queries[index],
+        dataframes[index],
+        plots[index](dataframes[index]),
+        index,
+    )
+
+
+def display_table_names(table_names):
+    return [table_names]
+
+
+def on_table_click(evt: gr.SelectData, table_names, sql_queries, dataframes, plots):
+    index = evt.index[1]
+    figure = plots[index](dataframes[index])
+    return (
+        sql_queries[index],
+        dataframes[index],
+        figure,
+        index,
+    )
+
+
+DETAILS_TEXT = """
+Hi, I'm **Talk to Drias**, designed to answer your questions using [**DRIAS - TRACC 2023**](https://www.drias-climat.fr/accompagnement/sections/401) data.  
+I'll answer by displaying a list of SQL queries, graphs and data most relevant to your question.
+
+❓ **How to use?**  
+You can ask me anything about these climate indicators: **temperature**, **precipitation** or **drought**.  
+You can specify **location** and/or **year**.  
+You can choose from a list of climate models. By default, we take the **average of each model**.
+
+For example, you can ask:  
+- 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?  
+
+**Example of indicators in the data**:  
+- Mean temperature (annual, winter, summer)  
+- Total precipitation (annual, winter, summer)  
+- Number of days with remarkable precipitations, with dry ground, with temperature above 30°C  
+
+⚠️ **Limitations**:  
+- You can't ask anything that isn't related to **DRIAS - TRACC 2023** data.  
+- You can only ask about **locations in France**.  
+- If you specify a year, there may be **no data for that year for some models**.  
+- You **cannot compare two models**.  
+
+🛈 **Information**  
+Please note that we **log your questions for meta-analysis purposes**, so avoid sharing any sensitive or personal information.
+"""
+
+
+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):
+        with gr.Accordion(label="Details"):
+            gr.Markdown(DETAILS_TEXT)
+
+        with gr.Row():
+            drias_direct_question = gr.Textbox(
+                label="Direct Question",
+                placeholder="You can write direct question here",
+                elem_id="direct-question",
+                interactive=True,
+            )
+
+        result_text = gr.Textbox(
+            label="", elem_id="no-result-label", interactive=False, visible=True
+        )
+
+        table_names_display = gr.DataFrame(
+            [], label="List of relevant indicators", headers=None, interactive=False, 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:
+            model_selection = gr.Dropdown(
+                label="Model", choices=DRIAS_MODELS, value="ALL", interactive=True
+            )
+            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")
+
+        pagination_display = gr.Markdown(
+            value="", visible=False, elem_id="pagination-display"
+        )
+
+        with gr.Row():
+            prev_button = gr.Button("Previous", visible=False)
+            next_button = gr.Button("Next", visible=False)
+
+        drias_direct_question.submit(
+            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],
+        )
+
+        model_selection.change(
+            filter_by_model,
+            inputs=[dataframes_state, plots_state, index_state, model_selection],
+            outputs=[drias_table, drias_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],
+        )
+
+        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],
+        )
+
+        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],
+
+        ).then(
+            update_pagination,
+            inputs=[index_state, sql_queries_state],
+            outputs=[pagination_display],
+        )

style.css

@@ -663,4 +663,11 @@ div#tab-vanna{
     text-align: center;
     font-weight: bold;
     font-size: 16px;
-}
+}
+
+#table-names table{
+    overflow: hidden;
+}
+#table-names thead{
+    display: none;
+}