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