From afc3b08351e572985f9e2121c16ac8093e6f8897 Mon Sep 17 00:00:00 2001 From: Game_Time <108236317+RayBytes@users.noreply.github.com> Date: Sat, 12 Aug 2023 20:56:21 +0500 Subject: [PATCH 1/4] some general changes --- api/core.py | 2 +- api/helpers/chat.py | 10 ++++++++++ api/helpers/exceptions.py | 2 -- api/helpers/tokens.py | 13 ++++++++++++- api/load_balancing.py | 2 +- api/transfer.py | 32 +++++++++++++------------------- 6 files changed, 37 insertions(+), 24 deletions(-) delete mode 100644 api/helpers/exceptions.py diff --git a/api/core.py b/api/core.py index f266147..b4622d0 100644 --- a/api/core.py +++ b/api/core.py @@ -51,7 +51,7 @@ async def create_user(incoming_request: fastapi.Request): auth_error = await check_core_auth(incoming_request) if auth_error: - return auth_error + return auth_error try: payload = await incoming_request.json() diff --git a/api/helpers/chat.py b/api/helpers/chat.py index a1002fc..bde2233 100644 --- a/api/helpers/chat.py +++ b/api/helpers/chat.py @@ -18,6 +18,16 @@ async def create_chat_id() -> str: return f'chatcmpl-{chat_id}' async def create_chat_chunk(chat_id: str, model: str, content=None) -> dict: + """Creates the chunk for streaming chat. + + Args: + chat_id (str): _description_ + model (str): _description_ + content (_type_, optional): _description_. Defaults to None. + + Returns: + dict: _description_ + """ content = content or {} delta = {} diff --git a/api/helpers/exceptions.py b/api/helpers/exceptions.py deleted file mode 100644 index 1fcff8b..0000000 --- a/api/helpers/exceptions.py +++ /dev/null @@ -1,2 +0,0 @@ -class Retry(Exception): - """The server should retry the request.""" diff --git a/api/helpers/tokens.py b/api/helpers/tokens.py index e1aede6..418bc6f 100644 --- a/api/helpers/tokens.py +++ b/api/helpers/tokens.py @@ -1,7 +1,18 @@ import tiktoken async def count_for_messages(messages: list, model: str='gpt-3.5-turbo-0613') -> int: - """Return the number of tokens used by a list of messages.""" + """Return the number of tokens used by a list of messages + + Args: + messages (list): _description_ + model (str, optional): _description_. Defaults to 'gpt-3.5-turbo-0613'. + + Raises: + NotImplementedError: _description_ + + Returns: + int: _description_ + """ try: encoding = tiktoken.encoding_for_model(model) diff --git a/api/load_balancing.py b/api/load_balancing.py index 3fb6f54..4d4f9fd 100644 --- a/api/load_balancing.py +++ b/api/load_balancing.py @@ -35,7 +35,7 @@ async def balance_chat_request(payload: dict) -> dict: return target async def balance_organic_request(request: dict) -> dict: - """Load balnace to non-chat completion request between other "organic" providers which respond in the desired format already.""" + """Load balance to non-chat completion request between other "organic" providers which respond in the desired format already.""" providers_available = [] diff --git a/api/transfer.py b/api/transfer.py index a4f83f6..acdc347 100644 --- a/api/transfer.py +++ b/api/transfer.py @@ -26,8 +26,7 @@ async def handle(incoming_request): # METHOD if incoming_request.method not in ['GET', 'POST', 'PUT', 'DELETE', 'PATCH']: - error = await errors.error(405, f'Method "{incoming_request.method}" is not allowed.', 'Change the request method to the correct one.') - return error + return await errors.error(405, f'Method "{incoming_request.method}" is not allowed.', 'Change the request method to the correct one.') # PAYLOAD try: @@ -35,42 +34,37 @@ async def handle(incoming_request): except json.decoder.JSONDecodeError: payload = {} - # TOKENS + # Tokenise w/ tiktoken try: input_tokens = await tokens.count_for_messages(payload['messages']) except (KeyError, TypeError): input_tokens = 0 - # AUTH + # Check user auth received_key = incoming_request.headers.get('Authorization') if not received_key: - error = await errors.error(401, 'No NovaAI API key given!', 'Add "Authorization: Bearer nv-..." to your request headers.') - return error + return await errors.error(401, 'No NovaAI API key given!', 'Add "Authorization: Bearer nv-..." to your request headers.') if received_key.startswith('Bearer '): received_key = received_key.split('Bearer ')[1] - # USER user = await users.by_api_key(received_key.strip()) if not user: - error = await errors.error(401, 'Invalid NovaAI API key!', 'Create a new NovaOSS API key.') - return error + return await errors.error(401, 'Invalid NovaAI API key!', 'Create a new NovaOSS API key.') ban_reason = user['status']['ban_reason'] if ban_reason: - error = await errors.error(403, f'Your NovaAI account has been banned. Reason: "{ban_reason}".', 'Contact the staff for an appeal.') - return error + return await errors.error(403, f'Your NovaAI account has been banned. Reason: "{ban_reason}".', 'Contact the staff for an appeal.') if not user['status']['active']: - error = await errors.error(418, 'Your NovaAI account is not active (paused).', 'Simply re-activate your account using a Discord command or the web panel.') - return error + return await errors.error(418, 'Your NovaAI account is not active (paused).', 'Simply re-activate your account using a Discord command or the web panel.') if '/models' in path: return fastapi.responses.JSONResponse(content=models_list) - # COST + # Calculate cost of tokens & check for nsfw prompts costs = credits_config['costs'] cost = costs['other'] @@ -94,17 +88,17 @@ async def handle(incoming_request): policy_violation = await moderation.is_policy_violated(inp) if policy_violation: - error = await errors.error(400, f'The request contains content which violates this model\'s policies for "{policy_violation}".', 'We currently don\'t support any NSFW models.') - return error + return await errors.error(400, f'The request contains content which violates this model\'s policies for "{policy_violation}".', 'We currently don\'t support any NSFW models.') + role_cost_multiplier = credits_config['bonuses'].get(user['role'], 1) cost = round(cost * role_cost_multiplier) if user['credits'] < cost: - error = await errors.error(429, 'Not enough credits.', 'Wait or earn more credits. Learn more on our website or Discord server.') - return error + return await errors.error(429, 'Not enough credits.', 'Wait or earn more credits. Learn more on our website or Discord server.') - # READY + + # Send the completion request if 'chat/completions' in path and not payload.get('stream') is True: payload['stream'] = False From f4a09515fec0ea5168947cb300ecd5954bebe8e1 Mon Sep 17 00:00:00 2001 From: Game_Time <108236317+RayBytes@users.noreply.github.com> Date: Sat, 12 Aug 2023 21:01:55 +0500 Subject: [PATCH 2/4] modified stuff --- api/core.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/api/core.py b/api/core.py index 02954e3..f266147 100644 --- a/api/core.py +++ b/api/core.py @@ -51,7 +51,7 @@ async def create_user(incoming_request: fastapi.Request): auth_error = await check_core_auth(incoming_request) if auth_error: - return auth_error + return auth_error try: payload = await incoming_request.json() @@ -63,3 +63,14 @@ async def create_user(incoming_request: fastapi.Request): await new_user_webhook(user) return user + +if __name__ == '__main__': + # new_user_webhook({ + # '_id': 'JUST_A_TEST_IGNORE_ME', + # 'auth': { + # 'discord': 123, + # 'github': 'abc' + # } + # }) + + pass From 1da39e5242ea7e44dd673933149a0a7b6302ddfe Mon Sep 17 00:00:00 2001 From: Game_Time <108236317+RayBytes@users.noreply.github.com> Date: Sat, 12 Aug 2023 21:05:53 +0500 Subject: [PATCH 3/4] a few more nice changes --- api/provider_auth.py | 13 ++++++++----- api/streaming.py | 6 ++---- 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/api/provider_auth.py b/api/provider_auth.py index bc499d7..5b19ec7 100644 --- a/api/provider_auth.py +++ b/api/provider_auth.py @@ -3,11 +3,14 @@ import asyncio async def invalidate_key(provider_and_key: str) -> none: - """Invalidates a key stored in the secret/ folder by storing it in the associated .invalid.txt file. -The schmea in which should be passed is: -, e.g. -closed4>sk-... -""" + """ + + Invalidates a key stored in the secret/ folder by storing it in the associated .invalid.txt file. + The schmea in which should be passed is: + , e.g. + closed4>sk-... + + """ if not provider_and_key: return diff --git a/api/streaming.py b/api/streaming.py index 295b06b..a1c7402 100644 --- a/api/streaming.py +++ b/api/streaming.py @@ -96,12 +96,11 @@ async def stream( webhook = dhooks.Webhook(os.getenv('DISCORD_WEBHOOK__API_ISSUE')) webhook.send(content=f'API Issue: **`{exc}`**\nhttps://i.imgflip.com/7uv122.jpg') - error = await errors.yield_error( + yield await errors.yield_error( 500, 'Sorry, the API has no working keys anymore.', 'The admins have been messaged automatically.' ) - yield error return for k, v in target_request.get('headers', {}).items(): @@ -180,12 +179,11 @@ async def stream( except Exception as exc: if 'Connection closed' in str(exc): - error = await errors.yield_error( + yield await errors.yield_error( 500, 'Sorry, there was an issue with the connection.', 'Please first check if the issue on your end. If this error repeats, please don\'t heistate to contact the staff!.' ) - yield error return break From d5e28465515e6a6536b2a8249a5119aafc87d3b4 Mon Sep 17 00:00:00 2001 From: Game_Time <108236317+RayBytes@users.noreply.github.com> Date: Sat, 12 Aug 2023 21:20:18 +0500 Subject: [PATCH 4/4] Another general cleanup commit --- api/core.py | 11 ----------- api/db/logs.py | 18 +++++++++++------- api/db/stats.py | 24 +++++++++++++++--------- api/db/users.py | 22 +++++++++++++--------- api/helpers/chat.py | 11 +---------- api/provider_auth.py | 4 ++-- 6 files changed, 42 insertions(+), 48 deletions(-) diff --git a/api/core.py b/api/core.py index f266147..fe16cd9 100644 --- a/api/core.py +++ b/api/core.py @@ -63,14 +63,3 @@ async def create_user(incoming_request: fastapi.Request): await new_user_webhook(user) return user - -if __name__ == '__main__': - # new_user_webhook({ - # '_id': 'JUST_A_TEST_IGNORE_ME', - # 'auth': { - # 'discord': 123, - # 'github': 'abc' - # } - # }) - - pass diff --git a/api/db/logs.py b/api/db/logs.py index caf2baf..9e6a9ea 100644 --- a/api/db/logs.py +++ b/api/db/logs.py @@ -16,8 +16,12 @@ UA_SIMPLIFY = { 'AppleWebKit/537.36 (KHTML, like Gecko)': 'K', } -async def _get_mongo(collection_name: str): - return AsyncIOMotorClient(os.getenv('MONGO_URI'))['nova-core'][collection_name] +## MONGODB Setup + +conn = AsyncIOMotorClient(os.getenv('MONGO_URI')) + +async def _get_collection(collection_name: str): + return conn['nova-core'][collection_name] async def replacer(text: str, dict_: dict) -> str: for k, v in dict_.items(): @@ -25,7 +29,7 @@ async def replacer(text: str, dict_: dict) -> str: return text async def log_api_request(user: dict, incoming_request, target_url: str): - db = await _get_mongo('logs') + db = await _get_collection('logs') payload = {} try: @@ -58,19 +62,19 @@ async def log_api_request(user: dict, incoming_request, target_url: str): return log_item async def by_id(log_id: str): - db = await _get_mongo('logs') + db = await _get_collection('logs') return await db.find_one({'_id': log_id}) async def by_user_id(user_id: str): - db = await _get_mongo('logs') + db = await _get_collection('logs') return await db.find({'user_id': user_id}) async def delete_by_id(log_id: str): - db = await _get_mongo('logs') + db = await _get_collection('logs') return await db.delete_one({'_id': log_id}) async def delete_by_user_id(user_id: str): - db = await _get_mongo('logs') + db = await _get_collection('logs') return await db.delete_many({'user_id': user_id}) if __name__ == '__main__': diff --git a/api/db/stats.py b/api/db/stats.py index 795b09f..d20ee4a 100644 --- a/api/db/stats.py +++ b/api/db/stats.py @@ -8,40 +8,46 @@ from motor.motor_asyncio import AsyncIOMotorClient load_dotenv() -async def _get_mongo(collection_name: str): - return AsyncIOMotorClient(os.getenv('MONGO_URI'))['nova-core'][collection_name] +## MONGODB Setup + +conn = AsyncIOMotorClient(os.getenv('MONGO_URI')) + +async def _get_collection(collection_name: str): + return conn['nova-core'][collection_name] + +## Statistics async def add_date(): date = datetime.datetime.now(pytz.timezone('GMT')).strftime('%Y.%m.%d') year, month, day = date.split('.') - db = await _get_mongo('stats') + db = await _get_collection('stats') await db.update_one({}, {'$inc': {f'dates.{year}.{month}.{day}': 1}}, upsert=True) async def add_ip_address(ip_address: str): ip_address = ip_address.replace('.', '_') - db = await _get_mongo('stats') + db = await _get_collection('stats') await db.update_one({}, {'$inc': {f'ips.{ip_address}': 1}}, upsert=True) async def add_target(url: str): - db = await _get_mongo('stats') + db = await _get_collection('stats') await db.update_one({}, {'$inc': {f'targets.{url}': 1}}, upsert=True) async def add_tokens(tokens: int, model: str): - db = await _get_mongo('stats') + db = await _get_collection('stats') await db.update_one({}, {'$inc': {f'tokens.{model}': tokens}}, upsert=True) async def add_model(model: str): - db = await _get_mongo('stats') + db = await _get_collection('stats') await db.update_one({}, {'$inc': {f'models.{model}': 1}}, upsert=True) async def add_path(path: str): path = path.replace('/', '_') - db = await _get_mongo('stats') + db = await _get_collection('stats') await db.update_one({}, {'$inc': {f'paths.{path}': 1}}, upsert=True) async def get_value(obj_filter): - db = await _get_mongo('stats') + db = await _get_collection('stats') return await db.find_one({obj_filter}) if __name__ == '__main__': diff --git a/api/db/users.py b/api/db/users.py index 580364a..8bafb21 100644 --- a/api/db/users.py +++ b/api/db/users.py @@ -12,8 +12,12 @@ load_dotenv() with open('config/credits.yml', encoding='utf8') as f: credits_config = yaml.safe_load(f) -async def _get_mongo(collection_name: str): - return AsyncIOMotorClient(os.getenv('MONGO_URI'))['nova-core'][collection_name] +## MONGODB Setup + +conn = AsyncIOMotorClient(os.getenv('MONGO_URI')) + +async def _get_collection(collection_name: str): + return conn['nova-core'][collection_name] async def create(discord_id: str='') -> dict: """Adds a new user to the MongoDB collection.""" @@ -41,33 +45,33 @@ async def create(discord_id: str='') -> dict: } } - db = await _get_mongo('users') + db = await _get_collection('users') await db.insert_one(new_user) user = await db.find_one({'api_key': new_api_key}) return user async def by_id(user_id: str): - db = await _get_mongo('users') + db = await _get_collection('users') return await db.find_one({'_id': user_id}) async def by_discord_id(discord_id: str): - db = await _get_mongo('users') + db = await _get_collection('users') return await db.find_one({'auth.discord': str(int(discord_id))}) async def by_api_key(key: str): - db = await _get_mongo('users') + db = await _get_collection('users') return await db.find_one({'api_key': key}) async def update_by_id(user_id: str, update): - db = await _get_mongo('users') + db = await _get_collection('users') return await db.update_one({'_id': user_id}, update) async def update_by_filter(obj_filter, update): - db = await _get_mongo('users') + db = await _get_collection('users') return await db.update_one(obj_filter, update) async def delete(user_id: str): - db = await _get_mongo('users') + db = await _get_collection('users') await db.delete_one({'_id': user_id}) async def demo(): diff --git a/api/helpers/chat.py b/api/helpers/chat.py index 2ab3f8e..6bca378 100644 --- a/api/helpers/chat.py +++ b/api/helpers/chat.py @@ -61,13 +61,4 @@ async def create_chat_chunk(chat_id: str, model: str, content=None) -> dict: ], } - return f'data: {json.dumps(chunk)}\n\n' - -if __name__ == '__main__': - demo_chat_id = asyncio.run(create_chat_id()) - print(demo_chat_id) - print(asyncio.run(create_chat_chunk( - model='gpt-4', - content='Hello', - chat_id=demo_chat_id, - ))) + return f'data: {json.dumps(chunk)}\n\n' \ No newline at end of file diff --git a/api/provider_auth.py b/api/provider_auth.py index 5b19ec7..a6aaed3 100644 --- a/api/provider_auth.py +++ b/api/provider_auth.py @@ -2,9 +2,9 @@ import asyncio -async def invalidate_key(provider_and_key: str) -> none: +async def invalidate_key(provider_and_key: str) -> None: """ - + Invalidates a key stored in the secret/ folder by storing it in the associated .invalid.txt file. The schmea in which should be passed is: , e.g.