Made the API return 404 if the model can't be found

This commit is contained in:
henceiusegentoo 2023-09-21 20:17:32 +02:00
parent 0d0453291f
commit 4673b23055

View file

@ -20,6 +20,7 @@ load_dotenv()
users = UserManager() users = UserManager()
models_list = json.load(open('models.json', encoding='utf8')) 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: with open('config/config.yml', encoding='utf8') as f:
config = yaml.safe_load(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' 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( return fastapi.responses.StreamingResponse(
content=responder.respond( content=responder.respond(
user=user, user=user,