feat: model filtering and UI upgrade for TTD

app.py

@@ -9,14 +9,14 @@ from climateqa.engine.embeddings import get_embeddings_function
 from climateqa.engine.llm import get_llm
 from climateqa.engine.vectorstore import get_pinecone_vectorstore
 from climateqa.engine.reranker import get_reranker
-from climateqa.engine.graph import make_graph_agent,make_graph_agent_poc
+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 import create_config_modal, cqa_tab, create_about_tab
+from front.tabs import MainTabPanel, ConfigPanel
 from front.utils import process_figures
 from gradio_modal import Modal
 
@@ -25,14 +25,14 @@ from utils import create_user_id
 import logging
 
 logging.basicConfig(level=logging.WARNING)
-os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'  # Suppresses INFO and WARNING logs
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"  # Suppresses INFO and WARNING logs
 logging.getLogger().setLevel(logging.WARNING)
 
 
-
 # Load environment variables in local mode
 try:
     from dotenv import load_dotenv
+
     load_dotenv()
 except Exception as e:
     pass
@@ -63,42 +63,105 @@ share_client = service.get_share_client(file_share_name)
 user_id = create_user_id()
 
 
-
 # Create vectorstore and retriever
 embeddings_function = get_embeddings_function()
-vectorstore = get_pinecone_vectorstore(embeddings_function, index_name=os.getenv("PINECONE_API_INDEX"))
-vectorstore_graphs = get_pinecone_vectorstore(embeddings_function, index_name=os.getenv("PINECONE_API_INDEX_OWID"), text_key="description")
-vectorstore_region = get_pinecone_vectorstore(embeddings_function, index_name=os.getenv("PINECONE_API_INDEX_LOCAL_V2"))
+vectorstore = get_pinecone_vectorstore(
+    embeddings_function, index_name=os.getenv("PINECONE_API_INDEX")
+)
+vectorstore_graphs = get_pinecone_vectorstore(
+    embeddings_function,
+    index_name=os.getenv("PINECONE_API_INDEX_OWID"),
+    text_key="description",
+)
+vectorstore_region = get_pinecone_vectorstore(
+    embeddings_function, index_name=os.getenv("PINECONE_API_INDEX_LOCAL_V2")
+)
 
-llm = get_llm(provider="openai",max_tokens = 1024,temperature = 0.0)
+llm = get_llm(provider="openai", max_tokens=1024, temperature=0.0)
 if os.environ["GRADIO_ENV"] == "local":
     reranker = get_reranker("nano")
-else :
+else:
     reranker = get_reranker("large")
 
-agent = make_graph_agent(llm=llm, vectorstore_ipcc=vectorstore, vectorstore_graphs=vectorstore_graphs, vectorstore_region = vectorstore_region, reranker=reranker, threshold_docs=0.2)
-agent_poc = make_graph_agent_poc(llm=llm, vectorstore_ipcc=vectorstore, vectorstore_graphs=vectorstore_graphs, vectorstore_region = vectorstore_region, reranker=reranker, threshold_docs=0, version="v4")#TODO put back default 0.2
-
-#Vanna object
-
-vn = MyVanna(config = {"temperature": 0, "api_key": os.getenv('THEO_API_KEY'), 'model': os.getenv('VANNA_MODEL'), 'pc_api_key': os.getenv('VANNA_PINECONE_API_KEY'), 'index_name': os.getenv('VANNA_INDEX_NAME'), "top_k" : 4})
-db_vanna_path = os.path.join(os.getcwd(), "data/drias/drias.db")
-vn.connect_to_sqlite(db_vanna_path)
+agent = make_graph_agent(
+    llm=llm,
+    vectorstore_ipcc=vectorstore,
+    vectorstore_graphs=vectorstore_graphs,
+    vectorstore_region=vectorstore_region,
+    reranker=reranker,
+    threshold_docs=0.2,
+)
+agent_poc = make_graph_agent_poc(
+    llm=llm,
+    vectorstore_ipcc=vectorstore,
+    vectorstore_graphs=vectorstore_graphs,
+    vectorstore_region=vectorstore_region,
+    reranker=reranker,
+    threshold_docs=0,
+    version="v4",
+)  # TODO put back default 0.2
+
+# Vanna object
+
+# vn = MyVanna(config = {"temperature": 0, "api_key": os.getenv('THEO_API_KEY'), 'model': os.getenv('VANNA_MODEL'), 'pc_api_key': os.getenv('VANNA_PINECONE_API_KEY'), 'index_name': os.getenv('VANNA_INDEX_NAME'), "top_k" : 4})
+# db_vanna_path = os.path.join(os.getcwd(), "data/drias/drias.db")
+# vn.connect_to_sqlite(db_vanna_path)
 
 # def ask_vanna_query(query):
 #     return ask_vanna(vn, db_vanna_path, query)
 
-def ask_drias_query(query: str, index_state: int, drias_model: str):
-    return ask_drias(db_vanna_path, query, index_state, drias_model)
 
-async def chat(query, history, audience, sources, reports, relevant_content_sources_selection, search_only):
+def ask_drias_query(query: str, index_state: int):
+    return ask_drias(query, index_state)
+
+
+async def chat(
+    query,
+    history,
+    audience,
+    sources,
+    reports,
+    relevant_content_sources_selection,
+    search_only,
+):
     print("chat cqa - message received")
-    async for event in chat_stream(agent, query, history, audience, sources, reports, relevant_content_sources_selection, search_only, share_client, user_id):
+    async for event in chat_stream(
+        agent,
+        query,
+        history,
+        audience,
+        sources,
+        reports,
+        relevant_content_sources_selection,
+        search_only,
+        share_client,
+        user_id,
+    ):
         yield event
-        
-async def chat_poc(query, history, audience, sources, reports, relevant_content_sources_selection, search_only):
+
+
+async def chat_poc(
+    query,
+    history,
+    audience,
+    sources,
+    reports,
+    relevant_content_sources_selection,
+    search_only,
+):
     print("chat poc - message received")
-    async for event in chat_stream(agent_poc, query, history, audience, sources, reports, relevant_content_sources_selection, search_only, share_client, user_id):
+    async for event in chat_stream(
+        agent_poc,
+        query,
+        history,
+        audience,
+        sources,
+        reports,
+        relevant_content_sources_selection,
+        search_only,
+        share_client,
+        user_id,
+    ):
         yield event
 
 
@@ -106,14 +169,17 @@ async def chat_poc(query, history, audience, sources, reports, relevant_content_
 # Gradio
 # --------------------------------------------------------------------
 
+
 # Function to update modal visibility
 def update_config_modal_visibility(config_open):
     print(config_open)
     new_config_visibility_status = not config_open
     return Modal(visible=new_config_visibility_status), new_config_visibility_status
-    
 
-def update_sources_number_display(sources_textbox, figures_cards, current_graphs, papers_html):
+
+def update_sources_number_display(
+    sources_textbox, figures_cards, current_graphs, papers_html
+):
     sources_number = sources_textbox.count("<h2>")
     figures_number = figures_cards.count("<h2>")
     graphs_number = current_graphs.count("<iframe")
@@ -122,9 +188,18 @@ def update_sources_number_display(sources_textbox, figures_cards, current_graphs
     figures_notif_label = f"Figures ({figures_number})"
     graphs_notif_label = f"Graphs ({graphs_number})"
     papers_notif_label = f"Papers ({papers_number})"
-    recommended_content_notif_label = f"Recommended content ({figures_number + graphs_number + papers_number})"
+    recommended_content_notif_label = (
+        f"Recommended content ({figures_number + graphs_number + papers_number})"
+    )
+
+    return (
+        gr.update(label=recommended_content_notif_label),
+        gr.update(label=sources_notif_label),
+        gr.update(label=figures_notif_label),
+        gr.update(label=graphs_notif_label),
+        gr.update(label=papers_notif_label),
+    )
 
-    return gr.update(label=recommended_content_notif_label), gr.update(label=sources_notif_label), gr.update(label=figures_notif_label), gr.update(label=graphs_notif_label), gr.update(label=papers_notif_label)
 
 # def create_drias_tab():
 #     with gr.Tab("Beta - Talk to DRIAS", elem_id="tab-vanna", id=6) as tab_vanna:
@@ -141,24 +216,112 @@ def update_sources_number_display(sources_textbox, figures_cards, current_graphs
 #         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)
-            model_selection = gr.Dropdown(label="Model", choices=DRIAS_MODELS ,elem_id="drias-model", value="ALL", interactive=True)
-        
-        with gr.Accordion(label="SQL Query Used"):
-            drias_sql_query = gr.Textbox(label="", elem_id="sql-query", interactive=False)
+            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='Data used', open=False):
-            drias_table = gr.DataFrame([], elem_id="vanna-table")
+        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"):
+        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")
-            next_button = gr.Button("Next")
+            prev_button = gr.Button("Previous", visible=False)
+            next_button = gr.Button("Next", visible=False)
 
         sql_queries_state = gr.State([])
         dataframes_state = gr.State([])
@@ -166,96 +329,104 @@ def create_drias_tab():
         index_state = gr.State(0)
 
         drias_direct_question.submit(
-            ask_drias_query, 
-            inputs=[drias_direct_question, index_state, model_selection], 
-            outputs=[drias_sql_query, drias_table, drias_display, sql_queries_state, dataframes_state, plots_state, index_state]
+            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(
-            ask_drias_query,
-            inputs=[drias_direct_question, index_state, model_selection],
-            outputs=[drias_sql_query, drias_table, drias_display, sql_queries_state, dataframes_state, plots_state, index_state]
+            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], index
+            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], index
+            return (
+                sql_queries[index],
+                dataframes[index],
+                plots[index](dataframes[index]),
+                index,
+            )
 
         prev_button.click(
-            show_previous, 
+            show_previous,
             inputs=[index_state, sql_queries_state, dataframes_state, plots_state],
-            outputs=[drias_sql_query, drias_table, drias_display, index_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, 
+            show_next,
             inputs=[index_state, sql_queries_state, dataframes_state, plots_state],
-            outputs=[drias_sql_query, drias_table, drias_display, index_state]
+            outputs=[drias_sql_query, drias_table, drias_display, index_state],
+        ).then(
+            update_pagination,
+            inputs=[index_state, sql_queries_state],
+            outputs=[pagination_display],
         )
-    
-# # UI Layout Components
-def cqa_tab(tab_name):
-    # State variables
-    current_graphs = gr.State([])
-    with gr.Tab(tab_name):
-        with gr.Row(elem_id="chatbot-row"):
-            # Left column - Chat interface
-            with gr.Column(scale=2):
-                chatbot, textbox, config_button = create_chat_interface(tab_name)
-
-            # Right column - Content panels
-            with gr.Column(scale=2, variant="panel", elem_id="right-panel"):
-                with gr.Tabs(elem_id="right_panel_tab") as tabs:
-                    # Examples tab
-                    with gr.TabItem("Examples", elem_id="tab-examples", id=0):
-                        examples_hidden = create_examples_tab(tab_name)
-
-                    # Sources tab
-                    with gr.Tab("Sources", elem_id="tab-sources", id=1) as tab_sources:
-                        sources_textbox = gr.HTML(show_label=False, elem_id="sources-textbox")
-
-
-                    # Recommended content tab
-                    with gr.Tab("Recommended content", elem_id="tab-recommended_content", id=2) as tab_recommended_content:
-                        with gr.Tabs(elem_id="group-subtabs") as tabs_recommended_content:
-                            # Figures subtab
-                            with gr.Tab("Figures", elem_id="tab-figures", id=3) as tab_figures:
-                                sources_raw, new_figures, used_figures, gallery_component, figures_cards, figure_modal = create_figures_tab()
-
-                            # Papers subtab
-                            with gr.Tab("Papers", elem_id="tab-citations", id=4) as tab_papers:
-                                papers_direct_search, papers_summary, papers_html, citations_network, papers_modal = create_papers_tab()
-
-                            # Graphs subtab
-                            with gr.Tab("Graphs", elem_id="tab-graphs", id=5) as tab_graphs:
-                                graphs_container = gr.HTML(
-                                    "<h2>There are no graphs to be displayed at the moment. Try asking another question.</h2>",
-                                    elem_id="graphs-container"
-                                )
-
-                                
-def config_event_handling(main_tabs_components : list[MainTabPanel], config_componenets : ConfigPanel):
+
+
+def config_event_handling(
+    main_tabs_components: list[MainTabPanel], config_componenets: ConfigPanel
+):
     config_open = config_componenets.config_open
     config_modal = config_componenets.config_modal
     close_config_modal = config_componenets.close_config_modal_button
-    
-    for button in [close_config_modal] + [main_tab_component.config_button for main_tab_component in main_tabs_components]:
+
+    for button in [close_config_modal] + [
+        main_tab_component.config_button for main_tab_component in main_tabs_components
+    ]:
         button.click(
             fn=update_config_modal_visibility,
             inputs=[config_open],
-            outputs=[config_modal, config_open]
-        ) 
-    
+            outputs=[config_modal, config_open],
+        )
+
+
 def event_handling(
-    main_tab_components : MainTabPanel,
-    config_components : ConfigPanel,
-    tab_name="ClimateQ&A"
+    main_tab_components: MainTabPanel,
+    config_components: ConfigPanel,
+    tab_name="ClimateQ&A",
 ):
     chatbot = main_tab_components.chatbot
     textbox = main_tab_components.textbox
@@ -279,7 +450,7 @@ def event_handling(
     graphs_container = main_tab_components.graph_container
     follow_up_examples = main_tab_components.follow_up_examples
     follow_up_examples_hidden = main_tab_components.follow_up_examples_hidden
-    
+
     dropdown_sources = config_components.dropdown_sources
     dropdown_reports = config_components.dropdown_reports
     dropdown_external_sources = config_components.dropdown_external_sources
@@ -288,91 +459,302 @@ def event_handling(
     after = config_components.after
     output_query = config_components.output_query
     output_language = config_components.output_language
-    
+
     new_sources_hmtl = gr.State([])
     ttd_data = gr.State([])
 
-    
     if tab_name == "ClimateQ&A":
         print("chat cqa - message sent")
 
         # Event for textbox
-        (textbox
-            .submit(start_chat, [textbox, chatbot, search_only], [textbox, tabs, chatbot, sources_raw], queue=False, api_name=f"start_chat_{textbox.elem_id}")
-            .then(chat, [textbox, chatbot, dropdown_audience, dropdown_sources, dropdown_reports, dropdown_external_sources, search_only], [chatbot, new_sources_hmtl, output_query, output_language, new_figures, current_graphs, follow_up_examples.dataset], concurrency_limit=8, api_name=f"chat_{textbox.elem_id}")
-            .then(finish_chat, None, [textbox], api_name=f"finish_chat_{textbox.elem_id}")
+        (
+            textbox.submit(
+                start_chat,
+                [textbox, chatbot, search_only],
+                [textbox, tabs, chatbot, sources_raw],
+                queue=False,
+                api_name=f"start_chat_{textbox.elem_id}",
+            )
+            .then(
+                chat,
+                [
+                    textbox,
+                    chatbot,
+                    dropdown_audience,
+                    dropdown_sources,
+                    dropdown_reports,
+                    dropdown_external_sources,
+                    search_only,
+                ],
+                [
+                    chatbot,
+                    new_sources_hmtl,
+                    output_query,
+                    output_language,
+                    new_figures,
+                    current_graphs,
+                    follow_up_examples.dataset,
+                ],
+                concurrency_limit=8,
+                api_name=f"chat_{textbox.elem_id}",
+            )
+            .then(
+                finish_chat, None, [textbox], api_name=f"finish_chat_{textbox.elem_id}"
+            )
         )
         # Event for examples_hidden
-        (examples_hidden
-            .change(start_chat, [examples_hidden, chatbot, search_only], [examples_hidden, tabs, chatbot, sources_raw], queue=False, api_name=f"start_chat_{examples_hidden.elem_id}")
-            .then(chat, [examples_hidden, chatbot, dropdown_audience, dropdown_sources, dropdown_reports, dropdown_external_sources, search_only], [chatbot, new_sources_hmtl, output_query, output_language, new_figures, current_graphs,follow_up_examples.dataset], concurrency_limit=8, api_name=f"chat_{examples_hidden.elem_id}")
-            .then(finish_chat, None, [textbox], api_name=f"finish_chat_{examples_hidden.elem_id}")
+        (
+            examples_hidden.change(
+                start_chat,
+                [examples_hidden, chatbot, search_only],
+                [examples_hidden, tabs, chatbot, sources_raw],
+                queue=False,
+                api_name=f"start_chat_{examples_hidden.elem_id}",
+            )
+            .then(
+                chat,
+                [
+                    examples_hidden,
+                    chatbot,
+                    dropdown_audience,
+                    dropdown_sources,
+                    dropdown_reports,
+                    dropdown_external_sources,
+                    search_only,
+                ],
+                [
+                    chatbot,
+                    new_sources_hmtl,
+                    output_query,
+                    output_language,
+                    new_figures,
+                    current_graphs,
+                    follow_up_examples.dataset,
+                ],
+                concurrency_limit=8,
+                api_name=f"chat_{examples_hidden.elem_id}",
+            )
+            .then(
+                finish_chat,
+                None,
+                [textbox],
+                api_name=f"finish_chat_{examples_hidden.elem_id}",
+            )
         )
-        (follow_up_examples_hidden
-            .change(start_chat, [follow_up_examples_hidden, chatbot, search_only], [follow_up_examples_hidden, tabs, chatbot, sources_raw], queue=False, api_name=f"start_chat_{examples_hidden.elem_id}")
-            .then(chat, [follow_up_examples_hidden, chatbot, dropdown_audience, dropdown_sources, dropdown_reports, dropdown_external_sources, search_only], [chatbot, new_sources_hmtl, output_query, output_language, new_figures, current_graphs,follow_up_examples.dataset], concurrency_limit=8, api_name=f"chat_{examples_hidden.elem_id}")
-            .then(finish_chat, None, [textbox], api_name=f"finish_chat_{follow_up_examples_hidden.elem_id}")
+        (
+            follow_up_examples_hidden.change(
+                start_chat,
+                [follow_up_examples_hidden, chatbot, search_only],
+                [follow_up_examples_hidden, tabs, chatbot, sources_raw],
+                queue=False,
+                api_name=f"start_chat_{examples_hidden.elem_id}",
+            )
+            .then(
+                chat,
+                [
+                    follow_up_examples_hidden,
+                    chatbot,
+                    dropdown_audience,
+                    dropdown_sources,
+                    dropdown_reports,
+                    dropdown_external_sources,
+                    search_only,
+                ],
+                [
+                    chatbot,
+                    new_sources_hmtl,
+                    output_query,
+                    output_language,
+                    new_figures,
+                    current_graphs,
+                    follow_up_examples.dataset,
+                ],
+                concurrency_limit=8,
+                api_name=f"chat_{examples_hidden.elem_id}",
+            )
+            .then(
+                finish_chat,
+                None,
+                [textbox],
+                api_name=f"finish_chat_{follow_up_examples_hidden.elem_id}",
+            )
         )
-    
+
     elif tab_name == "Beta - POC Adapt'Action":
         print("chat poc - message sent")
         # Event for textbox
-        (textbox
-            .submit(start_chat, [textbox, chatbot, search_only], [textbox, tabs, chatbot, sources_raw], queue=False, api_name=f"start_chat_{textbox.elem_id}")
-            .then(chat_poc, [textbox, chatbot, dropdown_audience, dropdown_sources, dropdown_reports, dropdown_external_sources, search_only], [chatbot, new_sources_hmtl, output_query, output_language, new_figures, current_graphs], concurrency_limit=8, api_name=f"chat_{textbox.elem_id}")
-            .then(finish_chat, None, [textbox], api_name=f"finish_chat_{textbox.elem_id}")
+        (
+            textbox.submit(
+                start_chat,
+                [textbox, chatbot, search_only],
+                [textbox, tabs, chatbot, sources_raw],
+                queue=False,
+                api_name=f"start_chat_{textbox.elem_id}",
+            )
+            .then(
+                chat_poc,
+                [
+                    textbox,
+                    chatbot,
+                    dropdown_audience,
+                    dropdown_sources,
+                    dropdown_reports,
+                    dropdown_external_sources,
+                    search_only,
+                ],
+                [
+                    chatbot,
+                    new_sources_hmtl,
+                    output_query,
+                    output_language,
+                    new_figures,
+                    current_graphs,
+                ],
+                concurrency_limit=8,
+                api_name=f"chat_{textbox.elem_id}",
+            )
+            .then(
+                finish_chat, None, [textbox], api_name=f"finish_chat_{textbox.elem_id}"
+            )
         )
         # Event for examples_hidden
-        (examples_hidden
-            .change(start_chat, [examples_hidden, chatbot, search_only], [examples_hidden, tabs, chatbot, sources_raw], queue=False, api_name=f"start_chat_{examples_hidden.elem_id}")
-            .then(chat_poc, [examples_hidden, chatbot, dropdown_audience, dropdown_sources, dropdown_reports, dropdown_external_sources, search_only], [chatbot, new_sources_hmtl, output_query, output_language, new_figures, current_graphs], concurrency_limit=8, api_name=f"chat_{examples_hidden.elem_id}")
-            .then(finish_chat, None, [textbox], api_name=f"finish_chat_{examples_hidden.elem_id}")
+        (
+            examples_hidden.change(
+                start_chat,
+                [examples_hidden, chatbot, search_only],
+                [examples_hidden, tabs, chatbot, sources_raw],
+                queue=False,
+                api_name=f"start_chat_{examples_hidden.elem_id}",
+            )
+            .then(
+                chat_poc,
+                [
+                    examples_hidden,
+                    chatbot,
+                    dropdown_audience,
+                    dropdown_sources,
+                    dropdown_reports,
+                    dropdown_external_sources,
+                    search_only,
+                ],
+                [
+                    chatbot,
+                    new_sources_hmtl,
+                    output_query,
+                    output_language,
+                    new_figures,
+                    current_graphs,
+                ],
+                concurrency_limit=8,
+                api_name=f"chat_{examples_hidden.elem_id}",
+            )
+            .then(
+                finish_chat,
+                None,
+                [textbox],
+                api_name=f"finish_chat_{examples_hidden.elem_id}",
+            )
         )
-        (follow_up_examples_hidden
-            .change(start_chat, [follow_up_examples_hidden, chatbot, search_only], [follow_up_examples_hidden, tabs, chatbot, sources_raw], queue=False, api_name=f"start_chat_{examples_hidden.elem_id}")
-            .then(chat, [follow_up_examples_hidden, chatbot, dropdown_audience, dropdown_sources, dropdown_reports, dropdown_external_sources, search_only], [chatbot, new_sources_hmtl, output_query, output_language, new_figures, current_graphs,follow_up_examples.dataset], concurrency_limit=8, api_name=f"chat_{examples_hidden.elem_id}")
-            .then(finish_chat, None, [textbox], api_name=f"finish_chat_{follow_up_examples_hidden.elem_id}")
+        (
+            follow_up_examples_hidden.change(
+                start_chat,
+                [follow_up_examples_hidden, chatbot, search_only],
+                [follow_up_examples_hidden, tabs, chatbot, sources_raw],
+                queue=False,
+                api_name=f"start_chat_{examples_hidden.elem_id}",
+            )
+            .then(
+                chat,
+                [
+                    follow_up_examples_hidden,
+                    chatbot,
+                    dropdown_audience,
+                    dropdown_sources,
+                    dropdown_reports,
+                    dropdown_external_sources,
+                    search_only,
+                ],
+                [
+                    chatbot,
+                    new_sources_hmtl,
+                    output_query,
+                    output_language,
+                    new_figures,
+                    current_graphs,
+                    follow_up_examples.dataset,
+                ],
+                concurrency_limit=8,
+                api_name=f"chat_{examples_hidden.elem_id}",
+            )
+            .then(
+                finish_chat,
+                None,
+                [textbox],
+                api_name=f"finish_chat_{follow_up_examples_hidden.elem_id}",
+            )
         )
-    
-    
-    new_sources_hmtl.change(lambda x : x, inputs = [new_sources_hmtl], outputs = [sources_textbox])
-    current_graphs.change(lambda x: x, inputs=[current_graphs], outputs=[graphs_container])
-    new_figures.change(process_figures, inputs=[sources_raw, new_figures], outputs=[sources_raw, figures_cards, gallery_component])
+
+    new_sources_hmtl.change(
+        lambda x: x, inputs=[new_sources_hmtl], outputs=[sources_textbox]
+    )
+    current_graphs.change(
+        lambda x: x, inputs=[current_graphs], outputs=[graphs_container]
+    )
+    new_figures.change(
+        process_figures,
+        inputs=[sources_raw, new_figures],
+        outputs=[sources_raw, figures_cards, gallery_component],
+    )
 
     # Update sources numbers
     for component in [sources_textbox, figures_cards, current_graphs, papers_html]:
-        component.change(update_sources_number_display, [sources_textbox, figures_cards, current_graphs, papers_html], [tab_recommended_content, tab_sources, tab_figures, tab_graphs, tab_papers])
-    
+        component.change(
+            update_sources_number_display,
+            [sources_textbox, figures_cards, current_graphs, papers_html],
+            [tab_recommended_content, tab_sources, tab_figures, tab_graphs, tab_papers],
+        )
+
     # Search for papers
     for component in [textbox, examples_hidden, papers_direct_search]:
-        component.submit(find_papers, [component, after, dropdown_external_sources], [papers_html, citations_network, papers_summary])
-        
+        component.submit(
+            find_papers,
+            [component, after, dropdown_external_sources],
+            [papers_html, citations_network, papers_summary],
+        )
 
     # if tab_name == "Beta - POC Adapt'Action": # Not untill results are good enough
     #     # Drias search
     #     textbox.submit(ask_vanna, [textbox], [vanna_sql_query ,vanna_table, vanna_display])
 
+
 def main_ui():
     # config_open = gr.State(True)
-    with gr.Blocks(title="Climate Q&A", css_paths=os.getcwd()+ "/style.css", theme=theme, elem_id="main-component") as demo: 
-        config_components = create_config_modal()  
-             
+    with gr.Blocks(
+        title="Climate Q&A",
+        css_paths=os.getcwd() + "/style.css",
+        theme=theme,
+        elem_id="main-component",
+    ) as demo:
+        config_components = create_config_modal()
+
         with gr.Tabs():
-            cqa_components = cqa_tab(tab_name = "ClimateQ&A")
-            local_cqa_components = cqa_tab(tab_name = "Beta - POC Adapt'Action")
+            cqa_components = cqa_tab(tab_name="ClimateQ&A")
+            local_cqa_components = cqa_tab(tab_name="Beta - POC Adapt'Action")
             create_drias_tab()
-        
+
             create_about_tab()
-        
-        event_handling(cqa_components, config_components, tab_name = 'ClimateQ&A')
-        event_handling(local_cqa_components, config_components, tab_name = "Beta - POC Adapt'Action")
-        
-        config_event_handling([cqa_components,local_cqa_components] ,config_components)
-        
+
+        event_handling(cqa_components, config_components, tab_name="ClimateQ&A")
+        event_handling(
+            local_cqa_components, config_components, tab_name="Beta - POC Adapt'Action"
+        )
+
+        config_event_handling([cqa_components, local_cqa_components], config_components)
+
         demo.queue()
-        
+
     return demo
 
-    
+
 demo = main_ui()
 demo.launch(ssr_mode=False)

climateqa/engine/talk_to_data/config.py

@@ -0,0 +1,33 @@
+DRIAS_TABLES = [
+    "total_winter_precipitation",
+    "total_summer_precipiation",
+    "total_annual_precipitation",
+    "total_remarkable_daily_precipitation",
+    "frequency_of_remarkable_daily_precipitation",
+    "extreme_precipitation_intensity",
+    "mean_winter_temperature",
+    "mean_summer_temperature",
+    "mean_annual_temperature",
+    "number_of_tropical_nights",
+    "maximum_summer_temperature",
+    "number_of_days_with_tx_above_30",
+    "number_of_days_with_tx_above_35",
+    "number_of_days_with_a_dry_ground",
+]
+
+INDICATOR_COLUMNS_PER_TABLE = {
+    "total_winter_precipitation": "total_winter_precipitation",
+    "total_summer_precipiation": "total_summer_precipitation",
+    "total_annual_precipitation": "total_annual_precipitation",
+    "total_remarkable_daily_precipitation": "total_remarkable_daily_precipitation",
+    "frequency_of_remarkable_daily_precipitation": "frequency_of_remarkable_daily_precipitation",
+    "extreme_precipitation_intensity": "extreme_precipitation_intensity",
+    "mean_winter_temperature": "mean_winter_temperature",
+    "mean_summer_temperature": "mean_summer_temperature",
+    "mean_annual_temperature": "mean_annual_temperature",
+    "number_of_tropical_nights": "number_tropical_nights",
+    "maximum_summer_temperature": "maximum_summer_temperature",
+    "number_of_days_with_tx_above_30": "number_of_days_with_tx_above_30",
+    "number_of_days_with_tx_above_35": "number_of_days_with_tx_above_35",
+    "number_of_days_with_a_dry_ground": "number_of_days_with_dry_ground"
+}

climateqa/engine/talk_to_data/main.py

@@ -13,8 +13,8 @@ def ask_llm_column_names(sql_query, llm):
     columns_list = ast.literal_eval(columns.strip("```python\n").strip())
     return columns_list
 
-def ask_drias(db_drias_path:str, query:str, index_state: int = 0, drias_model: str = "ALL"):
-    final_state = drias_workflow(db_drias_path, query, drias_model)
+def ask_drias(query:str, index_state: int = 0):
+    final_state = drias_workflow(query)
     sql_queries = []
     result_dataframes = []
     figures = []
@@ -28,10 +28,15 @@ def ask_drias(db_drias_path:str, query:str, index_state: int = 0, drias_model: s
                 if 'dataframe' in table_state and table_state['dataframe'] is not None:
                     result_dataframes.append(table_state['dataframe'])
                     if 'figure' in table_state and table_state['figure'] is not None:
-                        figures.append(table_state['figure'](table_state['dataframe']))
+                        figures.append(table_state['figure'])
 
-                
-    return sql_queries[index_state], result_dataframes[index_state], figures[index_state], sql_queries, result_dataframes, figures, index_state
+    if "error" in final_state and final_state["error"] != "":
+        return None, None, None, [], [], [], 0, final_state["error"]
+
+    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, ""
 
 DRIAS_MODELS = [
     'ALL',

climateqa/engine/talk_to_data/plot.py

@@ -29,7 +29,6 @@ def plot_indicator_evolution_at_location(params: dict) -> Callable[..., Figure]:
         Callable[..., Figure]: Function which can be call to create the figure with the associated dataframe
     """
     indicator = params["indicator_column"]
-    model = params["model"]
     location = params["location"]
     indicator_label = " ".join([word.capitalize() for word in indicator.split("_")])
 
@@ -43,7 +42,7 @@ def plot_indicator_evolution_at_location(params: dict) -> Callable[..., Figure]:
             Figure: Plotly figure
         """
         fig = go.Figure()
-        if model == "ALL":
+        if df['model'].nunique() != 1:
             df_avg = df.groupby("year", as_index=False)[indicator].mean()
 
             # Transform to list to avoid pandas encoding
@@ -58,8 +57,10 @@ def plot_indicator_evolution_at_location(params: dict) -> Callable[..., Figure]:
                 .astype(float)
                 .tolist()
             )
+            model_label = "Model Average"
+
         else:
-            df_model = df[df["model"] == model]
+            df_model = df
 
             # Transform to list to avoid pandas encoding
             indicators = df_model[indicator].astype(float).tolist()
@@ -73,6 +74,8 @@ def plot_indicator_evolution_at_location(params: dict) -> Callable[..., Figure]:
                 .astype(float)
                 .tolist()
             )
+            model_label = f"Model : {df['model'].unique()[0]}"
+
 
         # Indicator per year plot
         fig.add_scatter(
@@ -93,7 +96,7 @@ def plot_indicator_evolution_at_location(params: dict) -> Callable[..., Figure]:
             marker=dict(color="#d62728"),
         )
         fig.update_layout(
-            title=f"Plot of {indicator_label} in {location} {'(Model Average)' if model == 'ALL' else '(Model : ' + model + ')'}",
+            title=f"Plot of {indicator_label} in {location} ({model_label})",
             xaxis_title="Year",
             yaxis_title=indicator_label,
             template="plotly_white",
@@ -125,7 +128,6 @@ def plot_indicator_number_of_days_per_year_at_location(
     """
 
     indicator = params["indicator_column"]
-    model = params["model"]
     location = params["location"]
 
     def plot_data(df: pd.DataFrame) -> Figure:
@@ -138,19 +140,21 @@ def plot_indicator_number_of_days_per_year_at_location(
             Figure: Plotly figure
         """
         fig = go.Figure()
-        if model == "ALL":
+        if df['model'].nunique() != 1:
             df_avg = df.groupby("year", as_index=False)[indicator].mean()
 
             # Transform to list to avoid pandas encoding
             indicators = df_avg[indicator].astype(float).tolist()
             years = df_avg["year"].astype(int).tolist()
+            model_label = "Model Average"
 
         else:
-            df_model = df[df["model"] == model]
-
+            df_model = df
             # Transform to list to avoid pandas encoding
             indicators = df_model[indicator].astype(float).tolist()
             years = df_model["year"].astype(int).tolist()
+            model_label = f"Model : {df['model'].unique()[0]}"
+
 
         # Bar plot
         fig.add_trace(
@@ -165,7 +169,7 @@ def plot_indicator_number_of_days_per_year_at_location(
         indicator_label = " ".join([word.capitalize() for word in indicator.split("_")])
 
         fig.update_layout(
-            title=f"{indicator_label} in {location} {'(Model Average)' if model == 'ALL' else '(Model : ' + model + ')'}",
+            title=f"{indicator_label} in {location} ({model_label})",
             xaxis_title="Year",
             yaxis_title=indicator,
             yaxis=dict(range=[0, max(indicators)]),
@@ -199,7 +203,6 @@ def plot_distribution_of_indicator_for_given_year(
         Callable[..., Figure]: Function which can be call to create the figure with the associated dataframe
     """
     indicator = params["indicator_column"]
-    model = params["model"]
     year = params["year"]
     indicator_label = " ".join([word.capitalize() for word in indicator.split("_")])
 
@@ -213,18 +216,22 @@ def plot_distribution_of_indicator_for_given_year(
             Figure: Plotly figure
         """
         fig = go.Figure()
-        if params["model"] == "ALL":
+        if df['model'].nunique() != 1:
             df_avg = df.groupby(["latitude", "longitude"], as_index=False)[
                 indicator
             ].mean()
 
             # Transform to list to avoid pandas encoding
             indicators = df_avg[indicator].astype(float).tolist()
+            model_label = "Model Average"
+
         else:
-            df_model = df[df["model"] == model]
+            df_model = df
 
             # Transform to list to avoid pandas encoding
             indicators = df_model[indicator].astype(float).tolist()
+            model_label = f"Model : {df['model'].unique()[0]}"
+
 
         fig.add_trace(
             go.Histogram(
@@ -236,7 +243,7 @@ def plot_distribution_of_indicator_for_given_year(
         )
 
         fig.update_layout(
-            title=f"Distribution of {indicator_label} in {year} {'(Model Average)' if model == 'ALL' else '(Model : ' + model + ')'}",
+            title=f"Distribution of {indicator_label} in {year} ({model_label})",
             xaxis_title=indicator_label,
             yaxis_title="Frequency",
             plot_bgcolor="rgba(0, 0, 0, 0)",
@@ -270,13 +277,12 @@ def plot_map_of_france_of_indicator_for_given_year(
     """
 
     indicator = params["indicator_column"]
-    model = params["model"]
     year = params["year"]
     indicator_label = " ".join([word.capitalize() for word in indicator.split("_")])
 
     def plot_data(df: pd.DataFrame) -> Figure:
         fig = go.Figure()
-        if model == "ALL":
+        if df['model'].nunique() != 1:
             df_avg = df.groupby(["latitude", "longitude"], as_index=False)[
                 indicator
             ].mean()
@@ -284,14 +290,17 @@ def plot_map_of_france_of_indicator_for_given_year(
             indicators = df_avg[indicator].astype(float).tolist()
             latitudes = df_avg["latitude"].astype(float).tolist()
             longitudes = df_avg["longitude"].astype(float).tolist()
+            model_label = "Model Average"
 
         else:
-            df_model = df[df["model"] == model]
+            df_model = df
 
             # Transform to list to avoid pandas encoding
             indicators = df_model[indicator].astype(float).tolist()
             latitudes = df_model["latitude"].astype(float).tolist()
             longitudes = df_model["longitude"].astype(float).tolist()
+            model_label = f"Model : {df['model'].unique()[0]}"
+
 
         fig.add_trace(
             go.Scattermapbox(
@@ -314,7 +323,7 @@ def plot_map_of_france_of_indicator_for_given_year(
             mapbox_zoom=3,
             mapbox_center={"lat": 46.6, "lon": 2.0},
             coloraxis_colorbar=dict(title=f"{indicator_label}"),  # Add legend
-            title=f"{indicator_label} in {year} in France {'(Model Average)' if model == 'ALL' else '(Model : ' + model + ')'} " # Title
+            title=f"{indicator_label} in {year} in France ({model_label}) " # Title
         )
         return fig
 

climateqa/engine/talk_to_data/sql_query.py

@@ -1,41 +1,23 @@
-import sqlite3
-from typing import Any, TypedDict
+from typing import TypedDict
+import duckdb
+import pandas as pd
 
-
-class SqlQueryOutput(TypedDict):
-    labels: list[str]
-    data: list[list[Any]]
-
-
-def execute_sql_query(db_path: str, sql_query: str) -> SqlQueryOutput:
+def execute_sql_query(sql_query: str) -> pd.DataFrame:
     """Execute the SQL Query on the sqlite database
 
     Args:
-        db_ (str): path to the sqlite database
         sql_query (str): sql query to execute
 
     Returns:
         SqlQueryOutput: labels of the selected column and fetched data
     """
 
-    # Connect to sqlite3 database
-    conn = sqlite3.connect(db_path)
-    cursor = conn.cursor()
 
     # Execute the query
-    cursor.execute(sql_query)
-
-    # Fetch labels of selected columns
-    labels = [desc[0] for desc in cursor.description]
+    results = duckdb.sql(sql_query)
 
-    # Fetch data
-    data = cursor.fetchall()
-    conn.close()
-
-    return {
-        "labels": labels,
-        "data": data,
-    }
+    # return fetched data
+    return results.fetchdf()
 
 
 class IndicatorPerYearAtLocationQueryParams(TypedDict, total=False):
@@ -60,15 +42,13 @@ def indicator_per_year_at_location_query(
     indicator_column = params.get("indicator_column")
     latitude = params.get("latitude")
     longitude = params.get("longitude")
-    model = params.get('model')
     
     if indicator_column is None or latitude is None or longitude is None: # If one parameter is missing, returns an empty query
         return ""
     
-    if model == 'ALL':
-        sql_query = f"SELECT year, {indicator_column}, model\nFROM {table}\nWHERE latitude = {latitude} \nAnd longitude = {longitude} \nOrder by Year"
-    else:
-        sql_query = f"SELECT year, {indicator_column}, model\nFROM {table}\nWHERE latitude = {latitude} \nAnd longitude = {longitude} \nAnd model = '{model}' \nOrder by Year"
+    table = f"'hf://datasets/timeki/drias_db/{table.lower()}.parquet'"
+
+    sql_query = f"SELECT year, {indicator_column}, model\nFROM {table}\nWHERE latitude = {latitude} \nAnd longitude = {longitude} \nOrder by Year"
 
     return sql_query
 
@@ -91,12 +71,10 @@ def indicator_for_given_year_query(
     """
     indicator_column = params.get("indicator_column")
     year = params.get('year')
-    model = params.get('model')
     if year is None or indicator_column is None: # If one parameter is missing, returns an empty query
         return ""
     
-    if model == 'ALL':
-        sql_query = f"Select {indicator_column}, latitude, longitude, model\nFrom {table}\nWhere year = {year}"
-    else: 
-        sql_query = f"Select {indicator_column}, latitude, longitude, model\nFrom {table}\nWhere year = {year}\nAnd model = '{model}'"
+    table = f"'hf://datasets/timeki/drias_db/{table.lower()}.parquet'"
+
+    sql_query = f"Select {indicator_column}, latitude, longitude, model\nFrom {table}\nWhere year = {year}"
     return sql_query    

climateqa/engine/talk_to_data/utils.py

@@ -1,11 +1,10 @@
 import re
 from typing import Annotated, TypedDict
-
-from sympy import use
+import duckdb
 from geopy.geocoders import Nominatim
-import sqlite3
 import ast
 from climateqa.engine.llm import get_llm
+from climateqa.engine.talk_to_data.config import DRIAS_TABLES
 from climateqa.engine.talk_to_data.plot import PLOTS, Plot
 from langchain_core.prompts import ChatPromptTemplate
 
@@ -35,7 +34,7 @@ class ArrayOutput(TypedDict):
 
     array: Annotated[str, ..., "Syntactically valid python array."]
 
-def detect_year_with_openai(sentence: str):
+def detect_year_with_openai(sentence: str) -> str:
     """
     Detects years in a sentence using OpenAI's API via LangChain.
     """
@@ -56,7 +55,7 @@ def detect_year_with_openai(sentence: str):
     if len(years_list) > 0:
         return years_list[0]
     else:
-        return None
+        return ""
 
 
 def detectTable(sql_query):
@@ -81,24 +80,26 @@ def coords2loc(coords: tuple):
         return "Unknown Location"
 
 
-def nearestNeighbourSQL(db: str, location: tuple, table: str):
-    conn = sqlite3.connect(db)
+def nearestNeighbourSQL(location: tuple, table: str) -> tuple[str, str]:
     long = round(location[1], 3)
     lat = round(location[0], 3)
-    cursor = conn.cursor()
-    cursor.execute(
+
+    table = f"'hf://datasets/timeki/drias_db/{table.lower()}.parquet'"
+
+    results = duckdb.sql(
         f"SELECT latitude, longitude FROM {table} WHERE latitude BETWEEN {lat - 0.3} AND {lat + 0.3} AND longitude BETWEEN {long - 0.3} AND {long + 0.3}"
-    )
+    ).fetchdf()
+
+    if len(results) == 0:
+        return "", ""
     # cursor.execute(f"SELECT latitude, longitude FROM {table} WHERE latitude BETWEEN {lat - 0.3} AND {lat + 0.3} AND longitude BETWEEN {long - 0.3} AND {long + 0.3}")
-    results = cursor.fetchall()
-    return results[0]
+    return results['latitude'].iloc[0], results['longitude'].iloc[0]
 
 
-def detect_relevant_tables(db: str, user_question: str, plot: Plot, llm) -> list[str]:
+def detect_relevant_tables(user_question: str, plot: Plot, llm) -> list[str]:
     """Detect relevant tables regarding the plot and the user input
 
     Args:
-        db (str): database path
         user_question (str): initial user input
         plot (Plot): plot object for which we wanna plot
         llm (_type_): LLM
@@ -106,19 +107,21 @@ def detect_relevant_tables(db: str, user_question: str, plot: Plot, llm) -> list
     Returns:
         list[str]: list of table names
     """
-    conn = sqlite3.connect(db)
-    cursor = conn.cursor()
 
     # Get all table names
-    cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
-    table_names_list = cursor.fetchall()
+    table_names_list = DRIAS_TABLES
 
     prompt = (
         f"You are helping to build a plot following this description : {plot['description']}."
+        f"You are given a list of tables and a user question."
         f"Based on the description of the plot, which table are appropriate for that kind of plot."
-        f"The different tables are {table_names_list}."
-        f"The user question is {user_question}. Write the relevant tables to use. Answer only a python list of table name."
+        f"Write the 3 most relevant tables to use. Answer only a python list of table name."
+        f"### List of tables : {table_names_list}"
+        f"### User question : {user_question}"
+        f"### List of table name : "
     )
+
+
     table_names = ast.literal_eval(
         llm.invoke(prompt).content.strip("```python\n").strip()
     )
@@ -141,17 +144,28 @@ def detect_relevant_plots(user_question: str, llm):
         plots_description += " - Description: " + plot["description"] + "\n"
 
     prompt = (
-        f"You are helping to answer a question with insightful visualizations. "
-        f"Given a list of plots with their name and description: "
-        f"{plots_description} "
-        f"The user question is: {user_question}. "
-        f"Choose the most relevant plots to answer the question. "
-        f"The answer must be a Python list with the names of the relevant plots, and nothing else. "
-        f"Ensure the response is in the exact format: ['PlotName1', 'PlotName2']."
+        f"You are helping to answer a quesiton with insightful visualizations."
+        f"You are given an user question and a list of plots with their name and description."
+        f"Based on the descriptions of the plots, which plot is appropriate to answer to this question."
+        f"Write the most relevant tables to use. Answer only a python list of plot name."
+        f"### Descriptions of the plots : {plots_description}"
+        f"### User question : {user_question}"
+        f"### Name of the plot : "
     )
-
-    response = llm.invoke(prompt).content
-    return eval(response)
+    # prompt = (
+    #     f"You are helping to answer a question with insightful visualizations. "
+    #     f"Given a list of plots with their name and description: "
+    #     f"{plots_description} "
+    #     f"The user question is: {user_question}. "
+    #     f"Choose the most relevant plots to answer the question. "
+    #     f"The answer must be a Python list with the names of the relevant plots, and nothing else. "
+    #     f"Ensure the response is in the exact format: ['PlotName1', 'PlotName2']."
+    # )
+
+    plot_names = ast.literal_eval(
+        llm.invoke(prompt).content.strip("```python\n").strip()
+    )
+    return plot_names
 
 
 # Next Version

climateqa/engine/talk_to_data/workflow.py

@@ -5,6 +5,7 @@ import pandas as pd
 
 from plotly.graph_objects import Figure
 from climateqa.engine.llm import get_llm
+from climateqa.engine.talk_to_data.config import INDICATOR_COLUMNS_PER_TABLE
 from climateqa.engine.talk_to_data.plot import PLOTS, Plot
 from climateqa.engine.talk_to_data.sql_query import execute_sql_query
 from climateqa.engine.talk_to_data.utils import (
@@ -37,12 +38,12 @@ class State(TypedDict):
     user_input: str
     plots: list[str]
     plot_states: dict[str, PlotState]
+    error: NotRequired[str]
     
-def drias_workflow(db_drias_path: str, user_input: str, model: str) -> State:
+def drias_workflow(user_input: str) -> State:
     """Performs the complete workflow of Talk To Drias : from user input to sql queries, dataframes and figures generated
 
     Args:
-        db_drias_path (str): path to the drias database
         user_input (str): initial user input
 
     Returns:
@@ -60,8 +61,12 @@ def drias_workflow(db_drias_path: str, user_input: str, model: str) -> State:
     state['plots'] = plots
 
     if not state['plots']:
+        state['error'] = 'There is no plot to answer to the question'
         return state
 
+    have_relevant_table = False
+    have_sql_query = False
+    have_dataframe = False
     for plot_name in state['plots']:
 
         plot = next((p for p in PLOTS if p['name'] == plot_name), None) # Find the associated plot object
@@ -76,21 +81,23 @@ def drias_workflow(db_drias_path: str, user_input: str, model: str) -> State:
 
         plot_state['plot_name'] = plot_name
 
-        relevant_tables = find_relevant_tables_per_plot(state, plot, db_drias_path, llm)
+        relevant_tables = find_relevant_tables_per_plot(state, plot, llm)
+        if len(relevant_tables) > 0 :
+            have_relevant_table = True
 
         plot_state['tables'] = relevant_tables
 
-        for table in plot_state['tables']:
+        for n, table in enumerate(plot_state['tables']):
+            if n > 2:
+                break
+
             table_state: TableState = {
                 'table_name': table,
                 'params': {},
                 'status': 'OK'
             } 
-            table_state['params'] = {
-                'model': model
-            }
             for param_name in plot['params']:
-                param = find_param(state, param_name, table, db_drias_path)
+                param = find_param(state, param_name, table)
                 if param:
                     table_state['params'].update(param)
             
@@ -99,17 +106,30 @@ def drias_workflow(db_drias_path: str, user_input: str, model: str) -> State:
             if sql_query == "":
                 table_state['status'] = 'ERROR'
                 continue
+            else : 
+                have_sql_query = True
 
             table_state['sql_query'] = sql_query
-            results = execute_sql_query(db_drias_path, sql_query)
+            df = execute_sql_query(sql_query)
+            
+            if len(df) > 0:
+                have_dataframe = True
 
-            df = pd.DataFrame(results['data'], columns=results['labels'])
             figure = plot['plot_function'](table_state['params'])
             table_state['dataframe'] = df
             table_state['figure'] = figure
             plot_state['table_states'][table] = table_state
 
         state['plot_states'][plot_name] = plot_state
+    
+    if not have_relevant_table:
+        state['error'] = "There is no relevant table in the our database to answer your question"
+    elif not have_sql_query:
+        state['error'] = "There is no relevant sql query on our database that can help to answer your question"
+    elif not have_dataframe:
+        state['error'] = "There is no data in our table that can answer to your question"
+    
+
     return state
 
 
@@ -118,26 +138,25 @@ def find_relevant_plots(state: State, llm) -> list[str]:
     relevant_plots = detect_relevant_plots(state['user_input'], llm)
     return relevant_plots
 
-def find_relevant_tables_per_plot(state: State, plot: Plot, db_path: str, llm) -> list[str]:
+def find_relevant_tables_per_plot(state: State, plot: Plot, llm) -> list[str]:
     print(f"---- Find relevant tables for {plot['name']} ----")
-    relevant_tables = detect_relevant_tables(db_path, state['user_input'], plot, llm)
+    relevant_tables = detect_relevant_tables(state['user_input'], plot, llm)
     return relevant_tables
 
 
-def find_param(state: State, param_name:str, table: str, db_path: str) -> dict[str, Any] | None:
+def find_param(state: State, param_name:str, table: str) -> dict[str, Any] | None:
     """Perform the good method to retrieve the desired parameter
 
     Args:
         state (State): state of the workflow
         param_name (str): name of the desired parameter
         table (str): name of the table
-        db_path (str): path to the databse
 
     Returns:
         dict[str, Any] | None: 
     """
     if param_name == 'location':
-        location = find_location(state['user_input'], table, db_path)
+        location = find_location(state['user_input'], table)
         return location
     if param_name == 'indicator_column':
         indicator_column = find_indicator_column(table)
@@ -153,13 +172,13 @@ class Location(TypedDict):
     latitude: NotRequired[str]
     longitude: NotRequired[str]
 
-def find_location(user_input: str, table: str, db_path: str) -> Location:
+def find_location(user_input: str, table: str) -> Location:
     print(f"---- Find location in table {table} ----")
     location = detect_location_with_openai(user_input)
     output: Location = {'location' : location}
     if location:
         coords = loc2coords(location)
-        neighbour = nearestNeighbourSQL(db_path, coords, table)
+        neighbour = nearestNeighbourSQL(coords, table)
         output.update({
             "latitude": neighbour[0],
             "longitude": neighbour[1],
@@ -182,23 +201,8 @@ def find_indicator_column(table: str) -> str:
     """
     
     print(f"---- Find indicator column in table {table} ----")
-    indicator_columns_per_table = {
-        "total_winter_precipitation": "total_winter_precipitation",
-        "total_summer_precipiation": "total_summer_precipitation",
-        "total_annual_precipitation": "total_annual_precipitation",
-        "total_remarkable_daily_precipitation": "total_remarkable_daily_precipitation",
-        "frequency_of_remarkable_daily_precipitation": "frequency_of_remarkable_daily_precipitation",
-        "extreme_precipitation_intensity": "extreme_precipitation_intensity",
-        "mean_winter_temperature": "mean_winter_temperature",
-        "mean_summer_temperature": "mean_summer_temperature",
-        "mean_annual_temperature": "mean_annual_temperature",
-        "number_of_tropical_nights": "number_tropical_nights",
-        "maximum_summer_temperature": "maximum_summer_temperature",
-        "number_of_days_with_tx_above_30": "number_of_days_with_tx_above_30",
-        "number_of_days_with_tx_above_35": "number_of_days_with_tx_above_35",
-        "number_of_days_with_a_dry_ground": "number_of_days_with_dry_ground"
-    }
-    return indicator_columns_per_table[table]
+
+    return INDICATOR_COLUMNS_PER_TABLE[table]
 
 
 # def make_write_query_node():

style.css

@@ -644,17 +644,23 @@ a {
     overflow-y:scroll;
 }
 
+#sql-query textarea{
+    min-height: 100px !important;
+}
+
 #sql-query span{
     display: none;
 }
 div#tab-vanna{
     max-height: 100¨vh;
-    overflow-y:scroll;
+    overflow-y: hidden;
 } 
 #vanna-plot{
     max-height:500px
 }
 
-#drias-model{
-    max-width: 25%;
+#pagination-display{
+    text-align: center;
+    font-weight: bold;
+    font-size: 16px;
 }