feat: added drias model choice and changed TTD UI

app.py

@@ -12,7 +12,7 @@ from climateqa.engine.reranker import get_reranker
 from climateqa.engine.graph import make_graph_agent,make_graph_agent_poc
 from climateqa.engine.chains.retrieve_papers import find_papers
 from climateqa.chat import start_chat, chat_stream, finish_chat
-from climateqa.engine.talk_to_data.main import ask_drias
+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, create_examples_tab, create_papers_tab, create_figures_tab, create_chat_interface, create_about_tab)
@@ -87,8 +87,8 @@ 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, index_state):
-    return ask_drias(db_vanna_path, query, index_state)
+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):
     print("chat cqa - message received")
@@ -139,27 +139,40 @@ 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 create_drias_tab():
     with gr.Tab("Beta - Talk to DRIAS", elem_id="tab-vanna", id=6):
-        drias_direct_question = gr.Textbox(label="Direct Question", placeholder="You can write direct question here", elem_id="direct-question", interactive=True)
+        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("Details", elem_id="vanna-details", open=False) as drias_details:
-            drias_sql_query = gr.Textbox(label="SQL Query Used", elem_id="sql-query", interactive=False)
+        with gr.Accordion(label="SQL Query Used"):
+            drias_sql_query = gr.Textbox(label="", elem_id="sql-query", interactive=False)
+
+        with gr.Accordion(label='Data used', open=False):
             drias_table = gr.DataFrame([], elem_id="vanna-table")
-            drias_display = gr.Plot()
 
-            # Navigation buttons
+        with gr.Accordion(label="Chart"):
+            drias_display = gr.Plot(elem_id="vanna-plot")
+        
+        with gr.Row():
             prev_button = gr.Button("Previous")
             next_button = gr.Button("Next")
 
-            sql_queries_state = gr.State([])
-            dataframes_state = gr.State([])
-            plots_state = gr.State([])
-            index_state = gr.State(0)  # To track the current position
+        sql_queries_state = gr.State([])
+        dataframes_state = gr.State([])
+        plots_state = gr.State([])
+        index_state = gr.State(0)
 
         drias_direct_question.submit(
             ask_drias_query, 
-            inputs=[drias_direct_question, index_state], 
+            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]
+        )
+
+        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]
         )
 
@@ -184,8 +197,6 @@ def create_drias_tab():
             inputs=[index_state, sql_queries_state, dataframes_state, plots_state],
             outputs=[drias_sql_query, drias_table, drias_display, index_state]
         )
-
- 
     
 # # UI Layout Components
 def cqa_tab(tab_name):

climateqa/engine/talk_to_data/main.py

@@ -13,13 +13,12 @@ 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):
-    final_state = drias_workflow(db_drias_path, query)
+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)
     sql_queries = []
     result_dataframes = []
     figures = []
 
-
     for plot_state in final_state['plot_states'].values():
         for table_state in plot_state['table_states'].values():
             if table_state['status'] == 'OK':
@@ -30,9 +29,29 @@ def ask_drias(db_drias_path:str, query:str , index_state: int):
                     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']))
+
                 
     return sql_queries[index_state], result_dataframes[index_state], figures[index_state], sql_queries, result_dataframes, figures, index_state
 
+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'
+]
 # def ask_vanna(vn,db_vanna_path, query):
 
 #     try :

climateqa/engine/talk_to_data/plot.py

@@ -53,7 +53,7 @@ def plot_indicator_evolution_at_location(params: dict) -> Callable[..., Figure]:
             # Compute the 10-year rolling average
             sliding_averages = (
                 df_avg[indicator]
-                .rolling(window=10, min_periods=5)
+                .rolling(window=10, min_periods=1)
                 .mean()
                 .astype(float)
                 .tolist()
@@ -68,7 +68,7 @@ def plot_indicator_evolution_at_location(params: dict) -> Callable[..., Figure]:
             # Compute the 10-year rolling average
             sliding_averages = (
                 df_model[indicator]
-                .rolling(window=10, min_periods=5)
+                .rolling(window=10, min_periods=1)
                 .mean()
                 .astype(float)
                 .tolist()
@@ -241,6 +241,7 @@ def plot_distribution_of_indicator_for_given_year(
             yaxis_title="Frequency",
             plot_bgcolor="rgba(0, 0, 0, 0)",
             showlegend=False,
+            pan=False
         )
 
         return fig
@@ -313,8 +314,8 @@ def plot_map_of_france_of_indicator_for_given_year(
             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}"),  # Add legend
-            title=f"{indicator_label} in {year} in France",  # Title
+            coloraxis_colorbar=dict(title=f"{indicator_label} {'(Model Average)' if model == 'ALL' else '(Model : ' + model + ')'}"),  # Add legend
+            title=f"{indicator_label} in {year} in France ",  # Title
         )
         return fig
 

climateqa/engine/talk_to_data/sql_query.py

@@ -60,10 +60,16 @@ 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 ""
-    sql_query = f"SELECT year, {indicator_column}, model\nFROM {table}\nWHERE latitude = {latitude} \nand longitude={longitude} \nOrder by Year"
+    
+    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"
+
     return sql_query
 
 class IndicatorForGivenYearQueryParams(TypedDict, total=False):
@@ -85,9 +91,12 @@ 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 ""
-
-    sql_query = f"Select {indicator_column}, latitude, longitude, model\nFrom {table}\nWhere year = {year}"
+    
+    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}'"
     return sql_query    

climateqa/engine/talk_to_data/utils.py

@@ -27,6 +27,31 @@ def detect_location_with_openai(sentence):
         return location_list[0]
     else:
         return ""
+    
+def detect_year_with_openai(sentence: str):
+    """
+    Detects years in a sentence using OpenAI's API via LangChain.
+    """
+    llm = get_llm()
+
+    prompt = f"""
+    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}"
+    """
+
+    response = llm.invoke(prompt)
+    if response is None:
+        return None
+    response_split = response.content.strip("```python\n").split('=')
+    years_list = []
+    if len(response_split) > 1:
+        years_list = ast.literal_eval(response_split[1])
+    if years_list and len(years_list) > 0:
+        return years_list[0]
+    else:
+        return None
 
 
 def detectTable(sql_query):
@@ -65,6 +90,17 @@ def nearestNeighbourSQL(db: str, location: tuple, table: str):
 
 
 def detect_relevant_tables(db: str, 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
+
+    Returns:
+        list[str]: list of table names
+    """
     conn = sqlite3.connect(db)
     cursor = conn.cursor()
 

climateqa/engine/talk_to_data/workflow.py

@@ -38,7 +38,7 @@ class State(TypedDict):
     plots: list[str]
     plot_states: dict[str, PlotState]
     
-def drias_workflow(db_drias_path: str, user_input: str) -> State:
+def drias_workflow(db_drias_path: str, user_input: str, model: str) -> State:
     """Performs the complete workflow of Talk To Drias : from user input to sql queries, dataframes and figures generated
 
     Args:
@@ -87,7 +87,7 @@ def drias_workflow(db_drias_path: str, user_input: str) -> State:
                 'status': 'OK'
             } 
             table_state['params'] = {
-                'model': 'ALL'
+                'model': model
             }
             for param_name in plot['params']:
                 param = find_param(state, param_name, table, db_drias_path)
@@ -99,6 +99,7 @@ def drias_workflow(db_drias_path: str, user_input: str) -> State:
             if sql_query == "":
                 table_state['status'] = 'ERROR'
                 continue
+            print(sql_query)
 
             table_state['sql_query'] = sql_query
             results = execute_sql_query(db_drias_path, sql_query)

style.css

@@ -487,7 +487,6 @@ a {
         height: calc(100vh - 190px) !important;
         overflow-y: scroll !important;
     }
-    div#tab-vanna,
     div#sources-figures,
     div#graphs-container,
     div#tab-citations {
@@ -607,14 +606,25 @@ 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;
+
+#sql-query span{
+    display: none;
+}
+div#tab-vanna{
+    max-height: 100¨vh;
     overflow-y:scroll;
+} 
+#vanna-plot{
+    max-height:500px
 }
+
+#drias-model{
+    max-width: 25%;
+}