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

Merge pull request #7 from RayBytes/main

Browse source 
General changes
This commit is contained in:
Game_Time 2023-08-12 21:25:23 +05:00 committed by GitHub
parent ea68a638e9 d5e2846551
commit bea0cecdd3
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
9 changed files with 85 additions and 67 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
api/db/logs.py
View file

@ -16,8 +16,12 @@ UA_SIMPLIFY = {
'AppleWebKit/537.36 (KHTML, like Gecko)': 'K',
}
async def _get_mongo(collection_name: str):
return AsyncIOMotorClient(os.getenv('MONGO_URI'))['nova-core'][collection_name]
## MONGODB Setup
conn = AsyncIOMotorClient(os.getenv('MONGO_URI'))
async def _get_collection(collection_name: str):
return conn['nova-core'][collection_name]
async def replacer(text: str, dict_: dict) -> str:
for k, v in dict_.items():
@ -25,7 +29,7 @@ async def replacer(text: str, dict_: dict) -> str:
return text
async def log_api_request(user: dict, incoming_request, target_url: str):
db = await _get_mongo('logs')
db = await _get_collection('logs')
payload = {}
try:
@ -58,19 +62,19 @@ async def log_api_request(user: dict, incoming_request, target_url: str):
return log_item
async def by_id(log_id: str):
db = await _get_mongo('logs')
db = await _get_collection('logs')
return await db.find_one({'_id': log_id})
async def by_user_id(user_id: str):
db = await _get_mongo('logs')
db = await _get_collection('logs')
return await db.find({'user_id': user_id})
async def delete_by_id(log_id: str):
db = await _get_mongo('logs')
db = await _get_collection('logs')
return await db.delete_one({'_id': log_id})
async def delete_by_user_id(user_id: str):
db = await _get_mongo('logs')
db = await _get_collection('logs')
return await db.delete_many({'user_id': user_id})
if __name__ == '__main__':

24
api/db/stats.py
View file

@ -8,40 +8,46 @@ from motor.motor_asyncio import AsyncIOMotorClient
load_dotenv()
async def _get_mongo(collection_name: str):
return AsyncIOMotorClient(os.getenv('MONGO_URI'))['nova-core'][collection_name]
## MONGODB Setup
conn = AsyncIOMotorClient(os.getenv('MONGO_URI'))
async def _get_collection(collection_name: str):
return conn['nova-core'][collection_name]
## Statistics
async def add_date():
date = datetime.datetime.now(pytz.timezone('GMT')).strftime('%Y.%m.%d')
year, month, day = date.split('.')
db = await _get_mongo('stats')
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_mongo('stats')
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_mongo('stats')
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_mongo('stats')
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_mongo('stats')
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_mongo('stats')
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_mongo('stats')
db = await _get_collection('stats')
return await db.find_one({obj_filter})
if __name__ == '__main__':

22
api/db/users.py
View file

@ -12,8 +12,12 @@ load_dotenv()
with open('config/credits.yml', encoding='utf8') as f:
credits_config = yaml.safe_load(f)
async def _get_mongo(collection_name: str):
return AsyncIOMotorClient(os.getenv('MONGO_URI'))['nova-core'][collection_name]
## MONGODB Setup
conn = AsyncIOMotorClient(os.getenv('MONGO_URI'))
async def _get_collection(collection_name: str):
return conn['nova-core'][collection_name]
async def create(discord_id: str='') -> dict:
"""Adds a new user to the MongoDB collection."""
@ -41,33 +45,33 @@ async def create(discord_id: str='') -> dict:
}
}
db = await _get_mongo('users')
db = await _get_collection('users')
await db.insert_one(new_user)
user = await db.find_one({'api_key': new_api_key})
return user
async def by_id(user_id: str):
db = await _get_mongo('users')
db = await _get_collection('users')
return await db.find_one({'_id': user_id})
async def by_discord_id(discord_id: str):
db = await _get_mongo('users')
db = await _get_collection('users')
return await db.find_one({'auth.discord': str(int(discord_id))})
async def by_api_key(key: str):
db = await _get_mongo('users')
db = await _get_collection('users')
return await db.find_one({'api_key': key})
async def update_by_id(user_id: str, update):
db = await _get_mongo('users')
db = await _get_collection('users')
return await db.update_one({'_id': user_id}, update)
async def update_by_filter(obj_filter, update):
db = await _get_mongo('users')
db = await _get_collection('users')
return await db.update_one(obj_filter, update)
async def delete(user_id: str):
db = await _get_mongo('users')
db = await _get_collection('users')
await db.delete_one({'_id': user_id})
async def demo():

21
api/helpers/chat.py
View file

@ -20,8 +20,16 @@ async def create_chat_id() -> str:
return f'chatcmpl-{chat_id}'
async def create_chat_chunk(chat_id: str, model: str, content=None) -> dict:
"""Creates a new chat chunk"""
"""Creates the chunk for streaming chat.
Args:
chat_id (str): _description_
model (str): _description_
content (_type_, optional): _description_. Defaults to None.
Returns:
dict: _description_
"""
content = content or {}
delta = {}
@ -53,13 +61,4 @@ async def create_chat_chunk(chat_id: str, model: str, content=None) -> dict:
],
}
return f'data: {json.dumps(chunk)}\n\n'
if __name__ == '__main__':
demo_chat_id = asyncio.run(create_chat_id())
print(demo_chat_id)
print(asyncio.run(create_chat_chunk(
model='gpt-4',
content='Hello',
chat_id=demo_chat_id,
)))
return f'data: {json.dumps(chunk)}\n\n'

13
api/helpers/tokens.py
View file

@ -1,7 +1,18 @@
import tiktoken
async def count_for_messages(messages: list, model: str='gpt-3.5-turbo-0613') -> int:
"""Return the number of tokens used by a list of messages."""
"""Return the number of tokens used by a list of messages
Args:
messages (list): _description_
model (str, optional): _description_. Defaults to 'gpt-3.5-turbo-0613'.
Raises:
NotImplementedError: _description_
Returns:
int: _description_
"""
try:
encoding = tiktoken.encoding_for_model(model)

1
api/load_balancing.py
View file

@ -39,7 +39,6 @@ async def balance_organic_request(request: dict) -> dict:
"""Load balnace to non-chat completion request between other "organic" providers which respond in the desired format already.
Organic providers are used for non-chat completions, such as moderation and other paths.
"""
providers_available = []
if not request.get('headers'):

15
api/provider_auth.py
View file

@ -2,12 +2,15 @@
import asyncio
async def invalidate_key(provider_and_key: str) -> none:
"""Invalidates a key stored in the secret/ folder by storing it in the associated .invalid.txt file.
The schmea in which <provider_and_key> should be passed is:
<provider_name><key>, e.g.
closed4>sk-...
"""
async def invalidate_key(provider_and_key: str) -> None:
"""
Invalidates a key stored in the secret/ folder by storing it in the associated .invalid.txt file.
The schmea in which <provider_and_key> should be passed is:
<provider_name><key>, e.g.
closed4>sk-...
"""
if not provider_and_key:
return

6
api/streaming.py
View file

@ -96,12 +96,11 @@ async def stream(
webhook = dhooks.Webhook(os.getenv('DISCORD_WEBHOOK__API_ISSUE'))
webhook.send(content=f'API Issue: **`{exc}`**\nhttps://i.imgflip.com/7uv122.jpg')
error = await errors.yield_error(
yield await errors.yield_error(
500,
'Sorry, the API has no working keys anymore.',
'The admins have been messaged automatically.'
)
yield error
return
for k, v in target_request.get('headers', {}).items():
@ -180,12 +179,11 @@ async def stream(
except Exception as exc:
if 'Connection closed' in str(exc):
error = await errors.yield_error(
yield await errors.yield_error(
500,
'Sorry, there was an issue with the connection.',
'Please first check if the issue on your end. If this error repeats, please don\'t heistate to contact the staff!.'
)
yield error
return
break

32
api/transfer.py
View file

@ -26,8 +26,7 @@ async def handle(incoming_request):
# METHOD
if incoming_request.method not in ['GET', 'POST', 'PUT', 'DELETE', 'PATCH']:
error = await errors.error(405, f'Method "{incoming_request.method}" is not allowed.', 'Change the request method to the correct one.')
return error
return await errors.error(405, f'Method "{incoming_request.method}" is not allowed.', 'Change the request method to the correct one.')
# PAYLOAD
try:
@ -35,42 +34,37 @@ async def handle(incoming_request):
except json.decoder.JSONDecodeError:
payload = {}
# TOKENS
# Tokenise w/ tiktoken
try:
input_tokens = await tokens.count_for_messages(payload['messages'])
except (KeyError, TypeError):
input_tokens = 0
# AUTH
# Check user auth
received_key = incoming_request.headers.get('Authorization')
if not received_key:
error = await errors.error(401, 'No NovaAI API key given!', 'Add "Authorization: Bearer nv-..." to your request headers.')
return error
return await errors.error(401, 'No NovaAI API key given!', 'Add "Authorization: Bearer nv-..." to your request headers.')
if received_key.startswith('Bearer '):
received_key = received_key.split('Bearer ')[1]
# USER
user = await users.by_api_key(received_key.strip())
if not user:
error = await errors.error(401, 'Invalid NovaAI API key!', 'Create a new NovaOSS API key.')
return error
return await errors.error(401, 'Invalid NovaAI API key!', 'Create a new NovaOSS API key.')
ban_reason = user['status']['ban_reason']
if ban_reason:
error = await errors.error(403, f'Your NovaAI account has been banned. Reason: "{ban_reason}".', 'Contact the staff for an appeal.')
return error
return await errors.error(403, f'Your NovaAI account has been banned. Reason: "{ban_reason}".', 'Contact the staff for an appeal.')
if not user['status']['active']:
error = await errors.error(418, 'Your NovaAI account is not active (paused).', 'Simply re-activate your account using a Discord command or the web panel.')
return error
return await errors.error(418, 'Your NovaAI account is not active (paused).', 'Simply re-activate your account using a Discord command or the web panel.')
if '/models' in path:
return fastapi.responses.JSONResponse(content=models_list)
# COST
# Calculate cost of tokens & check for nsfw prompts
costs = credits_config['costs']
cost = costs['other']
@ -94,17 +88,17 @@ async def handle(incoming_request):
policy_violation = await moderation.is_policy_violated(inp)
if policy_violation:
error = await errors.error(400, f'The request contains content which violates this model\'s policies for "{policy_violation}".', 'We currently don\'t support any NSFW models.')
return error
return await errors.error(400, f'The request contains content which violates this model\'s policies for "{policy_violation}".', 'We currently don\'t support any NSFW models.')
role_cost_multiplier = credits_config['bonuses'].get(user['role'], 1)
cost = round(cost * role_cost_multiplier)
if user['credits'] < cost:
error = await errors.error(429, 'Not enough credits.', 'Wait or earn more credits. Learn more on our website or Discord server.')
return error
return await errors.error(429, 'Not enough credits.', 'Wait or earn more credits. Learn more on our website or Discord server.')
# READY
# Send the completion request
if 'chat/completions' in path and not payload.get('stream') is True:
payload['stream'] = False