From 7084544c13504e7e1f964a495cae531bec8ea4ea Mon Sep 17 00:00:00 2001 From: Game_Time <108236317+RayBytes@users.noreply.github.com> Date: Sat, 24 Jun 2023 18:53:02 +0500 Subject: [PATCH 1/8] Add files via upload --- api/keys.py | 86 +++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 86 insertions(+) create mode 100644 api/keys.py diff --git a/api/keys.py b/api/keys.py new file mode 100644 index 0000000..9dcf126 --- /dev/null +++ b/api/keys.py @@ -0,0 +1,86 @@ +import json +import os +import random +import sqlite3 + +with open('./config.json', 'r') as file: + config = json.load(file) + +class Keys: + # --- START OF CONFIG --- + + DB_PATH = os.getenv('DB_PATH') or config.get('DB_PATH') + + # --- END OF CONFIG --- + + locked_keys = set() + cache = set() + + def __init__(self, key: str): + self.key = key + if not Keys.cache: + self._load_keys() + + def _connect(self): + conn = sqlite3.connect(self.DB_PATH) + return conn + + def _load_keys(self) -> None: + conn = self._connect() + cursor = conn.cursor() + cursor.execute("SELECT key_value FROM keys") + rows = cursor.fetchall() + for row in rows: + key_value = row[0] + self.cache.add(key_value) + conn.commit() + + def lock(self) -> None: + self.locked_keys.add(self.key) + + def unlock(self) -> None: + self.locked_keys.remove(self.key) + + def is_locked(self) -> bool: + return self.key in self.locked_keys + + def get() -> str: + key_candidates = list(Keys.cache) + random.shuffle(key_candidates) + + for key_candidate in key_candidates: + key = Keys(key_candidate) + + if not key.is_locked(): + key.lock() + return key.key + + print("[WARN] No unlocked keys found in get keys request!") + + def delete(self) -> None: + conn = self._connect() + cursor = conn.cursor() + + cursor.execute("DELETE FROM keys WHERE key_value=?", (self.key,)) + conn.commit() + # Update cache + try: + Keys.cache.remove(self.key) + except KeyError: + print("[WARN] Tried to remove a key from cache which was not present: " + self.key) + + def save(self) -> None: + conn = self._connect() + cursor = conn.cursor() + cursor.execute("INSERT INTO keys (key_value) VALUES (?)", (self.key,)) + conn.commit() + # Update cache + Keys.cache.add(self.key) + + +# Usage example: +# os.environ['DB_PATH'] = "keys.db" +# key_instance = Keys("example_ky") +# key_instance.save() +# key_value = Keys.get() + From c19297fe0e093f58ebac25504ab8867b1cdffd5b Mon Sep 17 00:00:00 2001 From: Game_Time <108236317+RayBytes@users.noreply.github.com> Date: Sat, 24 Jun 2023 18:53:52 +0500 Subject: [PATCH 2/8] Add files via upload --- config.json | 3 +++ 1 file changed, 3 insertions(+) create mode 100644 config.json diff --git a/config.json b/config.json new file mode 100644 index 0000000..6971d41 --- /dev/null +++ b/config.json @@ -0,0 +1,3 @@ +{ + "DB_PATH": "keys.db" +} \ No newline at end of file From 5bf7d9fc5dc80369ae74e29bd350d9dd8f32fe25 Mon Sep 17 00:00:00 2001 From: Game_Time <108236317+RayBytes@users.noreply.github.com> Date: Sat, 24 Jun 2023 18:55:35 +0500 Subject: [PATCH 3/8] Update setup.py --- setup.py | 55 +++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 55 insertions(+) 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() + + From 6ca84147f24df0fb40e10bb491efd5d7f91afdaf Mon Sep 17 00:00:00 2001 From: Game_Time <108236317+RayBytes@users.noreply.github.com> Date: Sat, 24 Jun 2023 19:18:17 +0500 Subject: [PATCH 4/8] Update keys.py --- api/keys.py | 44 +++++++++++++++----------------------------- 1 file changed, 15 insertions(+), 29 deletions(-) diff --git a/api/keys.py b/api/keys.py index 9dcf126..ace05b1 100644 --- a/api/keys.py +++ b/api/keys.py @@ -1,39 +1,34 @@ import json import os import random -import sqlite3 +from pymongo import MongoClient with open('./config.json', 'r') as file: config = json.load(file) class Keys: # --- START OF CONFIG --- - - DB_PATH = os.getenv('DB_PATH') or config.get('DB_PATH') - + MONGO_URI = os.getenv('MONGO_URI') or config.get('MONGO_URI') # --- END OF CONFIG --- locked_keys = set() cache = set() + # Initialize MongoDB + client = MongoClient(MONGO_URI) + db = client.get_database('keys_db') + collection = db['keys'] + def __init__(self, key: str): self.key = key if not Keys.cache: self._load_keys() - def _connect(self): - conn = sqlite3.connect(self.DB_PATH) - return conn - def _load_keys(self) -> None: - conn = self._connect() - cursor = conn.cursor() - cursor.execute("SELECT key_value FROM keys") - rows = cursor.fetchall() - for row in rows: - key_value = row[0] + cursor = self.collection.find({}, {'_id': 0, 'key_value': 1}) + for doc in cursor: + key_value = doc['key_value'] self.cache.add(key_value) - conn.commit() def lock(self) -> None: self.locked_keys.add(self.key) @@ -56,13 +51,9 @@ class Keys: return key.key print("[WARN] No unlocked keys found in get keys request!") - - def delete(self) -> None: - conn = self._connect() - cursor = conn.cursor() - cursor.execute("DELETE FROM keys WHERE key_value=?", (self.key,)) - conn.commit() + def delete(self) -> None: + self.collection.delete_one({'key_value': self.key}) # Update cache try: Keys.cache.remove(self.key) @@ -70,17 +61,12 @@ class Keys: print("[WARN] Tried to remove a key from cache which was not present: " + self.key) def save(self) -> None: - conn = self._connect() - cursor = conn.cursor() - cursor.execute("INSERT INTO keys (key_value) VALUES (?)", (self.key,)) - conn.commit() + self.collection.insert_one({'key_value': self.key}) # Update cache Keys.cache.add(self.key) - # Usage example: -# os.environ['DB_PATH'] = "keys.db" -# key_instance = Keys("example_ky") +# os.environ['MONGO_URI'] = "mongodb://localhost:27017" +# key_instance = Keys("example_key") # key_instance.save() # key_value = Keys.get() - From 280e5d9a581231c09d7f1c1ce899a3735864a5f4 Mon Sep 17 00:00:00 2001 From: Game_Time <108236317+RayBytes@users.noreply.github.com> Date: Sat, 24 Jun 2023 22:26:36 +0500 Subject: [PATCH 5/8] Add model support --- api/keys.py | 33 ++++++++++++++++++--------------- 1 file changed, 18 insertions(+), 15 deletions(-) diff --git a/api/keys.py b/api/keys.py index ace05b1..5eb978b 100644 --- a/api/keys.py +++ b/api/keys.py @@ -12,23 +12,25 @@ class Keys: # --- END OF CONFIG --- locked_keys = set() - cache = set() + cache = {} # Initialize MongoDB client = MongoClient(MONGO_URI) db = client.get_database('keys_db') collection = db['keys'] - def __init__(self, key: str): + def __init__(self, key: str, model: str): self.key = key + self.model = model if not Keys.cache: self._load_keys() def _load_keys(self) -> None: - cursor = self.collection.find({}, {'_id': 0, 'key_value': 1}) + cursor = Keys.collection.find({}, {'_id': 0, 'key_value': 1, 'model': 1}) for doc in cursor: key_value = doc['key_value'] - self.cache.add(key_value) + model = doc['model'] + Keys.cache.setdefault(model, set()).add(key_value) def lock(self) -> None: self.locked_keys.add(self.key) @@ -39,34 +41,35 @@ class Keys: def is_locked(self) -> bool: return self.key in self.locked_keys - def get() -> str: - key_candidates = list(Keys.cache) + @staticmethod + def get(model: str) -> str: + key_candidates = list(Keys.cache.get(model, set())) random.shuffle(key_candidates) for key_candidate in key_candidates: - key = Keys(key_candidate) + key = Keys(key_candidate, model) if not key.is_locked(): key.lock() return key.key - print("[WARN] No unlocked keys found in get keys request!") + print(f"[WARN] No unlocked keys found for model '{model}' in get keys request!") def delete(self) -> None: - self.collection.delete_one({'key_value': self.key}) + Keys.collection.delete_one({'key_value': self.key, 'model': self.model}) # Update cache try: - Keys.cache.remove(self.key) + Keys.cache[self.model].remove(self.key) except KeyError: - print("[WARN] Tried to remove a key from cache which was not present: " + self.key) + print(f"[WARN] Tried to remove a key from cache which was not present: {self.key}") def save(self) -> None: - self.collection.insert_one({'key_value': self.key}) + Keys.collection.insert_one({'key_value': self.key, 'model': self.model}) # Update cache - Keys.cache.add(self.key) + Keys.cache.setdefault(self.model, set()).add(self.key) # Usage example: # os.environ['MONGO_URI'] = "mongodb://localhost:27017" -# key_instance = Keys("example_key") +# key_instance = Keys("example_key", "example_model") # key_instance.save() -# key_value = Keys.get() +# key_value = Keys.get("example_model") From 04413a3c6ac664005541ba4dd8064a8902ce0039 Mon Sep 17 00:00:00 2001 From: Game_Time <108236317+RayBytes@users.noreply.github.com> Date: Sat, 24 Jun 2023 22:30:18 +0500 Subject: [PATCH 6/8] Update config.json for mongodb --- config.json | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/config.json b/config.json index 6971d41..fb89592 100644 --- a/config.json +++ b/config.json @@ -1,3 +1,3 @@ { - "DB_PATH": "keys.db" -} \ No newline at end of file + "MONGO_URI": "" +} From 3ef5746b1a363d4ddf5f293ca56f2b792e4e4e68 Mon Sep 17 00:00:00 2001 From: Game_Time <108236317+RayBytes@users.noreply.github.com> Date: Sat, 24 Jun 2023 23:56:56 +0500 Subject: [PATCH 7/8] merge main.py with new keys system --- api/main.py | 21 ++++++++++++++++----- 1 file changed, 16 insertions(+), 5 deletions(-) diff --git a/api/main.py b/api/main.py index c50a0a0..ed04c66 100644 --- a/api/main.py +++ b/api/main.py @@ -1,6 +1,7 @@ import os import httpx import fastapi +from keys import Keys from starlette.requests import Request from starlette.responses import StreamingResponse @@ -31,33 +32,43 @@ async def startup_event(): security.enable_proxy() security.ip_protection_check() + # Setup key cache + Keys() + @app.get('/') async def root(): """Returns the root endpoint.""" return { 'status': 'ok', - 'discord': 'https://discord.gg/mX9BYdFeQF', + 'discord': 'https://discord.gg/85gdcd57Xr', 'github': 'https://github.com/Luna-OSS' } async def _reverse_proxy(request: Request): target_url = f'https://api.openai.com/v1/{request.url.path}' - + key = Keys.get(request.body()['model']) + if not key: + return fastapi.responses.JSONResponse( + status_code=400, + content={ + 'error': 'No API Key for model given, please try again with a valid model.' + } + ) request_to_api = target_api_client.build_request( method=request.method, url=target_url, headers={ - 'Authorization': 'Bearer ' + os.getenv('OPENAI_KEY'), + 'Authorization': 'Bearer ' + key, 'Content-Type': 'application/json' }, content=await request.body(), ) api_response = await target_api_client.send(request_to_api, stream=True) - + print(f'[{request.method}] {request.url.path} {api_response.status_code}') - + Keys(key).unlock() return StreamingResponse( api_response.aiter_raw(), status_code=api_response.status_code, From 9c41ac77be2a58e9386e6d49cf9fa6ce155a7960 Mon Sep 17 00:00:00 2001 From: Game_Time <108236317+RayBytes@users.noreply.github.com> Date: Sun, 25 Jun 2023 00:03:45 +0500 Subject: [PATCH 8/8] revert setup.py --- setup.py | 56 -------------------------------------------------------- 1 file changed, 56 deletions(-) diff --git a/setup.py b/setup.py index 0f40caa..8c27dd1 100644 --- a/setup.py +++ b/setup.py @@ -1,8 +1,4 @@ import setuptools -import os -import sqlite3 -import argparse -import json with open('README.md', 'r') as fh: long_description = fh.read() @@ -22,55 +18,3 @@ setuptools.setup( 'Operating System :: OS Independent', ] ) - - -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() - -