Merge branch 'dev' of https://bitbucket.org/ekimetrics/climate_qa into dev

app.py

@@ -9,14 +9,13 @@ 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_vanna
-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.tabs.tab_drias import create_drias_tab
 from front.utils import process_figures
 from gradio_modal import Modal
 
@@ -25,14 +24,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,39 +62,103 @@ 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
+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)
 
-#Vanna object
+# def ask_vanna_query(query):
+#     return ask_vanna(vn, db_vanna_path, query)
 
-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)
 
-async def chat(query, history, audience, sources, reports, relevant_content_sources_selection, search_only):
+
+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
 
 
@@ -103,14 +166,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")
@@ -119,42 +185,40 @@ 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})"
-
-    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:
-        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_vanna_query, [vanna_direct_question], [vanna_sql_query ,vanna_table, vanna_display])
-  
-                                
-def config_event_handling(main_tabs_components : list[MainTabPanel], config_componenets : ConfigPanel):
+    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),
+    )
+
+
+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
@@ -178,7 +242,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
@@ -187,91 +251,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,99 @@
+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"
+}
+
+DRIAS_MODELS = [
+    'ALL',
+    'RegCM4-6_MPI-ESM-LR',
+    'RACMO22E_EC-EARTH',
+    'RegCM4-6_HadGEM2-ES',
+    'HadREM3-GA7_EC-EARTH',
+    'HadREM3-GA7_CNRM-CM5',
+    'REMO2015_NorESM1-M',
+    'SMHI-RCA4_EC-EARTH',
+    'WRF381P_NorESM1-M',
+    'ALADIN63_CNRM-CM5',
+    'CCLM4-8-17_MPI-ESM-LR',
+    'HIRHAM5_IPSL-CM5A-MR',
+    'HadREM3-GA7_HadGEM2-ES',
+    'SMHI-RCA4_IPSL-CM5A-MR',
+    'HIRHAM5_NorESM1-M',
+    'REMO2009_MPI-ESM-LR',
+    'CCLM4-8-17_HadGEM2-ES'
+]
+# Mapping between indicator columns and their units
+INDICATOR_TO_UNIT = {
+    "total_winter_precipitation": "mm",
+    "total_summer_precipitation": "mm",
+    "total_annual_precipitation": "mm",
+    "total_remarkable_daily_precipitation": "mm",
+    "frequency_of_remarkable_daily_precipitation": "days",
+    "extreme_precipitation_intensity": "mm",
+    "mean_winter_temperature": "°C",
+    "mean_summer_temperature": "°C",
+    "mean_annual_temperature": "°C",
+    "number_tropical_nights": "days",
+    "maximum_summer_temperature": "°C",
+    "number_of_days_with_tx_above_30": "days",
+    "number_of_days_with_tx_above_35": "days",
+    "number_of_days_with_dry_ground": "days"
+}
+
+DRIAS_UI_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.
+"""

climateqa/engine/talk_to_data/main.py

@@ -1,47 +1,115 @@
-from climateqa.engine.talk_to_data.myVanna import MyVanna
-from climateqa.engine.talk_to_data.utils import loc2coords, detect_location_with_openai, detectTable, nearestNeighbourSQL, detect_relevant_tables, replace_coordonates
-import sqlite3
-import os
-import pandas as pd
+from climateqa.engine.talk_to_data.workflow import drias_workflow
 from climateqa.engine.llm import get_llm
 import ast
 
-
-
 llm = get_llm(provider="openai")
 
-def ask_llm_to_add_table_names(sql_query, llm):
+def ask_llm_to_add_table_names(sql_query: str, llm) -> str:
+    """Adds table names to the SQL query result rows using LLM.
+    
+    This function modifies the SQL query to include the source table name in each row
+    of the result set, making it easier to track which data comes from which table.
+    
+    Args:
+        sql_query (str): The original SQL query to modify
+        llm: The language model instance to use for generating the modified query
+        
+    Returns:
+        str: The modified SQL query with table names included in the result rows
+    """
     sql_with_table_names = llm.invoke(f"Make the following sql query display the source table in the rows {sql_query}. Just answer the query. The answer should not include ```sql\n").content
     return sql_with_table_names
 
-def ask_llm_column_names(sql_query, llm):
+def ask_llm_column_names(sql_query: str, llm) -> list[str]:
+    """Extracts column names from a SQL query using LLM.
+    
+    This function analyzes a SQL query to identify which columns are being selected
+    in the result set.
+    
+    Args:
+        sql_query (str): The SQL query to analyze
+        llm: The language model instance to use for column extraction
+        
+    Returns:
+        list[str]: A list of column names being selected in the query
+    """
     columns = llm.invoke(f"From the given sql query, list the columns that are being selected. The answer should only be a python list. Just answer the list. The SQL query : {sql_query}").content
     columns_list = ast.literal_eval(columns.strip("```python\n").strip())
     return columns_list
 
-def ask_vanna(vn,db_vanna_path, query):
+async def ask_drias(query: str, index_state: int = 0) -> tuple:
+    """Main function to process a DRIAS query and return results.
+    
+    This function orchestrates the DRIAS workflow, processing a user query to generate
+    SQL queries, dataframes, and visualizations. It handles multiple results and allows
+    pagination through them.
+    
+    Args:
+        query (str): The user's question about climate data
+        index_state (int, optional): The index of the result to return. Defaults to 0.
+        
+    Returns:
+        tuple: A tuple containing:
+            - sql_query (str): The SQL query used
+            - dataframe (pd.DataFrame): The resulting data
+            - figure (Callable): Function to generate the visualization
+            - sql_queries (list): All generated SQL queries
+            - result_dataframes (list): All resulting dataframes
+            - figures (list): All figure generation functions
+            - index_state (int): Current result index
+            - table_list (list): List of table names used
+            - error (str): Error message if any
+    """
+    final_state = await drias_workflow(query)
+    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'])
+                
+                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'])
+
+    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)
 
-    try :
-        location = detect_location_with_openai(query)
-        if location:
+    return sql_query, dataframe, figure, sql_queries, result_dataframes, figures, index_state, table_list, ""
 
-            coords = loc2coords(location)
-            user_input = query.lower().replace(location.lower(), f"lat, long : {coords}")
+# def ask_vanna(vn,db_vanna_path, query):
+
+#     try :
+#         location = detect_location_with_openai(query)
+#         if location:
+
+#             coords = loc2coords(location)
+#             user_input = query.lower().replace(location.lower(), f"lat, long : {coords}")
             
-            relevant_tables = detect_relevant_tables(user_input, llm)
-            coords_tables = [nearestNeighbourSQL(db_vanna_path, coords, relevant_tables[i]) for i in range(len(relevant_tables))]
-            user_input_with_coords = replace_coordonates(coords, user_input, coords_tables)
-
-            sql_query, result_dataframe, figure = vn.ask(user_input_with_coords, print_results=False, allow_llm_to_see_data=True, auto_train=False)
-
-            return sql_query, result_dataframe, figure
-
-        else : 
-            empty_df = pd.DataFrame()
-            empty_fig = None
-            return "", empty_df, empty_fig
-    except Exception as e:
-        print(f"Error: {e}")
-        empty_df = pd.DataFrame()
-        empty_fig = None
-        return "", empty_df, empty_fig
+#             relevant_tables = detect_relevant_tables(db_vanna_path, user_input, llm)
+#             coords_tables = [nearestNeighbourSQL(db_vanna_path, coords, relevant_tables[i]) for i in range(len(relevant_tables))]
+#             user_input_with_coords = replace_coordonates(coords, user_input, coords_tables)
+
+            # sql_query, result_dataframe, figure = vn.ask(user_input_with_coords, print_results=False, allow_llm_to_see_data=True, auto_train=False)
+
+#             return sql_query, result_dataframe, figure
+#         else : 
+#             empty_df = pd.DataFrame()
+#             empty_fig = None
+#             return "", empty_df, empty_fig
+#     except Exception as e:
+#         print(f"Error: {e}")
+#         empty_df = pd.DataFrame()
+#         empty_fig = None
+#         return "", empty_df, empty_fig

climateqa/engine/talk_to_data/plot.py

@@ -0,0 +1,402 @@
+from typing import Callable, TypedDict
+from matplotlib.figure import figaspect
+import pandas as pd
+from plotly.graph_objects import Figure
+import plotly.graph_objects as go
+import plotly.express as px
+
+from climateqa.engine.talk_to_data.sql_query import (
+    indicator_for_given_year_query,
+    indicator_per_year_at_location_query,
+)
+from climateqa.engine.talk_to_data.config import INDICATOR_TO_UNIT
+
+
+
+
+class Plot(TypedDict):
+    """Represents a plot configuration in the DRIAS system.
+    
+    This class defines the structure for configuring different types of plots
+    that can be generated from climate data.
+    
+    Attributes:
+        name (str): The name of the plot type
+        description (str): A description of what the plot shows
+        params (list[str]): List of required parameters for the plot
+        plot_function (Callable[..., Callable[..., Figure]]): Function to generate the plot
+        sql_query (Callable[..., str]): Function to generate the SQL query for the plot
+    """
+    name: str
+    description: str
+    params: list[str]
+    plot_function: Callable[..., Callable[..., Figure]]
+    sql_query: Callable[..., str]
+
+
+def plot_indicator_evolution_at_location(params: dict) -> Callable[..., Figure]:
+    """Generates a function to plot indicator evolution over time at a location.
+    
+    This function creates a line plot showing how a climate indicator changes
+    over time at a specific location. It handles temperature, precipitation,
+    and other climate indicators.
+    
+    Args:
+        params (dict): Dictionary containing:
+            - indicator_column (str): The column name for the indicator
+            - location (str): The location to plot
+            - model (str): The climate model to use
+            
+    Returns:
+        Callable[..., Figure]: A function that takes a DataFrame and returns a plotly Figure
+        
+    Example:
+        >>> plot_func = plot_indicator_evolution_at_location({
+        ...     'indicator_column': 'mean_temperature',
+        ...     'location': 'Paris',
+        ...     'model': 'ALL'
+        ... })
+        >>> fig = plot_func(df)
+    """
+    indicator = params["indicator_column"]
+    location = params["location"]
+    indicator_label = " ".join([word.capitalize() for word in indicator.split("_")])
+    unit = INDICATOR_TO_UNIT.get(indicator, "")
+
+    def plot_data(df: pd.DataFrame) -> Figure:
+        """Generates the actual plot from the data.
+        
+        Args:
+            df (pd.DataFrame): DataFrame containing the data to plot
+            
+        Returns:
+            Figure: A plotly Figure object showing the indicator evolution
+        """
+        fig = go.Figure()
+        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()
+
+            # Compute the 10-year rolling average
+            sliding_averages = (
+                df_avg[indicator]
+                .rolling(window=10, min_periods=1)
+                .mean()
+                .astype(float)
+                .tolist()
+            )
+            model_label = "Model Average"
+
+        else:
+            df_model = df
+
+            # Transform to list to avoid pandas encoding
+            indicators = df_model[indicator].astype(float).tolist()
+            years = df_model["year"].astype(int).tolist()
+
+            # Compute the 10-year rolling average
+            sliding_averages = (
+                df_model[indicator]
+                .rolling(window=10, min_periods=1)
+                .mean()
+                .astype(float)
+                .tolist()
+            )
+            model_label = f"Model : {df['model'].unique()[0]}"
+
+
+        # Indicator per year plot
+        fig.add_scatter(
+            x=years,
+            y=indicators,
+            name=f"Yearly {indicator_label}",
+            mode="lines",
+            marker=dict(color="#1f77b4"),
+            hovertemplate=f"{indicator_label}: %{{y:.2f}} {unit}<br>Year: %{{x}}<extra></extra>"
+        )
+
+        # Sliding average dashed line
+        fig.add_scatter(
+            x=years,
+            y=sliding_averages,
+            mode="lines",
+            name="10 years rolling average",
+            line=dict(dash="dash"),
+            marker=dict(color="#d62728"),
+            hovertemplate=f"10-year average: %{{y:.2f}} {unit}<br>Year: %{{x}}<extra></extra>"
+        )
+        fig.update_layout(
+            title=f"Plot of {indicator_label} in {location} ({model_label})",
+            xaxis_title="Year",
+            yaxis_title=f"{indicator_label} ({unit})",
+            template="plotly_white",
+        )
+        return fig
+
+    return plot_data
+
+
+indicator_evolution_at_location: Plot = {
+    "name": "Indicator evolution at location",
+    "description": "Plot an evolution of the indicator at a certain location",
+    "params": ["indicator_column", "location", "model"],
+    "plot_function": plot_indicator_evolution_at_location,
+    "sql_query": indicator_per_year_at_location_query,
+}
+
+
+def plot_indicator_number_of_days_per_year_at_location(
+    params: dict,
+) -> Callable[..., Figure]:
+    """Generates a function to plot the number of days per year for an indicator.
+    
+    This function creates a bar chart showing the frequency of certain climate
+    events (like days above a temperature threshold) per year at a specific location.
+    
+    Args:
+        params (dict): Dictionary containing:
+            - indicator_column (str): The column name for the indicator
+            - location (str): The location to plot
+            - model (str): The climate model to use
+            
+    Returns:
+        Callable[..., Figure]: A function that takes a DataFrame and returns a plotly Figure
+    """
+    indicator = params["indicator_column"]
+    location = params["location"]
+    indicator_label = " ".join([word.capitalize() for word in indicator.split("_")])
+    unit = INDICATOR_TO_UNIT.get(indicator, "")
+
+    def plot_data(df: pd.DataFrame) -> Figure:
+        """Generate the figure thanks to the dataframe
+
+        Args:
+            df (pd.DataFrame): pandas dataframe with the required data
+
+        Returns:
+            Figure: Plotly figure
+        """
+        fig = go.Figure()
+        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
+            # 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(
+            go.Bar(
+                x=years,
+                y=indicators,
+                width=0.5,
+                marker=dict(color="#1f77b4"),
+                hovertemplate=f"{indicator_label}: %{{y:.2f}} {unit}<br>Year: %{{x}}<extra></extra>"
+            )
+        )
+
+        fig.update_layout(
+            title=f"{indicator_label} in {location} ({model_label})",
+            xaxis_title="Year",
+            yaxis_title=f"{indicator_label} ({unit})",
+            yaxis=dict(range=[0, max(indicators)]),
+            bargap=0.5,
+            template="plotly_white",
+        )
+
+        return fig
+
+    return plot_data
+
+
+indicator_number_of_days_per_year_at_location: Plot = {
+    "name": "Indicator number of days per year at location",
+    "description": "Plot a barchart of the number of days per year of a certain indicator at a certain location. It is appropriate for frequency indicator.",
+    "params": ["indicator_column", "location", "model"],
+    "plot_function": plot_indicator_number_of_days_per_year_at_location,
+    "sql_query": indicator_per_year_at_location_query,
+}
+
+
+def plot_distribution_of_indicator_for_given_year(
+    params: dict,
+) -> Callable[..., Figure]:
+    """Generates a function to plot the distribution of an indicator for a year.
+    
+    This function creates a histogram showing the distribution of a climate
+    indicator across different locations for a specific year.
+    
+    Args:
+        params (dict): Dictionary containing:
+            - indicator_column (str): The column name for the indicator
+            - year (str): The year to plot
+            - model (str): The climate model to use
+            
+    Returns:
+        Callable[..., Figure]: A function that takes a DataFrame and returns a plotly Figure
+    """
+    indicator = params["indicator_column"]
+    year = params["year"]
+    indicator_label = " ".join([word.capitalize() for word in indicator.split("_")])
+    unit = INDICATOR_TO_UNIT.get(indicator, "")
+
+    def plot_data(df: pd.DataFrame) -> Figure:
+        """Generate the figure thanks to the dataframe
+
+        Args:
+            df (pd.DataFrame): pandas dataframe with the required data
+
+        Returns:
+            Figure: Plotly figure
+        """
+        fig = go.Figure()
+        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
+
+            # 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(
+                x=indicators,
+                opacity=0.8,
+                histnorm="percent",
+                marker=dict(color="#1f77b4"),
+                hovertemplate=f"{indicator_label}: %{{x:.2f}} {unit}<br>Frequency: %{{y:.2f}}%<extra></extra>"
+            )
+        )
+
+        fig.update_layout(
+            title=f"Distribution of {indicator_label} in {year} ({model_label})",
+            xaxis_title=f"{indicator_label} ({unit})",
+            yaxis_title="Frequency (%)",
+            plot_bgcolor="rgba(0, 0, 0, 0)",
+            showlegend=False,
+        )
+
+        return fig
+
+    return plot_data
+
+
+distribution_of_indicator_for_given_year: Plot = {
+    "name": "Distribution of an indicator for a given year",
+    "description": "Plot an histogram of the distribution for a given year of the values of an indicator",
+    "params": ["indicator_column", "model", "year"],
+    "plot_function": plot_distribution_of_indicator_for_given_year,
+    "sql_query": indicator_for_given_year_query,
+}
+
+
+def plot_map_of_france_of_indicator_for_given_year(
+    params: dict,
+) -> Callable[..., Figure]:
+    """Generates a function to plot a map of France for an indicator.
+    
+    This function creates a choropleth map of France showing the spatial
+    distribution of a climate indicator for a specific year.
+    
+    Args:
+        params (dict): Dictionary containing:
+            - indicator_column (str): The column name for the indicator
+            - year (str): The year to plot
+            - model (str): The climate model to use
+            
+    Returns:
+        Callable[..., Figure]: A function that takes a DataFrame and returns a plotly Figure
+    """
+    indicator = params["indicator_column"]
+    year = params["year"]
+    indicator_label = " ".join([word.capitalize() for word in indicator.split("_")])
+    unit = INDICATOR_TO_UNIT.get(indicator, "")
+
+    def plot_data(df: pd.DataFrame) -> Figure:
+        fig = go.Figure()
+        if df['model'].nunique() != 1:
+            df_avg = df.groupby(["latitude", "longitude"], as_index=False)[
+                indicator
+            ].mean()
+
+            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
+
+            # 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(
+                lat=latitudes,
+                lon=longitudes,
+                mode="markers",
+                marker=dict(
+                    size=10,
+                    color=indicators,  # Color mapped to values
+                    colorscale="Turbo",  # Color scale (can be 'Plasma', 'Jet', etc.)
+                    cmin=min(indicators),  # Minimum color range
+                    cmax=max(indicators),  # Maximum color range
+                    showscale=True,  # Show colorbar
+                ),
+                text=[f"{indicator_label}: {value:.2f} {unit}" for value in indicators],  # Add hover text showing the indicator value
+                hoverinfo="text"  # Only show the custom text on hover
+            )
+        )
+
+        fig.update_layout(
+            mapbox_style="open-street-map",  # Use OpenStreetMap
+            mapbox_zoom=3,
+            mapbox_center={"lat": 46.6, "lon": 2.0},
+            coloraxis_colorbar=dict(title=f"{indicator_label} ({unit})"),  # Add legend
+            title=f"{indicator_label} in {year} in France ({model_label}) " # Title
+        )
+        return fig
+
+    return plot_data
+
+
+map_of_france_of_indicator_for_given_year: Plot = {
+    "name": "Map of France of an indicator for a given year",
+    "description": "Heatmap on the map of France of the values of an in indicator for a given year",
+    "params": ["indicator_column", "year", "model"],
+    "plot_function": plot_map_of_france_of_indicator_for_given_year,
+    "sql_query": indicator_for_given_year_query,
+}
+
+
+PLOTS = [
+    indicator_evolution_at_location,
+    indicator_number_of_days_per_year_at_location,
+    distribution_of_indicator_for_given_year,
+    map_of_france_of_indicator_for_given_year,
+]

climateqa/engine/talk_to_data/sql_query.py

@@ -0,0 +1,113 @@
+import asyncio
+from concurrent.futures import ThreadPoolExecutor
+from typing import TypedDict
+import duckdb
+import pandas as pd
+
+async def execute_sql_query(sql_query: str) -> pd.DataFrame:
+    """Executes a SQL query on the DRIAS database and returns the results.
+    
+    This function connects to the DuckDB database containing DRIAS climate data
+    and executes the provided SQL query. It handles the database connection and
+    returns the results as a pandas DataFrame.
+    
+    Args:
+        sql_query (str): The SQL query to execute
+        
+    Returns:
+        pd.DataFrame: A DataFrame containing the query results
+        
+    Raises:
+        duckdb.Error: If there is an error executing the SQL query
+    """
+    def _execute_query():
+        # Execute the query
+        results = duckdb.sql(sql_query)
+        # return fetched data
+        return results.fetchdf()
+
+    # Run the query in a thread pool to avoid blocking
+    loop = asyncio.get_event_loop()
+    with ThreadPoolExecutor() as executor:
+        return await loop.run_in_executor(executor, _execute_query)
+
+
+class IndicatorPerYearAtLocationQueryParams(TypedDict, total=False):
+    """Parameters for querying an indicator's values over time at a location.
+    
+    This class defines the parameters needed to query climate indicator data
+    for a specific location over multiple years.
+    
+    Attributes:
+        indicator_column (str): The column name for the climate indicator
+        latitude (str): The latitude coordinate of the location
+        longitude (str): The longitude coordinate of the location
+        model (str): The climate model to use (optional)
+    """
+    indicator_column: str
+    latitude: str
+    longitude: str
+    model: str
+
+
+def indicator_per_year_at_location_query(
+    table: str, params: IndicatorPerYearAtLocationQueryParams
+) -> str:
+    """SQL Query to get the evolution of an indicator per year at a certain location
+
+    Args:
+        table (str): sql table of the indicator
+        params (IndicatorPerYearAtLocationQueryParams) : dictionary with the required params for the query
+
+    Returns:
+        str: the sql query
+    """
+    indicator_column = params.get("indicator_column")
+    latitude = params.get("latitude")
+    longitude = params.get("longitude")
+    
+    if indicator_column is None or latitude is None or longitude is None: # If one parameter is missing, returns an empty query
+        return ""
+    
+    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
+
+class IndicatorForGivenYearQueryParams(TypedDict, total=False):
+    """Parameters for querying an indicator's values across locations for a year.
+    
+    This class defines the parameters needed to query climate indicator data
+    across different locations for a specific year.
+    
+    Attributes:
+        indicator_column (str): The column name for the climate indicator
+        year (str): The year to query
+        model (str): The climate model to use (optional)
+    """
+    indicator_column: str
+    year: str
+    model: str
+
+def indicator_for_given_year_query(
+        table:str, params: IndicatorForGivenYearQueryParams 
+) -> str:
+    """SQL Query to get the values of an indicator with their latitudes, longitudes and models for a given year
+
+    Args:
+        table (str): sql table of the indicator
+        params (IndicatorForGivenYearQueryParams): dictionarry with the required params for the query
+
+    Returns:
+        str: the sql query
+    """
+    indicator_column = params.get("indicator_column")
+    year = params.get('year')
+    if year is None or indicator_column is None: # If one parameter is missing, returns an empty query
+        return ""
+    
+    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,12 +1,15 @@
 import re
-import openai
-import pandas as pd
+from typing import Annotated, TypedDict
+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
 
-def detect_location_with_openai(sentence):
+
+async def detect_location_with_openai(sentence):
     """
     Detects locations in a sentence using OpenAI's API via LangChain.
     """
@@ -19,74 +22,260 @@ def detect_location_with_openai(sentence):
     Sentence: "{sentence}"
     """
 
-    response = llm.invoke(prompt)
+    response = await llm.ainvoke(prompt)
     location_list = ast.literal_eval(response.content.strip("```python\n").strip())
     if location_list:
         return location_list[0]
     else:
         return ""
 
-def detectTable(sql_query):
+class ArrayOutput(TypedDict):
+    """Represents the output of a function that returns an array.
+    
+    This class is used to type-hint functions that return arrays,
+    ensuring consistent return types across the codebase.
+    
+    Attributes:
+        array (str): A syntactically valid Python array string
+    """
+    array: Annotated[str, "Syntactically valid python array."]
+
+async def detect_year_with_openai(sentence: str) -> str:
+    """
+    Detects years in a sentence using OpenAI's API via LangChain.
+    """
+    llm = get_llm()
+
+    prompt = """
+    Extract all years mentioned in the following sentence.
+    Return the result as a Python list. If no year are mentioned, return an empty list.
+    
+    Sentence: "{sentence}"
+    """
+
+    prompt = ChatPromptTemplate.from_template(prompt)
+    structured_llm = llm.with_structured_output(ArrayOutput)
+    chain = prompt | structured_llm
+    response: ArrayOutput = await chain.ainvoke({"sentence": sentence})
+    years_list = eval(response['array'])
+    if len(years_list) > 0:
+        return years_list[0]
+    else:
+        return ""
+
+
+def detectTable(sql_query: str) -> list[str]:
+    """Extracts table names from a SQL query.
+    
+    This function uses regular expressions to find all table names
+    referenced in a SQL query's FROM clause.
+    
+    Args:
+        sql_query (str): The SQL query to analyze
+        
+    Returns:
+        list[str]: A list of table names found in the query
+        
+    Example:
+        >>> detectTable("SELECT * FROM temperature_data WHERE year > 2000")
+        ['temperature_data']
+    """
     pattern = r'(?i)\bFROM\s+((?:`[^`]+`|"[^"]+"|\'[^\']+\'|\w+)(?:\.(?:`[^`]+`|"[^"]+"|\'[^\']+\'|\w+))*)'
     matches = re.findall(pattern, sql_query)
     return matches
 
 
-
-def loc2coords(location : str):
+def loc2coords(location: str) -> tuple[float, float]:
+    """Converts a location name to geographic coordinates.
+    
+    This function uses the Nominatim geocoding service to convert
+    a location name (e.g., city name) to its latitude and longitude.
+    
+    Args:
+        location (str): The name of the location to geocode
+        
+    Returns:
+        tuple[float, float]: A tuple containing (latitude, longitude)
+        
+    Raises:
+        AttributeError: If the location cannot be found
+    """
     geolocator = Nominatim(user_agent="city_to_latlong")
-    location = geolocator.geocode(location)
-    return (location.latitude, location.longitude)
+    coords = geolocator.geocode(location)
+    return (coords.latitude, coords.longitude)
 
 
-def coords2loc(coords : tuple):
+def coords2loc(coords: tuple[float, float]) -> str:
+    """Converts geographic coordinates to a location name.
+    
+    This function uses the Nominatim reverse geocoding service to convert
+    latitude and longitude coordinates to a human-readable location name.
+    
+    Args:
+        coords (tuple[float, float]): A tuple containing (latitude, longitude)
+        
+    Returns:
+        str: The address of the location, or "Unknown Location" if not found
+        
+    Example:
+        >>> coords2loc((48.8566, 2.3522))
+        'Paris, France'
+    """
     geolocator = Nominatim(user_agent="coords_to_city")
     try:
         location = geolocator.reverse(coords)
         return location.address
     except Exception as e:
         print(f"Error: {e}")
-        return "Unknown Location"  
+        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(f"SELECT lat, lon FROM {table} WHERE lat BETWEEN {lat - 0.3} AND {lat + 0.3} AND lon BETWEEN {long - 0.3} AND {long + 0.3}")
-    results = cursor.fetchall()
-    return results[0]
-
-def detect_relevant_tables(user_question, llm):
-    table_names_list = [
-        "Frequency_of_rainy_days_index",
-        "Winter_precipitation_total",
-        "Summer_precipitation_total",
-        "Annual_precipitation_total",
-        # "Remarkable_daily_precipitation_total_(Q99)",
-        "Frequency_of_remarkable_daily_precipitation",
-        "Extreme_precipitation_intensity",
-        "Mean_winter_temperature",
-        "Mean_summer_temperature",
-        "Number_of_tropical_nights",
-        "Maximum_summer_temperature",
-        "Number_of_days_with_Tx_above_30C",
-        "Number_of_days_with_Tx_above_35C",
-        "Drought_index"
-    ]
+
+    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}")
+    return results['latitude'].iloc[0], results['longitude'].iloc[0]
+
+
+async def detect_relevant_tables(user_question: str, plot: Plot, llm) -> list[str]:
+    """Identifies relevant tables for a plot based on user input.
+    
+    This function uses an LLM to analyze the user's question and the plot
+    description to determine which tables in the DRIAS database would be
+    most relevant for generating the requested visualization.
+    
+    Args:
+        user_question (str): The user's question about climate data
+        plot (Plot): The plot configuration object
+        llm: The language model instance to use for analysis
+        
+    Returns:
+        list[str]: A list of table names that are relevant for the plot
+        
+    Example:
+        >>> detect_relevant_tables(
+        ...     "What will the temperature be like in Paris?",
+        ...     indicator_evolution_at_location,
+        ...     llm
+        ... )
+        ['mean_annual_temperature', 'mean_summer_temperature']
+    """
+    # Get all table names
+    table_names_list = DRIAS_TABLES
+
     prompt = (
-        f"You are helping to build a sql query to retrieve relevant data for a user question."
-        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"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"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(
+        (await llm.ainvoke(prompt)).content.strip("```python\n").strip()
     )
-    table_names = ast.literal_eval(llm.invoke(prompt).content.strip("```python\n").strip())
     return table_names
 
+
 def replace_coordonates(coords, query, coords_tables):
     n = query.count(str(coords[0]))
 
     for i in range(n):
-        query = query.replace(str(coords[0]), str(coords_tables[i][0]),1)
-        query = query.replace(str(coords[1]), str(coords_tables[i][1]),1)
-    return query
+        query = query.replace(str(coords[0]), str(coords_tables[i][0]), 1)
+        query = query.replace(str(coords[1]), str(coords_tables[i][1]), 1)
+    return query
+
+
+async def detect_relevant_plots(user_question: str, llm):
+    plots_description = ""
+    for plot in PLOTS:
+        plots_description += "Name: " + plot["name"]
+        plots_description += " - Description: " + plot["description"] + "\n"
+
+    prompt = (
+        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 : "
+    )
+    # 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(
+        (await llm.ainvoke(prompt)).content.strip("```python\n").strip()
+    )
+    return plot_names
+
+
+# Next Version
+# class QueryOutput(TypedDict):
+#     """Generated SQL query."""
+
+#     query: Annotated[str, ..., "Syntactically valid SQL query."]
+
+
+# class PlotlyCodeOutput(TypedDict):
+#     """Generated Plotly code"""
+
+#     code: Annotated[str, ..., "Synatically valid Plotly python code."]
+# def write_sql_query(user_input: str, db: SQLDatabase, relevant_tables: list[str], llm):
+#     """Generate SQL query to fetch information."""
+#     prompt_params = {
+#         "dialect": db.dialect,
+#         "table_info": db.get_table_info(),
+#         "input": user_input,
+#         "relevant_tables": relevant_tables,
+#         "model": "ALADIN63_CNRM-CM5",
+#     }
+
+#     prompt = ChatPromptTemplate.from_template(query_prompt_template)
+#     structured_llm = llm.with_structured_output(QueryOutput)
+#     chain = prompt | structured_llm
+#     result = chain.invoke(prompt_params)
+
+#     return result["query"]
+
+
+# def fetch_data_from_sql_query(db: str, sql_query: str):
+#     conn = sqlite3.connect(db)
+#     cursor = conn.cursor()
+#     cursor.execute(sql_query)
+#     column_names = [desc[0] for desc in cursor.description]
+#     values = cursor.fetchall()
+#     return {"column_names": column_names, "data": values}
+
+
+# def generate_chart_code(user_input: str, sql_query: list[str], llm):
+#     """ "Generate plotly python code for the chart based on the sql query and the user question"""
+
+#     class PlotlyCodeOutput(TypedDict):
+#         """Generated Plotly code"""
+
+#         code: Annotated[str, ..., "Synatically valid Plotly python code."]
+
+#     prompt = ChatPromptTemplate.from_template(plot_prompt_template)
+#     structured_llm = llm.with_structured_output(PlotlyCodeOutput)
+#     chain = prompt | structured_llm
+#     result = chain.invoke({"input": user_input, "sql_query": sql_query})
+#     return result["code"]

climateqa/engine/talk_to_data/workflow.py

@@ -0,0 +1,287 @@
+import os
+
+from typing import Any, Callable, NotRequired, TypedDict
+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 (
+    detect_relevant_plots,
+    detect_year_with_openai,
+    loc2coords,
+    detect_location_with_openai,
+    nearestNeighbourSQL,
+    detect_relevant_tables,
+)
+
+ROOT_PATH = os.path.dirname(os.path.dirname(os.getcwd()))
+
+class TableState(TypedDict):
+    """Represents the state of a table in the DRIAS workflow.
+    
+    This class defines the structure for tracking the state of a table during the
+    data processing workflow, including its name, parameters, SQL query, and results.
+    
+    Attributes:
+        table_name (str): The name of the table in the database
+        params (dict[str, Any]): Parameters used for querying the table
+        sql_query (str, optional): The SQL query used to fetch data
+        dataframe (pd.DataFrame | None, optional): The resulting data
+        figure (Callable[..., Figure], optional): Function to generate visualization
+        status (str): The current status of the table processing ('OK' or 'ERROR')
+    """
+    table_name: str
+    params: dict[str, Any]
+    sql_query: NotRequired[str]
+    dataframe: NotRequired[pd.DataFrame | None]
+    figure: NotRequired[Callable[..., Figure]]
+    status: str
+
+class PlotState(TypedDict):
+    """Represents the state of a plot in the DRIAS workflow.
+    
+    This class defines the structure for tracking the state of a plot during the
+    data processing workflow, including its name and associated tables.
+    
+    Attributes:
+        plot_name (str): The name of the plot
+        tables (list[str]): List of tables used in the plot
+        table_states (dict[str, TableState]): States of the tables used in the plot
+    """
+    plot_name: str
+    tables: list[str]
+    table_states: dict[str, TableState]
+
+class State(TypedDict):
+    user_input: str
+    plots: list[str]
+    plot_states: dict[str, PlotState]
+    error: NotRequired[str]
+    
+async 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:
+        user_input (str): initial user input
+
+    Returns:
+        State: Final state with all the results
+    """
+    state: State = {
+        'user_input': user_input,
+        'plots': [],
+        'plot_states': {}
+    }
+
+    llm = get_llm(provider="openai")
+
+    plots = await find_relevant_plots(state, llm)
+    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
+        if plot is None:
+            continue
+        
+        plot_state: PlotState = {
+            'plot_name': plot_name,
+            'tables': [],
+            'table_states': {}
+        }
+
+        plot_state['plot_name'] = plot_name
+
+        relevant_tables = await find_relevant_tables_per_plot(state, plot, llm)
+        if len(relevant_tables) > 0 :
+            have_relevant_table = True
+
+        plot_state['tables'] = relevant_tables
+
+        params = {}
+        for param_name in plot['params']:
+            param = await 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,
+                'status': 'OK'
+            } 
+
+            table_state["params"]['indicator_column'] = find_indicator_column(table)
+            
+            sql_query = plot['sql_query'](table, table_state['params'])
+
+            if sql_query == "":
+                table_state['status'] = 'ERROR'
+                continue
+            else : 
+                have_sql_query = True
+
+            table_state['sql_query'] = sql_query
+            df = await execute_sql_query(sql_query)
+
+            if len(df) > 0:
+                have_dataframe = True
+
+            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
+
+async def find_relevant_plots(state: State, llm) -> list[str]:
+    print("---- Find relevant plots ----")
+    relevant_plots = await detect_relevant_plots(state['user_input'], llm)
+    return relevant_plots
+
+async def find_relevant_tables_per_plot(state: State, plot: Plot, llm) -> list[str]:
+    print(f"---- Find relevant tables for {plot['name']} ----")
+    relevant_tables = await detect_relevant_tables(state['user_input'], plot, llm)
+    return relevant_tables
+
+async 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
+
+    Returns:
+        dict[str, Any] | None: 
+    """
+    if param_name == 'location':
+        location = await find_location(state['user_input'], table)
+        return location
+    if param_name == 'year':
+        year = await find_year(state['user_input'])
+        return {'year': year}
+    return None
+
+class Location(TypedDict):
+    location: str
+    latitude: NotRequired[str]
+    longitude: NotRequired[str]
+
+async def find_location(user_input: str, table: str) -> Location:
+    print(f"---- Find location in table {table} ----")
+    location = await detect_location_with_openai(user_input)
+    output: Location = {'location' : location}
+    if location:
+        coords = loc2coords(location)
+        neighbour = nearestNeighbourSQL(coords, table)
+        output.update({
+            "latitude": neighbour[0],
+            "longitude": neighbour[1],
+        })
+    return output
+
+async def find_year(user_input: str) -> str:
+    """Extracts year information from user input using LLM.
+    
+    This function uses an LLM to identify and extract year information from the
+    user's query, which is used to filter data in subsequent queries.
+    
+    Args:
+        user_input (str): The user's query text
+        
+    Returns:
+        str: The extracted year, or empty string if no year found
+    """
+    print(f"---- Find year ---")
+    year = await detect_year_with_openai(user_input)
+    return year
+
+def find_indicator_column(table: str) -> str:
+    """Retrieves the name of the indicator column within a table.
+    
+    This function maps table names to their corresponding indicator columns
+    using the predefined mapping in INDICATOR_COLUMNS_PER_TABLE.
+    
+    Args:
+        table (str): Name of the table in the database
+        
+    Returns:
+        str: Name of the indicator column for the specified table
+        
+    Raises:
+        KeyError: If the table name is not found in the mapping
+    """
+    print(f"---- Find indicator column in table {table} ----")
+    return INDICATOR_COLUMNS_PER_TABLE[table]
+
+
+# def make_write_query_node():
+
+#     def write_query(state):
+#         print("---- Write query ----")
+#         for table in state["tables"]:
+#             sql_query = QUERIES[state[table]['query_type']](
+#                 table=table,
+#                 indicator_column=state[table]["columns"],
+#                 longitude=state[table]["longitude"],
+#                 latitude=state[table]["latitude"],
+#             )
+#             state[table].update({"sql_query": sql_query})
+
+#         return state
+
+#     return write_query
+
+# def make_fetch_data_node(db_path):
+
+#     def fetch_data(state):
+#         print("---- Fetch data ----")
+#         for table in state["tables"]:
+#             results = execute_sql_query(db_path, state[table]['sql_query'])
+#             state[table].update(results)
+
+#         return state
+        
+#     return fetch_data
+
+
+
+## V2
+
+
+# def make_fetch_data_node(db_path: str, llm):
+#     def fetch_data(state):
+#         print("---- Fetch data ----")
+#         db = SQLDatabase.from_uri(f"sqlite:///{db_path}")
+#         output = {}
+#         sql_query = write_sql_query(state["query"], db, state["tables"], llm)
+#         # TO DO : Add query checker
+#         print(f"SQL query  : {sql_query}")
+#         output["sql_query"] = sql_query
+#         output.update(fetch_data_from_sql_query(db_path, sql_query))
+#         return output
+
+#     return fetch_data

front/tabs/tab_drias.py

@@ -0,0 +1,332 @@
+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)
+
+
+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 df.empty:
+        return df, None
+    if "model" not in df.columns:
+        return df, figures[index_state](df)
+    if model_selection != "ALL":
+        df = df[df["model"] == model_selection]
+        if df.empty:
+            return df, None
+    figure = figures[index_state](df)
+    return df, figure
+
+
+def 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,
+    )
+
+
+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
+        examples_hidden = gr.Textbox(visible=False, elem_id="drias-examples-hidden")
+        examples = gr.Examples(
+            examples=[
+                ["What will the temperature be like in Paris?"],
+                ["What will be the total rainfall in France in 2030?"],
+                ["How frequent will extreme events be in Lyon?"],
+                ["Comment va évoluer la température en France entre 2030 et 2050 ?"]
+            ],
+            label="Example Questions",
+            inputs=[examples_hidden],
+            outputs=[examples_hidden],
+        )
+        
+        with gr.Row():
+            drias_direct_question = gr.Textbox(
+                label="Direct Question",
+                placeholder="You can write direct question here",
+                elem_id="direct-question",
+                interactive=True,
+            )
+
+        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)
+
+        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
+        )
+
+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([])
+
+    # 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"]],
+    )
+
+    # 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"]],
+    )
+
+    # Handle model selection change
+    ui_elements["model_selection"].change(
+        filter_by_model,
+        inputs=[dataframes_state, plots_state, index_state, ui_elements["model_selection"]],
+        outputs=[ui_elements["drias_table"], ui_elements["drias_display"]],
+    )
+
+    # Handle 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"]],
+    )
+
+    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)
+    
+        

style.css

@@ -520,7 +520,6 @@ a {
         height: calc(100vh - 190px) !important;
         overflow-y: scroll !important;
     }
-    div#tab-vanna,
     div#sources-figures,
     div#graphs-container,
     div#tab-citations {
@@ -653,14 +652,61 @@ a {
 }
 
 #vanna-display {
-    max-height: 300px;
+    max-height: 200px;
     /* overflow-y: scroll; */
 }
 #sql-query{
-    max-height: 100px;
+    max-height: 300px;
     overflow-y:scroll;
 }
-#vanna-details{
-    max-height: 500px;
-    overflow-y:scroll;
+
+#sql-query textarea{
+    min-height: 100px !important;
+}
+
+#sql-query span{
+    display: none;
+}
+div#tab-vanna{
+    max-height: 100¨vh;
+    overflow-y: hidden;
+} 
+#vanna-plot{
+    max-height:500px
+}
+
+#pagination-display{
+    text-align: center;
+    font-weight: bold;
+    font-size: 16px;
+}
+
+#table-names table{
+    overflow: hidden;
+}
+#table-names thead{
+    display: none;
+}
+
+/* DRIAS Data Table Styles */
+#vanna-table {
+    height: 400px !important;
+    overflow-y: auto !important;
+}
+
+#vanna-table > div[class*="table"] {
+    height: 400px !important;
+    overflow-y: None !important;
+}
+
+#vanna-table .table-wrap {
+    height: 400px !important;
+    overflow-y: None !important;
+}
+
+#vanna-table thead {
+    position: sticky;
+    top: 0;
+    background: white;
+    z-index: 1;
 }