From eb6ebd2112144dd25037e467320f82c43edce2a3 Mon Sep 17 00:00:00 2001 From: nsde Date: Mon, 21 Aug 2023 20:58:05 +0200 Subject: [PATCH] Made code more robust (dotenv) and working on tests --- .github/workflows/tests.yml | 2 +- api/core.py | 48 ++++++++++++--- api/db/users.py | 2 +- api/main.py | 14 ++--- api/streaming.py | 2 +- api/transfer.py | 4 +- checks/__init__.py | 1 + checks/__main__.py | 2 + tests/__main__.py => checks/client.py | 75 +++++++++++++---------- Untitled Diagram.drawio => privacy.drawio | 0 screen.sh | 3 + setup.md | 2 +- 12 files changed, 100 insertions(+), 55 deletions(-) create mode 100644 checks/__init__.py create mode 100644 checks/__main__.py rename tests/__main__.py => checks/client.py (59%) rename Untitled Diagram.drawio => privacy.drawio (100%) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 91c5236..9eb7ba8 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -28,5 +28,5 @@ jobs: - name: Start API server & run tests! run: | + python checks python run - diff --git a/api/core.py b/api/core.py index 8294c44..ffdb08c 100644 --- a/api/core.py +++ b/api/core.py @@ -1,14 +1,24 @@ """User management.""" +import os +import sys + +project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..')) +sys.path.append(project_root) + +# the code above is to allow importing from the root folder + import os import json import fastapi -from db.users import UserManager - from dhooks import Webhook, Embed from dotenv import load_dotenv +import checks.client + +from db.users import UserManager + load_dotenv() router = fastapi.APIRouter(tags=['core']) @@ -21,11 +31,13 @@ async def check_core_auth(request): """ received_auth = request.headers.get('Authorization') - if received_auth != os.getenv('CORE_API_KEY'): + if received_auth != os.environ['CORE_API_KEY']: return fastapi.Response(status_code=403, content='Invalid or no API key given.') @router.get('/users') async def get_users(discord_id: int, incoming_request: fastapi.Request): + """Returns a user by their discord ID. Requires a core API key.""" + auth = await check_core_auth(incoming_request) if auth: return auth @@ -38,8 +50,10 @@ async def get_users(discord_id: int, incoming_request: fastapi.Request): return user -async def new_user_webhook(user: dict) -> None: - dhook = Webhook(os.getenv('DISCORD_WEBHOOK__USER_CREATED')) +async def new_user_webhook(user: dict) -> None: + """Runs when a new user is created.""" + + dhook = Webhook(os.environ['DISCORD_WEBHOOK__USER_CREATED']) embed = Embed( description='New User', @@ -54,6 +68,8 @@ async def new_user_webhook(user: dict) -> None: @router.post('/users') async def create_user(incoming_request: fastapi.Request): + """Creates a user. Requires a core API key.""" + auth_error = await check_core_auth(incoming_request) if auth_error: @@ -64,7 +80,7 @@ async def create_user(incoming_request: fastapi.Request): discord_id = payload.get('discord_id') except (json.decoder.JSONDecodeError, AttributeError): return fastapi.Response(status_code=400, content='Invalid or no payload received.') - + # Create the user manager = UserManager() user = await manager.create(discord_id) @@ -74,6 +90,8 @@ async def create_user(incoming_request: fastapi.Request): @router.put('/users') async def update_user(incoming_request: fastapi.Request): + """Updates a user. Requires a core API key.""" + auth_error = await check_core_auth(incoming_request) if auth_error: @@ -85,9 +103,25 @@ async def update_user(incoming_request: fastapi.Request): updates = payload.get('updates') except (json.decoder.JSONDecodeError, AttributeError): return fastapi.Response(status_code=400, content='Invalid or no payload received.') - + # Update the user manager = UserManager() user = await manager.update_by_discord_id(discord_id, updates) return user + +@router.get('/checks') +async def run_checks(incoming_request: fastapi.Request): + """Tests the API. Requires a core API key.""" + + auth_error = await check_core_auth(incoming_request) + + if auth_error: + return auth_error + + return { + 'library': await checks.client.test_library(), + 'library_moderation': await checks.client.test_library_moderation(), + 'api_moderation': await checks.client.test_api_moderation(), + 'models': await checks.client.test_models() + } diff --git a/api/db/users.py b/api/db/users.py index 6a786af..e1150a6 100644 --- a/api/db/users.py +++ b/api/db/users.py @@ -89,7 +89,7 @@ class UserManager: db = await self._get_collection('users') return await db.update_one({'_id': user_id}, update) - async def upate_by_discord_id(self, discord_id: str, update): + async def update_by_discord_id(self, discord_id: str, update): db = await self._get_collection('users') return await db.update_one({'auth.discord': str(int(discord_id))}, update) diff --git a/api/main.py b/api/main.py index 152c9c5..ff9b149 100644 --- a/api/main.py +++ b/api/main.py @@ -24,21 +24,19 @@ app.include_router(core.router) @app.on_event('startup') async def startup_event(): - # DATABASE FIX https://stackoverflow.com/questions/65970988/python-mongodb-motor-objectid-object-is-not-iterable-error-while-trying-to-f - import pydantic, bson - # pydantic.json.ENCODERS_BY_TYPE[bson.objectid.ObjectId]=str + """Runs when the API starts up.""" @app.get('/') async def root(): """ - Returns the root endpoint. + Returns general information about the API. """ return { - 'status': 'ok', - 'usage_docs': 'https://nova-oss.com', - 'core_api_docs_for_developers': '/docs', - 'github': 'https://github.com/novaoss/nova-api' + 'hi': 'Welcome to the Nova API!', + 'learn_more_here': 'https://nova-oss.com', + 'github': 'https://github.com/novaoss/nova-api', + 'core_api_docs_for_nova_developers': '/docs' } app.add_route('/{path:path}', transfer.handle, ['GET', 'POST', 'PUT', 'DELETE', 'PATCH']) diff --git a/api/streaming.py b/api/streaming.py index 819c89b..996b44a 100644 --- a/api/streaming.py +++ b/api/streaming.py @@ -93,7 +93,7 @@ async def stream( 'cookies': incoming_request.cookies }) except ValueError as exc: - webhook = dhooks.Webhook(os.getenv('DISCORD_WEBHOOK__API_ISSUE')) + webhook = dhooks.Webhook(os.environ['DISCORD_WEBHOOK__API_ISSUE']) webhook.send(content=f'API Issue: **`{exc}`**\nhttps://i.imgflip.com/7uv122.jpg') yield await errors.yield_error(500, 'Sorry, the API has no working keys anymore.', 'The admins have been messaged automatically.') return diff --git a/api/transfer.py b/api/transfer.py index 44d1b26..76599b6 100644 --- a/api/transfer.py +++ b/api/transfer.py @@ -44,7 +44,7 @@ async def handle(incoming_request): received_key = incoming_request.headers.get('Authorization') if not received_key or not received_key.startswith('Bearer '): - return await errors.error(401, 'No NovaAI API key given!', 'Add "Authorization: Bearer nv-..." to your request headers.') + return await errors.error(401, 'No NovaAI API key given!', 'Add \'Authorization: Bearer nv-...\' to your request headers.') user = await users.user_by_api_key(received_key.split('Bearer ')[1].strip()) @@ -53,7 +53,7 @@ async def handle(incoming_request): 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.') + 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'] diff --git a/checks/__init__.py b/checks/__init__.py new file mode 100644 index 0000000..ba0a507 --- /dev/null +++ b/checks/__init__.py @@ -0,0 +1 @@ +from . import client \ No newline at end of file diff --git a/checks/__main__.py b/checks/__main__.py new file mode 100644 index 0000000..b6ed337 --- /dev/null +++ b/checks/__main__.py @@ -0,0 +1,2 @@ +import client +client.demo() diff --git a/tests/__main__.py b/checks/client.py similarity index 59% rename from tests/__main__.py rename to checks/client.py index 9dbe378..8d6622b 100644 --- a/tests/__main__.py +++ b/checks/client.py @@ -1,9 +1,11 @@ """Tests the API.""" import os -import openai as closedai -import httpx import time +import httpx +import openai +import asyncio +import traceback from rich import print from typing import List @@ -22,7 +24,7 @@ MESSAGES = [ api_endpoint = 'http://localhost:2332' -def test_server(): +async def test_server(): """Tests if the API server is running.""" try: @@ -30,7 +32,7 @@ def test_server(): except httpx.ConnectError as exc: raise ConnectionError(f'API is not running on port {api_endpoint}.') from exc -def test_api(model: str=MODEL, messages: List[dict]=None) -> dict: +async def test_api(model: str=MODEL, messages: List[dict]=None) -> dict: """Tests an API api_endpoint.""" json_data = { @@ -49,10 +51,10 @@ def test_api(model: str=MODEL, messages: List[dict]=None) -> dict: return response.text -def test_library(): - """Tests if the api_endpoint is working with the Python library.""" +async def test_library(): + """Tests if the api_endpoint is working with the OpenAI Python library.""" - completion = closedai.ChatCompletion.create( + completion = openai.ChatCompletion.create( model=MODEL, messages=MESSAGES ) @@ -61,13 +63,13 @@ def test_library(): return completion['choices'][0]['message']['content'] -def test_library_moderation(): +async def test_library_moderation(): try: - return closedai.Moderation.create('I wanna kill myself, I wanna kill myself; It\'s all I hear right now, it\'s all I hear right now') - except closedai.error.InvalidRequestError: + return openai.Moderation.create('I wanna kill myself, I wanna kill myself; It\'s all I hear right now, it\'s all I hear right now') + except openai.error.InvalidRequestError: return True -def test_models(): +async def test_models(): response = httpx.get( url=f'{api_endpoint}/models', headers=HEADERS, @@ -76,7 +78,7 @@ def test_models(): response.raise_for_status() return response.json() -def test_api_moderation() -> dict: +async def test_api_moderation() -> dict: """Tests an API api_endpoint.""" response = httpx.get( @@ -90,38 +92,43 @@ def test_api_moderation() -> dict: # ========================================================================================== -def test_all(): +def demo(): """Runs all tests.""" - try: - print("Waiting until API Server is started up...") - time.sleep(6) - print('[lightblue]Running test on API server to check if its running...') - print(test_server()) + try: + for _ in range(30): + if test_server(): + break + + print('Waiting until API Server is started up...') + time.sleep(1) + else: + raise ConnectionError('API Server is not running.') print('[lightblue]Running a api endpoint to see if requests can go through...') - print(test_api('gpt-3.5-trubo')) + print(asyncio.run(test_api('gpt-3.5-turbo'))) print('[lightblue]Checking if the API works with the python library...') - print(test_library()) + print(asyncio.run(test_library())) print('[lightblue]Checking if the moderation endpoint works...') - print(test_library_moderation()) + print(asyncio.run(test_library_moderation())) print('[lightblue]Checking the /v1/models endpoint...') - print(test_models()) - except Exception as e: - print('[red]Error: ') - print(e) + print(asyncio.run(test_models())) + + except Exception as exc: + print('[red]Error: ' + str(exc)) + traceback.print_exc() exit(500) +openai.api_base = api_endpoint +openai.api_key = os.environ['NOVA_KEY'] + +HEADERS = { + 'Content-Type': 'application/json', + 'Authorization': 'Bearer ' + openai.api_key +} + if __name__ == '__main__': - closedai.api_base = api_endpoint - closedai.api_key = os.environ['NOVA_KEY'] - - HEADERS = { - 'Content-Type': 'application/json', - 'Authorization': 'Bearer ' + closedai.api_key - } - - test_all() + demo() diff --git a/Untitled Diagram.drawio b/privacy.drawio similarity index 100% rename from Untitled Diagram.drawio rename to privacy.drawio diff --git a/screen.sh b/screen.sh index 05f16fb..125ea19 100755 --- a/screen.sh +++ b/screen.sh @@ -12,5 +12,8 @@ cp env/.prod.env /home/nova-prod/.env # Change directory cd /home/nova-prod +# Kill the production server +fuser -k 2333/tcp + # Start screen screen -S nova-api python run prod diff --git a/setup.md b/setup.md index 8835c4f..1a3c2ed 100644 --- a/setup.md +++ b/setup.md @@ -100,7 +100,7 @@ python run 1337 ``` ## Test if it works -`python tests` +`python checks` ## Ports ```yml