Update queryHelperManagerCoT.py

queryHelperManagerCoT.py

@@ -57,13 +57,16 @@ class QueryHelperChainOfThought:
     return extracted_json
   
   def getQueryForUserInputCoT(self, userInput):
+    selectedTablesAndCols = self.metadataLayout.getSelectedTablesAndCols()
+    systemPromptTail = self.getSystemPromptTailForCoTStep1(selectedTablesAndCols)
     #1. Is the input complete to create a query, or ask user to reask with more detailed input
     systemPromptForInputClarification = """Given an input text, user want to generate sql query. Please answer if the user input is complete or user needs to ask in more detailed way. Answer in following format. 'Yes' ; if yes, break the userinput into smaller subtask for query generation. Formatted into 
         {
         "Task 1": "task 1 description",
         "Task 2": "task 2 description"
         }
-'No' ; if no, then Reason- please be more detailed about customer details; if more modification needed"""
+'No' ; if no, then Reason- please be more detailed about customer details; if more modification needed."""
+    systemPromptForInputClarification = systemPromptForInputClarification + '\n' + systemPromptTail
     cotStep1 = self.gptInstance.getResponseForUserInput(userInput, systemPromptForInputClarification)
     if "yes" in cotStep1.lower()[:5]:
       print("User input sufficient")
@@ -139,4 +142,14 @@ ORDER BY chandelier_count DESC"""
     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", "    ")
+      
+  def getSystemPromptTailForCoTStep1(self, selectedTablesAndCols):
+    schemaName = self.schemaName
+    platform = self.platform
+
+    prompt = f"""schema name is {schemaName}. And sql platform is {platform}. and table info are below.\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", "    ")