From d72a0863a30d2fcb9f5f347acde5b4e1b90386d6 Mon Sep 17 00:00:00 2001
From: crennis <gitea@crennis.com>
Date: Thu, 4 May 2023 11:33:42 +0200
Subject: [PATCH] Added a SQL-Validity Check

---
 nogui/modules/aisql.py | 26 +++++++++++++++++---------
 1 file changed, 17 insertions(+), 9 deletions(-)

diff --git a/nogui/modules/aisql.py b/nogui/modules/aisql.py
index af26a4d..2c77e54 100644
--- a/nogui/modules/aisql.py
+++ b/nogui/modules/aisql.py
@@ -6,6 +6,7 @@ class AI:
         self.api_key = api_key
         openai.api_key = self.api_key
         self.convertlog = []
+        self.response = None
 
     def humantosql(self, text: str, dbtype: str, tableschema: list) -> str:
         if not self.convertlog:
@@ -13,15 +14,22 @@ class AI:
 
         prompt = {"role": "user", "content": text}
         self.convertlog.append(prompt)
-        response = openai.ChatCompletion.create(
+        self.response = openai.ChatCompletion.create(
             model="gpt-4",
             messages=self.convertlog
         )
-        response = response['choices'][0]['message']['content']
-        self.convertlog.append({"role": "system", "content": response})
+        self.response = self.response['choices'][0]['message']['content']
+        self.convertlog.append({"role": "assistant", "content": self.response})
 
-        if self.validate(response):
-            return response
+        for i in self.convertlog:
+            print(i)
+
+        # FIXME: After the first validation failed it will always return None
+
+        print(self.response)
+
+        if self.validate(self.response) is True:
+            return self.response
         else:
             self.humantosql("Only answer as a Valid SQL-Statement don't add any additional text", dbtype, tableschema)
 
@@ -40,15 +48,15 @@ class AI:
         return response
 
     def validate(self, sql: str) -> bool:
-        prompt = [{"role": "system", "content": "You check if the User Input is a valid SQL-Statement and only return True or False. Most important: ONLY ANSWER True OR False."},
+        prompt = [{"role": "system", "content": "You check if the User Input is a valid SQL-Statement. You also return False if there is any kind of text before and after the statement. You only return True or False. Most important: ONLY ANSWER True OR False."},
                   {"role": "user", "content": sql}]
 
-        response = openai.ChatCompletion.create(
+        valid = openai.ChatCompletion.create(
             model="gpt-4",
             messages=prompt
         )
-        response = response['choices'][0]['message']['content']
-        if response == "True":
+        valid = valid['choices'][0]['message']['content']
+        if valid == "True":
             return True
         else:
             return False