33 lines
1.4 KiB
Python
33 lines
1.4 KiB
Python
import openai
|
|
|
|
|
|
class AI:
|
|
def __init__(self, api_key, tableschema):
|
|
self.api_key = api_key
|
|
openai.api_key = self.api_key
|
|
self.convertlog = [{"role": "system", "content": f"You convert human language to SQL. You do not add any adidtional information. You are indirectly connected to a Database, so you can Query, Execute etc, just make the SQL statement. For better context you get the tables and columns from the database: {tableschema}"}]
|
|
|
|
def humantosql(self, text) -> str:
|
|
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:
|
|
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
|