app.py

from openai import OpenAI
import pandas as pd
import psycopg2
import time
import gradio as gr
import sqlparse
import re
import os
import warnings


from persistStorage import saveLog, getAllLogFilesPaths
from config import *
from constants import *
from utils import *
from gptManager import ChatgptManager
from queryHelperManager import QueryHelper

logsDir = os.getenv("HF_HOME", "/data")


pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', 10)

# Filter out all warning messages
warnings.filterwarnings("ignore")

dbCreds = DataWrapper(DB_CREDS_DATA)
dbEngine = DbEngine(dbCreds)

tablesAndCols = getAllTablesInfo(dbEngine, SCHEMA_NAME)
metadataLayout = MetaDataLayout(schemaName=SCHEMA_NAME, allTablesAndCols=tablesAndCols)
metadataLayout.setSelection(DEFAULT_TABLES_COLS)

selectedTablesAndCols = metadataLayout.getSelectedTablesAndCols()
  

openAIClient = OpenAI(api_key=OPENAI_API_KEY)
gptInstance = ChatgptManager(openAIClient, model=GPT_MODEL)
queryHelper = QueryHelper(gptInstance=gptInstance,
                          schemaName=SCHEMA_NAME,platform=PLATFORM, 
                          metadataLayout=metadataLayout, 
                          sampleDataRows=SAMPLE_ROW_MAX, 
                          gptSampleRows=GPT_SAMPLE_ROWS,
                         dbEngine=dbEngine,
                         getSampleDataForTablesAndCols=getSampleDataForTablesAndCols)

def checkAuth(username, password):
  global ADMIN, PASSWD
  if username == ADMIN and password == PASSWD:
      return True
  return False



# Function to save history of chat
def respond(message, chatHistory):
  """gpt response handler for gradio ui"""
  global queryHelper
  try:
      if "modify" in message:
        botMessage, prospectTablesAndCols  = queryHelper.getQueryForUserInput(message, chatHistory)
      else:
         botMessage, prospectTablesAndCols  = 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, "tablesColsSelectedByGpt":str(prospectTablesAndCols) , "queryGenerated":queryGenerated, "completeGptResponse":botMessage}
  saveLog(logMessage)
  chatHistory.append((message, botMessage))
  time.sleep(2)
  return "", chatHistory



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')
  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:
    dbEngine.connect()
    conn = dbEngine.getConnection()
    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')
      
      print(f"Error occured during running the query {sql}.\n and the error is {str(e)}")
        
      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}"


def onSelectedTablesChange(tablesSelected):
  #Updates tables visible and allow selecting columns for them
  global queryHelper
  print(f"Selected tables : {tablesSelected}")
  metadataLayout = queryHelper.getMetadata()
  allTables = list(metadataLayout.getAllTablesCols())
  tableBoxes = []
  for i in range(len(allTables)):
    if allTables[i] in tablesSelected:
      tableBoxes.append(gr.Textbox(f"Textbox {allTables[i]}", visible=True, label=f"{allTables[i]}"))
    else:
      tableBoxes.append(gr.Textbox(f"Textbox {allTables[i]}", visible=False, label=f"{allTables[i]}"))
  return tableBoxes

def onSelectedColumnsChange(*tableBoxes):
  #update selection of columns and tables (include new tables and cols in gpts context)
  global queryHelper
  metadataLayout = queryHelper.getMetadata()
  allTablesList = list(metadataLayout.getAllTablesCols().keys())
  tablesAndCols = {}
  result = ''
  print("Getting selected tables and columns from gradio")
  for tableBox, table in zip(tableBoxes, allTablesList):
    if isinstance(tableBox, list):
      if len(tableBox)!=0:
        tablesAndCols[table] = tableBox
    else:
      pass
  
  metadataLayout.setSelection(tablesAndCols=tablesAndCols)
  print("metadata updated")
  print("Updating queryHelper state, and sample data")
  queryHelper.updateMetadata(metadataLayout)
  return "Columns udpated"

def onResetToDefaultSelection():
  global queryHelper
  tablesSelected = list(DEFAULT_TABLES_COLS.keys())
  tableBoxes = []
  allTablesList = list(metadataLayout.getAllTablesCols().keys())
  for i in range(len(allTablesList)):
    if allTablesList[i] in tablesSelected:
      tableBoxes.append(gr.Textbox(f"Textbox {allTablesList[i]}", visible=True, label=f"{allTablesList[i]}"))
    else:
      tableBoxes.append(gr.Textbox(f"Textbox {allTablesList[i]}", visible=False, label=f"{allTablesList[i]}"))

  metadataLayout.resetSelection()
  metadataLayout.setSelection(DEFAULT_TABLES_COLS)
  queryHelper.updateMetadata(metadataLayout)

  return tableBoxes

def onSyncLogsWithDataDir():
    downloadableFilesPaths = getAllLogFilesPaths()
    fileComponent = gr.File(downloadableFilesPaths, file_count='multiple')
    return fileComponent


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])

    # screen 2 : To run sql query against database
    with gr.Tab("Run Query"):
        gr.Markdown("""<h1><center> Run Query </center></h1>""")
        text_input = gr.Textbox(label = 'Input SQL Query', placeholder="Write your SQL query here ...")
        text_output = gr.Textbox(label = 'Result')
        text_button = gr.Button("RUN QUERY")
        clear = gr.ClearButton([text_input, text_output])
        text_button.click(testSQL, inputs=text_input, outputs=text_output)
     # screen 3 : To set creds, schema, tables and columns
    with gr.Tab("Setup"):
        gr.Markdown("""<h1><center> Run Query </center></h1>""")
        text_input = gr.Textbox(label = 'schema name', value= SCHEMA_NAME)
        allTablesAndCols = queryHelper.getMetadata().getAllTablesCols()
        selectedTablesAndCols = queryHelper.getMetadata().getSelectedTablesAndCols()
        allTablesList = list(allTablesAndCols.keys()) 
        selectedTablesList = list(selectedTablesAndCols.keys())
        
        dropDown = gr.Dropdown(
             allTablesList, value=selectedTablesList, multiselect=True, label="Selected Tables", info="Select Tables from available tables of the schema"
        )

        refreshTables = gr.Button("Refresh selected tables")

        tableBoxes = []
        
        for i in range(len(allTablesList)):
            if allTablesList[i] in selectedTablesList:
              columnsDropDown = gr.Dropdown(
             allTablesAndCols[allTablesList[i]],visible=True,value=selectedTablesAndCols.get(allTablesList[i],None), multiselect=True, label=allTablesList[i], info="Select columns of a table"
        )
              #tableBoxes[allTables[i]] = columnsDropDown
              tableBoxes.append(columnsDropDown)
            else:
              columnsDropDown = gr.Dropdown(
             allTablesAndCols[allTablesList[i]], visible=False, value=None, multiselect=True, label=allTablesList[i], info="Select columns of a table"
        )
              #tableBoxes[allTables[i]] = columnsDropDown
              tableBoxes.append(columnsDropDown)

        refreshTables.click(onSelectedTablesChange, inputs=dropDown, outputs=tableBoxes)




        columnsTextBox = gr.Textbox(label = 'Result')
        refreshColumns = gr.Button("Refresh selected columns and Reload Data")
        refreshColumns.click(onSelectedColumnsChange, inputs=tableBoxes, outputs=columnsTextBox)

        resetToDefaultSelection = gr.Button("Reset to Default")
        resetToDefaultSelection.click(onResetToDefaultSelection, inputs=None, outputs=tableBoxes)
    
    #screen 4 for downloading logs
    with gr.Tab("Log-files"):
        downloadableFilesPaths = getAllLogFilesPaths()
        fileComponent = gr.File(downloadableFilesPaths, file_count='multiple')
        refreshLogs = gr.Button("Sync Log files from /data")
        refreshLogs.click(onSyncLogsWithDataDir, inputs=None, outputs=fileComponent)

demo.launch(share=True, debug=True, ssl_verify=False, auth=checkAuth)
dbEngine.connect()