tjdify/api/controllers/console/intention.py

import json
import logging
import os
import zipfile

from flask import request, send_file
from flask_restful import Resource, marshal, marshal_with, reqparse
from werkzeug.exceptions import Forbidden, NotFound

import services
from controllers.console import api
from controllers.console.error import FileTooLargeError, UnsupportedFileTypeError
from controllers.console.wraps import account_initialization_required, setup_required
from fields.intention_fields import (
    intention_corpus_detail_fields,
    intention_corpus_similarity_question_fields,
    intention_detail_fields,
    intention_keyword_detail_fields,
    intention_keyword_fields,
    intention_page_fields,
    intention_train_file_binding_fields,
    intention_train_file_fields,
    intention_train_task_fields,
    intention_type_detail_fields,
    intention_type_page_fields,
)
from libs.login import current_user, login_required
from models import UploadFile
from models.intention import IntentionTrainTask
from services.errors.intention import IntentionTrainFileDuplicateError
from services.file_service import FileService
from services.intention_service import (
    IntentionCorpusService,
    IntentionCorpusSimilarityQuestionService,
    IntentionKeywordService,
    IntentionService,
    IntentionTrainFileBindingService,
    IntentionTrainFileService,
    IntentionTrainTaskService,
    IntentionTypeService,
)
from services.upload_file_service import UploadFileService


class IntentionListApi(Resource):

    @setup_required
    @login_required
    @account_initialization_required
    def get(self):
        page = request.args.get("page", default=1, type=int)
        limit = request.args.get("limit", default=20, type=int)
        type_id = request.args.get("type_id", default=None, type=str)
        name_search = request.args.get("name_search", default=None, type=str)
        intentions, total = IntentionService.get_intentions(page, limit, type_id, name_search)
        data = marshal(intentions, intention_page_fields)
        response = {"data": data, "has_more": len(intentions) == limit, "limit": limit,
                    "total": total, "page": page}
        return response, 200

    @setup_required
    @login_required
    @account_initialization_required
    def post(self):
        parser = reqparse.RequestParser()
        parser.add_argument(
            "name",
            nullable=False,
            required=True,
            help="type is required. Name must be between 1 to 40 characters.",
        )
        parser.add_argument(
            "type_id",
            nullable=False,
            required=True,
            help="type is required.",
        )
        args = parser.parse_args()
        intention = IntentionService.save_intention(args)
        response = marshal(intention, intention_detail_fields)
        return response, 200

class IntentionApi(Resource):
    @setup_required
    @login_required
    @account_initialization_required
    def get(self, intention_id):
        intention = IntentionService.get_intention(intention_id)
        return marshal(intention, intention_detail_fields)

    @setup_required
    @login_required
    @account_initialization_required
    def patch(self, intention_id):
        parser = reqparse.RequestParser()
        parser.add_argument(
            "name",
            nullable=False,
            required=True,
            help="type is required. Name must be between 1 to 40 characters.",
        )
        parser.add_argument(
            "type_id",
            nullable=False,
            required=True,
            help="type is required.",
        )
        args = parser.parse_args()
        intention = IntentionService.update_intention(intention_id, args)
        response = marshal(intention, intention_detail_fields)
        return response, 200

    @setup_required
    @login_required
    @account_initialization_required
    def delete(self, intention_id):
        IntentionService.delete_intention(intention_id)
        return 200

class IntentionTypeListApi(Resource):

    @setup_required
    @login_required
    @account_initialization_required
    def get(self):
        page = request.args.get("page", default=1, type=int)
        limit = request.args.get("limit", default=20, type=int)
        search = request.args.get("search", default=None, type=str)
        intention_types, total = IntentionTypeService.get_intention_types(page, limit, search)
        data = marshal(intention_types, intention_type_page_fields)
        response = {"data": data, "has_more": len(intention_types) == limit, "limit": limit,
                    "total": total, "page": page}
        return response, 200

    @setup_required
    @login_required
    @account_initialization_required
    def post(self):
        parser = reqparse.RequestParser()
        parser.add_argument(
            "name",
            nullable=False,
            required=True,
            help="type is required. Name must be between 1 to 40 characters.",
        )
        args = parser.parse_args()
        intention_type = IntentionTypeService.save_intention_type(args)
        response = marshal(intention_type, intention_type_detail_fields)
        return response, 200

class IntentionTypeApi(Resource):
    @setup_required
    @login_required
    @account_initialization_required
    def get(self, intention_type_id):
        intention_type = IntentionTypeService.get_intention_type(intention_type_id)
        return marshal(intention_type, intention_type_detail_fields)

    @setup_required
    @login_required
    @account_initialization_required
    def patch(self, intention_type_id):
        parser = reqparse.RequestParser()
        parser.add_argument(
            "name",
            nullable=False,
            required=True,
            help="type is required. Name must be between 1 to 40 characters.",
        )
        args = parser.parse_args()
        intention_type = IntentionTypeService.update_intention_type(intention_type_id, args)
        return marshal(intention_type, intention_type_detail_fields), 200

    @setup_required
    @login_required
    @account_initialization_required
    def delete(self, intention_type_id):
        IntentionTypeService.delete_intention_type(intention_type_id)
        return 200

class IntentionKeywordListApi(Resource):
    @setup_required
    @login_required
    @account_initialization_required
    @marshal_with(intention_keyword_fields)
    def get(self, intention_id):
        search = request.args.get("search", default=None, type=str)
        intention = IntentionService.get_intention(intention_id)
        if not intention:
            raise NotFound("Intention not found")
        intention_keywords = IntentionKeywordService.get_intention_keywords(intention_id, search)
        return intention_keywords, 200

    @setup_required
    @login_required
    @account_initialization_required
    @marshal_with(intention_keyword_detail_fields)
    def post(self, intention_id):
        parser = reqparse.RequestParser()
        parser.add_argument(
            "name",
            nullable=False,
            required=True,
            help="type is required. Name must be between 1 to 40 characters.",
        )
        args = parser.parse_args()
        intention = IntentionService.get_intention(intention_id)
        if not intention:
            raise NotFound("Intention not found")
        intention_keyword = IntentionKeywordService.save_intention_keyword(intention_id, args)
        return intention_keyword, 200

    @setup_required
    @login_required
    @account_initialization_required
    def delete(self, intention_id):
        intention = IntentionService.get_intention(intention_id)
        if not intention:
            raise NotFound("Intention not found")

        IntentionKeywordService.delete_intention_keywords_by_intention_id(intention_id)
        return 200

class IntentionKeywordApi(Resource):

    @setup_required
    @login_required
    @account_initialization_required
    def get(self, intention_keyword_id):
        intention_keyword = IntentionKeywordService.get_intention_keyword(intention_keyword_id)
        if not intention_keyword:
            return {}, 200
        return marshal(intention_keyword, intention_keyword_detail_fields), 200

    @setup_required
    @login_required
    @account_initialization_required
    @marshal_with(intention_keyword_detail_fields)
    def patch(self, intention_keyword_id):
        parser = reqparse.RequestParser()
        parser.add_argument(
            "name",
            nullable=False,
            required=True,
            help="type is required. Name must be between 1 to 40 characters.",
        )
        parser.add_argument(
            "intention_id",
            nullable=False,
            required=True,
            help="type is required.",
        )
        args = parser.parse_args()
        intention_keyword = IntentionKeywordService.update_intention_keyword(intention_keyword_id, args)
        return intention_keyword, 200

    @setup_required
    @login_required
    @account_initialization_required
    def delete(self, intention_keyword_id):
        IntentionKeywordService.delete_intention_keyword(intention_keyword_id)
        return 200

class IntentionKeywordBatchApi(Resource):

    @setup_required
    @login_required
    @account_initialization_required
    def post(self):
        parser = reqparse.RequestParser()
        parser.add_argument(
            "method",
            nullable=False,
            required=True,
            help="method is required.",
            choices=["create", "update", "delete"],
            type=str,
            location="json",
        )
        parser.add_argument(
            "delete_data",
            nullable=False,
            required=True,
            help="delete_data is required.",
            type=list,
            location="json",
        )
        args = parser.parse_args()
        logging.info(args)

        method = args["method"]
        if method == "delete":
            intention_keyword_ids = args["delete_data"]
            IntentionKeywordService.delete_intention_keywords(intention_keyword_ids)
            return 200
        else:
            raise NotFound(f"method with name {method} not found")

class IntentionCorpusListApi(Resource):
    @setup_required
    @login_required
    @account_initialization_required
    def get(self):
        page = request.args.get("page", default=1, type=int)
        limit = request.args.get("limit", default=20, type=int)
        question_search = request.args.get("question_search", default=None, type=str)
        intention_id = request.args.get("intention_id", default=None, type=str)
        intention_corpus, total = IntentionCorpusService.get_page_intention_corpus(
            page, limit, question_search, intention_id)
        data = marshal(intention_corpus, intention_corpus_detail_fields)
        response = {"data": data, "has_more": len(intention_corpus) == limit, "limit": limit,
                    "total": total, "page": page}
        return response, 200

    @setup_required
    @login_required
    @account_initialization_required
    def post(self):
        parser = reqparse.RequestParser()
        parser.add_argument(
            "question",
            nullable=False,
            required=True,
            help="type is required. Question must be between 1 to 40 characters.",
        )
        parser.add_argument(
            "question_config",
            nullable=True,
            required=False,
            location="json",
        )
        parser.add_argument(
            "intention_id",
            nullable=False,
            required=True,
            help="type is required.",
        )
        args = parser.parse_args()
        intention_corpus = IntentionCorpusService.save_intention_corpus(args)
        return marshal(intention_corpus, intention_corpus_detail_fields), 200

class IntentionCorpusApi(Resource):
    @setup_required
    @login_required
    @account_initialization_required
    def get(self, corpus_id):
        intention_corpus = IntentionCorpusService.get_intention_corpus(corpus_id)
        if not intention_corpus:
            raise NotFound(f"IntentionCorpus with id {corpus_id} not found")
        return marshal(intention_corpus, intention_corpus_detail_fields), 200

    @setup_required
    @login_required
    @account_initialization_required
    def patch(self, corpus_id):
        parser = reqparse.RequestParser()
        parser.add_argument(
            "question",
            nullable=True,
            required=False,
            type=str,
            location="json",
        )
        parser.add_argument(
            "question_config",
            nullable=True,
            required=False,
            location="json",
        )
        parser.add_argument(
            "intention_id",
            nullable=True,
            required=False,
            type=str,
            location="json",
        )
        args = parser.parse_args()
        intention_corpus = IntentionCorpusService.update_intention_corpus(corpus_id, args)
        return marshal(intention_corpus, intention_corpus_detail_fields), 200

    @setup_required
    @login_required
    @account_initialization_required
    def delete(self, corpus_id):
        intention_corpus = IntentionCorpusService.get_intention_corpus(corpus_id)
        if not intention_corpus:
            raise NotFound(f"未发现Id未{corpus_id}的训练语料")

        similarity_questions = intention_corpus.similarity_questions
        if similarity_questions:
            raise Forbidden(f"存在与其关联的相似问题，无法删除Id为{corpus_id}训练语料")

        IntentionCorpusService.delete_intention_corpus(intention_corpus)
        return 200

class IntentionCorpusSimilarityQuestionApi(Resource):
    @setup_required
    @login_required
    @account_initialization_required
    def get(self, corpus_id):
        search = request.args.get("search", default=None, type=str)
        similarity_questions = (
            IntentionCorpusSimilarityQuestionService
            .get_similarity_questions_by_corpus_id_like_question(corpus_id, search)
        )
        return marshal(similarity_questions, intention_corpus_similarity_question_fields), 200

    @setup_required
    @login_required
    @account_initialization_required
    def post(self, corpus_id):
        parser = reqparse.RequestParser()
        parser.add_argument(
            "question",
            nullable=False,
            required=True,
            help="type is required. Question must be between 1 to 40 characters.",
            location="json",
        )
        parser.add_argument(
            "question_config",
            nullable=True,
            required=False,
            location="json",
        )
        args = parser.parse_args()
        intention_corpus_similarity_question = (
            IntentionCorpusSimilarityQuestionService.save_similarity_question(corpus_id, args)
        )
        return marshal(intention_corpus_similarity_question, intention_corpus_similarity_question_fields), 200

    @setup_required
    @login_required
    @account_initialization_required
    def delete(self, corpus_id):
        IntentionCorpusSimilarityQuestionService.delete_similarity_question_by_corpus_id(corpus_id)
        return 200

class IntentionCorpusSimilarityQuestionUpdateAndDeleteApi(Resource):
    @setup_required
    @login_required
    @account_initialization_required
    def patch(self, similarity_question_id):
        parser = reqparse.RequestParser()
        parser.add_argument(
            "question",
            nullable=True,
            required=False,
            help="type is required. Question must be between 1 to 40 characters.",
            location="json",
        )
        parser.add_argument(
            "question_config",
            nullable=True,
            required=False,
            location="json",
        )
        parser.add_argument(
            "corpus_id",
            nullable=True,
            required=False,
            location="json",
        )
        args = parser.parse_args()

        similarity_question = (
            IntentionCorpusSimilarityQuestionService.update_similarity_question(similarity_question_id, args)
        )
        return marshal(similarity_question, intention_corpus_similarity_question_fields), 200

    @setup_required
    @login_required
    @account_initialization_required
    def delete(self, similarity_question_id):
        IntentionCorpusSimilarityQuestionService.delete_similarity_question_by_id(similarity_question_id)
        return 200

class IntentionCorpusSimilarityQuestionBatchApi(Resource):
    @setup_required
    @login_required
    @account_initialization_required
    def post(self):
        parser = reqparse.RequestParser()
        parser.add_argument(
            "method",
            nullable=False,
            required=True,
            help="method is required.",
            choices=["create", "update", "delete"],
            type=str,
            location="json",
        )
        parser.add_argument(
            "data",
            nullable=False,
            required=True,
            help="data is required.",
            type=list,
            location="json",
        )
        args = parser.parse_args()
        logging.info(args)

        method = args["method"]
        if method == "delete":
            similarity_question_ids = args["data"]
            IntentionCorpusSimilarityQuestionService.delete_similarity_questions_by_ids(similarity_question_ids)
            return 200
        else:
            raise NotFound(f"method with name {method} not found")

class IntentionTrainTaskListApi(Resource):
    @setup_required
    @login_required
    @account_initialization_required
    def get(self):
        page = request.args.get("page", default=1, type=int)
        limit = request.args.get("limit", default=20, type=int)
        search = request.args.get("search", default=None, type=str)
        intention_train_tasks, total = IntentionTrainTaskService.get_page_intention_train_tasks(
            page, limit, search)
        data = marshal(intention_train_tasks, intention_train_task_fields)
        response = {"data": data, "has_more": len(intention_train_tasks) == limit, "limit": limit,
                    "total": total, "page": page}
        return response, 200

    @setup_required
    @login_required
    @account_initialization_required
    def post(self):
        parser = reqparse.RequestParser()
        parser.add_argument(
            "name",
            nullable=False,
            required=True,
            help="name is required.",
            location="json",
        )
        parser.add_argument(
            "status",
            nullable=False,
            required=True,
            help="status is required.",
            choices=IntentionTrainTask.STATUS_LIST,
            location="json",
        )
        args = parser.parse_args()
        train_task = IntentionTrainTaskService.save_train_task(args)
        return marshal(train_task, intention_train_task_fields), 200

class IntentionTrainTaskApi(Resource):
    @setup_required
    @login_required
    @account_initialization_required
    def patch(self, task_id):
        parser = reqparse.RequestParser()
        parser.add_argument(
            "name",
            nullable=False,
            required=True,
            help="name is required.",
            location="json",
        )
        parser.add_argument(
            "status",
            nullable=False,
            required=True,
            help="status is required.",
            choices=IntentionTrainTask.STATUS_LIST,
            location="json",
        )
        args = parser.parse_args()

        train_task = IntentionTrainTaskService.update_train_task(task_id, args)
        return marshal(train_task, intention_train_task_fields), 200

class IntentionTrainTaskDownloadApi(Resource):
    @setup_required
    @login_required
    @account_initialization_required
    def get(self, task_id):
        train_task = IntentionTrainTaskService.get_train_task(task_id)
        if train_task.status != "COMPLETED":
            raise Forbidden(f"Task with id {task_id} not completed")

        dataset_info = train_task.dataset_info
        dataset_source_info = json.loads(dataset_info.data_source_info)
        dataset_file_id = dataset_source_info["upload_file_id"]
        dataset_upload_file: UploadFile = UploadFileService.get_upload_file(dataset_file_id)

        model_info = train_task.model_info
        model_source_info = json.loads(model_info.data_source_info)
        model_file_id = model_source_info["upload_file_id"]
        model_upload_file: UploadFile = UploadFileService.get_upload_file(model_file_id)

        def file2zip(zip_filename: str, upload_files: list[UploadFile]):
            with zipfile.ZipFile(zip_filename, "w", compression=zipfile.ZIP_DEFLATED) as zip_file:
                for upload_file in upload_files:
                    filename = f"storage/{dataset_upload_file.key}"
                    zip_file.write(filename, arcname=upload_file.name)

        # 生成待下载的zip包
        zip_filename = f"{train_task.name}.zip"
        upload_files: list[UploadFile] = [dataset_upload_file, model_upload_file]
        file2zip(zip_filename, upload_files)

        # 下载zip包
        response = send_file(zip_filename, as_attachment=True, download_name=zip_filename)
        os.remove(zip_filename)
        return response

class IntentionTrainFileApi(Resource):
    @setup_required
    @login_required
    @account_initialization_required
    def get(self):
        name = request.args.get("name", default=None, type=str)
        version = request.args.get("version", default=None, type=str)
        type = request.args.get("type", default=None, type=str)
        train_files = IntentionTrainFileService.get_train_files(name, version, type)
        return marshal(train_files, intention_train_file_fields), 200

    @setup_required
    @login_required
    @account_initialization_required
    def post(self):
        name = request.form.get("name")
        version = request.form.get("version")
        type = request.form.get("type")
        train_file = IntentionTrainFileService.get_train_file(name, version, type)
        if train_file:
            raise IntentionTrainFileDuplicateError(f"IntentionTrainFile with name-version-type "
                                                   f"{name}-{version}-{type} already exists.")

        data_source_type = request.form.get("data_source_type")

        # get file from request
        file = request.files["file"]
        filename = file.filename
        mimetype = file.mimetype
        if not filename or not mimetype:
            raise Forbidden("Invalid request.")

        try:
            upload_file = FileService.upload_file(
                filename=filename,
                content=file.read(),
                mimetype=mimetype,
                user=current_user,
                source=None,
            )

            args = {
                "name": name,
                "version": version,
                "type": type,
                "data_source_type": data_source_type,
                "data_source_info": {
                    "upload_file_id": upload_file.id
                }
            }
            intention_train_file = IntentionTrainFileService.save_train_file(args)
            return marshal(intention_train_file, intention_train_file_fields), 200
        except services.errors.file.FileTooLargeError as file_too_large_error:
            raise FileTooLargeError(file_too_large_error.description)
        except services.errors.file.UnsupportedFileTypeError:
            raise UnsupportedFileTypeError()

class IntentionTrainFileBindingApi(Resource):
    @setup_required
    @login_required
    @account_initialization_required
    def post(self):
        parser = reqparse.RequestParser()
        parser.add_argument(
            "file_id",
            nullable=False,
            required=True,
            help="file_id is required.",
            location="json",
        )
        parser.add_argument(
            "task_id",
            nullable=False,
            required=True,
            help="task_id is required.",
            location="json",
        )
        args = parser.parse_args()

        train_file_binding = IntentionTrainFileBindingService.save_train_file_binding(args)
        return marshal(train_file_binding, intention_train_file_binding_fields), 200

api.add_resource(IntentionListApi, "/intentions")
api.add_resource(IntentionApi, "/intentions/<uuid:intention_id>")
api.add_resource(IntentionTypeListApi, "/intentions/types")
api.add_resource(IntentionTypeApi, "/intentions/types/<uuid:intention_type_id>")
api.add_resource(IntentionKeywordListApi, "/intentions/<uuid:intention_id>/keywords")
api.add_resource(IntentionKeywordApi, "/intentions/keywords/<uuid:intention_keyword_id>")
api.add_resource(IntentionKeywordBatchApi, "/intentions/keywords/batch")

api.add_resource(IntentionCorpusListApi, "/intentions/corpus")
api.add_resource(IntentionCorpusApi, "/intentions/corpus/<uuid:corpus_id>")
api.add_resource(IntentionCorpusSimilarityQuestionApi, "/intentions/corpus/<uuid:corpus_id>/similarity_questions")

api.add_resource(IntentionCorpusSimilarityQuestionUpdateAndDeleteApi,
                 "/intentions/similarity_questions/<uuid:similarity_question_id>")
api.add_resource(IntentionCorpusSimilarityQuestionBatchApi, "/intentions/similarity_questions/batch")

api.add_resource(IntentionTrainTaskListApi, "/intentions/train_tasks")
api.add_resource(IntentionTrainTaskApi, "/intentions/train_tasks/<uuid:task_id>")
api.add_resource(IntentionTrainTaskDownloadApi, "/intentions/train_tasks/download/<uuid:task_id>")

api.add_resource(IntentionTrainFileApi, "/intentions/train_files")

api.add_resource(IntentionTrainFileBindingApi, "/intentions/train_file_bindings")