Underground-Digital/Workflow-Engine

1import flask_restful2from flask import request3from flask_login import current_user4from flask_restful import Resource, marshal, marshal_with, reqparse5from werkzeug.exceptions import Forbidden, NotFound6 7import services8from configs import dify_config9from controllers.console import api10from controllers.console.apikey import api_key_fields, api_key_list11from controllers.console.app.error import ProviderNotInitializeError12from controllers.console.datasets.error import DatasetInUseError, DatasetNameDuplicateError, IndexingEstimateError13from controllers.console.wraps import account_initialization_required, setup_required14from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError15from core.indexing_runner import IndexingRunner16from core.model_runtime.entities.model_entities import ModelType17from core.provider_manager import ProviderManager18from core.rag.datasource.vdb.vector_type import VectorType19from core.rag.extractor.entity.extract_setting import ExtractSetting20from core.rag.retrieval.retrieval_methods import RetrievalMethod21from extensions.ext_database import db22from fields.app_fields import related_app_list23from fields.dataset_fields import dataset_detail_fields, dataset_query_detail_fields24from fields.document_fields import document_status_fields25from libs.login import login_required26from models import ApiToken, Dataset, Document, DocumentSegment, UploadFile27from models.dataset import DatasetPermissionEnum28from services.dataset_service import DatasetPermissionService, DatasetService, DocumentService29 30 31def _validate_name(name):32    if not name or len(name) < 1 or len(name) > 40:33        raise ValueError("Name must be between 1 to 40 characters.")34    return name35 36 37def _validate_description_length(description):38    if len(description) > 400:39        raise ValueError("Description cannot exceed 400 characters.")40    return description41 42 43class DatasetListApi(Resource):44    @setup_required45    @login_required46    @account_initialization_required47    def get(self):48        page = request.args.get("page", default=1, type=int)49        limit = request.args.get("limit", default=20, type=int)50        ids = request.args.getlist("ids")51        # provider = request.args.get("provider", default="vendor")52        search = request.args.get("keyword", default=None, type=str)53        tag_ids = request.args.getlist("tag_ids")54 55        if ids:56            datasets, total = DatasetService.get_datasets_by_ids(ids, current_user.current_tenant_id)57        else:58            datasets, total = DatasetService.get_datasets(59                page, limit, current_user.current_tenant_id, current_user, search, tag_ids60            )61 62        # check embedding setting63        provider_manager = ProviderManager()64        configurations = provider_manager.get_configurations(tenant_id=current_user.current_tenant_id)65 66        embedding_models = configurations.get_models(model_type=ModelType.TEXT_EMBEDDING, only_active=True)67 68        model_names = []69        for embedding_model in embedding_models:70            model_names.append(f"{embedding_model.model}:{embedding_model.provider.provider}")71 72        data = marshal(datasets, dataset_detail_fields)73        for item in data:74            if item["indexing_technique"] == "high_quality":75                item_model = f"{item['embedding_model']}:{item['embedding_model_provider']}"76                if item_model in model_names:77                    item["embedding_available"] = True78                else:79                    item["embedding_available"] = False80            else:81                item["embedding_available"] = True82 83            if item.get("permission") == "partial_members":84                part_users_list = DatasetPermissionService.get_dataset_partial_member_list(item["id"])85                item.update({"partial_member_list": part_users_list})86            else:87                item.update({"partial_member_list": []})88 89        response = {"data": data, "has_more": len(datasets) == limit, "limit": limit, "total": total, "page": page}90        return response, 20091 92    @setup_required93    @login_required94    @account_initialization_required95    def post(self):96        parser = reqparse.RequestParser()97        parser.add_argument(98            "name",99            nullable=False,100            required=True,101            help="type is required. Name must be between 1 to 40 characters.",102            type=_validate_name,103        )104        parser.add_argument(105            "description",106            type=str,107            nullable=True,108            required=False,109            default="",110        )111        parser.add_argument(112            "indexing_technique",113            type=str,114            location="json",115            choices=Dataset.INDEXING_TECHNIQUE_LIST,116            nullable=True,117            help="Invalid indexing technique.",118        )119        parser.add_argument(120            "external_knowledge_api_id",121            type=str,122            nullable=True,123            required=False,124        )125        parser.add_argument(126            "provider",127            type=str,128            nullable=True,129            choices=Dataset.PROVIDER_LIST,130            required=False,131            default="vendor",132        )133        parser.add_argument(134            "external_knowledge_id",135            type=str,136            nullable=True,137            required=False,138        )139        args = parser.parse_args()140 141        # The role of the current user in the ta table must be admin, owner, or editor, or dataset_operator142        if not current_user.is_dataset_editor:143            raise Forbidden()144 145        try:146            dataset = DatasetService.create_empty_dataset(147                tenant_id=current_user.current_tenant_id,148                name=args["name"],149                description=args["description"],150                indexing_technique=args["indexing_technique"],151                account=current_user,152                permission=DatasetPermissionEnum.ONLY_ME,153                provider=args["provider"],154                external_knowledge_api_id=args["external_knowledge_api_id"],155                external_knowledge_id=args["external_knowledge_id"],156            )157        except services.errors.dataset.DatasetNameDuplicateError:158            raise DatasetNameDuplicateError()159 160        return marshal(dataset, dataset_detail_fields), 201161 162 163class DatasetApi(Resource):164    @setup_required165    @login_required166    @account_initialization_required167    def get(self, dataset_id):168        dataset_id_str = str(dataset_id)169        dataset = DatasetService.get_dataset(dataset_id_str)170        if dataset is None:171            raise NotFound("Dataset not found.")172        try:173            DatasetService.check_dataset_permission(dataset, current_user)174        except services.errors.account.NoPermissionError as e:175            raise Forbidden(str(e))176        data = marshal(dataset, dataset_detail_fields)177        if data.get("permission") == "partial_members":178            part_users_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str)179            data.update({"partial_member_list": part_users_list})180 181        # check embedding setting182        provider_manager = ProviderManager()183        configurations = provider_manager.get_configurations(tenant_id=current_user.current_tenant_id)184 185        embedding_models = configurations.get_models(model_type=ModelType.TEXT_EMBEDDING, only_active=True)186 187        model_names = []188        for embedding_model in embedding_models:189            model_names.append(f"{embedding_model.model}:{embedding_model.provider.provider}")190 191        if data["indexing_technique"] == "high_quality":192            item_model = f"{data['embedding_model']}:{data['embedding_model_provider']}"193            if item_model in model_names:194                data["embedding_available"] = True195            else:196                data["embedding_available"] = False197        else:198            data["embedding_available"] = True199 200        if data.get("permission") == "partial_members":201            part_users_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str)202            data.update({"partial_member_list": part_users_list})203 204        return data, 200205 206    @setup_required207    @login_required208    @account_initialization_required209    def patch(self, dataset_id):210        dataset_id_str = str(dataset_id)211        dataset = DatasetService.get_dataset(dataset_id_str)212        if dataset is None:213            raise NotFound("Dataset not found.")214 215        parser = reqparse.RequestParser()216        parser.add_argument(217            "name",218            nullable=False,219            help="type is required. Name must be between 1 to 40 characters.",220            type=_validate_name,221        )222        parser.add_argument("description", location="json", store_missing=False, type=_validate_description_length)223        parser.add_argument(224            "indexing_technique",225            type=str,226            location="json",227            choices=Dataset.INDEXING_TECHNIQUE_LIST,228            nullable=True,229            help="Invalid indexing technique.",230        )231        parser.add_argument(232            "permission",233            type=str,234            location="json",235            choices=(DatasetPermissionEnum.ONLY_ME, DatasetPermissionEnum.ALL_TEAM, DatasetPermissionEnum.PARTIAL_TEAM),236            help="Invalid permission.",237        )238        parser.add_argument("embedding_model", type=str, location="json", help="Invalid embedding model.")239        parser.add_argument(240            "embedding_model_provider", type=str, location="json", help="Invalid embedding model provider."241        )242        parser.add_argument("retrieval_model", type=dict, location="json", help="Invalid retrieval model.")243        parser.add_argument("partial_member_list", type=list, location="json", help="Invalid parent user list.")244 245        parser.add_argument(246            "external_retrieval_model",247            type=dict,248            required=False,249            nullable=True,250            location="json",251            help="Invalid external retrieval model.",252        )253 254        parser.add_argument(255            "external_knowledge_id",256            type=str,257            required=False,258            nullable=True,259            location="json",260            help="Invalid external knowledge id.",261        )262 263        parser.add_argument(264            "external_knowledge_api_id",265            type=str,266            required=False,267            nullable=True,268            location="json",269            help="Invalid external knowledge api id.",270        )271        args = parser.parse_args()272        data = request.get_json()273 274        # check embedding model setting275        if data.get("indexing_technique") == "high_quality":276            DatasetService.check_embedding_model_setting(277                dataset.tenant_id, data.get("embedding_model_provider"), data.get("embedding_model")278            )279 280        # The role of the current user in the ta table must be admin, owner, editor, or dataset_operator281        DatasetPermissionService.check_permission(282            current_user, dataset, data.get("permission"), data.get("partial_member_list")283        )284 285        dataset = DatasetService.update_dataset(dataset_id_str, args, current_user)286 287        if dataset is None:288            raise NotFound("Dataset not found.")289 290        result_data = marshal(dataset, dataset_detail_fields)291        tenant_id = current_user.current_tenant_id292 293        if data.get("partial_member_list") and data.get("permission") == "partial_members":294            DatasetPermissionService.update_partial_member_list(295                tenant_id, dataset_id_str, data.get("partial_member_list")296            )297        # clear partial member list when permission is only_me or all_team_members298        elif (299            data.get("permission") == DatasetPermissionEnum.ONLY_ME300            or data.get("permission") == DatasetPermissionEnum.ALL_TEAM301        ):302            DatasetPermissionService.clear_partial_member_list(dataset_id_str)303 304        partial_member_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str)305        result_data.update({"partial_member_list": partial_member_list})306 307        return result_data, 200308 309    @setup_required310    @login_required311    @account_initialization_required312    def delete(self, dataset_id):313        dataset_id_str = str(dataset_id)314 315        # The role of the current user in the ta table must be admin, owner, or editor316        if not current_user.is_editor or current_user.is_dataset_operator:317            raise Forbidden()318 319        try:320            if DatasetService.delete_dataset(dataset_id_str, current_user):321                DatasetPermissionService.clear_partial_member_list(dataset_id_str)322                return {"result": "success"}, 204323            else:324                raise NotFound("Dataset not found.")325        except services.errors.dataset.DatasetInUseError:326            raise DatasetInUseError()327 328 329class DatasetUseCheckApi(Resource):330    @setup_required331    @login_required332    @account_initialization_required333    def get(self, dataset_id):334        dataset_id_str = str(dataset_id)335 336        dataset_is_using = DatasetService.dataset_use_check(dataset_id_str)337        return {"is_using": dataset_is_using}, 200338 339 340class DatasetQueryApi(Resource):341    @setup_required342    @login_required343    @account_initialization_required344    def get(self, dataset_id):345        dataset_id_str = str(dataset_id)346        dataset = DatasetService.get_dataset(dataset_id_str)347        if dataset is None:348            raise NotFound("Dataset not found.")349 350        try:351            DatasetService.check_dataset_permission(dataset, current_user)352        except services.errors.account.NoPermissionError as e:353            raise Forbidden(str(e))354 355        page = request.args.get("page", default=1, type=int)356        limit = request.args.get("limit", default=20, type=int)357 358        dataset_queries, total = DatasetService.get_dataset_queries(dataset_id=dataset.id, page=page, per_page=limit)359 360        response = {361            "data": marshal(dataset_queries, dataset_query_detail_fields),362            "has_more": len(dataset_queries) == limit,363            "limit": limit,364            "total": total,365            "page": page,366        }367        return response, 200368 369 370class DatasetIndexingEstimateApi(Resource):371    @setup_required372    @login_required373    @account_initialization_required374    def post(self):375        parser = reqparse.RequestParser()376        parser.add_argument("info_list", type=dict, required=True, nullable=True, location="json")377        parser.add_argument("process_rule", type=dict, required=True, nullable=True, location="json")378        parser.add_argument(379            "indexing_technique",380            type=str,381            required=True,382            choices=Dataset.INDEXING_TECHNIQUE_LIST,383            nullable=True,384            location="json",385        )386        parser.add_argument("doc_form", type=str, default="text_model", required=False, nullable=False, location="json")387        parser.add_argument("dataset_id", type=str, required=False, nullable=False, location="json")388        parser.add_argument(389            "doc_language", type=str, default="English", required=False, nullable=False, location="json"390        )391        args = parser.parse_args()392        # validate args393        DocumentService.estimate_args_validate(args)394        extract_settings = []395        if args["info_list"]["data_source_type"] == "upload_file":396            file_ids = args["info_list"]["file_info_list"]["file_ids"]397            file_details = (398                db.session.query(UploadFile)399                .filter(UploadFile.tenant_id == current_user.current_tenant_id, UploadFile.id.in_(file_ids))400                .all()401            )402 403            if file_details is None:404                raise NotFound("File not found.")405 406            if file_details:407                for file_detail in file_details:408                    extract_setting = ExtractSetting(409                        datasource_type="upload_file", upload_file=file_detail, document_model=args["doc_form"]410                    )411                    extract_settings.append(extract_setting)412        elif args["info_list"]["data_source_type"] == "notion_import":413            notion_info_list = args["info_list"]["notion_info_list"]414            for notion_info in notion_info_list:415                workspace_id = notion_info["workspace_id"]416                for page in notion_info["pages"]:417                    extract_setting = ExtractSetting(418                        datasource_type="notion_import",419                        notion_info={420                            "notion_workspace_id": workspace_id,421                            "notion_obj_id": page["page_id"],422                            "notion_page_type": page["type"],423                            "tenant_id": current_user.current_tenant_id,424                        },425                        document_model=args["doc_form"],426                    )427                    extract_settings.append(extract_setting)428        elif args["info_list"]["data_source_type"] == "website_crawl":429            website_info_list = args["info_list"]["website_info_list"]430            for url in website_info_list["urls"]:431                extract_setting = ExtractSetting(432                    datasource_type="website_crawl",433                    website_info={434                        "provider": website_info_list["provider"],435                        "job_id": website_info_list["job_id"],436                        "url": url,437                        "tenant_id": current_user.current_tenant_id,438                        "mode": "crawl",439                        "only_main_content": website_info_list["only_main_content"],440                    },441                    document_model=args["doc_form"],442                )443                extract_settings.append(extract_setting)444        else:445            raise ValueError("Data source type not support")446        indexing_runner = IndexingRunner()447        try:448            response = indexing_runner.indexing_estimate(449                current_user.current_tenant_id,450                extract_settings,451                args["process_rule"],452                args["doc_form"],453                args["doc_language"],454                args["dataset_id"],455                args["indexing_technique"],456            )457        except LLMBadRequestError:458            raise ProviderNotInitializeError(459                "No Embedding Model available. Please configure a valid provider " "in the Settings -> Model Provider."460            )461        except ProviderTokenNotInitError as ex:462            raise ProviderNotInitializeError(ex.description)463        except Exception as e:464            raise IndexingEstimateError(str(e))465 466        return response, 200467 468 469class DatasetRelatedAppListApi(Resource):470    @setup_required471    @login_required472    @account_initialization_required473    @marshal_with(related_app_list)474    def get(self, dataset_id):475        dataset_id_str = str(dataset_id)476        dataset = DatasetService.get_dataset(dataset_id_str)477        if dataset is None:478            raise NotFound("Dataset not found.")479 480        try:481            DatasetService.check_dataset_permission(dataset, current_user)482        except services.errors.account.NoPermissionError as e:483            raise Forbidden(str(e))484 485        app_dataset_joins = DatasetService.get_related_apps(dataset.id)486 487        related_apps = []488        for app_dataset_join in app_dataset_joins:489            app_model = app_dataset_join.app490            if app_model:491                related_apps.append(app_model)492 493        return {"data": related_apps, "total": len(related_apps)}, 200494 495 496class DatasetIndexingStatusApi(Resource):497    @setup_required498    @login_required499    @account_initialization_required500    def get(self, dataset_id):501        dataset_id = str(dataset_id)502        documents = (503            db.session.query(Document)504            .filter(Document.dataset_id == dataset_id, Document.tenant_id == current_user.current_tenant_id)505            .all()506        )507        documents_status = []508        for document in documents:509            completed_segments = DocumentSegment.query.filter(510                DocumentSegment.completed_at.isnot(None),511                DocumentSegment.document_id == str(document.id),512                DocumentSegment.status != "re_segment",513            ).count()514            total_segments = DocumentSegment.query.filter(515                DocumentSegment.document_id == str(document.id), DocumentSegment.status != "re_segment"516            ).count()517            document.completed_segments = completed_segments518            document.total_segments = total_segments519            documents_status.append(marshal(document, document_status_fields))520        data = {"data": documents_status}521        return data522 523 524class DatasetApiKeyApi(Resource):525    max_keys = 10526    token_prefix = "dataset-"527    resource_type = "dataset"528 529    @setup_required530    @login_required531    @account_initialization_required532    @marshal_with(api_key_list)533    def get(self):534        keys = (535            db.session.query(ApiToken)536            .filter(ApiToken.type == self.resource_type, ApiToken.tenant_id == current_user.current_tenant_id)537            .all()538        )539        return {"items": keys}540 541    @setup_required542    @login_required543    @account_initialization_required544    @marshal_with(api_key_fields)545    def post(self):546        # The role of the current user in the ta table must be admin or owner547        if not current_user.is_admin_or_owner:548            raise Forbidden()549 550        current_key_count = (551            db.session.query(ApiToken)552            .filter(ApiToken.type == self.resource_type, ApiToken.tenant_id == current_user.current_tenant_id)553            .count()554        )555 556        if current_key_count >= self.max_keys:557            flask_restful.abort(558                400,559                message=f"Cannot create more than {self.max_keys} API keys for this resource type.",560                code="max_keys_exceeded",561            )562 563        key = ApiToken.generate_api_key(self.token_prefix, 24)564        api_token = ApiToken()565        api_token.tenant_id = current_user.current_tenant_id566        api_token.token = key567        api_token.type = self.resource_type568        db.session.add(api_token)569        db.session.commit()570        return api_token, 200571 572 573class DatasetApiDeleteApi(Resource):574    resource_type = "dataset"575 576    @setup_required577    @login_required578    @account_initialization_required579    def delete(self, api_key_id):580        api_key_id = str(api_key_id)581 582        # The role of the current user in the ta table must be admin or owner583        if not current_user.is_admin_or_owner:584            raise Forbidden()585 586        key = (587            db.session.query(ApiToken)588            .filter(589                ApiToken.tenant_id == current_user.current_tenant_id,590                ApiToken.type == self.resource_type,591                ApiToken.id == api_key_id,592            )593            .first()594        )595 596        if key is None:597            flask_restful.abort(404, message="API key not found")598 599        db.session.query(ApiToken).filter(ApiToken.id == api_key_id).delete()600        db.session.commit()601 602        return {"result": "success"}, 204603 604 605class DatasetApiBaseUrlApi(Resource):606    @setup_required607    @login_required608    @account_initialization_required609    def get(self):610        return {"api_base_url": (dify_config.SERVICE_API_URL or request.host_url.rstrip("/")) + "/v1"}611 612 613class DatasetRetrievalSettingApi(Resource):614    @setup_required615    @login_required616    @account_initialization_required617    def get(self):618        vector_type = dify_config.VECTOR_STORE619        match vector_type:620            case (621                VectorType.MILVUS622                | VectorType.RELYT623                | VectorType.PGVECTOR624                | VectorType.TIDB_VECTOR625                | VectorType.CHROMA626                | VectorType.TENCENT627                | VectorType.PGVECTO_RS628                | VectorType.BAIDU629                | VectorType.VIKINGDB630                | VectorType.UPSTASH631                | VectorType.OCEANBASE632            ):633                return {"retrieval_method": [RetrievalMethod.SEMANTIC_SEARCH.value]}634            case (635                VectorType.QDRANT636                | VectorType.WEAVIATE637                | VectorType.OPENSEARCH638                | VectorType.ANALYTICDB639                | VectorType.MYSCALE640                | VectorType.ORACLE641                | VectorType.ELASTICSEARCH642                | VectorType.PGVECTOR643                | VectorType.TIDB_ON_QDRANT644                | VectorType.LINDORM645                | VectorType.COUCHBASE646            ):647                return {648                    "retrieval_method": [649                        RetrievalMethod.SEMANTIC_SEARCH.value,650                        RetrievalMethod.FULL_TEXT_SEARCH.value,651                        RetrievalMethod.HYBRID_SEARCH.value,652                    ]653                }654            case _:655                raise ValueError(f"Unsupported vector db type {vector_type}.")656 657 658class DatasetRetrievalSettingMockApi(Resource):659    @setup_required660    @login_required661    @account_initialization_required662    def get(self, vector_type):663        match vector_type:664            case (665                VectorType.MILVUS666                | VectorType.RELYT667                | VectorType.TIDB_VECTOR668                | VectorType.CHROMA669                | VectorType.TENCENT670                | VectorType.PGVECTO_RS671                | VectorType.BAIDU672                | VectorType.VIKINGDB673                | VectorType.UPSTASH674                | VectorType.OCEANBASE675            ):676                return {"retrieval_method": [RetrievalMethod.SEMANTIC_SEARCH.value]}677            case (678                VectorType.QDRANT679                | VectorType.WEAVIATE680                | VectorType.OPENSEARCH681                | VectorType.ANALYTICDB682                | VectorType.MYSCALE683                | VectorType.ORACLE684                | VectorType.ELASTICSEARCH685                | VectorType.COUCHBASE686                | VectorType.PGVECTOR687                | VectorType.LINDORM688            ):689                return {690                    "retrieval_method": [691                        RetrievalMethod.SEMANTIC_SEARCH.value,692                        RetrievalMethod.FULL_TEXT_SEARCH.value,693                        RetrievalMethod.HYBRID_SEARCH.value,694                    ]695                }696            case _:697                raise ValueError(f"Unsupported vector db type {vector_type}.")698 699 700class DatasetErrorDocs(Resource):701    @setup_required702    @login_required703    @account_initialization_required704    def get(self, dataset_id):705        dataset_id_str = str(dataset_id)706        dataset = DatasetService.get_dataset(dataset_id_str)707        if dataset is None:708            raise NotFound("Dataset not found.")709        results = DocumentService.get_error_documents_by_dataset_id(dataset_id_str)710 711        return {"data": [marshal(item, document_status_fields) for item in results], "total": len(results)}, 200712 713 714class DatasetPermissionUserListApi(Resource):715    @setup_required716    @login_required717    @account_initialization_required718    def get(self, dataset_id):719        dataset_id_str = str(dataset_id)720        dataset = DatasetService.get_dataset(dataset_id_str)721        if dataset is None:722            raise NotFound("Dataset not found.")723        try:724            DatasetService.check_dataset_permission(dataset, current_user)725        except services.errors.account.NoPermissionError as e:726            raise Forbidden(str(e))727 728        partial_members_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str)729 730        return {731            "data": partial_members_list,732        }, 200733 734 735api.add_resource(DatasetListApi, "/datasets")736api.add_resource(DatasetApi, "/datasets/<uuid:dataset_id>")737api.add_resource(DatasetUseCheckApi, "/datasets/<uuid:dataset_id>/use-check")738api.add_resource(DatasetQueryApi, "/datasets/<uuid:dataset_id>/queries")739api.add_resource(DatasetErrorDocs, "/datasets/<uuid:dataset_id>/error-docs")740api.add_resource(DatasetIndexingEstimateApi, "/datasets/indexing-estimate")741api.add_resource(DatasetRelatedAppListApi, "/datasets/<uuid:dataset_id>/related-apps")742api.add_resource(DatasetIndexingStatusApi, "/datasets/<uuid:dataset_id>/indexing-status")743api.add_resource(DatasetApiKeyApi, "/datasets/api-keys")744api.add_resource(DatasetApiDeleteApi, "/datasets/api-keys/<uuid:api_key_id>")745api.add_resource(DatasetApiBaseUrlApi, "/datasets/api-base-info")746api.add_resource(DatasetRetrievalSettingApi, "/datasets/retrieval-setting")747api.add_resource(DatasetRetrievalSettingMockApi, "/datasets/retrieval-setting/<string:vector_type>")748api.add_resource(DatasetPermissionUserListApi, "/datasets/<uuid:dataset_id>/permission-part-users")749