diff --git a/setup.py b/setup.py index 72e5d4c..0f40caa 100644 --- a/setup.py +++ b/setup.py @@ -1,4 +1,8 @@ import setuptools +import os +import sqlite3 +import argparse +import json with open('README.md', 'r') as fh: long_description = fh.read() @@ -19,3 +23,54 @@ setuptools.setup( ] ) + +with open('config.json', 'r') as file: + config = json.load(file) + +DB_PATH = os.getenv('DB_PATH') or config.get('DB_PATH') + +def create_database(): + conn = sqlite3.connect(DB_PATH) + cursor = conn.cursor() + + cursor.execute('''CREATE TABLE IF NOT EXISTS keys ( + key_value TEXT NOT NULL + );''') + conn.commit() + print("Database created!") + +def add_key(key): + conn = sqlite3.connect(DB_PATH) + cursor = conn.cursor() + + cursor.execute("INSERT INTO keys (key_value) VALUES (?)", (key,)) + conn.commit() + print(f"Added key: {key}") + +def remove_key(key_id): + conn = sqlite3.connect(DB_PATH) + cursor = conn.cursor() + + cursor.execute("DELETE FROM keys WHERE key_value=?", (key_id,)) + conn.commit() + print(f"Removed key with key_id: {key_id}") + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="NovaGPT Key Setup Utility") + + parser.add_argument('--create-db', action="store_true", help="Create the keys database") + parser.add_argument('--add-key', type=str, help="Add a key to the keys database") + parser.add_argument('--remove-key', type=int, help="Remove a key from the keys database by key_id") + + args = parser.parse_args() + + if args.create_db: + create_database() + elif args.add_key: + add_key(args.add_key) + elif args.remove_key: + remove_key(args.remove_key) + else: + parser.print_help() + +