Fixed some stuff, needs testing
This commit is contained in:
crennis
parent
ca09ec3b5b
commit
d03e7fa086
|
|
@ -1,3 +1,2 @@
|
|||
# openai-sql
|
||||
|
||||
An GUI SQL Table viewer, with translation from Human Language to SQL-Statement
|
||||
|
|
@ -17,52 +17,56 @@ class Postgres:
|
|||
database=self.db_name
|
||||
)
|
||||
self.tables = []
|
||||
fetch = self.fetchall("SELECT table_name FROM information_schema.tables WHERE table_schema = 'public';")
|
||||
for item in fetch:
|
||||
self.tables.append(item[0])
|
||||
|
||||
self.tableschema = {}
|
||||
for table in self.tables:
|
||||
fetch = self.fetchall(f"SELECT column_name, data_type FROM information_schema.columns WHERE table_name = '{table}';")
|
||||
self.tableschema.update({table: fetch})
|
||||
self.get_schema()
|
||||
|
||||
def __str__(self):
|
||||
return f'PostgreSQL Server: {self.db_ip}:{self.db_port} as {self.db_username} on database {self.db_name}'
|
||||
|
||||
def fetchall(self, query):
|
||||
def get_schema(self) -> dict:
|
||||
fetch = self.fetchall("SELECT table_name FROM information_schema.tables WHERE table_schema = 'public';")
|
||||
for item in fetch:
|
||||
self.tables.append(item[0])
|
||||
|
||||
for table in self.tables:
|
||||
fetch = self.fetchall(f"SELECT column_name, data_type FROM information_schema.columns WHERE table_name = '{table}';")
|
||||
self.tableschema.update({table: fetch})
|
||||
|
||||
return self.tableschema
|
||||
|
||||
def fetchall(self, query: str):
|
||||
cur = self.conn.cursor()
|
||||
cur.execute(query)
|
||||
result = cur.fetchall()
|
||||
cur.close()
|
||||
return result
|
||||
|
||||
def fetchone(self, query):
|
||||
|
||||
def fetchone(self, query: str):
|
||||
cur = self.conn.cursor()
|
||||
cur.execute(query)
|
||||
result = cur.fetchone()
|
||||
cur.close()
|
||||
return result
|
||||
|
||||
def fetchmany(self, query, size):
|
||||
def fetchmany(self, query: str, size: int):
|
||||
cur = self.conn.cursor()
|
||||
cur.execute(query)
|
||||
result = cur.fetchmany(size)
|
||||
cur.close()
|
||||
return result
|
||||
|
||||
def execute(self, query):
|
||||
def execute(self, query: str):
|
||||
cur = self.conn.cursor()
|
||||
cur.execute(query)
|
||||
self.conn.commit()
|
||||
cur.close()
|
||||
|
||||
def executemany(self, query, values):
|
||||
def executemany(self, query: str, values):
|
||||
cur = self.conn.cursor()
|
||||
cur.executemany(query, values)
|
||||
self.conn.commit()
|
||||
cur.close()
|
||||
|
||||
def closeconnection(self):
|
||||
def close(self):
|
||||
self.conn.close()
|
||||
|
||||
|
|
|
|||
|
|
@ -1,12 +1,15 @@
|
|||
import getpass
|
||||
import modules.aisql as aisql
|
||||
|
||||
#TODO: Clear screen after every option
|
||||
|
||||
|
||||
class Main:
|
||||
def __init__(self):
|
||||
self.tableschema = None
|
||||
self.dbtype = 'postgresql'
|
||||
self.ip = 'localhost'
|
||||
self.port = '32771'
|
||||
self.port = '32768'
|
||||
self.username = 'root'
|
||||
self.password = 'wv6G*Uny#4^CUA9E'
|
||||
self.database = 'discord'
|
||||
|
|
@ -17,46 +20,49 @@ class Main:
|
|||
print(f'Database type')
|
||||
print("postgresql - PostgreSQL")
|
||||
print("mysql - MySQL (WIP)")
|
||||
self.dbtype = input('Database type: ')
|
||||
self.dbtype = input(f'Database type (current: {self.dbtype}): ') or self.dbtype
|
||||
|
||||
self.ip = input(f'IP: (current: {self.ip})') or self.ip
|
||||
self.port = input(f'Port: (current: {self.port})') or self.port
|
||||
self.username = input(f'Username: (current: {self.username})') or self.username
|
||||
self.password = getpass.getpass(f'Password: (current: {self.password})') or self.password
|
||||
self.database = input(f'Database: (current: {self.database})') or self.database
|
||||
self.ip = input(f'IP (current: {self.ip}): ') or self.ip
|
||||
self.port = input(f'Port(current: {self.port}): ') or self.port
|
||||
self.username = input(f'Username (current: {self.username}): ') or self.username
|
||||
self.password = getpass.getpass(f'Password (current: {self.password}): ') or self.password
|
||||
self.database = input(f'Database (current: {self.database}): ') or self.database
|
||||
|
||||
if self.dbtype == 'postgresql':
|
||||
# TODO: Add more Database languages
|
||||
import database.postgresql as pg
|
||||
db = pg.Postgres(self.ip, self.port, self.username, self.password, self.database)
|
||||
return db, db.tableschema
|
||||
self.db = pg.Postgres(self.ip, self.port, self.username, self.password, self.database)
|
||||
return self.db
|
||||
|
||||
def select(self):
|
||||
print('Select an option to continue')
|
||||
print('1. Convert human language to SQL')
|
||||
print('2. Run SQL query')
|
||||
print('3. Configure database')
|
||||
print('4. Exit')
|
||||
print('4. Get Tables Schema')
|
||||
print('5. Exit')
|
||||
option = input('Option: ')
|
||||
if option == '1':
|
||||
self.generate_sql()
|
||||
elif option == '2':
|
||||
query = input('Query: ')
|
||||
print(self.db.fetchall(query))
|
||||
self.run_sql(query)
|
||||
elif option == '3':
|
||||
self.db, self.tableschema = self.config()
|
||||
elif option == '4':
|
||||
print(self.db.get_schema())
|
||||
elif option == '5':
|
||||
exit()
|
||||
else:
|
||||
print('Invalid option')
|
||||
|
||||
def generate_sql(self):
|
||||
ai = aisql.AI(self.apikey, self.tableschema)
|
||||
ai = aisql.AI(self.apikey)
|
||||
generate = True
|
||||
while generate:
|
||||
text = input('Text: ')
|
||||
sql = ai.humantosql(text)
|
||||
print(f'SQL:\n{sql}')
|
||||
sql = ai.humantosql(text, self.db.get_schema())
|
||||
print(f'SQL:\n{sql}\n___')
|
||||
print('Would you like to run it?')
|
||||
print('1. Yes')
|
||||
print('2. No')
|
||||
|
|
@ -73,27 +79,43 @@ class Main:
|
|||
print('Invalid option')
|
||||
|
||||
def run_sql(self, query):
|
||||
ai = aisql.AI(self.apikey, self.tableschema)
|
||||
function = ai.decide(query)
|
||||
print(f'Using {function}')
|
||||
if "fetchall".lower() in function:
|
||||
self.db.fetchall(query)
|
||||
elif "fetchone".lower() in function:
|
||||
self.db.fetchone(query)
|
||||
elif "fetchmany".lower() in function:
|
||||
size = function.split('=')[1].strip(']')
|
||||
self.db.fetchmany(query, size)
|
||||
elif "execute".lower() in function:
|
||||
self.db.execute(query)
|
||||
elif "executemany".lower() in function:
|
||||
size = function.split('=')[1].strip(']')
|
||||
self.db.executemany(query, size)
|
||||
else:
|
||||
print('Invalid option')
|
||||
if self.db is None:
|
||||
print('You have to configure the database first')
|
||||
print('1. Configure')
|
||||
print('2. Exit')
|
||||
choice = input('Choice: ')
|
||||
if choice == '1':
|
||||
self.db, self.tableschema = self.config()
|
||||
self.run_sql(query)
|
||||
elif choice == '2':
|
||||
exit()
|
||||
else:
|
||||
print('Invalid option')
|
||||
elif self.db is not None:
|
||||
ai = aisql.AI(self.apikey)
|
||||
function = ai.decide(query)
|
||||
print(f'Using {function}')
|
||||
try:
|
||||
if "fetchall".casefold() in function.casefold():
|
||||
self.db.fetchall(query)
|
||||
elif "fetchone".casefold() in function.casefold():
|
||||
self.db.fetchone(query)
|
||||
elif "fetchmany".casefold() in function.casefold():
|
||||
size = function.split('=')[1].strip(']')
|
||||
self.db.fetchmany(query, size)
|
||||
elif "execute".casefold() in function.casefold():
|
||||
self.db.execute(query)
|
||||
elif "executemany".casefold() in function.casefold():
|
||||
size = function.split('=')[1].strip(']')
|
||||
self.db.executemany(query, size)
|
||||
else:
|
||||
print('Invalid option')
|
||||
except Exception as e:
|
||||
print('Something went wrong:')
|
||||
print(e)
|
||||
|
||||
if __name__ == '__main__':
|
||||
main = Main()
|
||||
|
||||
|
||||
while True:
|
||||
main.select()
|
||||
|
|
@ -2,12 +2,13 @@ import openai
|
|||
|
||||
|
||||
class AI:
|
||||
def __init__(self, api_key, tableschema):
|
||||
def __init__(self, api_key):
|
||||
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}"}]
|
||||
self.convertlog = []
|
||||
|
||||
def humantosql(self, text) -> str:
|
||||
def humantosql(self, text: str, tableschema: list) -> str:
|
||||
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}"}]
|
||||
prompt = {"role": "user", "content": text}
|
||||
self.convertlog.append(prompt)
|
||||
response = openai.ChatCompletion.create(
|
||||
|
|
@ -18,9 +19,9 @@ class AI:
|
|||
self.convertlog.append({"role": "system", "content": response})
|
||||
return response
|
||||
|
||||
def decide(self, sql) -> str:
|
||||
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}]
|
||||
{"role": "user", "content": sql}]
|
||||
|
||||
response = openai.ChatCompletion.create(
|
||||
model="gpt-4",
|
||||
|
|
|
|||
Loading…
Reference in New Issue
Block a user