Save history/cache QnA pair wrt to each table

Author: pramitchoudhary

sidekick/memory.py

@@ -31,8 +31,6 @@ def save_context(self, info: str, extract_context: bool = True) -> Dict:
         split_token = ";"
         query = " ".join(info.partition(":")[2].split(split_token)[0].strip().split())
         response = " ".join(info.partition(":")[2].split(split_token)[1].partition(":")[2].strip().split())
-        # TODO add additional guardrails to check if the response is a valid response.
-        # At-least syntactically correct SQL.
 
         # Check if entity extraction is enabled
         # Add logic for entity extraction
@@ -66,7 +64,8 @@ def save_context(self, info: str, extract_context: bool = True) -> Dict:
 
         # Persist added information locally
         if chat_history:
-            with open(f"{self.path}/var/lib/tmp/data/history.jsonl", "a") as outfile:
+            # TODO: Persist history for each user. This flow is currently only affects openai models.
+            with open(f"{self.path}/var/lib/tmp/.cache/history.jsonl", "a") as outfile:
                 json.dump(chat_history, outfile)
                 outfile.write("\n")
             if extract_context:

sidekick/query.py

@@ -21,7 +21,7 @@
 from sidekick.logger import logger
 from sidekick.utils import (_check_file_info, filter_samples, is_resource_low,
                             load_causal_lm_model, load_embedding_model,
-                            read_sample_pairs, remove_duplicates)
+                            make_dir, read_sample_pairs, remove_duplicates)
 from sqlalchemy import create_engine
 
 
@@ -119,19 +119,21 @@ def build_index(self, persist: bool = True):
 
     def update_context_queries(self):
         # Check if seed samples were provided
+        cache_path = f"{self.path}/var/lib/tmp/.cache/{self.table_name}/"
         new_context_queries = []
         if self.sample_queries_path is not None and Path(self.sample_queries_path).exists():
             logger.info(f"Using QnA samples from path {self.sample_queries_path}")
             new_context_queries = read_sample_pairs(self.sample_queries_path, "h2ogpt-sql")
             # cache the samples for future use
-            with open(f"{self.path}/var/lib/tmp/data/queries_cache.json", "w") as f:
+            make_dir(cache_path)
+            with open(f"{cache_path}/queries_cache.json", "w") as f:
                 json.dump(new_context_queries, f, indent=2)
-        elif self.sample_queries_path is None and Path(f"{self.path}/var/lib/tmp/data/queries_cache.json").exists():
+        elif self.sample_queries_path is None and Path(f"{cache_path}/queries_cache.json").exists():
             logger.info(f"Using samples from cache")
-            with open(f"{self.path}/var/lib/tmp/data/queries_cache.json", "r") as f:
+            with open(f"{cache_path}/queries_cache.json", "r") as f:
                 new_context_queries = json.load(f)
         # Read the history file and update the context queries
-        history_file = f"{self.path}/var/lib/tmp/data/{self.table_name}/history.jsonl"
+        history_file = f"{self.path}/var/lib/tmp/.cache/{self.table_name}/history.jsonl"
         try:
             if Path(history_file).exists():
                 with open(history_file, "r") as in_file:

sidekick/utils.py

@@ -1,3 +1,4 @@
+import errno
 import glob
 import json
 import os
@@ -134,14 +135,19 @@ def remove_duplicates(
     return res
 
 
-def save_query(output_path: str, query, response, extracted_entity: Optional[dict] = ""):
+def save_query(output_path: str, table_name: str, query, response, extracted_entity: Optional[dict] = ""):
     _response = response
     # Probably need to find a better way to extra the info rather than depending on key phrases
     if response and "Generated response for question,".lower() in response.lower():
-        _response = response.split("**Generated response for question,**")[1].split("\n")[3].strip()
+        _response = (
+            response.split("**Generated response for question,**")[1].split("``` sql")[1].split("```")[0].strip()
+        )
     chat_history = {"Query": query, "Answer": _response, "Entity": extracted_entity}
 
-    with open(f"{output_path}/var/lib/tmp/data/history.jsonl", "a") as outfile:
+    # Persist history for contextual reference wrt to the table.
+    dir_name = f"{output_path}/var/lib/tmp/.cache/{table_name}"
+    make_dir(dir_name)
+    with open(f"{dir_name}/history.jsonl", "a") as outfile:
         json.dump(chat_history, outfile)
         outfile.write("\n")
 
@@ -378,7 +384,7 @@ def make_dir(path: str):
     try:
         os.makedirs(path)
     except OSError as exc:
-        if exc.errno == errno.EXIST and os.path.isdir(path):
+        if exc.errno == errno.EEXIST and os.path.isdir(path):
             pass
         else:
             raise Exception("Error reported while creating default directory path.")

ui/app.py

@@ -578,6 +578,7 @@ async def on_event(q: Q):
         question = q.client.query
         _val = q.client.llm_response
         # Currently, any manual input by the user is a Question by default
+        table_name = q.user.table_name if q.user.table_name else "default"
         if (
             question is not None
             and "SELECT" in question
@@ -586,11 +587,11 @@ async def on_event(q: Q):
             _q = question.lower().split("q:")[1].split("r:")[0].strip()
             _r = question.lower().split("r:")[1].strip()
             logging.info(f"Saving conversation for question: {_q} and response: {_r}")
-            save_query(base_path, query=_q, response=_r)
+            save_query(base_path, table_name, query=_q, response=_r)
             _msg = "Conversation saved successfully!"
         elif question is not None and _val is not None and _val.strip() != "":
             logging.info(f"Saving conversation for question: {question} and response: {_val}")
-            save_query(base_path, query=question, response=_val)
+            save_query(base_path, table_name, query=question, response=_val)
             _msg = "Conversation saved successfully!"
         else:
             _msg = "Sorry, try generating a conversation to save."