dify/api/controllers/console/datasets/metadata.py

from typing import Literal

from flask_restx import Resource
from werkzeug.exceptions import NotFound

from controllers.common.controller_schemas import MetadataUpdatePayload
from controllers.common.schema import register_response_schema_models, register_schema_models
from controllers.console import console_ns
from controllers.console.wraps import account_initialization_required, enterprise_license_required, setup_required
from fields.dataset_fields import (
    DatasetMetadataBuiltInFieldsResponse,
    DatasetMetadataListResponse,
    DatasetMetadataResponse,
)
from libs.helper import dump_response
from libs.login import current_account_with_tenant, login_required
from services.dataset_service import DatasetService
from services.entities.knowledge_entities.knowledge_entities import (
    DocumentMetadataOperation,
    MetadataArgs,
    MetadataDetail,
    MetadataOperationData,
)
from services.metadata_service import MetadataService

register_schema_models(
    console_ns, MetadataArgs, MetadataOperationData, MetadataUpdatePayload, DocumentMetadataOperation, MetadataDetail
)
register_response_schema_models(
    console_ns,
    DatasetMetadataBuiltInFieldsResponse,
    DatasetMetadataListResponse,
    DatasetMetadataResponse,
)


@console_ns.route("/datasets/<uuid:dataset_id>/metadata")
class DatasetMetadataCreateApi(Resource):
    @setup_required
    @login_required
    @account_initialization_required
    @enterprise_license_required
    @console_ns.response(201, "Metadata created successfully", console_ns.models[DatasetMetadataResponse.__name__])
    @console_ns.expect(console_ns.models[MetadataArgs.__name__])
    def post(self, dataset_id):
        current_user, _ = current_account_with_tenant()
        metadata_args = MetadataArgs.model_validate(console_ns.payload or {})

        dataset_id_str = str(dataset_id)
        dataset = DatasetService.get_dataset(dataset_id_str)
        if dataset is None:
            raise NotFound("Dataset not found.")
        DatasetService.check_dataset_permission(dataset, current_user)

        metadata = MetadataService.create_metadata(dataset_id_str, metadata_args)
        return dump_response(DatasetMetadataResponse, metadata), 201

    @setup_required
    @login_required
    @account_initialization_required
    @enterprise_license_required
    @console_ns.response(
        200, "Metadata retrieved successfully", console_ns.models[DatasetMetadataListResponse.__name__]
    )
    def get(self, dataset_id):
        dataset_id_str = str(dataset_id)
        dataset = DatasetService.get_dataset(dataset_id_str)
        if dataset is None:
            raise NotFound("Dataset not found.")
        metadata = MetadataService.get_dataset_metadatas(dataset)
        return dump_response(DatasetMetadataListResponse, metadata), 200


@console_ns.route("/datasets/<uuid:dataset_id>/metadata/<uuid:metadata_id>")
class DatasetMetadataApi(Resource):
    @setup_required
    @login_required
    @account_initialization_required
    @enterprise_license_required
    @console_ns.response(200, "Metadata updated successfully", console_ns.models[DatasetMetadataResponse.__name__])
    @console_ns.expect(console_ns.models[MetadataUpdatePayload.__name__])
    def patch(self, dataset_id, metadata_id):
        current_user, _ = current_account_with_tenant()
        payload = MetadataUpdatePayload.model_validate(console_ns.payload or {})
        name = payload.name

        dataset_id_str = str(dataset_id)
        metadata_id_str = str(metadata_id)
        dataset = DatasetService.get_dataset(dataset_id_str)
        if dataset is None:
            raise NotFound("Dataset not found.")
        DatasetService.check_dataset_permission(dataset, current_user)

        metadata = MetadataService.update_metadata_name(dataset_id_str, metadata_id_str, name)
        return dump_response(DatasetMetadataResponse, metadata), 200

    @setup_required
    @login_required
    @account_initialization_required
    @enterprise_license_required
    @console_ns.response(204, "Metadata deleted successfully")
    def delete(self, dataset_id, metadata_id):
        current_user, _ = current_account_with_tenant()
        dataset_id_str = str(dataset_id)
        metadata_id_str = str(metadata_id)
        dataset = DatasetService.get_dataset(dataset_id_str)
        if dataset is None:
            raise NotFound("Dataset not found.")
        DatasetService.check_dataset_permission(dataset, current_user)

        MetadataService.delete_metadata(dataset_id_str, metadata_id_str)
        # Frontend callers only await success and invalidate metadata caches; no response body is consumed.
        return "", 204


@console_ns.route("/datasets/metadata/built-in")
class DatasetMetadataBuiltInFieldApi(Resource):
    @setup_required
    @login_required
    @account_initialization_required
    @enterprise_license_required
    @console_ns.response(
        200,
        "Built-in fields retrieved successfully",
        console_ns.models[DatasetMetadataBuiltInFieldsResponse.__name__],
    )
    def get(self):
        built_in_fields = MetadataService.get_built_in_fields()
        return dump_response(DatasetMetadataBuiltInFieldsResponse, {"fields": built_in_fields}), 200


@console_ns.route("/datasets/<uuid:dataset_id>/metadata/built-in/<string:action>")
class DatasetMetadataBuiltInFieldActionApi(Resource):
    @setup_required
    @login_required
    @account_initialization_required
    @enterprise_license_required
    @console_ns.response(204, "Action completed successfully")
    def post(self, dataset_id, action: Literal["enable", "disable"]):
        current_user, _ = current_account_with_tenant()
        dataset_id_str = str(dataset_id)
        dataset = DatasetService.get_dataset(dataset_id_str)
        if dataset is None:
            raise NotFound("Dataset not found.")
        DatasetService.check_dataset_permission(dataset, current_user)

        match action:
            case "enable":
                MetadataService.enable_built_in_field(dataset)
            case "disable":
                MetadataService.disable_built_in_field(dataset)
        # Frontend callers only await success and invalidate metadata caches; no response body is consumed.
        return "", 204


@console_ns.route("/datasets/<uuid:dataset_id>/documents/metadata")
class DocumentMetadataEditApi(Resource):
    @setup_required
    @login_required
    @account_initialization_required
    @enterprise_license_required
    @console_ns.expect(console_ns.models[MetadataOperationData.__name__])
    @console_ns.response(
        204,
        "Documents metadata updated successfully",
    )
    def post(self, dataset_id):
        current_user, _ = current_account_with_tenant()
        dataset_id_str = str(dataset_id)
        dataset = DatasetService.get_dataset(dataset_id_str)
        if dataset is None:
            raise NotFound("Dataset not found.")
        DatasetService.check_dataset_permission(dataset, current_user)

        metadata_args = MetadataOperationData.model_validate(console_ns.payload or {})

        MetadataService.update_documents_metadata(dataset, metadata_args)

        # Frontend callers only await success and invalidate caches; no response body is consumed.
        return "", 204