from flask import request from flask_login import current_user from flask_restful import Resource, marshal, marshal_with, reqparse from werkzeug.exceptions import Forbidden from controllers.console import api from controllers.console.app.error import NoFileUploadedError from controllers.console.datasets.error import TooManyFilesError from controllers.console.wraps import ( account_initialization_required, cloud_edition_billing_resource_check, setup_required, ) from extensions.ext_redis import redis_client from fields.annotation_fields import ( annotation_fields, annotation_hit_history_fields, ) from libs.login import login_required from services.annotation_service import AppAnnotationService class AnnotationReplyActionApi(Resource): @setup_required @login_required @account_initialization_required @cloud_edition_billing_resource_check("annotation") def post(self, app_id, action): if not current_user.is_editor: raise Forbidden() app_id = str(app_id) parser = reqparse.RequestParser() parser.add_argument("score_threshold", required=True, type=float, location="json") parser.add_argument("embedding_provider_name", required=True, type=str, location="json") parser.add_argument("embedding_model_name", required=True, type=str, location="json") args = parser.parse_args() if action == "enable": result = AppAnnotationService.enable_app_annotation(args, app_id) elif action == "disable": result = AppAnnotationService.disable_app_annotation(app_id) else: raise ValueError("Unsupported annotation reply action") return result, 200 class AppAnnotationSettingDetailApi(Resource): @setup_required @login_required @account_initialization_required def get(self, app_id): if not current_user.is_editor: raise Forbidden() app_id = str(app_id) result = AppAnnotationService.get_app_annotation_setting_by_app_id(app_id) return result, 200 class AppAnnotationSettingUpdateApi(Resource): @setup_required @login_required @account_initialization_required def post(self, app_id, annotation_setting_id): if not current_user.is_editor: raise Forbidden() app_id = str(app_id) annotation_setting_id = str(annotation_setting_id) parser = reqparse.RequestParser() parser.add_argument("score_threshold", required=True, type=float, location="json") args = parser.parse_args() result = AppAnnotationService.update_app_annotation_setting(app_id, annotation_setting_id, args) return result, 200 class AnnotationReplyActionStatusApi(Resource): @setup_required @login_required @account_initialization_required @cloud_edition_billing_resource_check("annotation") def get(self, app_id, job_id, action): if not current_user.is_editor: raise Forbidden() job_id = str(job_id) app_annotation_job_key = "{}_app_annotation_job_{}".format(action, str(job_id)) cache_result = redis_client.get(app_annotation_job_key) if cache_result is None: raise ValueError("The job is not exist.") job_status = cache_result.decode() error_msg = "" if job_status == "error": app_annotation_error_key = "{}_app_annotation_error_{}".format(action, str(job_id)) error_msg = redis_client.get(app_annotation_error_key).decode() return {"job_id": job_id, "job_status": job_status, "error_msg": error_msg}, 200 class AnnotationListApi(Resource): @setup_required @login_required @account_initialization_required def get(self, app_id): if not current_user.is_editor: raise Forbidden() page = request.args.get("page", default=1, type=int) limit = request.args.get("limit", default=20, type=int) keyword = request.args.get("keyword", default=None, type=str) app_id = str(app_id) annotation_list, total = AppAnnotationService.get_annotation_list_by_app_id(app_id, page, limit, keyword) response = { "data": marshal(annotation_list, annotation_fields), "has_more": len(annotation_list) == limit, "limit": limit, "total": total, "page": page, } return response, 200 class AnnotationExportApi(Resource): @setup_required @login_required @account_initialization_required def get(self, app_id): if not current_user.is_editor: raise Forbidden() app_id = str(app_id) annotation_list = AppAnnotationService.export_annotation_list_by_app_id(app_id) response = {"data": marshal(annotation_list, annotation_fields)} return response, 200 class AnnotationCreateApi(Resource): @setup_required @login_required @account_initialization_required @cloud_edition_billing_resource_check("annotation") @marshal_with(annotation_fields) def post(self, app_id): if not current_user.is_editor: raise Forbidden() app_id = str(app_id) parser = reqparse.RequestParser() parser.add_argument("question", required=True, type=str, location="json") parser.add_argument("answer", required=True, type=str, location="json") args = parser.parse_args() annotation = AppAnnotationService.insert_app_annotation_directly(args, app_id) return annotation class AnnotationUpdateDeleteApi(Resource): @setup_required @login_required @account_initialization_required @cloud_edition_billing_resource_check("annotation") @marshal_with(annotation_fields) def post(self, app_id, annotation_id): if not current_user.is_editor: raise Forbidden() app_id = str(app_id) annotation_id = str(annotation_id) parser = reqparse.RequestParser() parser.add_argument("question", required=True, type=str, location="json") parser.add_argument("answer", required=True, type=str, location="json") args = parser.parse_args() annotation = AppAnnotationService.update_app_annotation_directly(args, app_id, annotation_id) return annotation @setup_required @login_required @account_initialization_required def delete(self, app_id, annotation_id): if not current_user.is_editor: raise Forbidden() app_id = str(app_id) annotation_id = str(annotation_id) AppAnnotationService.delete_app_annotation(app_id, annotation_id) return {"result": "success"}, 200 class AnnotationBatchImportApi(Resource): @setup_required @login_required @account_initialization_required @cloud_edition_billing_resource_check("annotation") def post(self, app_id): if not current_user.is_editor: raise Forbidden() app_id = str(app_id) # get file from request file = request.files["file"] # check file if "file" not in request.files: raise NoFileUploadedError() if len(request.files) > 1: raise TooManyFilesError() # check file type if not file.filename.endswith(".csv"): raise ValueError("Invalid file type. Only CSV files are allowed") return AppAnnotationService.batch_import_app_annotations(app_id, file) class AnnotationBatchImportStatusApi(Resource): @setup_required @login_required @account_initialization_required @cloud_edition_billing_resource_check("annotation") def get(self, app_id, job_id): if not current_user.is_editor: raise Forbidden() job_id = str(job_id) indexing_cache_key = "app_annotation_batch_import_{}".format(str(job_id)) cache_result = redis_client.get(indexing_cache_key) if cache_result is None: raise ValueError("The job is not exist.") job_status = cache_result.decode() error_msg = "" if job_status == "error": indexing_error_msg_key = "app_annotation_batch_import_error_msg_{}".format(str(job_id)) error_msg = redis_client.get(indexing_error_msg_key).decode() return {"job_id": job_id, "job_status": job_status, "error_msg": error_msg}, 200 class AnnotationHitHistoryListApi(Resource): @setup_required @login_required @account_initialization_required def get(self, app_id, annotation_id): if not current_user.is_editor: raise Forbidden() page = request.args.get("page", default=1, type=int) limit = request.args.get("limit", default=20, type=int) app_id = str(app_id) annotation_id = str(annotation_id) annotation_hit_history_list, total = AppAnnotationService.get_annotation_hit_histories( app_id, annotation_id, page, limit ) response = { "data": marshal(annotation_hit_history_list, annotation_hit_history_fields), "has_more": len(annotation_hit_history_list) == limit, "limit": limit, "total": total, "page": page, } return response api.add_resource(AnnotationReplyActionApi, "/apps//annotation-reply/") api.add_resource( AnnotationReplyActionStatusApi, "/apps//annotation-reply//status/" ) api.add_resource(AnnotationListApi, "/apps//annotations") api.add_resource(AnnotationExportApi, "/apps//annotations/export") api.add_resource(AnnotationUpdateDeleteApi, "/apps//annotations/") api.add_resource(AnnotationBatchImportApi, "/apps//annotations/batch-import") api.add_resource(AnnotationBatchImportStatusApi, "/apps//annotations/batch-import-status/") api.add_resource(AnnotationHitHistoryListApi, "/apps//annotations//hit-histories") api.add_resource(AppAnnotationSettingDetailApi, "/apps//annotation-setting") api.add_resource(AppAnnotationSettingUpdateApi, "/apps//annotation-settings/")