68 lines
2.7 KiB
Python
68 lines
2.7 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)
|
|
self.response = openai.ChatCompletion.create(
|
|
model="gpt-4",
|
|
messages=self.convertlog
|
|
)
|
|
self.response = self.response['choices'][0]['message']['content']
|
|
self.convertlog.append({"role": "assistant", "content": self.response})
|
|
|
|
if self.validate(self.response) is True:
|
|
return self.response
|
|
else:
|
|
valid = self.humantosql("Only answer as a Valid SQL-Statement don't add any additional text", dbtype,
|
|
tableschema)
|
|
return valid
|
|
|
|
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. 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}]
|
|
|
|
valid = openai.ChatCompletion.create(
|
|
model="gpt-4",
|
|
messages=prompt
|
|
)
|
|
valid = valid['choices'][0]['message']['content']
|
|
if valid == "True":
|
|
return True
|
|
else:
|
|
return False
|
|
|
|
def test_key(self) -> bool:
|
|
openai.api_key = str(self.api_key)
|
|
try:
|
|
openai.Engine.list()
|
|
return True
|
|
except Exception as e:
|
|
print(e)
|
|
return False
|