added bug fixes, small features : see description

In run query: tell user that it is limit 5 there, also show total number of results (Like pandas)
bug fix: always add schema to table names in query
Improve loggings : log complete gpt response and separate sql.
Manage History without sending systemPrompt each time.

app.py

@@ -9,18 +9,18 @@ import os
 import warnings
 
 
-from persistStorage import saveLog
+from persistStorage import saveLog, getAllLogFilesPaths
 from config import *
 from constants import *
 from utils import *
 from gptManager import ChatgptManager
-# from queryHelper import QueryHelper
+from queryHelperManager import QueryHelper
 
 logsDir = os.getenv("HF_HOME", "/data")
 
 
 pd.set_option('display.max_columns', None)
-pd.set_option('display.max_rows', None)
+pd.set_option('display.max_rows', 10)
 
 # Filter out all warning messages
 warnings.filterwarnings("ignore")
@@ -47,96 +47,6 @@ def getSampleDataForTablesAndCols(dbEngine, schemaName, tablesAndCols, maxRows):
         print(f"couldn't read table data. Table: {table}")
     return data
   
-class QueryHelper:
-  def __init__(self, gptInstance, dbEngine, schemaName, 
-               platform, metadataLayout, sampleDataRows, 
-               gptSampleRows, getSampleDataForTablesAndCols):
-    self.gptInstance = gptInstance
-    self.schemaName = schemaName
-    self.platform = platform
-    self.metadataLayout = metadataLayout
-    self.sampleDataRows = sampleDataRows
-    self.gptSampleRows = gptSampleRows
-    self.getSampleDataForTablesAndCols = getSampleDataForTablesAndCols
-    self.dbEngine = dbEngine
-    self._onMetadataChange()
-  
-  def _onMetadataChange(self):
-    metadataLayout = self.metadataLayout
-    sampleDataRows = self.sampleDataRows
-    dbEngine = self.dbEngine
-    schemaName = self.schemaName
-    
-    selectedTablesAndCols = metadataLayout.getSelectedTablesAndCols()
-    self.sampleData = self.getSampleDataForTablesAndCols(dbEngine=dbEngine,schemaName=schemaName,
-                                                         tablesAndCols=selectedTablesAndCols, maxRows=sampleDataRows)
-  
-  def getMetadata(self):
-    return self.metadataLayout
-  
-  def updateMetadata(self, metadataLayout):
-    self.metadataLayout = metadataLayout
-    self._onMetadataChange()
-
-  def modifySqlQueryEnteredByUser(self, userSqlQuery):
-    platform = self.platform
-    userPrompt = f"Please correct the following sql query, also it has to be run on {platform}. sql query is \n {userSqlQuery}."
-    systemPrompt = ""
-    modifiedSql = self.gptInstance.getResponseForUserInput(userPrompt, systemPrompt)
-    return modifiedSql
-  
-  def filteredSampleDataForProspects(self, prospectTablesAndCols):
-    sampleData = self.sampleData
-    filteredData = {}
-    for table in prospectTablesAndCols.keys():
-      # filteredData[table] = sampleData[table][prospectTablesAndCols[table]]
-      #take all columns of prospects
-      filteredData[table] = sampleData[table]
-    return filteredData
-  
-  def getQueryForUserInput(self, userInput):
-    gptSampleRows = self.gptSampleRows
-    selectedTablesAndCols = self.metadataLayout.getSelectedTablesAndCols()
-    prospectTablesAndCols = self.getProspectiveTablesAndCols(userInput, selectedTablesAndCols)
-    print("getting prospects", prospectTablesAndCols)
-    prospectTablesData = self.filteredSampleDataForProspects(prospectTablesAndCols)
-    systemPromptForQueryGeneration = self.getSystemPromptForQueryGeneration(prospectTablesData, gptSampleRows=gptSampleRows)
-    queryByGpt = self.gptInstance.getResponseForUserInput(userInput, systemPromptForQueryGeneration)
-    return queryByGpt
-
-  def getProspectiveTablesAndCols(self, userInput, selectedTablesAndCols):
-    schemaName = self.schemaName
-
-    systemPromptForProspectColumns = self.getSystemPromptForProspectColumns(selectedTablesAndCols)
-    prospectiveTablesColsText = self.gptInstance.getResponseForUserInput(userInput, systemPromptForProspectColumns)
-    prospectTablesAndCols = {}
-    for table in selectedTablesAndCols.keys():
-      if table in prospectiveTablesColsText:
-        prospectTablesAndCols[table] = []
-        for column in selectedTablesAndCols[table]:
-          if column in prospectiveTablesColsText:
-            prospectTablesAndCols[table].append(column)
-    return prospectTablesAndCols
-
-  def getSystemPromptForQueryGeneration(self, prospectTablesData, gptSampleRows):
-    schemaName = self.schemaName
-    platform = self.platform
-    prompt = f"""Given an input text, generate the corresponding SQL query for given details. Schema Name is {schemaName}. And sql platform is {platform}.\n following is sample data"""
-    for idx, tableName in enumerate(prospectTablesData.keys(), start=1):
-        prompt += f"table name is {tableName}, table data is {prospectTablesData[tableName].head(gptSampleRows)}"
-    prompt += "XXXX"
-    return prompt.replace("\n"," ").replace("\\"," ").replace("  "," ").replace("XXXX", "    ")
-
-  def getSystemPromptForProspectColumns(self, selectedTablesAndCols):
-    schemaName = self.schemaName
-    platform = self.platform
-
-    prompt = f"""Given an input text, User wants to know which all tables and columns would be possibily to have the desired data. Output them as json. Schema Name is {schemaName}. And sql platform is {platform}.\n"""
-    for idx, tableName in enumerate(selectedTablesAndCols.keys(), start=1):
-        prompt += f"table name {tableName} {', '.join(selectedTablesAndCols[tableName])}"
-    prompt += "XXXX"
-    return prompt.replace("\n"," ").replace("\\"," ").replace("  "," ").replace("XXXX", "    ")
-  
 
 openAIClient = OpenAI(api_key=OPENAI_API_KEY)
 gptInstance = ChatgptManager(openAIClient, model=GPT_MODEL)
@@ -155,49 +65,51 @@ 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)
+      botMessage, prospectTablesAndCols  = queryHelper.getQueryForUserInput(message, chatHistory)
   except Exception as e:
       errorMessage = {"function":"queryHelper.getQueryForUserInput","error":str(e), "userInput":message}
       saveLog(errorMessage, 'error')
-      
-  logMessage = {"userInput":message, "queryGenerated":botMessage}
+  queryGenerated = extractSqlFromGptResponse(botMessage)
+  logMessage = {"userInput":message, "tablesColsSelectedByGpt":str(prospectTablesAndCols) , "queryGenerated":queryGenerated, "completeGptResponse":botMessage}
   saveLog(logMessage)
   chatHistory.append((message, botMessage))
   time.sleep(2)
   return "", chatHistory
 
-# Function to test the generated sql query
-def isDataQuery(sql_query):
-    upper_query = sql_query.upper()
-
-    dml_keywords = ['INSERT', 'UPDATE', 'DELETE', 'MERGE']
-    for keyword in dml_keywords:
-        if re.search(fr'\b{keyword}\b', upper_query):
-            return False  # Found a DML keyword, indicating modification
-
-    # If no DML keywords are found, it's likely a data query
-    return True
-
-def testSQL(sql):
-  global dbEngine, queryHelper
+def preProcessGptQueryReponse(gptResponse, metadataLayout: MetaDataLayout):
+   schemaName = metadataLayout.schemaName
+   tablesList = metadataLayout.getAllTablesCols().keys()
+   gptResponse = addSchemaToTableInSQL(gptResponse, schemaName=schemaName, tablesList=tablesList)
+   return gptResponse
 
+def preProcessSQL(sql):
   sql=sql.replace(';', '')
+  disclaimerOutputStripping = ""
   if ('limit' in sql[-15:].lower())==False:
     sql = sql + ' ' + 'limit 5'
+    disclaimerOutputStripping = """Results are stripped to show only top 5 rows.
+Please add your custom limit to get extend result.
+eg\n select * from schema.table limit 20\n\n"""
   sql = str(sql)
   sql = sqlparse.format(sql, reindent=True, keyword_case='upper')
-  print(sql)
+  return sql, disclaimerOutputStripping
+
+def testSQL(sql):
+  global dbEngine, queryHelper
+
+  sql, disclaimerOutputStripping = preProcessSQL(sql=sql)
   if not isDataQuery(sql):
     return "Sorry not allowed to run. As the query modifies the data."
   try:
     conn = dbEngine.connection
-    df = pd.read_sql_query(sql, con=conn)
-    return pd.DataFrame(df)
+    df = pd.read_sql_query(sql, con=conn)       
+    return disclaimerOutputStripping + str(pd.DataFrame(df))
   except Exception as e:
       errorMessage = {"function":"testSQL","error":str(e), "userInput":sql}
       saveLog(errorMessage, 'error')
@@ -206,6 +118,8 @@ def testSQL(sql):
         
       prompt = f"Please correct the following sql query, also it has to be run on {PLATFORM}. sql query is \n {sql}. the error occured is {str(e)}."
       modifiedSql = queryHelper.modifySqlQueryEnteredByUser(prompt)
+      logMessage = {"function":"queryHelper.modifySqlQueryEnteredByUser", "sqlQuery":sql, "modifiedSQLQuery":modifiedSql}
+      saveLog(logMessage, 'info')
       return f"The query you entered throws some error. Here is modified version. Please try this.\n {modifiedSql}"
 
 
@@ -246,7 +160,7 @@ def onSelectedColumnsChange(*tableBoxes):
 
 def onResetToDefaultSelection():
   global queryHelper
-  tablesSelected = list(DefaultTablesAndCols.keys())
+  tablesSelected = list(DEFAULT_TABLES_COLS.keys())
   tableBoxes = []
   allTablesList = list(metadataLayout.getAllTablesCols().keys())
   for i in range(len(allTablesList)):
@@ -256,21 +170,11 @@ def onResetToDefaultSelection():
       tableBoxes.append(gr.Textbox(f"Textbox {allTablesList[i]}", visible=False, label=f"{allTablesList[i]}"))
 
   metadataLayout.resetSelection()
-  metadataLayout.setSelection(DefaultTablesAndCols)
+  metadataLayout.setSelection(DEFAULT_TABLES_COLS)
   queryHelper.updateMetadata(metadataLayout)
 
   return tableBoxes
 
-def getAllLogFilesPaths():
-    global logsDir
-    # Save processed data to temporary file
-    
-    logFiles = [file for file in os.listdir(logsDir) if 'log' in file.lower()]
-    print(logFiles,"avaiable logs")
-    
-    downloadableFilesPaths = [os.path.join(os.path.abspath(logsDir), logFilePath) for logFilePath in logFiles]
-    return downloadableFilesPaths
-
 def onSyncLogsWithDataDir():
     downloadableFilesPaths = getAllLogFilesPaths()
     fileComponent = gr.File(downloadableFilesPaths, file_count='multiple')

gptManager.py

@@ -5,20 +5,34 @@ class ChatgptManager:
     self.client = openAIClient
     self.tokenLimit = tokenLimit
     self.model = model
+  
+  def _chatHistoryToGptMessages(self, chatHistory=[]):
+    messages = []
+    for i in range(len(chatHistory)):
+      if i%2==0:
+        message = {"role":"user", "content":chatHistory[i]}
+      else:
+        message = {"role":"assistant", "content": chatHistory[i]}
+      messages.append(message)
+    return messages
 
-  def getResponseForUserInput(self, userInput, systemPrompt):
-    self.messages = []
+  def getResponseForUserInput(self, userInput, systemPrompt, chatHistory=[]):
+    self.messages = self._chatHistoryToGptMessages(chatHistory[:])
     newMessage = {"role":"system", "content":systemPrompt}
     if not self.isTokeLimitExceeding(newMessage):
       self.messages.append(newMessage)
     else:
-      raise ValueError("System Prompt Too long.")
+      if chatHistory==[]:
+        raise ValueError("System Prompt Too long.")
+      return self.getResponseForUserInput(userInput=userInput, systemPrompt=systemPrompt)
 
     userMessage = {"role":"user", "content":userInput}
     if not self.isTokeLimitExceeding(userMessage):
       self.messages.append(userMessage)
     else:
-      raise ValueError("Token Limit exceeding. With user input")
+      if chatHistory==[]:
+        raise ValueError("Token Limit exceeding. With user input")
+      return self.getResponseForUserInput(userInput=userInput, systemPrompt=systemPrompt)
 
     # completion = self.client.chat.completions.create(
     #   model="gpt-3.5-turbo-1106",

persistStorage.py

@@ -6,7 +6,7 @@ import os
 from config import HUGGING_FACE_TOKEN
 import csv 
 
-logs_dir = os.getenv("HF_HOME", "/data")
+logsDir = os.getenv("HF_HOME", "/data")
 
 # # Create a new file
 # with open(os.path.join(data_dir, "my_data.txt"), "a") as f:
@@ -34,9 +34,24 @@ def append_dict_to_csv(file_path, row_data):
         csv_writer.writerow(row_data)
 
 def saveLog(message, level='info') -> None:
-    global logs_dir
+    global logsDir
+    if not os.path.isdir(logsDir):
+        print("Log directory/Data Directory not available.")
+        return
     current_time = datetime.now(TIMEZONE_OBJ)
     message = str(message)
-    log_file_path = os.path.join(logs_dir, f"{current_time.strftime('%Y-%m')}-log.csv")
+    log_file_path = os.path.join(logsDir, f"{current_time.strftime('%Y-%m')}-log.csv")
     data_dict = {"time":str(current_time),  "level": level, "message": message}
     append_dict_to_csv(log_file_path, data_dict)
+
+def getAllLogFilesPaths():
+    global logsDir
+    # Save processed data to temporary file
+    if not os.path.isdir(logsDir):
+        print("Log directory/Data Directory not available.")
+        return []
+    logFiles = [file for file in os.listdir(logsDir) if 'log' in file.lower()]
+    print(logFiles,"avaiable logs")
+    
+    downloadableFilesPaths = [os.path.join(os.path.abspath(logsDir), logFilePath) for logFilePath in logFiles]
+    return downloadableFilesPaths

queryHelperManager.py

@@ -0,0 +1,95 @@
+from gptManager import ChatgptManager
+from utils import MetaDataLayout
+
+class QueryHelper:
+  def __init__(self, gptInstance: ChatgptManager, dbEngine, schemaName, 
+               platform, metadataLayout: MetaDataLayout, sampleDataRows, 
+               gptSampleRows, getSampleDataForTablesAndCols):
+    self.gptInstance = gptInstance
+    self.schemaName = schemaName
+    self.platform = platform
+    self.metadataLayout = metadataLayout
+    self.sampleDataRows = sampleDataRows
+    self.gptSampleRows = gptSampleRows
+    self.getSampleDataForTablesAndCols = getSampleDataForTablesAndCols
+    self.dbEngine = dbEngine
+    self._onMetadataChange()
+  
+  def _onMetadataChange(self):
+    metadataLayout = self.metadataLayout
+    sampleDataRows = self.sampleDataRows
+    dbEngine = self.dbEngine
+    schemaName = self.schemaName
+    
+    selectedTablesAndCols = metadataLayout.getSelectedTablesAndCols()
+    self.sampleData = self.getSampleDataForTablesAndCols(dbEngine=dbEngine,schemaName=schemaName,
+                                                         tablesAndCols=selectedTablesAndCols, maxRows=sampleDataRows)
+  
+  def getMetadata(self) -> MetaDataLayout :
+    return self.metadataLayout
+  
+  def updateMetadata(self, metadataLayout):
+    self.metadataLayout = metadataLayout
+    self._onMetadataChange()
+
+  def modifySqlQueryEnteredByUser(self, userSqlQuery):
+    platform = self.platform
+    userPrompt = f"Please correct the following sql query, also it has to be run on {platform}. sql query is \n {userSqlQuery}."
+    systemPrompt = ""
+    modifiedSql = self.gptInstance.getResponseForUserInput(userPrompt, systemPrompt)
+    return modifiedSql
+  
+  def filteredSampleDataForProspects(self, prospectTablesAndCols):
+    sampleData = self.sampleData
+    filteredData = {}
+    for table in prospectTablesAndCols.keys():
+      # filteredData[table] = sampleData[table][prospectTablesAndCols[table]]
+      #take all columns of prospects
+      filteredData[table] = sampleData[table]
+    return filteredData
+  
+  def getQueryForUserInput(self, userInput, chatHistory=[]):
+    gptSampleRows = self.gptSampleRows
+    selectedTablesAndCols = self.metadataLayout.getSelectedTablesAndCols()
+    prospectTablesAndCols = self.getProspectiveTablesAndCols(userInput, selectedTablesAndCols, chatHistory)
+    print("getting prospects", prospectTablesAndCols)
+    prospectTablesData = self.filteredSampleDataForProspects(prospectTablesAndCols)
+    systemPromptForQueryGeneration = self.getSystemPromptForQueryGeneration(prospectTablesData, gptSampleRows=gptSampleRows)
+    queryByGpt = self.gptInstance.getResponseForUserInput(userInput, systemPromptForQueryGeneration, chatHistory)
+    
+    queryByGpt = preProcessGptQueryReponse(queryByGpt, metadataLayout=metadataLayout)
+    return queryByGpt, prospectTablesAndCols
+
+  def getProspectiveTablesAndCols(self, userInput, selectedTablesAndCols, chatHistory=[]):
+    schemaName = self.schemaName
+
+    systemPromptForProspectColumns = self.getSystemPromptForProspectColumns(selectedTablesAndCols)
+    prospectiveTablesColsText = self.gptInstance.getResponseForUserInput(userInput, systemPromptForProspectColumns, chatHistory)
+    prospectTablesAndCols = {}
+    for table in selectedTablesAndCols.keys():
+      if table in prospectiveTablesColsText:
+        prospectTablesAndCols[table] = []
+        for column in selectedTablesAndCols[table]:
+          if column in prospectiveTablesColsText:
+            prospectTablesAndCols[table].append(column)
+    return prospectTablesAndCols
+
+  def getSystemPromptForQueryGeneration(self, prospectTablesData, gptSampleRows):
+    schemaName = self.schemaName
+    platform = self.platform
+    prompt = f"""Given an input text, generate the corresponding SQL query for given details. Schema Name is {schemaName}. And sql platform is {platform}.\n following is sample data"""
+    for idx, tableName in enumerate(prospectTablesData.keys(), start=1):
+        prompt += f"table name is {tableName}, table data is {prospectTablesData[tableName].head(gptSampleRows)}"
+    prompt += "XXXX"
+    return prompt.replace("\n"," ").replace("\\"," ").replace("  "," ").replace("XXXX", "    ")
+
+  def getSystemPromptForProspectColumns(self, selectedTablesAndCols):
+    schemaName = self.schemaName
+    platform = self.platform
+
+    prompt = f"""Given an input text, User wants to know which all tables and columns would be possibily to have the desired data. Output them as json. Schema Name is {schemaName}. And sql platform is {platform}.\n"""
+    for idx, tableName in enumerate(selectedTablesAndCols.keys(), start=1):
+        prompt += f"table name {tableName} {', '.join(selectedTablesAndCols[tableName])}"
+    prompt += "XXXX"
+    return prompt.replace("\n"," ").replace("\\"," ").replace("  "," ").replace("XXXX", "    ")
+  

utils.py

@@ -1,5 +1,5 @@
 import psycopg2
-
+import re
 
 class DataWrapper:
   def __init__(self, data):
@@ -108,4 +108,37 @@ def getSampleDataForTablesAndCols(dbEngine, schemaName, tablesAndCols, maxRows):
         data[table] = pd.read_sql_query(sqlQuery, con=conn)
       except:
         print(f"couldn't read table data. Table: {table}")
-    return data
+    return data
+
+# Function to test the generated sql query
+def isDataQuery(sql_query):
+    upper_query = sql_query.upper()
+
+    dml_keywords = ['INSERT', 'UPDATE', 'DELETE', 'MERGE']
+    for keyword in dml_keywords:
+        if re.search(fr'\b{keyword}\b', upper_query):
+            return False  # Found a DML keyword, indicating modification
+
+    # If no DML keywords are found, it's likely a data query
+    return True
+
+def extractSqlFromGptResponse(gptReponse):
+  sqlPattern = re.compile(r"```sql\n(.*?)```", re.DOTALL)
+
+  # Find the match in the text
+  match = re.search(sqlPattern, gptReponse)
+
+  # Extract the SQL query if a match is found
+  if match:
+    sqlQuery = match.group(1)
+    return sqlQuery
+  else:
+    return ""
+  
+def addSchemaToTableInSQL(sqlQuery, schemaName, tablesList):
+
+  for table in tablesList:
+    pattern = re.compile(rf'(?<![a-zA-Z0-9_]){re.escape(table)}(?![a-zA-Z0-9_])', re.IGNORECASE)
+    replacement = f'{schemaName}.{table}'
+    sqlQuery = re.sub(pattern, replacement, sqlQuery)
+  return sqlQuery