change logs repo

climateqa/logging.py

@@ -124,22 +124,37 @@ def serialize_docs(docs:list)->list:
     
 ## HUGGING FACE LOGGING
 
-def log_on_huggingface(log_filename, logs):
+def log_on_huggingface(log_filename, logs, log_type="chat"):
     """Log data to Hugging Face dataset repository.
     
     Args:
         log_filename (str): Name of the file to store logs
         logs (dict): Log data to store
+        log_type (str): Type of log to store
     """
     try:
-        # Get Hugging Face token from environment
-        hf_token = os.getenv("HF_LOGS_TOKEN")
-        if not hf_token:
-            print("HF_LOGS_TOKEN not found in environment variables")
-            return
+        if log_type =="chat":   
+            # Get Hugging Face token from environment
+            hf_token = os.getenv("HF_LOGS_TOKEN")
+            if not hf_token:
+                print("HF_LOGS_TOKEN not found in environment variables")
+                return
 
-        # Get repository name from environment or use default
-        repo_id = os.getenv("HF_DATASET_REPO", "Ekimetrics/climateqa_logs")
+            # Get repository name from environment or use default
+            repo_id = os.getenv("HF_DATASET_REPO", "Ekimetrics/climateqa_logs")
+        
+        elif log_type =="drias":
+            # Get Hugging Face token from environment
+            hf_token = os.getenv("HF_LOGS_DRIAS_TOKEN")
+            if not hf_token:
+                print("HF_LOGS_DRIAS_TOKEN not found in environment variables")
+                return
+
+            # Get repository name from environment or use default
+            repo_id = os.getenv("HF_DATASET_REPO_DRIAS", "Ekimetrics/climateqa_logs_talk_to_data")
+
+        else:
+            raise ValueError(f"Invalid log type: {log_type}")
         
         # Initialize HfApi
         api = HfApi(token=hf_token)
@@ -189,7 +204,7 @@ def log_interaction_to_huggingface(history, output_query, sources, docs, share_c
                 "time": timestamp,
             }
             # Log to Hugging Face
-            log_on_huggingface(f"chat/{timestamp}.json", logs)
+            log_on_huggingface(f"chat/{timestamp}.json", logs, log_type="chat")
             print(f"Logged interaction to Hugging Face")
         else:
             print("Did not log to Hugging Face because GRADIO_ENV is local")
@@ -216,7 +231,7 @@ def log_drias_interaction_to_huggingface(query, sql_query, user_id):
                 "sql_query": sql_query,
                 "time": timestamp,
             }
-            log_on_huggingface(f"drias/drias_{timestamp}.json", logs)
+            log_on_huggingface(f"drias/drias_{timestamp}.json", logs, log_type="drias")
             print(f"Logged Drias interaction to Hugging Face: {logs}")
         else:
             print("share_client or user_id is None, or GRADIO_ENV is local")