From bb1e9de563b61e8073364a0fec32e6d93841c0a3 Mon Sep 17 00:00:00 2001
From: Game_Time <108236317+RayBytes@users.noreply.github.com>
Date: Sun, 13 Aug 2023 21:42:38 +0500
Subject: [PATCH] Update streaming to work & change stats to class

---
 api/db/stats.py  | 65 ++++++++++++++++++++++++++++--------------------
 api/streaming.py | 20 +++++++--------
 2 files changed, 48 insertions(+), 37 deletions(-)

diff --git a/api/db/stats.py b/api/db/stats.py
index d20ee4a..f3df4ec 100644
--- a/api/db/stats.py
+++ b/api/db/stats.py
@@ -17,39 +17,50 @@ async def _get_collection(collection_name: str):
 
 ## Statistics
 
-async def add_date():
-    date = datetime.datetime.now(pytz.timezone('GMT')).strftime('%Y.%m.%d')
-    year, month, day = date.split('.')
+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)
+        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_collection('stats')
-    await db.update_one({}, {'$inc': {f'ips.{ip_address}': 1}}, upsert=True)
+    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):
-    db = await _get_collection('stats')
-    await db.update_one({}, {'$inc': {f'targets.{url}': 1}}, upsert=True)
+    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):
-    db = await _get_collection('stats')
-    await db.update_one({}, {'$inc': {f'tokens.{model}': tokens}}, upsert=True)
+    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):
-    db = await _get_collection('stats')
-    await db.update_one({}, {'$inc': {f'models.{model}': 1}}, upsert=True)
+    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):
-    path = path.replace('/', '_')
-    db = await _get_collection('stats')
-    await db.update_one({}, {'$inc': {f'paths.{path}': 1}}, upsert=True)
+    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):
-    db = await _get_collection('stats')
-    return await db.find_one({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'))
diff --git a/api/streaming.py b/api/streaming.py
index 3aa0e53..885eba9 100644
--- a/api/streaming.py
+++ b/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())