parse final result from gpt response for CoT response

queryHelperManagerCoT.py

@@ -1,6 +1,7 @@
 from gptManager import ChatgptManager
 from utils import *
 import json 
+import sqlparse
 from constants import TABLE_RELATIONS
 
 class QueryHelperChainOfThought:
@@ -40,7 +41,16 @@ class QueryHelperChainOfThought:
     prompt = self.getPromptForCot()
     self.gptInstanceForCoT.setSystemPrompt(prompt)
     gptResponse = self.gptInstanceForCoT.getResponseForUserInput(userInput)
-    return gptResponse
+    try:
+      txt = response.split("```json")[-1].split('```')[0].replace('\n', '')
+      sqlResult = json.loads(txt)['finalResult']
+      formattedSql = sqlparse.format(sqlResult, reindent=True)
+      print(formattedSql)
+      responseToReturn = formattedSql
+    except:
+      print("Couldn't parse desired result from gpt response")
+      responseToReturn = gptResponse
+    return responseToReturn
   
   def getPromptForCot(self):
     schemaName = self.schemaName