From 4673b2305523308a8d9deb621b30f8c4a4af616e Mon Sep 17 00:00:00 2001 From: henceiusegentoo Date: Thu, 21 Sep 2023 20:17:32 +0200 Subject: [PATCH] Made the API return 404 if the model can't be found --- api/handler.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/api/handler.py b/api/handler.py index 101d2d9..5c988d5 100644 --- a/api/handler.py +++ b/api/handler.py @@ -20,6 +20,7 @@ load_dotenv() users = UserManager() models_list = json.load(open('models.json', encoding='utf8')) +models = [model['id'] for model in models_list['data']] with open('config/config.yml', encoding='utf8') as f: config = yaml.safe_load(f) @@ -149,6 +150,9 @@ async def handle(incoming_request: fastapi.Request): media_type = 'text/event-stream' if payload.get('stream', False) else 'application/json' + if payload['model'] not in models: + return await errors.error(404, 'Model not found.', 'Check the model name and try again.') + return fastapi.responses.StreamingResponse( content=responder.respond( user=user,