Update streaming to work & change stats to class

This commit is contained in:
Game_Time 2023-08-13 21:42:38 +05:00
parent 885da2a27e
commit bb1e9de563
2 changed files with 48 additions and 37 deletions

View file

@ -17,6 +17,17 @@ async def _get_collection(collection_name: str):
## Statistics ## Statistics
class Stats:
"""
### The manager for all statistics tracking
Stats tracked:
- Dates
- IPs
- Target URLs
- Tokens
- Models
- URL Paths
"""
async def add_date(): async def add_date():
date = datetime.datetime.now(pytz.timezone('GMT')).strftime('%Y.%m.%d') date = datetime.datetime.now(pytz.timezone('GMT')).strftime('%Y.%m.%d')
year, month, day = date.split('.') year, month, day = date.split('.')
@ -51,5 +62,5 @@ async def get_value(obj_filter):
return await db.find_one({obj_filter}) return await db.find_one({obj_filter})
if __name__ == '__main__': if __name__ == '__main__':
asyncio.run(add_date()) asyncio.run(Stats.add_date())
asyncio.run(add_path('/__demo/test')) asyncio.run(Stats.add_path('/__demo/test'))

View file

@ -15,7 +15,8 @@ import proxies
import provider_auth import provider_auth
import load_balancing 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 from helpers import network, chat, errors
load_dotenv() 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 """Proccesses chunks from streaming
Args: Args:
@ -154,7 +155,7 @@ async def stream(
if 'Too Many Requests' in str(exc): if 'Too Many Requests' in str(exc):
continue 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 yield chunk
break break
@ -177,14 +178,13 @@ async def stream(
await users.update_by_id(user['_id'], {'$inc': {'credits': -credits_cost}}) await users.update_by_id(user['_id'], {'$inc': {'credits': -credits_cost}})
ip_address = await network.get_ip(incoming_request) ip_address = await network.get_ip(incoming_request)
await stats.add_date() await Stats.add_date()
await stats.add_ip_address(ip_address) await Stats.add_ip_address(ip_address)
await stats.add_path(path) await Stats.add_path(path)
await stats.add_target(target_request['url']) await Stats.add_target(target_request['url'])
if is_chat: if is_chat:
await stats.add_model(model) await Stats.add_model(model)
await stats.add_tokens(input_tokens, model) await Stats.add_tokens(input_tokens, model)
if __name__ == '__main__': if __name__ == '__main__':
asyncio.run(stream()) asyncio.run(stream())