alex/nova-api
1
0
Fork 0
mirror of https://github.com/NovaOSS/nova-api.git synced 2024-11-25 16:23:57 +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
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'))

20
api/streaming.py
View file

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