diff --git a/api/transfer.py b/api/transfer.py index cb813f7..614f1e8 100644 --- a/api/transfer.py +++ b/api/transfer.py @@ -2,13 +2,12 @@ import os import json +import tokens import logging import starlette -import netclient - from dotenv import load_dotenv -from helpers import requesting +from helpers import requesting, tokens, errors load_dotenv() @@ -21,27 +20,45 @@ logging.basicConfig( logging.info('API started') -async def handle_api_request(incoming_request, target_endpoint: str=''): +DEFAULT_ENDPOINT = os.getenv('CLOSEDAI_ENDPOINT') + +async def handle_api_request(incoming_request, target_endpoint: str=DEFAULT_ENDPOINT): """Transfer a streaming response from the incoming request to the target endpoint""" - if not target_endpoint: - target_endpoint = os.getenv('CLOSEDAI_ENDPOINT') target_url = f'{target_endpoint}{incoming_request.url.path}' - logging.info('TRANSFER %s -> %s', incoming_request.url.path, target_url) - if target_url.endswith('/v1'): - return starlette.responses.Response(status_code=200, content='200. /v1 is the API endpoint root path.') - - if incoming_request.headers.get('Authorization') != f'Bearer {os.getenv("DEMO_AUTH")}': - return starlette.responses.Response(status_code=403, content='Invalid API key') + if incoming_request.method not in ['GET', 'POST', 'PUT', 'DELETE', 'PATCH']: + return errors.error(405, f'Method "{incoming_request.method}" is not allowed.', 'Change the request method to the correct one.') try: payload = await incoming_request.json() except json.decoder.JSONDecodeError: payload = {} - if 'temperature' in payload or 'functions' in payload: - target_provider = 'closed' + try: + input_tokens = tokens.count_for_messages(payload['messages']) + except: + input_tokens = 0 + + auth_header = incoming_request.headers.get('Authorization') + + if not auth_header: + return errors.error(401, 'No NovaOSS API key given!', 'Add "Authorization: Bearer nv-..." to your request headers.') + + received_key = auth_header + + if auth_header.startswith('Bearer '): + received_key = auth_header.split('Bearer ')[1] + + user = users.get_user(by_api_key=received_key) + + if not user: + return errors.error(401, 'Invalid NovaOSS API key!', 'Create a new NovaOSS API key.') + + if not user['active']: + return errors.error(403, 'Your account is not active.', 'Activate your account.') + + logging.info(f'[%s] %s -> %s', incoming_request.method, incoming_request.url.path, target_url) request = requesting.Request( url=target_url,