mirror of
https://github.com/NovaOSS/nova-api.git
synced 2024-11-25 20:13:57 +01:00
118 lines
2.9 KiB
Python
118 lines
2.9 KiB
Python
import os
|
|
import uuid
|
|
import time
|
|
import string
|
|
import random
|
|
import asyncio
|
|
|
|
from dotenv import load_dotenv
|
|
from motor.motor_asyncio import AsyncIOMotorClient
|
|
|
|
from rich import print
|
|
|
|
load_dotenv()
|
|
|
|
MONGO_URI = os.getenv('MONGO_URI')
|
|
MONGO_DB_NAME = 'users'
|
|
|
|
def get_mongo(collection_name):
|
|
client = AsyncIOMotorClient(MONGO_URI)
|
|
db = client[MONGO_DB_NAME]
|
|
return db[collection_name]
|
|
|
|
async def prepare() -> None:
|
|
"""Create the MongoDB collection."""
|
|
|
|
collection = get_mongo('users')
|
|
|
|
await collection.create_index('id', unique=True)
|
|
await collection.create_index('discord_id', unique=True)
|
|
await collection.create_index('api_key', unique=True)
|
|
|
|
async def add_user(discord_id: int = 0, tags: list = None) -> dict:
|
|
"""Adds a new user to the MongoDB collection."""
|
|
|
|
chars = string.ascii_letters + string.digits
|
|
|
|
infix = os.getenv('KEYGEN_INFIX')
|
|
suffix = ''.join(random.choices(chars, k=20))
|
|
prefix = ''.join(random.choices(chars, k=20))
|
|
|
|
key = f'nv-{prefix}{infix}{suffix}'
|
|
|
|
tags = tags or []
|
|
new_user = {
|
|
'id': str(uuid.uuid4()),
|
|
'api_key': key,
|
|
'created_at': int(time.time()),
|
|
'ban_reason': '',
|
|
'active': True,
|
|
'discord_id': discord_id,
|
|
'credit': 0,
|
|
'tags': '/'.join(tags),
|
|
'usage': {
|
|
'events': [],
|
|
'num_tokens': 0
|
|
}
|
|
}
|
|
|
|
collection = get_mongo('users')
|
|
await collection.insert_one(new_user)
|
|
|
|
return new_user
|
|
|
|
async def get_user(by_id: str = '', by_discord_id: int = 0, by_api_key: str = ''):
|
|
"""Retrieve a user from the MongoDB collection."""
|
|
|
|
collection = get_mongo('users')
|
|
query = {
|
|
'$or': [
|
|
{'id': by_id},
|
|
{'discord_id': by_discord_id},
|
|
{'api_key': by_api_key},
|
|
]
|
|
}
|
|
return await collection.find_one(query)
|
|
|
|
async def get_all_users():
|
|
"""Retrieve all users from the MongoDB collection."""
|
|
|
|
collection = get_mongo('users')
|
|
return list(await collection.find())
|
|
|
|
async def user_used_api(user_id: str, num_tokens: int = 0, model='', ip_address: str = '', user_agent: str = '') -> None:
|
|
"""Update the stats of a user."""
|
|
|
|
collection = get_mongo('users')
|
|
user = await get_user(by_id=user_id)
|
|
|
|
if not user:
|
|
raise ValueError('User not found.')
|
|
|
|
usage = user['usage']
|
|
usage['events'].append({
|
|
'timestamp': time.time(),
|
|
'ip_address': ip_address,
|
|
'user_agent': user_agent,
|
|
'model': model,
|
|
'num_tokens': num_tokens
|
|
})
|
|
|
|
usage['num_tokens'] += num_tokens
|
|
|
|
await collection.update_one({'id': user_id}, {'$set': {'usage': usage}})
|
|
|
|
async def demo():
|
|
await prepare()
|
|
|
|
example_id = 133769420
|
|
user = await get_user(by_discord_id=example_id)
|
|
print(user)
|
|
uid = await user['id']
|
|
|
|
await user_used_api(uid, model='gpt-5', num_tokens=42, ip_address='9.9.9.9', user_agent='Mozilla/5.0')
|
|
# print(user)
|
|
|
|
if __name__ == '__main__':
|
|
asyncio.run(demo())
|