model-service/model_service/views/root_view.py

74 lines
2.1 KiB
Python

#!/usr/bin/env python3
from flask import jsonify, abort, request
from flask_classful import FlaskView, route
from marshmallow import ValidationError
from utils import json_required
from model import db, AIModel, AIModelType, Default
from schemas import AIModelSchema, DefaultSchema
class RootView(FlaskView):
route_base = '/'
aimodels_schema = AIModelSchema(many=True, exclude=['timestamp', 'details', 'target_class_name'])
aimodel_schema = AIModelSchema(many=False)
default_schema = DefaultSchema(many=False)
## Shared stuff goes here
def index(self):
models = AIModel.query.all()
return jsonify(self.aimodels_schema.dump(models))
@route('/<type_>')
def get_models(self, type_: str):
try:
aimodel_type = AIModelType[type_]
except KeyError:
abort(404, "Unknown type")
models = AIModel.query.filter_by(type=aimodel_type).all()
return jsonify(self.aimodels_schema.dump(models)), 200
@route('/<type_>/<id_>')
def get_model(self, type_: str, id_: str):
try:
aimodel_type = AIModelType[type_]
except KeyError:
abort(404, "Unknown type")
if id_ == "$default":
default = Default.query.filter_by(type=aimodel_type).first_or_404()
m = default.aimodel
else:
m = AIModel.query.filter_by(type=aimodel_type, id=id_).first_or_404()
return jsonify(self.aimodel_schema.dump(m))
@json_required
@route('/<type_>/$default', methods=['PUT'])
def put_default(self, type_: str):
try:
aimodel_type = AIModelType[type_]
except KeyError:
abort(404, "Unknown type")
try:
req = self.default_schema.load(request.json)
except ValidationError as e:
abort(400, str(e))
m = AIModel.query.filter_by(type=aimodel_type, id=req['id']).first_or_404()
Default.query.filter_by(type=aimodel_type).delete()
new_default = Default(type=aimodel_type, aimodel=m)
db.session.add(new_default)
db.session.commit()
return '', 204