From d03e7fa08651614a5e9c89bb7386c7c2f61f36da Mon Sep 17 00:00:00 2001 From: crennis Date: Thu, 4 May 2023 09:04:34 +0200 Subject: [PATCH] Fixed some stuff, needs testing --- README.md | 1 - nogui/database/postgresql.py | 32 ++++++++------ nogui/main.py | 86 ++++++++++++++++++++++-------------- nogui/modules/aisql.py | 11 ++--- 4 files changed, 78 insertions(+), 52 deletions(-) diff --git a/README.md b/README.md index d72ef30..4a11cf3 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,2 @@ # openai-sql -An GUI SQL Table viewer, with translation from Human Language to SQL-Statement \ No newline at end of file diff --git a/nogui/database/postgresql.py b/nogui/database/postgresql.py index 8b42c62..2aae6c6 100644 --- a/nogui/database/postgresql.py +++ b/nogui/database/postgresql.py @@ -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() diff --git a/nogui/main.py b/nogui/main.py index 9a64414..4a621c1 100644 --- a/nogui/main.py +++ b/nogui/main.py @@ -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() \ No newline at end of file diff --git a/nogui/modules/aisql.py b/nogui/modules/aisql.py index ef3fa03..ca45085 100644 --- a/nogui/modules/aisql.py +++ b/nogui/modules/aisql.py @@ -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",