From 6ecc5f59ce4b818757b7db90c7427cac651429f7 Mon Sep 17 00:00:00 2001
From: Game_Time <108236317+RayBytes@users.noreply.github.com>
Date: Sun, 13 Aug 2023 20:12:35 +0500
Subject: [PATCH] Codebase changes + a lot of commenting

---
 api/db/logs.py        | 13 +++++++++++++
 api/db/users.py       |  9 ++++++++-
 api/load_balancing.py | 14 +++++++++-----
 api/main.py           |  4 +++-
 api/proxies.py        |  6 +++++-
 api/streaming.py      | 10 ++++++++++
 api/transfer.py       |  6 +++++-
 7 files changed, 53 insertions(+), 9 deletions(-)

diff --git a/api/db/logs.py b/api/db/logs.py
index 9e6a9ea..7ecee9f 100644
--- a/api/db/logs.py
+++ b/api/db/logs.py
@@ -24,11 +24,24 @@ async def _get_collection(collection_name: str):
     return conn['nova-core'][collection_name]
 
 async def replacer(text: str, dict_: dict) -> str:
+    # This seems to exist for a very specific and dumb purpose :D
     for k, v in dict_.items():
         text = text.replace(k, v)
     return text
 
 async def log_api_request(user: dict, incoming_request, target_url: str):
+    """Logs the API Request into the database.
+    No input prompt is logged, however data such as IP & useragent is noted.
+    This would be useful for security reasons. Other minor data is also collected.
+
+    Args:
+        user (dict): User dict object
+        incoming_request (_type_): Request
+        target_url (str): The URL the api request was targetted to.
+
+    Returns:
+        _type_: _description_
+    """
     db = await _get_collection('logs')
     payload = {}
 
diff --git a/api/db/users.py b/api/db/users.py
index 8bafb21..ab4a4c7 100644
--- a/api/db/users.py
+++ b/api/db/users.py
@@ -20,7 +20,14 @@ 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."""
+    """Add a user to the mongodb
+
+    Args:
+        discord_id (str): Defaults to ''.
+
+    Returns:
+        dict: The user object
+    """
 
     chars = string.ascii_letters + string.digits
 
diff --git a/api/load_balancing.py b/api/load_balancing.py
index 84adb8a..bc781ea 100644
--- a/api/load_balancing.py
+++ b/api/load_balancing.py
@@ -10,8 +10,10 @@ async def _get_module_name(module) -> str:
     return name
 
 async def balance_chat_request(payload: dict) -> dict:
-    """Load balance the chat completion request between chat providers.
-"""
+    """
+    ### Load balance the chat completion request between chat providers.
+    Providers are sorted by streaming and models. Target (provider.chat_completion) is returned
+    """
 
     providers_available = []
 
@@ -36,9 +38,11 @@ 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.
-Organic providers are used for non-chat completions, such as moderation and other paths.    
-"""
+    """
+    ### Load balance non-chat completion request
+    Balances between other "organic" providers which respond in the desired format already.
+    Organic providers are used for non-chat completions, such as moderation and other paths.    
+    """
     providers_available = []
 
     if not request.get('headers'):
diff --git a/api/main.py b/api/main.py
index b177efa..8bc8810 100644
--- a/api/main.py
+++ b/api/main.py
@@ -31,7 +31,9 @@ async def startup_event():
 
 @app.get('/')
 async def root():
-    """Returns the root endpoint."""
+    """
+    Returns the root endpoint.
+    """
 
     return {
         'status': 'ok',
diff --git a/api/proxies.py b/api/proxies.py
index e692384..0dd84b2 100644
--- a/api/proxies.py
+++ b/api/proxies.py
@@ -63,7 +63,11 @@ class Proxy:
 
     @property
     def connector(self):
-        """Returns an aiohttp_socks.ProxyConnector object. Which can be used in aiohttp.ClientSession."""
+        """
+        ### Returns a proxy connector
+        Returns an aiohttp_socks.ProxyConnector object. 
+        This can be used in aiohttp.ClientSession.
+        """
 
         proxy_types = {
             'http': aiohttp_socks.ProxyType.HTTP,
diff --git a/api/streaming.py b/api/streaming.py
index a1c7402..9b8e4f6 100644
--- a/api/streaming.py
+++ b/api/streaming.py
@@ -38,6 +38,16 @@ async def stream(
     input_tokens: int=0,
     incoming_request: starlette.requests.Request=None,
 ):
+    """Stream the completions request. Sends data in chunks
+
+    Args:
+        path (str, optional): URL Path. Defaults to '/v1/chat/completions'.
+        user (dict, optional): User object (dict) Defaults to None.
+        payload (dict, optional): Payload. Defaults to None.
+        credits_cost (int, optional): Cost of the credits of the request. Defaults to 0.
+        input_tokens (int, optional): Total tokens calculated with tokenizer. Defaults to 0.
+        incoming_request (starlette.requests.Request, optional): Incoming request. Defaults to None.
+    """
     is_chat = False
     is_stream = payload.get('stream', False)
 
diff --git a/api/transfer.py b/api/transfer.py
index c7931c3..157c68a 100644
--- a/api/transfer.py
+++ b/api/transfer.py
@@ -20,7 +20,11 @@ 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"""
+    """
+    ### Transfer a streaming response 
+    Takes the request from the incoming request to the target endpoint.
+    Checks method, token amount, auth and cost along with if request is NSFW.
+    """
 
     path = incoming_request.url.path.replace('v1/v1/', 'v1/')