Alotta changes.

Setup managers and modified other files, cleaning up codebase.
Created user.py class for future type usage.

.vscode/settings.json

@@ -14,5 +14,9 @@
     "hide-files.files": [
         "tests/__pycache__",
         "api/__pycache__"
+    ],
+    "python.analysis.extraPaths": [
+        ".",
+        "./api/db"
     ]
 }

api/core.py

@@ -4,7 +4,7 @@ import os
 import json
 import fastapi
 
-from db import users
+from users import UserManager
 
 from dhooks import Webhook, Embed
 from dotenv import load_dotenv
@@ -29,10 +29,12 @@ async def check_core_auth(request):
 async def get_users(discord_id: int, incoming_request: fastapi.Request):
     auth = await check_core_auth(incoming_request)
     if auth:
-        return auth_error
+        return auth
 
     # Get user by discord ID
-    if not await users.by_discord_id(discord_id):
+    manager = UserManager()
+    user = await manager.user_by_discord_id(discord_id)
+    if not user:
         return fastapi.Response(status_code=404, content='User not found.')
 
     return user
@@ -64,7 +66,9 @@ async def create_user(incoming_request: fastapi.Request):
     except (json.decoder.JSONDecodeError, AttributeError):
         return fastapi.Response(status_code=400, content='Invalid or no payload received.')
     
-    user = await users.create(discord_id)
+    # Create the user 
+    manager = UserManager()
+    user = await manager.create(discord_id)
     await new_user_webhook(user)
 
     return user

api/db/stats.py

@@ -8,16 +8,9 @@ from motor.motor_asyncio import AsyncIOMotorClient
 
 load_dotenv()
 
-## MONGODB Setup
-
-conn = AsyncIOMotorClient(os.getenv('MONGO_URI'))
-
-async def _get_collection(collection_name: str):
-    return conn['nova-core'][collection_name]
-
 ## Statistics
 
-class Stats:
+class StatsManager:
     """
     ### The manager for all statistics tracking
     Stats tracked:
@@ -28,37 +21,44 @@ class Stats:
     - Models
     - URL Paths
     """
-    async def add_date():
+
+    def __init__(self):
+        self.conn = AsyncIOMotorClient(os.getenv('MONGO_URI'))
+
+    async def _get_collection(self, collection_name: str):
+        return self.conn['nova-core'][collection_name]
+    
+    async def add_date(self):
         date = datetime.datetime.now(pytz.timezone('GMT')).strftime('%Y.%m.%d')
         year, month, day = date.split('.')
 
-        db = await _get_collection('stats')
+        db = await self._get_collection('stats')
         await db.update_one({}, {'$inc': {f'dates.{year}.{month}.{day}': 1}}, upsert=True)
 
-    async def add_ip_address(ip_address: str):
+    async def add_ip_address(self, ip_address: str):
         ip_address = ip_address.replace('.', '_')
-        db = await _get_collection('stats')
+        db = await self._get_collection('stats')
         await db.update_one({}, {'$inc': {f'ips.{ip_address}': 1}}, upsert=True)
 
-    async def add_target(url: str):
-        db = await _get_collection('stats')
+    async def add_target(self, url: str):
+        db = await self._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_collection('stats')
+    async def add_tokens(self, tokens: int, model: str):
+        db = await self._get_collection('stats')
         await db.update_one({}, {'$inc': {f'tokens.{model}': tokens}}, upsert=True)
 
-    async def add_model(model: str):
-        db = await _get_collection('stats')
+    async def add_model(self, model: str):
+        db = await self._get_collection('stats')
         await db.update_one({}, {'$inc': {f'models.{model}': 1}}, upsert=True)
 
-    async def add_path(path: str):
+    async def add_path(self, path: str):
         path = path.replace('/', '_')
-        db = await _get_collection('stats')
+        db = await self._get_collection('stats')
         await db.update_one({}, {'$inc': {f'paths.{path}': 1}}, upsert=True)
 
-    async def get_value(obj_filter):
-        db = await _get_collection('stats')
+    async def get_value(self, obj_filter):
+        db = await self._get_collection('stats')
         return await db.find_one({obj_filter})
 
 if __name__ == '__main__':

api/db/users.py

@@ -14,75 +14,85 @@ with open('config/credits.yml', encoding='utf8') as f:
 
 ## MONGODB Setup
 
-conn = AsyncIOMotorClient(os.getenv('MONGO_URI'))
+class UserManager:
+    """
+    ### Manager of all users in the database.
+    Following methods are available:
 
-async def _get_collection(collection_name: str):
-    return conn['nova-core'][collection_name]
-
-async def create(discord_id: str='') -> dict:
-    """Add a user to the mongodb
-
-    Args:
-        discord_id (str): Defaults to ''.
-
-    Returns:
-        dict: The user object
+    - `_get_collection(collection_name)`
+    - `create(discord_id)`
+    - `user_by_id(user_id)`
+    - `user_by_discord_id(discord_id)`
+    - `user_by_api_key(api_key)`
+    - `update_by_id(user_id, new_obj)`
+    - `update_by_filter(filter_object, new_obj)`
+    - `delete(user_id)`
     """
 
-    chars = string.ascii_letters + string.digits
+    def __init__(self):
+        self.conn = AsyncIOMotorClient(os.getenv('MONGO_URI'))
 
-    infix = os.getenv('KEYGEN_INFIX')
-    suffix = ''.join(random.choices(chars, k=20))
-    prefix = ''.join(random.choices(chars, k=20))
+    async def _get_collection(self, collection_name: str):
+        return self.conn['nova-core'][collection_name]
     
-    new_api_key = f'nv-{prefix}{infix}{suffix}'
+    async def get_all_users(self):
+        return self.conn['nova-core']['users']
 
-    new_user = {
-        'api_key': new_api_key,
-        'credits': credits_config['start-credits'],
-        'role': '',
-        'level': '',
-        'status': {
-            'active': True,
-            'ban_reason': '',
-        },
-        'auth': {
-            'discord': str(discord_id),
-            'github': None
+    async def create(self, discord_id: str = '') -> dict:
+        chars = string.ascii_letters + string.digits
+
+        infix = os.getenv('KEYGEN_INFIX')
+        suffix = ''.join(random.choices(chars, k=20))
+        prefix = ''.join(random.choices(chars, k=20))
+
+        new_api_key = f'nv-{prefix}{infix}{suffix}'
+
+        new_user = {
+            'api_key': new_api_key,
+            'credits': credits_config['start-credits'],
+            'role': '',
+            'level': '',
+            'status': {
+                'active': True,
+                'ban_reason': '',
+            },
+            'auth': {
+                'discord': str(discord_id),
+                'github': None
+            }
         }
-    }
 
-    db = await _get_collection('users')
-    await db.insert_one(new_user)
-    user = await db.find_one({'api_key': new_api_key})	
-    return user
+        db = await self._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_collection('users')
-    return await db.find_one({'_id': user_id})
+    async def user_by_id(self, user_id: str):
+        db = await self._get_collection('users')
+        return await db.find_one({'_id': user_id})
 
-async def by_discord_id(discord_id: str):
-    db = await _get_collection('users')
-    return await db.find_one({'auth.discord': str(int(discord_id))})
+    async def user_by_discord_id(self, discord_id: str):
+        db = await self._get_collection('users')
+        return await db.find_one({'auth.discord': str(int(discord_id))})
 
-async def by_api_key(key: str):
-    db = await _get_collection('users')
-    return await db.find_one({'api_key': key})
+    async def user_by_api_key(self, key: str):
+        db = await self._get_collection('users')
+        return await db.find_one({'api_key': key})
 
-async def update_by_id(user_id: str, update):
-    db = await _get_collection('users')
-    return await db.update_one({'_id': user_id}, update)
+    async def update_by_id(self, user_id: str, update):
+        db = await self._get_collection('users')
+        return await db.update_one({'_id': user_id}, update)
 
-async def update_by_filter(obj_filter, update):
-    db = await _get_collection('users')
-    return await db.update_one(obj_filter, update)
+    async def update_by_filter(self, obj_filter, update):
+        db = await self._get_collection('users')
+        return await db.update_one(obj_filter, update)
 
-async def delete(user_id: str):
-    db = await _get_collection('users')
-    await db.delete_one({'_id': user_id})
+    async def delete(self, user_id: str):
+        db = await self._get_collection('users')
+        await db.delete_one({'_id': user_id})
 
 async def demo():
-    user = await create(69420)
+    user = await UserManager().create(69420)
     print(user)
 
 if __name__ == '__main__':

api/streaming.py

@@ -15,8 +15,9 @@ import proxies
 import provider_auth
 import load_balancing
 
-from db import logs, users
-from db.stats import Stats
+from db import logs
+from db.users import UserManager
+from db.stats import StatsManager
 from helpers import network, chat, errors
 
 load_dotenv()
@@ -76,7 +77,8 @@ async def stream(
         input_tokens (int, optional): Total tokens calculated with tokenizer. Defaults to 0.
         incoming_request (starlette.requests.Request, optional): Incoming request. Defaults to None.
     """
-
+    db = UserManager()
+    stats = StatsManager()
     is_chat = False
     is_stream = payload.get('stream', False)
 
@@ -175,16 +177,16 @@ async def stream(
         await logs.log_api_request(user=user, incoming_request=incoming_request, target_url=target_request['url'])
 
     if credits_cost and user:
-        await users.update_by_id(user['_id'], {'$inc': {'credits': -credits_cost}})
+        await db.update_by_id(user['_id'], {'$inc': {'credits': -credits_cost}})
 
     ip_address = await network.get_ip(incoming_request)
-    await Stats.add_date()
-    await Stats.add_ip_address(ip_address)
-    await Stats.add_path(path)
-    await Stats.add_target(target_request['url'])
+    await stats.add_date()
+    await stats.add_ip_address(ip_address)
+    await stats.add_path(path)
+    await stats.add_target(target_request['url'])
     if is_chat:
-        await Stats.add_model(model)
-        await Stats.add_tokens(input_tokens, model)
+        await stats.add_model(model)
+        await stats.add_tokens(input_tokens, model)
 
 if __name__ == '__main__':
     asyncio.run(stream())

api/transfer.py

@@ -9,7 +9,7 @@ from dotenv import load_dotenv
 import streaming
 import moderation
 
-from db import users
+from users import UserManager
 from helpers import tokens, errors
 
 load_dotenv()
@@ -25,7 +25,7 @@ async def handle(incoming_request):
     Takes the request from the incoming request to the target endpoint.
     Checks method, token amount, auth and cost along with if request is NSFW.
     """
-    
+    users = UserManager()
     path = incoming_request.url.path.replace('v1/v1/', 'v1/')
 
     allowed_methods = {'GET', 'POST', 'PUT', 'DELETE', 'PATCH'}
@@ -46,7 +46,7 @@ async def handle(incoming_request):
     if not received_key or not received_key.startswith('Bearer '):
         return await errors.error(401, 'No NovaAI API key given!', 'Add "Authorization: Bearer nv-..." to your request headers.')
 
-    user = await users.by_api_key(received_key.split('Bearer ')[1].strip())
+    user = await users.user_by_api_key(received_key.split('Bearer ')[1].strip())
 
     if not user or not user['status']['active']:
         return await errors.error(401, 'Invalid or inactive NovaAI API key!', 'Create a new NovaOSS API key or reactivate your account.')

api/types/user.py

@@ -0,0 +1 @@
+## TODO: Create user type out of JSON object.

rewards/autocredits.py

@@ -1,10 +1,9 @@
-async def get_all_users(client):
-    users = client['nova-core']['users']
-    return users
+from users import UserManager
 
 
 async def update_credits(pymongo_client, settings=None):
-    users = await get_all_users(pymongo_client)
+    manager = UserManager()
+    users = await manager.get_all_users(pymongo_client)
 
     if not settings:
         users.update_many({}, {'$inc': {'credits': 2500}})