Several improvements, about to change DB

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,