alex/nova-api
1
0
Fork 0
mirror of https://github.com/NovaOSS/nova-api.git synced 2024-11-25 21:33:58 +01:00
Code Issues Packages Projects Releases Wiki Activity

Update streaming to work & change stats to class

Browse source
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
Show stats Download patch file Download diff file Expand all files Collapse all files

65
api/db/stats.py
View file

@ -17,39 +17,50 @@ async def _get_collection(collection_name: str):
## Statistics ## Statistics
async def add_date(): class Stats:
date = datetime.datetime.now(pytz.timezone('GMT')).strftime('%Y.%m.%d') """
year, month, day = date.split('.') ### 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') db = await _get_collection('stats')
await db.update_one({}, {'$inc': {f'dates.{year}.{month}.{day}': 1}}, upsert=True) 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('.', '_') ip_address = ip_address.replace('.', '_')
db = await _get_collection('stats') db = await _get_collection('stats')
await db.update_one({}, {'$inc': {f'ips.{ip_address}': 1}}, upsert=True) 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') db = await _get_collection('stats')
await db.update_one({}, {'$inc': {f'targets.{url}': 1}}, upsert=True) 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') db = await _get_collection('stats')
await db.update_one({}, {'$inc': {f'tokens.{model}': tokens}}, upsert=True) 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') db = await _get_collection('stats')
await db.update_one({}, {'$inc': {f'models.{model}': 1}}, upsert=True) 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('/', '_') path = path.replace('/', '_')
db = await _get_collection('stats') db = await _get_collection('stats')
await db.update_one({}, {'$inc': {f'paths.{path}': 1}}, upsert=True) 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') db = await _get_collection('stats')
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'))

20
api/streaming.py
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())