Underground-Digital/Workflow-Engine

1from flask import request2from flask_login import current_user3from flask_restful import Resource, marshal, marshal_with, reqparse4from werkzeug.exceptions import Forbidden5 6from controllers.console import api7from controllers.console.app.error import NoFileUploadedError8from controllers.console.datasets.error import TooManyFilesError9from controllers.console.wraps import (10    account_initialization_required,11    cloud_edition_billing_resource_check,12    setup_required,13)14from extensions.ext_redis import redis_client15from fields.annotation_fields import (16    annotation_fields,17    annotation_hit_history_fields,18)19from libs.login import login_required20from services.annotation_service import AppAnnotationService21 22 23class AnnotationReplyActionApi(Resource):24    @setup_required25    @login_required26    @account_initialization_required27    @cloud_edition_billing_resource_check("annotation")28    def post(self, app_id, action):29        if not current_user.is_editor:30            raise Forbidden()31 32        app_id = str(app_id)33        parser = reqparse.RequestParser()34        parser.add_argument("score_threshold", required=True, type=float, location="json")35        parser.add_argument("embedding_provider_name", required=True, type=str, location="json")36        parser.add_argument("embedding_model_name", required=True, type=str, location="json")37        args = parser.parse_args()38        if action == "enable":39            result = AppAnnotationService.enable_app_annotation(args, app_id)40        elif action == "disable":41            result = AppAnnotationService.disable_app_annotation(app_id)42        else:43            raise ValueError("Unsupported annotation reply action")44        return result, 20045 46 47class AppAnnotationSettingDetailApi(Resource):48    @setup_required49    @login_required50    @account_initialization_required51    def get(self, app_id):52        if not current_user.is_editor:53            raise Forbidden()54 55        app_id = str(app_id)56        result = AppAnnotationService.get_app_annotation_setting_by_app_id(app_id)57        return result, 20058 59 60class AppAnnotationSettingUpdateApi(Resource):61    @setup_required62    @login_required63    @account_initialization_required64    def post(self, app_id, annotation_setting_id):65        if not current_user.is_editor:66            raise Forbidden()67 68        app_id = str(app_id)69        annotation_setting_id = str(annotation_setting_id)70 71        parser = reqparse.RequestParser()72        parser.add_argument("score_threshold", required=True, type=float, location="json")73        args = parser.parse_args()74 75        result = AppAnnotationService.update_app_annotation_setting(app_id, annotation_setting_id, args)76        return result, 20077 78 79class AnnotationReplyActionStatusApi(Resource):80    @setup_required81    @login_required82    @account_initialization_required83    @cloud_edition_billing_resource_check("annotation")84    def get(self, app_id, job_id, action):85        if not current_user.is_editor:86            raise Forbidden()87 88        job_id = str(job_id)89        app_annotation_job_key = "{}_app_annotation_job_{}".format(action, str(job_id))90        cache_result = redis_client.get(app_annotation_job_key)91        if cache_result is None:92            raise ValueError("The job is not exist.")93 94        job_status = cache_result.decode()95        error_msg = ""96        if job_status == "error":97            app_annotation_error_key = "{}_app_annotation_error_{}".format(action, str(job_id))98            error_msg = redis_client.get(app_annotation_error_key).decode()99 100        return {"job_id": job_id, "job_status": job_status, "error_msg": error_msg}, 200101 102 103class AnnotationListApi(Resource):104    @setup_required105    @login_required106    @account_initialization_required107    def get(self, app_id):108        if not current_user.is_editor:109            raise Forbidden()110 111        page = request.args.get("page", default=1, type=int)112        limit = request.args.get("limit", default=20, type=int)113        keyword = request.args.get("keyword", default=None, type=str)114 115        app_id = str(app_id)116        annotation_list, total = AppAnnotationService.get_annotation_list_by_app_id(app_id, page, limit, keyword)117        response = {118            "data": marshal(annotation_list, annotation_fields),119            "has_more": len(annotation_list) == limit,120            "limit": limit,121            "total": total,122            "page": page,123        }124        return response, 200125 126 127class AnnotationExportApi(Resource):128    @setup_required129    @login_required130    @account_initialization_required131    def get(self, app_id):132        if not current_user.is_editor:133            raise Forbidden()134 135        app_id = str(app_id)136        annotation_list = AppAnnotationService.export_annotation_list_by_app_id(app_id)137        response = {"data": marshal(annotation_list, annotation_fields)}138        return response, 200139 140 141class AnnotationCreateApi(Resource):142    @setup_required143    @login_required144    @account_initialization_required145    @cloud_edition_billing_resource_check("annotation")146    @marshal_with(annotation_fields)147    def post(self, app_id):148        if not current_user.is_editor:149            raise Forbidden()150 151        app_id = str(app_id)152        parser = reqparse.RequestParser()153        parser.add_argument("question", required=True, type=str, location="json")154        parser.add_argument("answer", required=True, type=str, location="json")155        args = parser.parse_args()156        annotation = AppAnnotationService.insert_app_annotation_directly(args, app_id)157        return annotation158 159 160class AnnotationUpdateDeleteApi(Resource):161    @setup_required162    @login_required163    @account_initialization_required164    @cloud_edition_billing_resource_check("annotation")165    @marshal_with(annotation_fields)166    def post(self, app_id, annotation_id):167        if not current_user.is_editor:168            raise Forbidden()169 170        app_id = str(app_id)171        annotation_id = str(annotation_id)172        parser = reqparse.RequestParser()173        parser.add_argument("question", required=True, type=str, location="json")174        parser.add_argument("answer", required=True, type=str, location="json")175        args = parser.parse_args()176        annotation = AppAnnotationService.update_app_annotation_directly(args, app_id, annotation_id)177        return annotation178 179    @setup_required180    @login_required181    @account_initialization_required182    def delete(self, app_id, annotation_id):183        if not current_user.is_editor:184            raise Forbidden()185 186        app_id = str(app_id)187        annotation_id = str(annotation_id)188        AppAnnotationService.delete_app_annotation(app_id, annotation_id)189        return {"result": "success"}, 200190 191 192class AnnotationBatchImportApi(Resource):193    @setup_required194    @login_required195    @account_initialization_required196    @cloud_edition_billing_resource_check("annotation")197    def post(self, app_id):198        if not current_user.is_editor:199            raise Forbidden()200 201        app_id = str(app_id)202        # get file from request203        file = request.files["file"]204        # check file205        if "file" not in request.files:206            raise NoFileUploadedError()207 208        if len(request.files) > 1:209            raise TooManyFilesError()210        # check file type211        if not file.filename.endswith(".csv"):212            raise ValueError("Invalid file type. Only CSV files are allowed")213        return AppAnnotationService.batch_import_app_annotations(app_id, file)214 215 216class AnnotationBatchImportStatusApi(Resource):217    @setup_required218    @login_required219    @account_initialization_required220    @cloud_edition_billing_resource_check("annotation")221    def get(self, app_id, job_id):222        if not current_user.is_editor:223            raise Forbidden()224 225        job_id = str(job_id)226        indexing_cache_key = "app_annotation_batch_import_{}".format(str(job_id))227        cache_result = redis_client.get(indexing_cache_key)228        if cache_result is None:229            raise ValueError("The job is not exist.")230        job_status = cache_result.decode()231        error_msg = ""232        if job_status == "error":233            indexing_error_msg_key = "app_annotation_batch_import_error_msg_{}".format(str(job_id))234            error_msg = redis_client.get(indexing_error_msg_key).decode()235 236        return {"job_id": job_id, "job_status": job_status, "error_msg": error_msg}, 200237 238 239class AnnotationHitHistoryListApi(Resource):240    @setup_required241    @login_required242    @account_initialization_required243    def get(self, app_id, annotation_id):244        if not current_user.is_editor:245            raise Forbidden()246 247        page = request.args.get("page", default=1, type=int)248        limit = request.args.get("limit", default=20, type=int)249        app_id = str(app_id)250        annotation_id = str(annotation_id)251        annotation_hit_history_list, total = AppAnnotationService.get_annotation_hit_histories(252            app_id, annotation_id, page, limit253        )254        response = {255            "data": marshal(annotation_hit_history_list, annotation_hit_history_fields),256            "has_more": len(annotation_hit_history_list) == limit,257            "limit": limit,258            "total": total,259            "page": page,260        }261        return response262 263 264api.add_resource(AnnotationReplyActionApi, "/apps/<uuid:app_id>/annotation-reply/<string:action>")265api.add_resource(266    AnnotationReplyActionStatusApi, "/apps/<uuid:app_id>/annotation-reply/<string:action>/status/<uuid:job_id>"267)268api.add_resource(AnnotationListApi, "/apps/<uuid:app_id>/annotations")269api.add_resource(AnnotationExportApi, "/apps/<uuid:app_id>/annotations/export")270api.add_resource(AnnotationUpdateDeleteApi, "/apps/<uuid:app_id>/annotations/<uuid:annotation_id>")271api.add_resource(AnnotationBatchImportApi, "/apps/<uuid:app_id>/annotations/batch-import")272api.add_resource(AnnotationBatchImportStatusApi, "/apps/<uuid:app_id>/annotations/batch-import-status/<uuid:job_id>")273api.add_resource(AnnotationHitHistoryListApi, "/apps/<uuid:app_id>/annotations/<uuid:annotation_id>/hit-histories")274api.add_resource(AppAnnotationSettingDetailApi, "/apps/<uuid:app_id>/annotation-setting")275api.add_resource(AppAnnotationSettingUpdateApi, "/apps/<uuid:app_id>/annotation-settings/<uuid:annotation_setting_id>")276