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,