diff --git a/nogui/modules/aisql.py b/nogui/modules/aisql.py index 07352d6..af26a4d 100644 --- a/nogui/modules/aisql.py +++ b/nogui/modules/aisql.py @@ -19,10 +19,12 @@ class AI: ) response = response['choices'][0]['message']['content'] self.convertlog.append({"role": "system", "content": response}) - for i in self.convertlog: - print(i) - return response + if self.validate(response): + return response + else: + self.humantosql("Only answer as a Valid SQL-Statement don't add any additional text", dbtype, tableschema) + def decide(self, sql: str) -> str: prompt = [{"role": "system", "content": "You have to decide which function it should use. Answer with [FETCHALL] to fetch all, [FETCHONE] to fetch only one, [FETCHMANY=N] to fetchmany with N being the range, [EXECUTE] to just execute, [EXECUTEMANY=N] to execute many with N being the range"}, @@ -36,3 +38,17 @@ class AI: response = response['choices'][0]['message']['content'] 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."}, + {"role": "user", "content": sql}] + + response = openai.ChatCompletion.create( + model="gpt-4", + messages=prompt + ) + response = response['choices'][0]['message']['content'] + if response == "True": + return True + else: + return False