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.