From 1a3e275a1c4ec0588861967d4ac5757dcb5acc1c Mon Sep 17 00:00:00 2001 From: nsde Date: Mon, 7 Aug 2023 23:28:24 +0200 Subject: [PATCH] Added auto-rewards and invalid key system --- api/load_balancing.py | 12 +----- api/moderation.py | 6 ++- api/provider_auth.py | 21 ++++++++++ api/transfer.py | 25 ++++------- rewards/__main__.py | 2 + {rewardsystem => rewards}/autocredits.py | 4 +- rewards/main.py | 53 ++++++++++++++++++++++++ {rewardsystem => rewards}/settings.py | 0 rewardsystem/main.py | 48 --------------------- tests/__main__.py | 2 +- 10 files changed, 93 insertions(+), 80 deletions(-) create mode 100644 api/provider_auth.py create mode 100644 rewards/__main__.py rename {rewardsystem => rewards}/autocredits.py (71%) create mode 100644 rewards/main.py rename {rewardsystem => rewards}/settings.py (100%) delete mode 100644 rewardsystem/main.py diff --git a/api/load_balancing.py b/api/load_balancing.py index 0755449..8673d6c 100644 --- a/api/load_balancing.py +++ b/api/load_balancing.py @@ -3,14 +3,6 @@ import asyncio import providers -provider_modules = [ - # providers.twa, - # providers.quantum, - providers.churchless, - providers.closed, - providers.closed4 -] - async def _get_module_name(module) -> str: name = module.__name__ if '.' in name: @@ -22,7 +14,7 @@ async def balance_chat_request(payload: dict) -> dict: providers_available = [] - for provider_module in provider_modules: + for provider_module in providers.MODULES: if payload['stream'] and not provider_module.STREAMING: continue @@ -52,7 +44,7 @@ async def balance_organic_request(request: dict) -> dict: 'Content-Type': 'application/json' } - for provider_module in provider_modules: + for provider_module in providers.MODULES: if not provider_module.ORGANIC: continue diff --git a/api/moderation.py b/api/moderation.py index be9add7..38b2d8b 100644 --- a/api/moderation.py +++ b/api/moderation.py @@ -1,7 +1,8 @@ import asyncio - import aiohttp + import proxies +import provider_auth import load_balancing async def is_safe(inp) -> bool: @@ -37,11 +38,12 @@ async def is_safe(inp) -> bool: timeout=aiohttp.ClientTimeout(total=5), ) as res: res.raise_for_status() - json_response = await res.json() return not json_response['results'][0]['flagged'] + except Exception as exc: + await provider_auth.invalidate_key(req.get('provider_auth')) print('[!] moderation error:', type(exc), exc) continue diff --git a/api/provider_auth.py b/api/provider_auth.py new file mode 100644 index 0000000..9393109 --- /dev/null +++ b/api/provider_auth.py @@ -0,0 +1,21 @@ +import asyncio + +async def invalidate_key(provider_and_key): + if not provider_and_key: + return + + provider = provider_and_key.split('>')[0] + provider_file = f'secret/{provider}.txt' + key = provider_and_key.split('>')[1] + + with open(provider_file, encoding='utf8') as f_in: + text = f_in.read() + + with open(provider_file, 'w', encoding='utf8') as f_out: + f_out.write(text.replace(key, '')) + + with open(f'secret/{provider}.invalid.txt', 'a', encoding='utf8') as f: + f.write(key + '\n') + +if __name__ == '__main__': + asyncio.run(invalidate_key('closed>sk-...')) diff --git a/api/transfer.py b/api/transfer.py index 3159941..13f80ac 100644 --- a/api/transfer.py +++ b/api/transfer.py @@ -1,36 +1,26 @@ """Module for transferring requests to ClosedAI API""" -import os import json import yaml -import logging import fastapi -import starlette from dotenv import load_dotenv import streaming import moderation -from db import logs, users -from helpers import tokens, errors, exceptions +from db import users +from helpers import tokens, errors load_dotenv() -# log to "api.log" file -logging.basicConfig( - filename='api.log', - level=logging.DEBUG, - format='%(asctime)s %(levelname)s %(name)s %(message)s' -) - with open('config/credits.yml', encoding='utf8') as f: credits_config = yaml.safe_load(f) async def handle(incoming_request): """Transfer a streaming response from the incoming request to the target endpoint""" - path = incoming_request.url.path + path = incoming_request.url.path.replace('v1/v1/', 'v1/') # METHOD if incoming_request.method not in ['GET', 'POST', 'PUT', 'DELETE', 'PATCH']: @@ -91,10 +81,11 @@ async def handle(incoming_request): else: inp = payload.get('input', payload.get('prompt')) - if inp and not '/moderations' in path: - is_safe = await moderation.is_safe(inp) - - if not is_safe: + if inp: + if len(inp) > 2 and not inp.isnumeric(): + is_safe = await moderation.is_safe(inp) + + if not is_safe and not '/moderations' in path: error = await errors.error(400, 'The request contains content which violates this model\'s policies.', 'We currently don\'t support any NSFW models.') return error diff --git a/rewards/__main__.py b/rewards/__main__.py new file mode 100644 index 0000000..8fea19b --- /dev/null +++ b/rewards/__main__.py @@ -0,0 +1,2 @@ +import main +main.launch() \ No newline at end of file diff --git a/rewardsystem/autocredits.py b/rewards/autocredits.py similarity index 71% rename from rewardsystem/autocredits.py rename to rewards/autocredits.py index 7971fab..6e870cb 100644 --- a/rewardsystem/autocredits.py +++ b/rewards/autocredits.py @@ -7,9 +7,9 @@ async def update_credits(pymongo_client, settings=None): users = await get_all_users(pymongo_client) if not settings: - users.update_many({}, {"$inc": {"credits": 2500}}) + users.update_many({}, {'$inc': {'credits': 2500}}) else: for key, value in settings.items(): users.update_many( - {'level': key}, {"$inc": {"credits": int(value)}}) + {'level': key}, {'$inc': {'credits': int(value)}}) diff --git a/rewards/main.py b/rewards/main.py new file mode 100644 index 0000000..ba0e87a --- /dev/null +++ b/rewards/main.py @@ -0,0 +1,53 @@ +import os +import time +import aiohttp +import pymongo +import asyncio +import autocredits + +from settings import roles +from dotenv import load_dotenv + +load_dotenv() +async def main(): + mongo = pymongo.MongoClient(os.getenv('MONGO_URI')) + + await update_roles(mongo) + await autocredits.update_credits(mongo, roles) + +async def update_roles(mongo): + async with aiohttp.ClientSession() as session: + try: + async with session.get('http://0.0.0.0:3224/get_roles') as response: + discord_users = await response.json() + except aiohttp.ClientError as e: + print(f'Error: {e}') + return + + level_role_names = [f'lvl{lvl}' for lvl in range(10, 110, 10)] + users = await autocredits.get_all_users(mongo) + + for user in users.find(): + discord = str(user['auth']['discord']) + + for user_id, role_names in discord_users.items(): + if user_id == discord: + for role in level_role_names: + if role in role_names: + users.update_one( + {'auth.discord': int(discord)}, + {'$set': {'level': role}} + ) + print(f'Updated {discord} to {role}') + + return users + +def launch(): + asyncio.run(main()) + + with open('rewards/last_update.txt', 'w') as f: + f.write(str(time.time())) + + +if __name__ == '__main__': + launch() diff --git a/rewardsystem/settings.py b/rewards/settings.py similarity index 100% rename from rewardsystem/settings.py rename to rewards/settings.py diff --git a/rewardsystem/main.py b/rewardsystem/main.py deleted file mode 100644 index f4b943a..0000000 --- a/rewardsystem/main.py +++ /dev/null @@ -1,48 +0,0 @@ -import asyncio -from settings import roles -import autocredits -import aiohttp -from dotenv import load_dotenv -import os -import pymongo - -load_dotenv() - -CONNECTION_STRING = os.getenv("MONGO_URI") - - -async def main(): - pymongo_client = pymongo.MongoClient(CONNECTION_STRING) - - await update_roles(pymongo_client) - await autocredits.update_credits(pymongo_client, roles) - - -async def update_roles(pymongo_client): - async with aiohttp.ClientSession() as session: - try: - async with session.get('http://0.0.0.0:3224/get_roles') as response: - data = await response.json() - except aiohttp.ClientError as e: - print(f"Error: {e}") - return - - lvlroles = [f"lvl{lvl}" for lvl in range(10, 110, 10)] - discord_users = data - users = await autocredits.get_all_users(pymongo_client) - - for user in users.find(): - discord = str(user['auth']['discord']) - - for id_, roles in discord_users.items(): - if id_ == discord: - for role in lvlroles: - if role in roles: - users.update_one({'auth.discord': int(discord)}, { - '$set': {'level': role}}) - print(f"Updated {discord} to {role}") - - return users - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/tests/__main__.py b/tests/__main__.py index 0a9f134..0033ee7 100644 --- a/tests/__main__.py +++ b/tests/__main__.py @@ -23,7 +23,7 @@ MODEL = 'gpt-3.5-turbo' MESSAGES = [ { 'role': 'user', - 'content': 'fuck you', + 'content': '1+1=', } ]