add logs of drias interactions

app.py

@@ -100,15 +100,6 @@ agent_poc = make_graph_agent_poc(
     version="v4",
 )  # TODO put back default 0.2
 
-# Vanna object
-
-# vn = MyVanna(config = {"temperature": 0, "api_key": os.getenv('THEO_API_KEY'), 'model': os.getenv('VANNA_MODEL'), 'pc_api_key': os.getenv('VANNA_PINECONE_API_KEY'), 'index_name': os.getenv('VANNA_INDEX_NAME'), "top_k" : 4})
-# db_vanna_path = os.path.join(os.getcwd(), "data/drias/drias.db")
-# vn.connect_to_sqlite(db_vanna_path)
-
-# def ask_vanna_query(query):
-#     return ask_vanna(vn, db_vanna_path, query)
-
 
 
 
@@ -532,7 +523,7 @@ def main_ui():
         with gr.Tabs():
             cqa_components = cqa_tab(tab_name="ClimateQ&A")
             local_cqa_components = cqa_tab(tab_name="Beta - POC Adapt'Action")
-            create_drias_tab()
+            create_drias_tab(share_client=share_client, user_id=user_id)
 
             create_about_tab()
 

climateqa/chat.py

@@ -61,6 +61,27 @@ def handle_numerical_data(event):
         return numerical_data, sql_query
     return None, None
     
+def log_drias_interaction_to_azure(query, sql_query, data, share_client, user_id):
+    try:
+        # Log interaction to Azure if not in local environment
+        if os.getenv("GRADIO_ENV") != "local":
+            timestamp = str(datetime.now().timestamp())
+            logs = {
+                "user_id": str(user_id),
+                "query": query,
+                "sql_query": sql_query,
+                # "data": data.to_dict() if data is not None else None,
+                "time": timestamp,
+            }
+            log_on_azure(f"drias_{timestamp}.json", logs, share_client)
+            print(f"Logged Drias interaction to Azure Blob Storage: {logs}")
+        else:
+            print("share_client or user_id is None, or GRADIO_ENV is local")
+    except Exception as e:
+        print(f"Error logging Drias interaction on Azure Blob Storage: {e}")
+        error_msg = f"Drias Error: {str(e)[:100]} - The error has been noted, try another question and if the error remains, you can contact us :)"
+        raise gr.Error(error_msg)
+
 # Main chat function
 async def chat_stream(
     agent : CompiledStateGraph,

front/tabs/tab_drias.py

@@ -1,8 +1,11 @@
 import gradio as gr
 from typing import TypedDict, List, Optional
+import os
+import pandas as pd
 
 from climateqa.engine.talk_to_data.main import ask_drias
 from climateqa.engine.talk_to_data.config import DRIAS_MODELS, DRIAS_UI_TEXT
+from climateqa.chat import log_drias_interaction_to_azure
 
 
 class DriasUIElements(TypedDict):
@@ -26,7 +29,8 @@ class DriasUIElements(TypedDict):
 
 
 async def ask_drias_query(query: str, index_state: int):
-    return await ask_drias(query, index_state)
+    result = await ask_drias(query, index_state)
+    return result
 
 
 def show_results(sql_queries_state, dataframes_state, plots_state):
@@ -192,7 +196,21 @@ def create_drias_ui() -> DriasUIElements:
             next_button=next_button
         )
 
-def setup_drias_events(ui_elements: DriasUIElements) -> None:
+def log_drias_to_azure(query: str, sql_query: str, data, share_client, user_id):
+    """Log Drias interaction to Azure storage."""
+    print("log_drias_to_azure")
+    if share_client is not None and user_id is not None:
+        log_drias_interaction_to_azure(
+            query=query,
+            sql_query=sql_query,
+            data=data,
+            share_client=share_client,
+            user_id=user_id
+        )
+    else:
+        print("share_client or user_id is None")
+
+def setup_drias_events(ui_elements: DriasUIElements, share_client=None, user_id=None) -> None:
     """Set up all event handlers for the DRIAS tab."""
     # Create state variables
     sql_queries_state = gr.State([])
@@ -200,6 +218,10 @@ def setup_drias_events(ui_elements: DriasUIElements) -> None:
     plots_state = gr.State([])
     index_state = gr.State(0)
     table_names_list = gr.State([])
+    
+    def log_drias_interaction(query: str, sql_query: str, data: pd.DataFrame):
+        log_drias_to_azure(query, sql_query, data, share_client, user_id)
+    
 
     # Handle example selection
     ui_elements["examples_hidden"].change(
@@ -220,6 +242,10 @@ def setup_drias_events(ui_elements: DriasUIElements) -> None:
             table_names_list,
             ui_elements["result_text"],
         ],
+    ).then(
+        log_drias_interaction,
+        inputs=[ui_elements["examples_hidden"], ui_elements["drias_sql_query"], ui_elements["drias_table"]],
+        outputs=[],
     ).then(
         show_results,
         inputs=[sql_queries_state, dataframes_state, plots_state],
@@ -262,6 +288,10 @@ def setup_drias_events(ui_elements: DriasUIElements) -> None:
             table_names_list,
             ui_elements["result_text"],
         ],
+    ).then(
+        log_drias_interaction,
+        inputs=[ui_elements["drias_direct_question"], ui_elements["drias_sql_query"], ui_elements["drias_table"]],
+        outputs=[],
     ).then(
         show_results,
         inputs=[sql_queries_state, dataframes_state, plots_state],
@@ -324,9 +354,9 @@ def setup_drias_events(ui_elements: DriasUIElements) -> None:
         outputs=[ui_elements["pagination_display"]],
     )
 
-def create_drias_tab():
-    """Main function to create the DRIAS tab with UI and event handling."""
+def create_drias_tab(share_client=None, user_id=None):
+    """Create the DRIAS tab with all its components and event handlers."""
     ui_elements = create_drias_ui()
-    setup_drias_events(ui_elements)
+    setup_drias_events(ui_elements, share_client=share_client, user_id=user_id)