mirror of
https://github.com/NovaOSS/nova-api.git
synced 2024-11-25 19:03:57 +01:00
143 lines
4.6 KiB
Python
143 lines
4.6 KiB
Python
"""Does quite a few checks and prepares the incoming request for the target endpoint, so it can be streamed"""
|
|
|
|
import json
|
|
import yaml
|
|
import time
|
|
import orjson
|
|
import fastapi
|
|
|
|
from dotenv import load_dotenv
|
|
|
|
import streaming
|
|
import moderation
|
|
|
|
from rich import print
|
|
from db.users import UserManager
|
|
from helpers import tokens, errors, network
|
|
|
|
load_dotenv()
|
|
|
|
users = UserManager()
|
|
models_list = json.load(open('models.json', encoding='utf8'))
|
|
|
|
with open('config/config.yml', encoding='utf8') as f:
|
|
config = yaml.safe_load(f)
|
|
|
|
async def handle(incoming_request: fastapi.Request):
|
|
"""
|
|
### Transfer a streaming response
|
|
Takes the request from the incoming request to the target endpoint.
|
|
Checks method, token amount, auth and cost along with if request is NSFW.
|
|
"""
|
|
path = incoming_request.url.path.replace('v1/v1', 'v1').replace('//', '/')
|
|
|
|
ip_address = await network.get_ip(incoming_request)
|
|
print(f'[bold green]>{ip_address}[/bold green]')
|
|
|
|
if '/models' in path:
|
|
return fastapi.responses.JSONResponse(content=models_list)
|
|
|
|
try:
|
|
payload = await incoming_request.json()
|
|
except json.decoder.JSONDecodeError:
|
|
payload = {}
|
|
|
|
received_key = incoming_request.headers.get('Authorization')
|
|
|
|
if not received_key or not received_key.startswith('Bearer '):
|
|
return await errors.error(403, 'No NovaAI API key given!', 'Add \'Authorization: Bearer nv-...\' to your request headers.')
|
|
|
|
if '#' in received_key:
|
|
key_tags = received_key.split('#')[1]
|
|
received_key = received_key.split('#')[0]
|
|
|
|
user = await users.user_by_api_key(received_key.split('Bearer ')[1].strip())
|
|
|
|
if not user or not user['status']['active']:
|
|
return await errors.error(403, 'Invalid or inactive NovaAI API key!', 'Create a new NovaOSS API key or reactivate your account.')
|
|
|
|
if user.get('auth', {}).get('discord'):
|
|
print(f'[bold green]>Discord[/bold green] {user["auth"]["discord"]}')
|
|
|
|
ban_reason = user['status']['ban_reason']
|
|
if ban_reason:
|
|
return await errors.error(403, f'Your NovaAI account has been banned. Reason: \'{ban_reason}\'.', 'Contact the staff for an appeal.')
|
|
|
|
costs = config['costs']
|
|
cost = costs['other']
|
|
|
|
if 'chat/completions' in path:
|
|
cost = costs['chat-models'].get(payload.get('model'), cost)
|
|
|
|
role = user.get('role', 'default')
|
|
|
|
try:
|
|
role_cost_multiplier = config['roles'][role]['bonus']
|
|
except KeyError:
|
|
role_cost_multiplier = 1
|
|
|
|
cost = round(cost * role_cost_multiplier)
|
|
|
|
if user['credits'] < cost:
|
|
return await errors.error(429, 'Not enough credits.', 'Wait or earn more credits. Learn more on our website or Discord server.')
|
|
|
|
payload_with_vars = json.dumps(payload)
|
|
|
|
replace_dict = {
|
|
'timestamp': str(int(time.time())),
|
|
'date': time.strftime('%Y-%m-%d'),
|
|
'time': time.strftime('%H:%M:%S'),
|
|
'datetime': time.strftime('%Y-%m-%d %H:%M:%S'),
|
|
'model': payload.get('model', 'unknown'),
|
|
}
|
|
|
|
if 'ALLOW_INSECURE_VARS' in key_tags:
|
|
replace_dict.update({
|
|
'my.ip': ip_address,
|
|
'my.id': str(user['_id']),
|
|
'my.role': user.get('role', 'default'),
|
|
'my.credits': str(user['credits']),
|
|
'my.discord': user.get('auth', {}).get('discord', ''),
|
|
})
|
|
|
|
for key, value in replace_dict.items():
|
|
payload_with_vars = payload_with_vars.replace(f'[[{key}]]', value)
|
|
|
|
payload = json.loads(payload_with_vars)
|
|
|
|
policy_violation = False
|
|
if '/moderations' not in path:
|
|
inp = ''
|
|
|
|
if 'input' in payload or 'prompt' in payload:
|
|
inp = payload.get('input', payload.get('prompt', ''))
|
|
|
|
if isinstance(payload.get('messages'), list):
|
|
inp = '\n'.join([message['content'] for message in payload['messages']])
|
|
|
|
if inp and len(inp) > 2 and not inp.isnumeric():
|
|
policy_violation = await moderation.is_policy_violated(inp)
|
|
|
|
if policy_violation:
|
|
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.'
|
|
)
|
|
|
|
if 'chat/completions' in path and not payload.get('stream', False):
|
|
payload['stream'] = False
|
|
|
|
media_type = 'text/event-stream' if payload.get('stream', False) else 'application/json'
|
|
|
|
return fastapi.responses.StreamingResponse(
|
|
content=streaming.stream(
|
|
user=user,
|
|
path=path,
|
|
payload=payload,
|
|
credits_cost=cost,
|
|
input_tokens=-1,
|
|
incoming_request=incoming_request,
|
|
),
|
|
media_type=media_type
|
|
)
|