34 lines
1.5 KiB
Python
34 lines
1.5 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:
|
|
self.convertlog = [{"role": "system", "content": f"You convert Human Language to SQL. Only answer as an SQL Statement as the output of you will be the direct input into the database. Always edit the old Statement and do not create a completly new one. You are using this Database: {dbtype} and 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})
|
|
return response
|
|
|
|
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
|