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

Enterprise fix

Browse source
This commit is contained in:
NovaOSS Admins 2023-10-15 20:50:13 +00:00
parent fa1200fba3
commit 45c6a14dc0
1 changed files with 11 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

19
api/handler.py
View file

@ -71,10 +71,10 @@ async def handle(incoming_request: fastapi.Request):
if ban_reason:
return await errors.error(403, f'Your NovaAI account has been banned. Reason: \'{ban_reason}\'.', 'Contact the staff for an appeal.')
# Checking for enterprise status
enterprise_keys = os.environ.get('ENTERPRISE_KEYS')
if path.startswith('/enterprise/v1') and user.get('api_key') not in enterprise_keys.split():
return await errors.error(403, 'Enterprise API is not available.', 'Contact the staff for an upgrade.')
is_enterprise_key = 'enterprise' in user.get('role', 'default')
if path.startswith('/enterprise/v1') and not is_enterprise_key:
return await errors.error(403, 'Enterprise API is not available for your API key.', 'Contact the staff for an upgrade.')
if 'account/credits' in path:
return fastapi.responses.JSONResponse({'credits': user['credits']})
@ -87,10 +87,13 @@ async def handle(incoming_request: fastapi.Request):
role = user.get('role', 'default')
try:
role_cost_multiplier = config['roles'][role]['bonus']
except KeyError:
role_cost_multiplier = 1
if 'enterprise' in role:
role_cost_multiplier = 0.1
else:
try:
role_cost_multiplier = config['roles'][role]['bonus']
except KeyError:
role_cost_multiplier = 1
cost = round(cost * role_cost_multiplier)