Fixed some stuff, needs testing

README.md

@@ -1,3 +1,2 @@
 # openai-sql
 
-An GUI SQL Table viewer, with translation from Human Language to SQL-Statement

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:
+        self.get_schema()
-            fetch = self.fetchall(f"SELECT column_name, data_type FROM information_schema.columns WHERE table_name = '{table}';")
-            self.tableschema.update({table: fetch})
 
     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()
 

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.ip = input(f'IP (current: {self.ip}): ') or self.ip
-        self.port = input(f'Port: (current: {self.port})') or self.port
+        self.port = input(f'Port(current: {self.port}): ') or self.port
-        self.username = input(f'Username: (current: {self.username})') or self.username
+        self.username = input(f'Username (current: {self.username}): ') or self.username
-        self.password = getpass.getpass(f'Password: (current: {self.password})') or self.password
+        self.password = getpass.getpass(f'Password (current: {self.password}): ') or self.password
-        self.database = input(f'Database: (current: {self.database})') or self.database
+        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)
+            self.db = pg.Postgres(self.ip, self.port, self.username, self.password, self.database)
-            return db, db.tableschema
+            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)
+            sql = ai.humantosql(text, self.db.get_schema())
-            print(f'SQL:\n{sql}')
+            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)
+        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}')
-        if "fetchall".lower() in function:
+            try:
+                if "fetchall".casefold() in function.casefold():
                     self.db.fetchall(query)
-        elif "fetchone".lower() in function:
+                elif "fetchone".casefold() in function.casefold():
                     self.db.fetchone(query)
-        elif "fetchmany".lower() in function:
+                elif "fetchmany".casefold() in function.casefold():
                     size = function.split('=')[1].strip(']')
                     self.db.fetchmany(query, size)
-        elif "execute".lower() in function:
+                elif "execute".casefold() in function.casefold():
                     self.db.execute(query)
-        elif "executemany".lower() in function:
+                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()

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,7 +19,7 @@ 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}]