combine all prospect tables and cols

queryHelperManagerCoT.py

@@ -70,9 +70,11 @@ class QueryHelperChainOfThought:
       tasks = self.extractSingleJson(cotStep1)
       print(f"tasks are {tasks}")
       taskQueries = {}
+      prospectTablesAndColsAll = []
       for key, task in tasks.items():
-        taskQuery = self.getQueryForUserInput(userInput)
+        taskQuery, prospectTablesAndCols = self.getQueryForUserInput(userInput)
         taskQueries[key] = {"task":task, "taskQuery":taskQuery}
+        prospectTablesAndColsAll.append(prospectTablesAndCols)
       print(f"tasks and their queries {taskQueries}")
             
       combiningSubtasksQueryPrompt = f"""Combine following subtask and their queries to generate sql query to answer the user input.\n """
@@ -81,8 +83,8 @@ class QueryHelperChainOfThought:
         task = taskQueries[key]["task"]
         query = taskQueries[key]["taskQuery"]
         userPrompt += f" task: {task}, task query: {query}"
-      return self.gptInstance.getResponseForUserInput(userPrompt, combiningSubtasksQueryPrompt)
-    return f"Please rephrase your query. {' '.join(cotStep1.split('Reason')[1:])}"
+      return self.gptInstance.getResponseForUserInput(userPrompt, combiningSubtasksQueryPrompt), prospectTablesAndColsAll
+    return f"Please rephrase your query. {' '.join(cotStep1.split('Reason')[1:])}", None
   
   def getQueryForUserInput(self, userInput, chatHistory=[]):
     gptSampleRows = self.gptSampleRows