cache for tablesAndCols and minor changes

Minor changes
1. Remove unused query helper tab. Keep only CoT
2. Generate csv result file (Rename the button)
3. Rename Run query (Setup tab) to schema and setup

app.py

@@ -26,32 +26,14 @@ warnings.filterwarnings("ignore")
 
 dbCreds = DataWrapper(DB_CREDS_DATA)
 dbEngine = DbEngine(dbCreds)
-
+print("getting tablesAndCols..")
 tablesAndCols = getAllTablesInfo(dbEngine, SCHEMA_NAME)
-##ToDo Resolve it and remove ittablesAndCols not getting flags table.
-# tablesAndCols['tbl_d_product_style_flags'] = ["product_id", "contemp_style_flag", "trad_style_flag", "country_style_flag", "trans_style_flag",
-#                              "mc_style_flag", "farm_style_flag", "wi_style_flag","iron_style_flag","crystal_style_flag","coast_style_flag","rustic_style_flag","ind_style_flag",
-#                              "glam_style_flag","ac_style_flag","kids_style_flag","asian_style_flag","tiff_style_flag","trop_style_flag","um_style_flag",
-#                              "sw_style_flag", "themed_style_flag", "west_style_flag", "style", "sku#"]
+print("Done.")
 
 metadataLayout = MetaDataLayout(schemaName=SCHEMA_NAME, allTablesAndCols=tablesAndCols)
 metadataLayout.setSelection(DEFAULT_TABLES_COLS)
 
 selectedTablesAndCols = metadataLayout.getSelectedTablesAndCols()
-  
-
-openAIClient = OpenAI(api_key=OPENAI_API_KEY)
-gptInstanceForTableCols = ChatgptManager(openAIClient, model=GPT_MODEL)
-gptInstanceForQuery = ChatgptManager(openAIClient, model=GPT_MODEL)
-queryHelper = QueryHelper(gptInstanceForTableCols=gptInstanceForTableCols,
-                          gptInstanceForQuery=gptInstanceForQuery,
-                          schemaName=SCHEMA_NAME,platform=PLATFORM, 
-                          metadataLayout=metadataLayout, 
-                          sampleDataRows=SAMPLE_ROW_MAX, 
-                          gptSampleRows=GPT_SAMPLE_ROWS,
-                         dbEngine=dbEngine,
-                         getSampleDataForTablesAndCols=getSampleDataForTablesAndCols)
-
 
 openAIClient2 = OpenAI(api_key=OPENAI_API_KEY)
 gptInstanceForCoT = ChatgptManager(openAIClient2, model=GPT_MODEL)
@@ -70,23 +52,6 @@ def checkAuth(username, password):
   return False
 
 
-
-# Function to save history of chat
-def respond(message, chatHistory):
-  """gpt response handler for gradio ui"""
-  global queryHelper
-  try:
-      botMessage  = queryHelper.getQueryForUserInput(message)
-  except Exception as e:
-      errorMessage = {"function":"queryHelper.getQueryForUserInput","error":str(e), "userInput":message}
-      saveLog(errorMessage, 'error')
-      raise ValueError(str(e))
-  queryGenerated = extractSqlFromGptResponse(botMessage)
-  logMessage = {"userInput":message, "queryGenerated":queryGenerated, "completeGptResponse":botMessage,  "function":"queryHelper.getQueryForUserInput"}
-  saveLog(logMessage)
-  chatHistory.append((message, botMessage))
-  return "", chatHistory
-
 # Function to save history of chat
 def respondCoT(message, chatHistory):
   """gpt response handler for gradio ui"""
@@ -115,7 +80,7 @@ eg\n select * from schema.table limit 200\n"""
   return sql, disclaimerOutputStripping
 
 def onGetResultCsvFile(sql):
-  global dbEngine, queryHelper
+  global dbEngine, queryHelperCot
   sql, disclaimerOutputStripping = preProcessSQL(sql=sql)
   if not isDataQuery(sql):
     return "Sorry not allowed to run. As the query modifies the data."
@@ -142,7 +107,7 @@ def onGetResultCsvFile(sql):
   return fileComponent
 
 def testSQL(sql):
-  global dbEngine, queryHelper
+  global dbEngine, queryHelperCot
 
   sql, disclaimerOutputStripping = preProcessSQL(sql=sql)
   if not isDataQuery(sql):
@@ -167,9 +132,9 @@ def testSQL(sql):
 
 def onSelectedTablesChange(tablesSelected):
   #Updates tables visible and allow selecting columns for them
-  global queryHelper
+  global queryHelperCot
   print(f"Selected tables : {tablesSelected}")
-  metadataLayout = queryHelper.getMetadata()
+  metadataLayout = queryHelperCot.getMetadata()
   allTablesAndCols = metadataLayout.getAllTablesCols()
   selectedTablesAndCols = metadataLayout.getSelectedTablesAndCols()
   allTablesList = list(allTablesAndCols.keys())
@@ -189,8 +154,8 @@ def onSelectedTablesChange(tablesSelected):
 
 def onSelectedColumnsChange(*tableBoxes):
   #update selection of columns and tables (include new tables and cols in gpts context)
-  global queryHelper
-  metadataLayout = queryHelper.getMetadata()
+  global queryHelperCot
+  metadataLayout = queryHelperCot.getMetadata()
   allTablesAndCols = metadataLayout.getAllTablesCols()
   allTablesList = list(allTablesAndCols.keys())
   tablesAndCols = {}
@@ -205,17 +170,17 @@ def onSelectedColumnsChange(*tableBoxes):
   
   metadataLayout.setSelection(tablesAndCols=tablesAndCols)
   print("metadata updated")
-  print("Updating queryHelper state, and sample data")
-  queryHelper.updateMetadata(metadataLayout)
+  print("Updating queryHelperCot state, and sample data")
+  queryHelperCot.updateMetadata(metadataLayout)
   return "Columns udpated"
 
 def onResetToDefaultSelection():
-  global queryHelper
-  metadataLayout = queryHelper.getMetadata()
+  global queryHelperCot
+  metadataLayout = queryHelperCot.getMetadata()
   metadataLayout.setSelection(tablesAndCols=tablesAndCols)
-  queryHelper.updateMetadata(metadataLayout)
+  queryHelperCot.updateMetadata(metadataLayout)
 
-  metadataLayout = queryHelper.getMetadata()
+  metadataLayout = queryHelperCot.getMetadata()
   allTablesAndCols = metadataLayout.getAllTablesCols()
   selectedTablesAndCols = metadataLayout.getSelectedTablesAndCols()
   allTablesList = list(allTablesAndCols.keys())
@@ -241,19 +206,12 @@ def onSyncLogsWithDataDir():
 
 
 with gr.Blocks() as demo:
-    # screen 1 : Chatbot for question answering to generate sql query from user input in english
+
     with gr.Tab("Query Helper"):
         gr.Markdown("""<h1><center> Query Helper</center></h1>""")
         chatbot = gr.Chatbot()
         msg = gr.Textbox()
         clear = gr.ClearButton([msg, chatbot])
-        msg.submit(respond, [msg, chatbot], [msg, chatbot])
-
-    with gr.Tab("Query Helper CoT"):
-        gr.Markdown("""<h1><center> Query Helper CoT</center></h1>""")
-        chatbot = gr.Chatbot()
-        msg = gr.Textbox()
-        clear = gr.ClearButton([msg, chatbot])
         msg.submit(respondCoT, [msg, chatbot], [msg, chatbot])
 
     # screen 2 : To run sql query against database
@@ -268,15 +226,15 @@ with gr.Blocks() as demo:
         text_button.click(testSQL, inputs=text_input, outputs=[text_output, table_output])
 
         csvFileComponent = gr.File([], file_count='multiple')
-        downloadCsv = gr.Button("Get result as csv")
+        downloadCsv = gr.Button("Generate csv result file")
         downloadCsv.click(onGetResultCsvFile, inputs=text_input, outputs=csvFileComponent)
         
      # screen 3 : To set creds, schema, tables and columns
     with gr.Tab("Setup"):
-        gr.Markdown("""<h1><center> Run Query </center></h1>""")
+        gr.Markdown("""<h1><center> Setup Tab </center></h1>""")
         text_input = gr.Textbox(label = 'schema name', value= SCHEMA_NAME)
-        allTablesAndCols = queryHelper.getMetadata().getAllTablesCols()
-        selectedTablesAndCols = queryHelper.getMetadata().getSelectedTablesAndCols()
+        allTablesAndCols = queryHelperCot.getMetadata().getAllTablesCols()
+        selectedTablesAndCols = queryHelperCot.getMetadata().getSelectedTablesAndCols()
         allTablesList = list(allTablesAndCols.keys()) 
         selectedTablesList = list(selectedTablesAndCols.keys())
         
@@ -320,5 +278,5 @@ with gr.Blocks() as demo:
         fileComponent = gr.File(downloadableFilesPaths, file_count='multiple')
         refreshLogs = gr.Button("Sync Log files from /data")
         refreshLogs.click(onSyncLogsWithDataDir, inputs=None, outputs=fileComponent)
-        
+print("Ready to launch...")
 demo.launch(share=True, debug=True, ssl_verify=False, auth=checkAuth)

configProd.py

@@ -26,4 +26,5 @@ STORAGE_DIR = os.getenv("HF_HOME", "/data")
 logsDir = STORAGE_DIR
 
 TABLES_DATA_DIR = os.path.join(STORAGE_DIR, "tablesData")
-RESULT_CSV_DIR = os.path.join(STORAGE_DIR, "csvResults")
+RESULT_CSV_DIR = os.path.join(STORAGE_DIR, "csvResults")
+SCHEMA_INFO_FILE_PATH = os.path.join(STORAGE_DIR, "schemaInfo.pickle")

utils.py

@@ -2,6 +2,9 @@ import psycopg2
 import re
 import pandas as pd
 from persistStorage import retrieveTablesDataFromLocalDb, saveTablesDataToLocalDB
+from config import SCHEMA_INFO_FILE_PATH
+import os
+import pickle
 
 class DataWrapper:
   def __init__(self, data):
@@ -105,8 +108,14 @@ def executeColumnsQuery(dbEngine, columnQuery):
 def closeDbEngine(dbEngine):
   dbEngine.disconnect()
 
-def getAllTablesInfo(dbEngine, schemaName):
+def getAllTablesInfo(dbEngine, schemaName, useCache=True):
+  if useCache:
+    if os.path.isfile(SCHEMA_INFO_FILE_PATH):
+      with open(SCHEMA_INFO_FILE_PATH,'rb') as fh:
+        tablesAndCols = pickle.load(fh)
+        return tablesAndCols
   tablesAndCols = {}
+  print("Getting tables Info, list of tables and columns...")
   allTablesQuery = f"""SELECT table_name FROM information_schema.tables
     WHERE table_schema = '{schemaName}'"""
   tables = executeQuery(dbEngine, allTablesQuery)
@@ -115,6 +124,9 @@ def getAllTablesInfo(dbEngine, schemaName):
     columnsQuery = f"""Select * FROM {schemaName}.{tableName} LIMIT 0"""
     columns = executeColumnsQuery(dbEngine, columnsQuery)
     tablesAndCols[tableName] = columns
+  
+  with open(SCHEMA_INFO_FILE_PATH, 'wb') as fh:
+    pickle.dump(tablesAndCols, fh)
   return tablesAndCols
 
 def getSampleDataForTablesAndCols(dbEngine, schemaName, tablesAndCols, maxRows):