55 lines
2.2 KiB
Python
55 lines
2.2 KiB
Python
import openai
|
|
|
|
|
|
class AI:
|
|
def __init__(self, api_key):
|
|
self.api_key = api_key
|
|
openai.api_key = self.api_key
|
|
self.convertlog = []
|
|
|
|
def humantosql(self, text: str, dbtype: str, tableschema: list) -> str:
|
|
if not self.convertlog:
|
|
self.convertlog = [{"role": "system", "content": f"You convert Human Language to SQL. Only answer as an Valid SQL Statement. Always modify your previous answer and do not create something new. You are using this Database: {dbtype}. For better context here are the tables and columns: {tableschema}"}]
|
|
|
|
prompt = {"role": "user", "content": text}
|
|
self.convertlog.append(prompt)
|
|
response = openai.ChatCompletion.create(
|
|
model="gpt-4",
|
|
messages=self.convertlog
|
|
)
|
|
response = response['choices'][0]['message']['content']
|
|
self.convertlog.append({"role": "system", "content": 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"},
|
|
{"role": "user", "content": sql}]
|
|
|
|
response = openai.ChatCompletion.create(
|
|
model="gpt-4",
|
|
messages=prompt
|
|
)
|
|
|
|
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
|