Update streaming to work & change stats to class

api/db/stats.py

@@ -17,39 +17,50 @@ async def _get_collection(collection_name: str):
 
 ## Statistics
 
-async def add_date():
+class Stats:
+    """
+    ### The manager for all statistics tracking
+    Stats tracked:
+    - Dates
+    - IPs
+    - Target URLs
+    - Tokens
+    - Models
+    - URL Paths
+    """
+    async def add_date():
         date = datetime.datetime.now(pytz.timezone('GMT')).strftime('%Y.%m.%d')
         year, month, day = date.split('.')
 
         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):
+    async def add_ip_address(ip_address: str):
         ip_address = ip_address.replace('.', '_')
         db = await _get_collection('stats')
         await db.update_one({}, {'$inc': {f'ips.{ip_address}': 1}}, upsert=True)
 
-async def add_target(url: str):
+    async def add_target(url: str):
         db = await _get_collection('stats')
         await db.update_one({}, {'$inc': {f'targets.{url}': 1}}, upsert=True)
 
-async def add_tokens(tokens: int, model: str):
+    async def add_tokens(tokens: int, model: str):
         db = await _get_collection('stats')
         await db.update_one({}, {'$inc': {f'tokens.{model}': tokens}}, upsert=True)
 
-async def add_model(model: str):
+    async def add_model(model: str):
         db = await _get_collection('stats')
         await db.update_one({}, {'$inc': {f'models.{model}': 1}}, upsert=True)
 
-async def add_path(path: str):
+    async def add_path(path: str):
         path = path.replace('/', '_')
         db = await _get_collection('stats')
         await db.update_one({}, {'$inc': {f'paths.{path}': 1}}, upsert=True)
 
-async def get_value(obj_filter):
+    async def get_value(obj_filter):
         db = await _get_collection('stats')
         return await db.find_one({obj_filter})
 
 if __name__ == '__main__':
-    asyncio.run(add_date())
-    asyncio.run(add_path('/__demo/test'))
+    asyncio.run(Stats.add_date())
+    asyncio.run(Stats.add_path('/__demo/test'))

api/streaming.py

@@ -15,7 +15,8 @@ import proxies
 import provider_auth
 import load_balancing
 
-from db import logs, users, stats
+from db import logs, users
+from db.stats import Stats
 from helpers import network, chat, errors
 
 load_dotenv()
@@ -30,7 +31,7 @@ DEMO_PAYLOAD = {
     ]
 }
 
-async def process_response(response, is_chat, chat_id, model):
+async def process_response(response, is_chat, chat_id, model, target_request):
     """Proccesses chunks from streaming
 
     Args:
@@ -154,7 +155,7 @@ async def stream(
                             if 'Too Many Requests' in str(exc):
                                 continue
 
-                        async for chunk in process_response(response, is_chat, chat_id, model):
+                        async for chunk in process_response(response, is_chat, chat_id, model, target_request):
                             yield chunk
 
                     break
@@ -177,14 +178,13 @@ async def stream(
         await users.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())