Another general cleanup commit

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

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__':

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__':

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():

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'

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 <provider_and_key> should be passed is:
     <provider_name><key>, e.g.