diff --git a/api/controllers/console/datasets/data_source.py b/api/controllers/console/datasets/data_source.py index 590fbdbb87d..9cf8420c251 100644 --- a/api/controllers/console/datasets/data_source.py +++ b/api/controllers/console/datasets/data_source.py @@ -36,6 +36,7 @@ from ..wraps import ( RBACPermission, RBACResourceScope, account_initialization_required, + model_validate, rbac_permission_required, setup_required, with_current_tenant_id, @@ -229,12 +230,18 @@ class DataSourceNotionListApi(Resource): @with_current_user @with_current_tenant_id @with_session(write=False) - def get(self, session: Session, current_tenant_id: str, current_user: Account) -> tuple[dict[str, Any], int]: - query = DataSourceNotionListQuery.model_validate(request.args.to_dict(flat=True)) + @model_validate(DataSourceNotionListQuery) + def get( + self, + req_data: DataSourceNotionListQuery, + session: Session, + current_tenant_id: str, + current_user: Account, + ) -> tuple[dict[str, Any], int]: datasource_provider_service = DatasourceProviderService() credential = datasource_provider_service.get_datasource_credentials( tenant_id=current_tenant_id, - credential_id=query.credential_id, + credential_id=req_data.credential_id, provider="notion_datasource", plugin_id="langgenius/notion_datasource", ) @@ -242,8 +249,8 @@ class DataSourceNotionListApi(Resource): raise NotFound("Credential not found.") exist_page_ids = [] # import notion in the exist dataset - if query.dataset_id: - dataset = DatasetService.get_dataset(query.dataset_id, session) + if req_data.dataset_id: + dataset = DatasetService.get_dataset(req_data.dataset_id, session) if not dataset: raise NotFound("Dataset not found.") if dataset.data_source_type != "notion_import": @@ -251,7 +258,7 @@ class DataSourceNotionListApi(Resource): documents = session.scalars( select(Document).where( - Document.dataset_id == query.dataset_id, + Document.dataset_id == req_data.dataset_id, Document.tenant_id == current_tenant_id, Document.data_source_type == "notion_import", Document.enabled.is_(True), @@ -318,13 +325,19 @@ class DataSourceNotionPreviewApi(Resource): @console_ns.doc(params=query_params_from_model(DataSourceNotionPreviewQuery)) @console_ns.response(200, "Success", console_ns.models[TextContentResponse.__name__]) @with_current_tenant_id - def get(self, current_tenant_id: str, page_id: UUID, page_type: str) -> tuple[dict[str, str], int]: - query = DataSourceNotionPreviewQuery.model_validate(request.args.to_dict(flat=True)) + @model_validate(DataSourceNotionPreviewQuery) + def get( + self, + req_data: DataSourceNotionPreviewQuery, + current_tenant_id: str, + page_id: UUID, + page_type: str, + ) -> tuple[dict[str, str], int]: datasource_provider_service = DatasourceProviderService() credential = datasource_provider_service.get_datasource_credentials( tenant_id=current_tenant_id, - credential_id=query.credential_id, + credential_id=req_data.credential_id, provider="notion_datasource", plugin_id="langgenius/notion_datasource", ) @@ -354,12 +367,17 @@ class DataSourceNotionIndexingEstimateApi(Resource): @console_ns.response(200, "Success", console_ns.models[IndexingEstimate.__name__]) @with_current_tenant_id @with_session - def post(self, session: Session, current_tenant_id: str) -> tuple[dict[str, Any], int]: - payload = NotionEstimatePayload.model_validate(console_ns.payload or {}) - args = payload.model_dump() + @model_validate(NotionEstimatePayload) + def post( + self, + req_data: NotionEstimatePayload, + session: Session, + current_tenant_id: str, + ) -> tuple[dict[str, Any], int]: + args = req_data.model_dump() # validate args DocumentService.estimate_args_validate(args) - notion_info_list = payload.notion_info_list + notion_info_list = req_data.notion_info_list extract_settings = [] for notion_info in notion_info_list: workspace_id = notion_info["workspace_id"] diff --git a/api/controllers/console/datasets/datasets.py b/api/controllers/console/datasets/datasets.py index 9d444978c01..ae3f74ce63d 100644 --- a/api/controllers/console/datasets/datasets.py +++ b/api/controllers/console/datasets/datasets.py @@ -26,6 +26,7 @@ from controllers.console.wraps import ( cloud_edition_billing_rate_limit_check, enterprise_license_required, is_admin_or_owner_required, + model_validate, rbac_permission_required, setup_required, with_current_tenant_id, @@ -572,9 +573,8 @@ class DatasetListApi(Resource): @with_current_user @with_current_tenant_id @with_session - def post(self, session: Session, current_tenant_id: str, current_user: Account): - payload = DatasetCreatePayload.model_validate(console_ns.payload or {}) - + @model_validate(DatasetCreatePayload) + def post(self, req_data: DatasetCreatePayload, session: Session, current_tenant_id: str, current_user: Account): # The role of the current user in the ta table must be admin, owner, or editor, or dataset_operator if not current_user.is_dataset_editor: raise Forbidden() @@ -582,20 +582,20 @@ class DatasetListApi(Resource): if dify_config.RBAC_ENABLED: permission = DatasetPermissionEnum.ALL_TEAM else: - permission = payload.permission or DatasetPermissionEnum.ONLY_ME + permission = req_data.permission or DatasetPermissionEnum.ONLY_ME try: dataset = DatasetService.create_empty_dataset( session=session, tenant_id=current_tenant_id, - name=payload.name, - description=payload.description, - indexing_technique=payload.indexing_technique, + name=req_data.name, + description=req_data.description, + indexing_technique=req_data.indexing_technique, account=current_user, permission=permission, - provider=payload.provider, - external_knowledge_api_id=payload.external_knowledge_api_id, - external_knowledge_id=payload.external_knowledge_id, + provider=req_data.provider, + external_knowledge_api_id=req_data.external_knowledge_api_id, + external_knowledge_id=req_data.external_knowledge_id, ) except services.errors.dataset.DatasetNameDuplicateError: raise DatasetNameDuplicateError() @@ -715,28 +715,35 @@ class DatasetApi(Resource): @with_current_tenant_id @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT) @with_session - def patch(self, session: Session, current_tenant_id: str, current_user: Account, dataset_id: UUID): + @model_validate(DatasetUpdatePayload) + def patch( + self, + req_data: DatasetUpdatePayload, + session: Session, + current_tenant_id: str, + current_user: Account, + dataset_id: UUID, + ): dataset_id_str = str(dataset_id) dataset = DatasetService.get_dataset(dataset_id_str, session) if dataset is None: raise NotFound("Dataset not found.") - payload = DatasetUpdatePayload.model_validate(console_ns.payload or {}) # check embedding model setting if ( - payload.indexing_technique == IndexTechniqueType.HIGH_QUALITY - and payload.embedding_model_provider is not None - and payload.embedding_model is not None + req_data.indexing_technique == IndexTechniqueType.HIGH_QUALITY + and req_data.embedding_model_provider is not None + and req_data.embedding_model is not None ): is_multimodal = DatasetService.check_is_multimodal_model( - dataset.tenant_id, payload.embedding_model_provider, payload.embedding_model + dataset.tenant_id, req_data.embedding_model_provider, req_data.embedding_model ) - payload.is_multimodal = is_multimodal - payload_data = payload.model_dump(exclude_unset=True) + req_data.is_multimodal = is_multimodal + payload_data = req_data.model_dump(exclude_unset=True) # The role of the current user in the ta table must be admin, owner, editor, or dataset_operator if not dify_config.RBAC_ENABLED: DatasetPermissionService.check_permission( - current_user, dataset, payload.permission, payload.partial_member_list, session=session + current_user, dataset, req_data.permission, req_data.partial_member_list, session=session ) dataset = DatasetService.update_dataset(dataset_id_str, payload_data, current_user, session=session) @@ -754,12 +761,12 @@ class DatasetApi(Resource): result_data["permission_keys"] = permission_keys_map.get(dataset_id_str, []) tenant_id = current_tenant_id - if payload.partial_member_list is not None and payload.permission == DatasetPermissionEnum.PARTIAL_TEAM: + if req_data.partial_member_list is not None and req_data.permission == DatasetPermissionEnum.PARTIAL_TEAM: DatasetPermissionService.update_partial_member_list( - tenant_id, dataset_id_str, payload.partial_member_list, session + tenant_id, dataset_id_str, req_data.partial_member_list, session ) # clear partial member list when permission is only_me or all_team_members - elif payload.permission in {DatasetPermissionEnum.ONLY_ME, DatasetPermissionEnum.ALL_TEAM}: + elif req_data.permission in {DatasetPermissionEnum.ONLY_ME, DatasetPermissionEnum.ALL_TEAM}: DatasetPermissionService.clear_partial_member_list(dataset_id_str, session) partial_member_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str, session) @@ -872,9 +879,9 @@ class DatasetIndexingEstimateApi(Resource): @console_ns.expect(console_ns.models[IndexingEstimatePayload.__name__]) @with_current_tenant_id @with_session - def post(self, session: Session, current_tenant_id: str): - payload = IndexingEstimatePayload.model_validate(console_ns.payload or {}) - args = payload.model_dump() + @model_validate(IndexingEstimatePayload) + def post(self, req_data: IndexingEstimatePayload, session: Session, current_tenant_id: str): + args = req_data.model_dump() # validate args DocumentService.estimate_args_validate(args) extract_settings = [] diff --git a/api/controllers/console/datasets/datasets_segments.py b/api/controllers/console/datasets/datasets_segments.py index 33a5a1f752a..d1fa14902e6 100644 --- a/api/controllers/console/datasets/datasets_segments.py +++ b/api/controllers/console/datasets/datasets_segments.py @@ -36,6 +36,7 @@ from controllers.console.wraps import ( cloud_edition_billing_knowledge_limit_check, cloud_edition_billing_rate_limit_check, cloud_edition_billing_resource_check, + model_validate, rbac_permission_required, setup_required, with_current_tenant_id, @@ -414,8 +415,10 @@ class DatasetDocumentSegmentAddApi(Resource): @with_current_tenant_id @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT) @with_session + @model_validate(SegmentCreatePayload) def post( self, + req_data: SegmentCreatePayload, session: Session, current_tenant_id: str, current_user: Account, @@ -455,8 +458,7 @@ class DatasetDocumentSegmentAddApi(Resource): except services.errors.account.NoPermissionError as e: raise Forbidden(str(e)) # validate args - payload = SegmentCreatePayload.model_validate(console_ns.payload or {}) - payload_dict = payload.model_dump(exclude_none=True) + payload_dict = req_data.model_dump(exclude_none=True) SegmentService.segment_create_args_validate(payload_dict, document) segment = type_cast(DocumentSegment, SegmentService.create_segment(payload_dict, document, dataset, session)) summary = SummaryIndexService.get_segment_summary( @@ -485,8 +487,10 @@ class DatasetDocumentSegmentUpdateApi(Resource): @with_current_tenant_id @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT) @with_session + @model_validate(SegmentUpdatePayload) def patch( self, + req_data: SegmentUpdatePayload, session: Session, current_tenant_id: str, current_user: Account, @@ -532,13 +536,12 @@ class DatasetDocumentSegmentUpdateApi(Resource): segment_id_str = str(segment_id) _, segment = _get_segment_for_document(session, dataset, document, segment_id_str) # validate args - payload = SegmentUpdatePayload.model_validate(console_ns.payload or {}) - payload_dict = payload.model_dump(exclude_none=True) + payload_dict = req_data.model_dump(exclude_none=True) SegmentService.segment_create_args_validate(payload_dict, document) # Update segment (summary update with change detection is handled in SegmentService.update_segment) segment = SegmentService.update_segment( - SegmentUpdateArgs.model_validate(payload.model_dump(exclude_none=True)), + SegmentUpdateArgs.model_validate(req_data.model_dump(exclude_none=True)), segment, document, dataset, @@ -616,8 +619,10 @@ class DatasetDocumentSegmentBatchImportApi(Resource): @with_current_tenant_id @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT) @with_session + @model_validate(BatchImportPayload) def post( self, + req_data: BatchImportPayload, session: Session, current_tenant_id: str, current_user: Account, @@ -635,8 +640,7 @@ class DatasetDocumentSegmentBatchImportApi(Resource): if not document: raise NotFound("Document not found.") - payload = BatchImportPayload.model_validate(console_ns.payload or {}) - upload_file_id = payload.upload_file_id + upload_file_id = req_data.upload_file_id upload_file = session.scalar(select(UploadFile).where(UploadFile.id == upload_file_id).limit(1)) if not upload_file: @@ -697,8 +701,10 @@ class ChildChunkAddApi(Resource): @with_current_tenant_id @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT) @with_session + @model_validate(ChildChunkCreatePayload) def post( self, + req_data: ChildChunkCreatePayload, session: Session, current_tenant_id: str, current_user: Account, @@ -742,8 +748,7 @@ class ChildChunkAddApi(Resource): _, segment = _get_segment_for_document(session, dataset, document, segment_id_str) # validate args try: - payload = ChildChunkCreatePayload.model_validate(console_ns.payload or {}) - child_chunk = SegmentService.create_child_chunk(payload.content, segment, document, dataset, session) + child_chunk = SegmentService.create_child_chunk(req_data.content, segment, document, dataset, session) except ChildChunkIndexingServiceError as e: raise ChildChunkIndexingError(str(e)) return dump_response(ChildChunkDetailResponse, {"data": child_chunk}), 200 @@ -812,8 +817,10 @@ class ChildChunkAddApi(Resource): @with_current_tenant_id @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT) @with_session + @model_validate(ChildChunkBatchUpdatePayload) def patch( self, + req_data: ChildChunkBatchUpdatePayload, session: Session, current_tenant_id: str, current_user: Account, @@ -843,9 +850,8 @@ class ChildChunkAddApi(Resource): segment_id_str = str(segment_id) _, segment = _get_segment_for_document(session, dataset, document, segment_id_str) # validate args - payload = ChildChunkBatchUpdatePayload.model_validate(console_ns.payload or {}) try: - child_chunks = SegmentService.update_child_chunks(payload.chunks, segment, document, dataset, session) + child_chunks = SegmentService.update_child_chunks(req_data.chunks, segment, document, dataset, session) except ChildChunkIndexingServiceError as e: raise ChildChunkIndexingError(str(e)) return dump_response(ChildChunkBatchUpdateResponse, {"data": child_chunks}), 200 @@ -918,8 +924,10 @@ class ChildChunkUpdateApi(Resource): @with_current_tenant_id @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT) @with_session + @model_validate(ChildChunkUpdatePayload) def patch( self, + req_data: ChildChunkUpdatePayload, session: Session, current_tenant_id: str, current_user: Account, @@ -955,9 +963,8 @@ class ChildChunkUpdateApi(Resource): raise NotFound("Child chunk not found.") # validate args try: - payload = ChildChunkUpdatePayload.model_validate(console_ns.payload or {}) child_chunk = SegmentService.update_child_chunk( - payload.content, child_chunk, segment, document, dataset, session + req_data.content, child_chunk, segment, document, dataset, session ) except ChildChunkIndexingServiceError as e: raise ChildChunkIndexingError(str(e)) diff --git a/api/controllers/console/datasets/external.py b/api/controllers/console/datasets/external.py index 34f6aecfcbd..26d958e15d3 100644 --- a/api/controllers/console/datasets/external.py +++ b/api/controllers/console/datasets/external.py @@ -5,7 +5,6 @@ from datetime import datetime from typing import Any from uuid import UUID -from flask import request from flask_restx import Resource from pydantic import AliasChoices, BaseModel, Field, field_validator from sqlalchemy.orm import Session @@ -26,6 +25,7 @@ from controllers.console.wraps import ( RBACResourceScope, account_initialization_required, edit_permission_required, + model_validate, rbac_permission_required, setup_required, with_current_tenant_id, @@ -186,18 +186,18 @@ class ExternalApiTemplateListApi(Resource): @with_current_tenant_id @account_initialization_required @with_session(write=False) - def get(self, session: Session, current_tenant_id: str): - query = ExternalApiTemplateListQuery.model_validate(request.args.to_dict()) + @model_validate(ExternalApiTemplateListQuery) + def get(self, req_data: ExternalApiTemplateListQuery, session: Session, current_tenant_id: str): external_knowledge_apis, total = ExternalDatasetService.get_external_knowledge_apis( - query.page, query.limit, current_tenant_id, query.keyword, session=session + req_data.page, req_data.limit, current_tenant_id, req_data.keyword, session=session ) return ExternalKnowledgeApiListResponse( data=[external_knowledge_api_response(item, session=session) for item in external_knowledge_apis], - has_more=len(external_knowledge_apis) == query.limit, - limit=query.limit, + has_more=len(external_knowledge_apis) == req_data.limit, + limit=req_data.limit, total=total, - page=query.page, + page=req_data.page, ).model_dump(mode="json"), 200 @console_ns.doc("create_external_api_template") @@ -215,10 +215,16 @@ class ExternalApiTemplateListApi(Resource): @with_current_user @with_current_tenant_id @with_session - def post(self, session: Session, current_tenant_id: str, current_user: Account): - payload = ExternalKnowledgeApiPayload.model_validate(console_ns.payload or {}) + @model_validate(ExternalKnowledgeApiPayload) + def post( + self, + req_data: ExternalKnowledgeApiPayload, + session: Session, + current_tenant_id: str, + current_user: Account, + ): - ExternalDatasetService.validate_api_list(payload.settings) + ExternalDatasetService.validate_api_list(req_data.settings) # The role of the current user in the ta table must be admin, owner, or editor, or dataset_operator if not current_user.is_dataset_editor: @@ -228,7 +234,7 @@ class ExternalApiTemplateListApi(Resource): external_knowledge_api = ExternalDatasetService.create_external_knowledge_api( tenant_id=current_tenant_id, user_id=current_user.id, - args=payload.model_dump(), + args=req_data.model_dump(), session=session, ) except services.errors.dataset.DatasetNameDuplicateError: @@ -279,17 +285,24 @@ class ExternalApiTemplateApi(Resource): @with_current_user @with_current_tenant_id @with_session - def patch(self, session: Session, current_tenant_id: str, current_user: Account, external_knowledge_api_id: UUID): + @model_validate(ExternalKnowledgeApiPayload) + def patch( + self, + req_data: ExternalKnowledgeApiPayload, + session: Session, + current_tenant_id: str, + current_user: Account, + external_knowledge_api_id: UUID, + ): external_knowledge_api_id_str = str(external_knowledge_api_id) - payload = ExternalKnowledgeApiPayload.model_validate(console_ns.payload or {}) - ExternalDatasetService.validate_api_list(payload.settings) + ExternalDatasetService.validate_api_list(req_data.settings) external_knowledge_api = ExternalDatasetService.update_external_knowledge_api( tenant_id=current_tenant_id, user_id=current_user.id, external_knowledge_api_id=external_knowledge_api_id_str, - args=payload.model_dump(), + args=req_data.model_dump(), session=session, ) @@ -354,9 +367,15 @@ class ExternalDatasetCreateApi(Resource): @with_current_user @with_current_tenant_id @with_session - def post(self, session: Session, current_tenant_id: str, current_user: Account): + @model_validate(ExternalDatasetCreatePayload) + def post( + self, + req_data: ExternalDatasetCreatePayload, + session: Session, + current_tenant_id: str, + current_user: Account, + ): # The role of the current user in the ta table must be admin, owner, or editor - payload = ExternalDatasetCreatePayload.model_validate(console_ns.payload or {}) # The role of the current user in the ta table must be admin, owner, or editor, or dataset_operator if not current_user.is_dataset_editor: @@ -366,7 +385,7 @@ class ExternalDatasetCreateApi(Resource): dataset = ExternalDatasetService.create_external_dataset( tenant_id=current_tenant_id, user_id=current_user.id, - args=payload, + args=req_data, session=session, ) except services.errors.dataset.DatasetNameDuplicateError: @@ -405,7 +424,8 @@ class ExternalKnowledgeHitTestingApi(Resource): @with_current_user @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_PIPELINE_TEST) @with_session - def post(self, session: Session, current_user: Account, dataset_id: UUID): + @model_validate(ExternalHitTestingPayload) + def post(self, req_data: ExternalHitTestingPayload, session: Session, current_user: Account, dataset_id: UUID): dataset_id_str = str(dataset_id) dataset = DatasetService.get_dataset(dataset_id_str, session) if dataset is None: @@ -416,17 +436,16 @@ class ExternalKnowledgeHitTestingApi(Resource): except services.errors.account.NoPermissionError as e: raise Forbidden(str(e)) - payload = ExternalHitTestingPayload.model_validate(console_ns.payload or {}) - HitTestingService.hit_testing_args_check(payload.model_dump()) + HitTestingService.hit_testing_args_check(req_data.model_dump()) try: response = HitTestingService.external_retrieve( session=session, dataset=dataset, - query=payload.query, + query=req_data.query, account=current_user, - external_retrieval_model=payload.external_retrieval_model, - metadata_filtering_conditions=payload.metadata_filtering_conditions, + external_retrieval_model=req_data.external_retrieval_model, + metadata_filtering_conditions=req_data.metadata_filtering_conditions, ) return dump_response(ExternalHitTestingResponse, response) @@ -441,11 +460,11 @@ class BedrockRetrievalApi(Resource): @console_ns.doc(description="Bedrock retrieval test (internal use only)") @console_ns.expect(console_ns.models[BedrockRetrievalPayload.__name__]) @console_ns.response(200, "Bedrock retrieval test completed", console_ns.models[BedrockRetrievalResponse.__name__]) - def post(self): - payload = BedrockRetrievalPayload.model_validate(console_ns.payload or {}) + @model_validate(BedrockRetrievalPayload) + def post(self, req_data: BedrockRetrievalPayload): # Call the knowledge retrieval service result = ExternalDatasetTestService.knowledge_retrieval( - payload.retrieval_setting, payload.query, payload.knowledge_id + req_data.retrieval_setting, req_data.query, req_data.knowledge_id ) return dump_response(BedrockRetrievalResponse, result), 200 diff --git a/api/controllers/console/datasets/metadata.py b/api/controllers/console/datasets/metadata.py index 32c4151a017..75aae6d50e0 100644 --- a/api/controllers/console/datasets/metadata.py +++ b/api/controllers/console/datasets/metadata.py @@ -14,6 +14,7 @@ from controllers.console.wraps import ( RBACResourceScope, account_initialization_required, enterprise_license_required, + model_validate, rbac_permission_required, setup_required, with_current_tenant_id, @@ -59,9 +60,15 @@ class DatasetMetadataCreateApi(Resource): @with_current_tenant_id @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT) @with_session - def post(self, session: Session, current_tenant_id: str, current_user: Account, dataset_id: UUID): - metadata_args = MetadataArgs.model_validate(console_ns.payload or {}) - + @model_validate(MetadataArgs) + def post( + self, + req_data: MetadataArgs, + session: Session, + current_tenant_id: str, + current_user: Account, + dataset_id: UUID, + ): dataset_id_str = str(dataset_id) dataset = DatasetService.get_dataset(dataset_id_str, session) if dataset is None: @@ -69,7 +76,7 @@ class DatasetMetadataCreateApi(Resource): DatasetService.check_dataset_permission(dataset, current_user, session) metadata = MetadataService.create_metadata( - dataset_id_str, metadata_args, current_user, current_tenant_id, session=session + dataset_id_str, req_data, current_user, current_tenant_id, session=session ) return dump_response(DatasetMetadataResponse, metadata), 201 @@ -103,16 +110,17 @@ class DatasetMetadataApi(Resource): @with_current_tenant_id @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT) @with_session + @model_validate(MetadataUpdatePayload) def patch( self, + req_data: MetadataUpdatePayload, session: Session, current_tenant_id: str, current_user: Account, dataset_id: UUID, metadata_id: UUID, ): - payload = MetadataUpdatePayload.model_validate(console_ns.payload or {}) - name = payload.name + name = req_data.name dataset_id_str = str(dataset_id) metadata_id_str = str(metadata_id) @@ -203,16 +211,15 @@ class DocumentMetadataEditApi(Resource): @with_current_user @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT) @with_session - def post(self, session: Session, current_user: Account, dataset_id: UUID): + @model_validate(MetadataOperationData) + def post(self, req_data: MetadataOperationData, session: Session, current_user: Account, dataset_id: UUID): dataset_id_str = str(dataset_id) dataset = DatasetService.get_dataset(dataset_id_str, session) if dataset is None: raise NotFound("Dataset not found.") DatasetService.check_dataset_permission(dataset, current_user, session) - metadata_args = MetadataOperationData.model_validate(console_ns.payload or {}) - - MetadataService.update_documents_metadata(dataset, metadata_args, current_user, session=session) + MetadataService.update_documents_metadata(dataset, req_data, current_user, session=session) # Frontend callers only await success and invalidate caches; no response body is consumed. return "", 204 diff --git a/api/tests/test_containers_integration_tests/controllers/console/datasets/test_data_source.py b/api/tests/test_containers_integration_tests/controllers/console/datasets/test_data_source.py index 814886b1772..bc0a205ecb1 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/datasets/test_data_source.py +++ b/api/tests/test_containers_integration_tests/controllers/console/datasets/test_data_source.py @@ -7,7 +7,7 @@ from uuid import uuid4 from flask import Flask from sqlalchemy.orm import Session -from controllers.console.datasets.data_source import DataSourceNotionListApi +from controllers.console.datasets.data_source import DataSourceNotionListApi, DataSourceNotionListQuery from models import Account from models.dataset import Document from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus @@ -78,7 +78,11 @@ def test_notion_page_is_marked_bound_from_persisted_document( ), ): response, status = unwrap(DataSourceNotionListApi().get)( - DataSourceNotionListApi(), db_session_with_containers, tenant_id, account + DataSourceNotionListApi(), + DataSourceNotionListQuery(credential_id="c1", dataset_id=dataset_id), + db_session_with_containers, + tenant_id, + account, ) assert status == 200 diff --git a/api/tests/unit_tests/controllers/console/datasets/test_data_source.py b/api/tests/unit_tests/controllers/console/datasets/test_data_source.py index 55ee174355c..c5339e166c2 100644 --- a/api/tests/unit_tests/controllers/console/datasets/test_data_source.py +++ b/api/tests/unit_tests/controllers/console/datasets/test_data_source.py @@ -14,7 +14,7 @@ from sqlalchemy.orm import Session from werkzeug.exceptions import NotFound from controllers.console.datasets import data_source as module -from controllers.console.datasets.data_source import DataSourceApi, DataSourceNotionListApi +from controllers.console.datasets.data_source import DataSourceApi, DataSourceNotionListApi, DataSourceNotionListQuery from models import Account, DataSourceOauthBinding from models.engine import db @@ -217,7 +217,11 @@ def test_notion_pre_import_pages_serializes_frontend_list_shape( patch("core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", return_value=runtime), ): response, status = unwrap(DataSourceNotionListApi().get)( - DataSourceNotionListApi(), sqlite_session, "tenant-1", current_user + DataSourceNotionListApi(), + DataSourceNotionListQuery(credential_id="credential-1"), + sqlite_session, + "tenant-1", + current_user, ) assert status == 200 @@ -255,7 +259,13 @@ def test_notion_pre_import_pages_rejects_missing_credential( patch.object(module.DatasourceProviderService, "get_datasource_credentials", return_value=None), pytest.raises(NotFound, match="Credential not found"), ): - unwrap(DataSourceNotionListApi().get)(DataSourceNotionListApi(), sqlite_session, TENANT_ID, current_user) + unwrap(DataSourceNotionListApi().get)( + DataSourceNotionListApi(), + DataSourceNotionListQuery(credential_id="credential-1"), + sqlite_session, + TENANT_ID, + current_user, + ) @pytest.mark.parametrize("sqlite_session", [()], indirect=True) @@ -276,4 +286,10 @@ def test_notion_pre_import_pages_rejects_non_notion_dataset( patch.object(module.DatasetService, "get_dataset", return_value=dataset), pytest.raises(ValueError, match="Dataset is not notion type"), ): - unwrap(DataSourceNotionListApi().get)(DataSourceNotionListApi(), sqlite_session, TENANT_ID, current_user) + unwrap(DataSourceNotionListApi().get)( + DataSourceNotionListApi(), + DataSourceNotionListQuery(credential_id="credential-1", dataset_id="dataset-1"), + sqlite_session, + TENANT_ID, + current_user, + ) diff --git a/api/tests/unit_tests/controllers/console/datasets/test_data_source_notion_apis.py b/api/tests/unit_tests/controllers/console/datasets/test_data_source_notion_apis.py index d86a1093cc8..ca3d18973cf 100644 --- a/api/tests/unit_tests/controllers/console/datasets/test_data_source_notion_apis.py +++ b/api/tests/unit_tests/controllers/console/datasets/test_data_source_notion_apis.py @@ -15,6 +15,8 @@ from controllers.console.datasets.data_source import ( DataSourceNotionDocumentSyncApi, DataSourceNotionIndexingEstimateApi, DataSourceNotionPreviewApi, + DataSourceNotionPreviewQuery, + NotionEstimatePayload, ) from core.rag.index_processor.constant.index_type import IndexStructureType from models import Account @@ -45,7 +47,7 @@ class TestDataSourceNotionPreviewApi: return_value=extractor, ), ): - response, status = method(api, "tenant-1", "p1", "page") + response, status = method(api, DataSourceNotionPreviewQuery(credential_id="c1"), "tenant-1", "p1", "page") assert status == 200 @@ -80,7 +82,7 @@ class TestDataSourceNotionIndexingEstimateApi: return_value=MagicMock(model_dump=lambda: {"total_pages": 1}), ), ): - response, status = method(api, sqlite_session, "tenant-1") + response, status = method(api, NotionEstimatePayload.model_validate(payload), sqlite_session, "tenant-1") assert status == 200 diff --git a/api/tests/unit_tests/controllers/console/datasets/test_datasets.py b/api/tests/unit_tests/controllers/console/datasets/test_datasets.py index 72c095e2356..43fa8f6513d 100644 --- a/api/tests/unit_tests/controllers/console/datasets/test_datasets.py +++ b/api/tests/unit_tests/controllers/console/datasets/test_datasets.py @@ -18,6 +18,7 @@ from controllers.console.datasets.datasets import ( DatasetApiDeleteApi, DatasetApiKeyApi, DatasetAutoDisableLogApi, + DatasetCreatePayload, DatasetEnableApiApi, DatasetErrorDocs, DatasetIndexingEstimateApi, @@ -28,7 +29,9 @@ from controllers.console.datasets.datasets import ( DatasetRelatedAppListApi, DatasetRetrievalSettingApi, DatasetRetrievalSettingMockApi, + DatasetUpdatePayload, DatasetUseCheckApi, + IndexingEstimatePayload, _get_retrieval_methods_by_vector_type, ) from controllers.console.datasets.error import DatasetInUseError, DatasetNameDuplicateError, IndexingEstimateError @@ -489,7 +492,7 @@ class TestDatasetListApiPost: patch.object(type(console_ns), "payload", payload), patch.object(DatasetService, "create_empty_dataset", return_value=dataset), ): - _, status = method(api, MagicMock(), "tenant-1", user) + _, status = method(api, DatasetCreatePayload(**payload), MagicMock(), "tenant-1", user) assert status == 201 def test_post_forbidden(self, app: Flask): @@ -499,7 +502,7 @@ class TestDatasetListApiPost: user = make_account(TenantAccountRole.NORMAL) with app.test_request_context("/datasets", json=payload), patch.object(type(console_ns), "payload", payload): with pytest.raises(Forbidden): - method(api, MagicMock(), "tenant-1", user) + method(api, DatasetCreatePayload(**payload), MagicMock(), "tenant-1", user) def test_post_duplicate_name(self, app: Flask): api = DatasetListApi() @@ -514,14 +517,14 @@ class TestDatasetListApiPost: ), ): with pytest.raises(DatasetNameDuplicateError): - method(api, MagicMock(), "tenant-1", user) + method(api, DatasetCreatePayload(**payload), MagicMock(), "tenant-1", user) def test_post_invalid_payload_missing_name(self, app: Flask): api = DatasetListApi() method = unwrap(api.post) with app.test_request_context("/datasets", json={}), patch.object(type(console_ns), "payload", {}): with pytest.raises(ValueError): - method(api, MagicMock(), "tenant-1", make_account()) + method(api, DatasetCreatePayload(), MagicMock(), "tenant-1", make_account()) def test_post_invalid_indexing_technique(self, app: Flask): api = DatasetListApi() @@ -529,7 +532,7 @@ class TestDatasetListApiPost: payload = {"name": "bad", "indexing_technique": "invalid-tech"} with app.test_request_context("/datasets", json=payload), patch.object(type(console_ns), "payload", payload): with pytest.raises(ValueError, match="Invalid indexing technique"): - method(api, MagicMock(), "tenant-1", make_account()) + method(api, DatasetCreatePayload(**payload), MagicMock(), "tenant-1", make_account()) def test_post_invalid_provider(self, app: Flask): api = DatasetListApi() @@ -537,7 +540,7 @@ class TestDatasetListApiPost: payload = {"name": "bad", "provider": "unknown"} with app.test_request_context("/datasets", json=payload), patch.object(type(console_ns), "payload", payload): with pytest.raises(ValueError, match="Invalid provider"): - method(api, MagicMock(), "tenant-1", make_account()) + method(api, DatasetCreatePayload(**payload), MagicMock(), "tenant-1", make_account()) class TestDatasetApiGet: @@ -692,7 +695,7 @@ class TestDatasetApiPatch: patch.object(DatasetService, "update_dataset", return_value=dataset), patch.object(DatasetPermissionService, "get_dataset_partial_member_list", return_value=[]), ): - result, status = method(api, MagicMock(), tenant_id, user, dataset_id) + result, status = method(api, DatasetUpdatePayload(), MagicMock(), tenant_id, user, dataset_id) assert status == 200 assert result["partial_member_list"] == [] @@ -704,7 +707,7 @@ class TestDatasetApiPatch: patch.object(DatasetService, "get_dataset", return_value=None), ): with pytest.raises(NotFound, match="Dataset not found"): - method(api, MagicMock(), "tenant-1", make_account(), "missing") + method(api, DatasetUpdatePayload(), MagicMock(), "tenant-1", make_account(), "missing") def test_patch_permission_denied(self, app: Flask): api = DatasetApi() @@ -719,7 +722,7 @@ class TestDatasetApiPatch: patch.object(DatasetPermissionService, "check_permission", side_effect=Forbidden("no permission")), ): with pytest.raises(Forbidden): - method(api, MagicMock(), "tenant", make_account(), dataset_id) + method(api, DatasetUpdatePayload(), MagicMock(), "tenant", make_account(), dataset_id) def test_patch_partial_members_update(self, app: Flask): api = DatasetApi() @@ -736,7 +739,7 @@ class TestDatasetApiPatch: patch.object(DatasetPermissionService, "update_partial_member_list", return_value=None), patch.object(DatasetPermissionService, "get_dataset_partial_member_list", return_value=["u1", "u2"]), ): - result, _ = method(api, MagicMock(), "tenant", make_account(), dataset_id) + result, _ = method(api, DatasetUpdatePayload(), MagicMock(), "tenant", make_account(), dataset_id) assert result["partial_member_list"] == ["u1", "u2"] def test_patch_clear_partial_members(self, app: Flask): @@ -754,7 +757,7 @@ class TestDatasetApiPatch: patch.object(DatasetPermissionService, "clear_partial_member_list", return_value=None), patch.object(DatasetPermissionService, "get_dataset_partial_member_list", return_value=[]), ): - result, _ = method(api, MagicMock(), "tenant", make_account(), dataset_id) + result, _ = method(api, DatasetUpdatePayload(), MagicMock(), "tenant", make_account(), dataset_id) assert result["partial_member_list"] == [] @@ -1014,7 +1017,12 @@ class TestDatasetIndexingEstimateApi: patch("controllers.console.datasets.datasets.DocumentService.estimate_args_validate", return_value=None), patch("controllers.console.datasets.datasets.IndexingRunner.indexing_estimate", return_value=mock_response), ): - response, status = method(api, session, "tenant-1") + response, status = method( + api, + IndexingEstimatePayload(**payload), + session, + "tenant-1", + ) assert status == 200 assert response == { "tokens": 0, @@ -1036,7 +1044,12 @@ class TestDatasetIndexingEstimateApi: patch("controllers.console.datasets.datasets.DocumentService.estimate_args_validate", return_value=None), ): with pytest.raises(NotFound): - method(api, session, "tenant-1") + method( + api, + IndexingEstimatePayload(**payload), + session, + "tenant-1", + ) def test_post_llm_bad_request_error(self, app: Flask): api = DatasetIndexingEstimateApi() @@ -1055,7 +1068,12 @@ class TestDatasetIndexingEstimateApi: ), ): with pytest.raises(ProviderNotInitializeError): - method(api, session, "tenant-1") + method( + api, + IndexingEstimatePayload(**payload), + session, + "tenant-1", + ) def test_post_provider_token_not_init(self, app: Flask): api = DatasetIndexingEstimateApi() @@ -1074,7 +1092,12 @@ class TestDatasetIndexingEstimateApi: ), ): with pytest.raises(ProviderNotInitializeError): - method(api, session, "tenant-1") + method( + api, + IndexingEstimatePayload(**payload), + session, + "tenant-1", + ) def test_post_generic_exception(self, app: Flask): api = DatasetIndexingEstimateApi() @@ -1092,7 +1115,12 @@ class TestDatasetIndexingEstimateApi: ), ): with pytest.raises(IndexingEstimateError): - method(api, session, "tenant-1") + method( + api, + IndexingEstimatePayload(**payload), + session, + "tenant-1", + ) class TestDatasetRelatedAppListApi: diff --git a/api/tests/unit_tests/controllers/console/datasets/test_datasets_segments.py b/api/tests/unit_tests/controllers/console/datasets/test_datasets_segments.py index 22597a21f9b..bbcda2fb372 100644 --- a/api/tests/unit_tests/controllers/console/datasets/test_datasets_segments.py +++ b/api/tests/unit_tests/controllers/console/datasets/test_datasets_segments.py @@ -9,9 +9,11 @@ from flask import Flask from werkzeug.exceptions import Forbidden, NotFound import services +from controllers.common.controller_schemas import ChildChunkCreatePayload, ChildChunkUpdatePayload from controllers.console import console_ns from controllers.console.app.error import ProviderNotInitializeError from controllers.console.datasets.datasets_segments import ( + BatchImportPayload, ChildChunkAddApi, ChildChunkBatchUpdatePayload, ChildChunkUpdateApi, @@ -20,6 +22,8 @@ from controllers.console.datasets.datasets_segments import ( DatasetDocumentSegmentBatchImportApi, DatasetDocumentSegmentListApi, DatasetDocumentSegmentUpdateApi, + SegmentCreatePayload, + SegmentUpdatePayload, ) from controllers.console.datasets.error import ChildChunkDeleteIndexError, ChildChunkIndexingError, InvalidActionError from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError @@ -50,6 +54,7 @@ def _segment(): status=SegmentStatus.COMPLETED, updated_by="u1", ) + segment.id = "seg-1" segment.created_at = naive_utc_now() segment.updated_at = naive_utc_now() @@ -354,7 +359,9 @@ class TestDatasetDocumentSegmentAddApi: return_value=None, ), ): - response, status = method(api, session, "tenant-1", user, "ds-1", "doc-1") + response, status = method( + api, SegmentCreatePayload(content="test content"), session, "tenant-1", user, "ds-1", "doc-1" + ) assert status == 200 assert response["data"]["id"] == "seg-1" @@ -378,7 +385,9 @@ class TestDatasetDocumentSegmentAddApi: ), ): with pytest.raises(ProviderNotInitializeError): - method(api, MagicMock(), "tenant-1", user, "ds-1", "doc-1") + method( + api, SegmentCreatePayload(content="test content"), MagicMock(), "tenant-1", user, "ds-1", "doc-1" + ) def test_post_provider_token_not_init(self, app: Flask): api = DatasetDocumentSegmentAddApi() @@ -400,7 +409,9 @@ class TestDatasetDocumentSegmentAddApi: ), ): with pytest.raises(ProviderNotInitializeError): - method(api, MagicMock(), "tenant-1", user, "ds-1", "doc-1") + method( + api, SegmentCreatePayload(content="test content"), MagicMock(), "tenant-1", user, "ds-1", "doc-1" + ) class TestDatasetDocumentSegmentUpdateApi: @@ -441,7 +452,9 @@ class TestDatasetDocumentSegmentUpdateApi: return_value=None, ), ): - response, status = method(api, session, "tenant-1", user, "ds-1", "doc-1", "seg-1") + response, status = method( + api, SegmentUpdatePayload(content="test content"), session, "tenant-1", user, "ds-1", "doc-1", "seg-1" + ) assert status == 200 assert "data" in response @@ -467,7 +480,16 @@ class TestDatasetDocumentSegmentUpdateApi: ), ): with pytest.raises(NotFound): - method(api, MagicMock(), "tenant-1", user, "ds-1", "doc-1", "seg-1") + method( + api, + SegmentUpdatePayload(content="test content"), + MagicMock(), + "tenant-1", + user, + "ds-1", + "doc-1", + "seg-1", + ) def test_patch_segment_not_found(self, app: Flask): api = DatasetDocumentSegmentUpdateApi() @@ -495,7 +517,16 @@ class TestDatasetDocumentSegmentUpdateApi: ), ): with pytest.raises(NotFound): - method(api, MagicMock(), "tenant-1", user, "ds-1", "doc-1", "seg-1") + method( + api, + SegmentUpdatePayload(content="test content"), + MagicMock(), + "tenant-1", + user, + "ds-1", + "doc-1", + "seg-1", + ) def test_patch_llm_bad_request(self, app: Flask): api = DatasetDocumentSegmentUpdateApi() @@ -525,7 +556,16 @@ class TestDatasetDocumentSegmentUpdateApi: ), ): with pytest.raises(ProviderNotInitializeError): - method(api, MagicMock(), "tenant-1", user, "ds-1", "doc-1", "seg-1") + method( + api, + SegmentUpdatePayload(content="test content"), + MagicMock(), + "tenant-1", + user, + "ds-1", + "doc-1", + "seg-1", + ) class TestDatasetDocumentSegmentBatchImportApi: @@ -564,7 +604,9 @@ class TestDatasetDocumentSegmentBatchImportApi: return_value=None, ), ): - response, status = method(api, session, "tenant-1", user, "ds-1", "doc-1") + response, status = method( + api, BatchImportPayload(upload_file_id="test-file-id"), session, "tenant-1", user, "ds-1", "doc-1" + ) assert status == 200 assert response["job_status"] == "waiting" @@ -581,7 +623,15 @@ class TestDatasetDocumentSegmentBatchImportApi: patch("controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=None), ): with pytest.raises(NotFound): - method(api, MagicMock(), "tenant-1", user, "ds-1", "doc-1") + method( + api, + BatchImportPayload(upload_file_id="test-file-id"), + MagicMock(), + "tenant-1", + user, + "ds-1", + "doc-1", + ) def test_post_document_not_found(self, app: Flask): api = DatasetDocumentSegmentBatchImportApi() @@ -599,7 +649,15 @@ class TestDatasetDocumentSegmentBatchImportApi: patch("controllers.console.datasets.datasets_segments.DocumentService.get_document", return_value=None), ): with pytest.raises(NotFound): - method(api, MagicMock(), "tenant-1", user, "ds-1", "doc-1") + method( + api, + BatchImportPayload(upload_file_id="test-file-id"), + MagicMock(), + "tenant-1", + user, + "ds-1", + "doc-1", + ) def test_post_upload_file_not_found(self, app: Flask): api = DatasetDocumentSegmentBatchImportApi() @@ -619,7 +677,9 @@ class TestDatasetDocumentSegmentBatchImportApi: ), ): with pytest.raises(NotFound): - method(api, session, "tenant-1", user, "ds-1", "doc-1") + method( + api, BatchImportPayload(upload_file_id="test-file-id"), session, "tenant-1", user, "ds-1", "doc-1" + ) def test_post_invalid_file_type(self, app: Flask): api = DatasetDocumentSegmentBatchImportApi() @@ -641,7 +701,9 @@ class TestDatasetDocumentSegmentBatchImportApi: ), ): with pytest.raises(ValueError): - method(api, session, "tenant-1", user, "ds-1", "doc-1") + method( + api, BatchImportPayload(upload_file_id="test-file-id"), session, "tenant-1", user, "ds-1", "doc-1" + ) def test_post_async_task_failure(self, app: Flask): api = DatasetDocumentSegmentBatchImportApi() @@ -665,7 +727,9 @@ class TestDatasetDocumentSegmentBatchImportApi: "controllers.console.datasets.datasets_segments.redis_client.setnx", side_effect=Exception("redis down") ), ): - response, status = method(api, session, "tenant-1", user, "ds-1", "doc-1") + response, status = method( + api, BatchImportPayload(upload_file_id="test-file-id"), session, "tenant-1", user, "ds-1", "doc-1" + ) assert status == 500 assert "error" in response @@ -747,7 +811,9 @@ class TestChildChunkAddApi: return_value=child_chunk, ), ): - response, status = method(api, MagicMock(), "tenant-1", user, "ds-1", "doc-1", "seg-1") + response, status = method( + api, ChildChunkCreatePayload(content="child"), MagicMock(), "tenant-1", user, "ds-1", "doc-1", "seg-1" + ) assert status == 200 assert response["data"]["id"] == "cc-1" @@ -778,7 +844,16 @@ class TestChildChunkAddApi: ), ): with pytest.raises(ChildChunkIndexingError): - method(api, MagicMock(), "tenant-1", user, "ds-1", "doc-1", "seg-1") + method( + api, + ChildChunkCreatePayload(content="child"), + MagicMock(), + "tenant-1", + user, + "ds-1", + "doc-1", + "seg-1", + ) def test_post_permission_denied(self, app: Flask): api = ChildChunkAddApi() @@ -799,7 +874,16 @@ class TestChildChunkAddApi: ), ): with pytest.raises(Forbidden): - method(api, MagicMock(), "tenant-1", user, "ds-1", "doc-1", "seg-1") + method( + api, + ChildChunkCreatePayload(content="child"), + MagicMock(), + "tenant-1", + user, + "ds-1", + "doc-1", + "seg-1", + ) class TestChildChunkUpdateApi: @@ -922,7 +1006,17 @@ class TestChildChunkUpdateApi: ), ): with pytest.raises(NotFound): - method(api, MagicMock(), "tenant-1", user, "ds-1", "doc-1", "seg-1", "cc-1") + method( + api, + ChildChunkUpdatePayload(content="updated child"), + MagicMock(), + "tenant-1", + user, + "ds-1", + "doc-1", + "seg-1", + "cc-1", + ) class TestSegmentListAdvancedCases: @@ -1047,7 +1141,9 @@ class TestSegmentOperationCases: ), ): with pytest.raises(ProviderTokenNotInitError): - method(api, MagicMock(), "tenant-1", user, "ds-1", "doc-1") + method( + api, SegmentCreatePayload(content="test content"), MagicMock(), "tenant-1", user, "ds-1", "doc-1" + ) def test_batch_import_with_document_not_found(self, app: Flask): """Test batch import with document not found""" @@ -1063,7 +1159,15 @@ class TestSegmentOperationCases: patch("controllers.console.datasets.datasets_segments.DocumentService.get_document", return_value=None), ): with pytest.raises(NotFound): - method(api, MagicMock(), "tenant-1", user, "ds-1", "doc-1") + method( + api, + BatchImportPayload(upload_file_id="test-file-id"), + MagicMock(), + "tenant-1", + user, + "ds-1", + "doc-1", + ) def test_batch_import_with_invalid_file(self, app: Flask): """Test batch import with invalid file type""" @@ -1083,7 +1187,9 @@ class TestSegmentOperationCases: patch("controllers.console.datasets.datasets_segments.DocumentService.get_document", return_value=document), ): with pytest.raises(NotFound): - method(api, session, "tenant-1", user, "ds-1", "doc-1") + method( + api, BatchImportPayload(upload_file_id="test-file-id"), session, "tenant-1", user, "ds-1", "doc-1" + ) def test_batch_import_with_async_task_failure(self, app: Flask): api = DatasetDocumentSegmentBatchImportApi() @@ -1122,7 +1228,9 @@ class TestSegmentOperationCases: side_effect=Exception("Task failed"), ), ): - response, status = method(api, session, "tenant-1", user, "ds-1", "doc-1") + response, status = method( + api, BatchImportPayload(upload_file_id="test-file-id"), session, "tenant-1", user, "ds-1", "doc-1" + ) assert status == 500 assert "error" in response diff --git a/api/tests/unit_tests/controllers/console/datasets/test_external.py b/api/tests/unit_tests/controllers/console/datasets/test_external.py index edcc2451131..305a9adbacc 100644 --- a/api/tests/unit_tests/controllers/console/datasets/test_external.py +++ b/api/tests/unit_tests/controllers/console/datasets/test_external.py @@ -12,10 +12,14 @@ from controllers.console import console_ns from controllers.console.datasets.error import DatasetNameDuplicateError from controllers.console.datasets.external import ( BedrockRetrievalApi, + BedrockRetrievalPayload, ExternalApiTemplateApi, ExternalApiTemplateListApi, + ExternalApiTemplateListQuery, ExternalApiUseCheckApi, ExternalDatasetCreateApi, + ExternalHitTestingPayload, + ExternalKnowledgeApiPayload, ExternalKnowledgeHitTestingApi, ) from models.account import Account, TenantAccountRole @@ -174,7 +178,9 @@ class TestExternalApiTemplateListApi: return_value=([api_item], 3), ) as get_external_knowledge_apis, ): - resp, status = method(api, session, "tenant-1") + resp, status = method( + api, ExternalApiTemplateListQuery(page=2, limit=1, keyword="vector"), session, "tenant-1" + ) assert status == 200 assert resp == { @@ -213,7 +219,9 @@ class TestExternalApiTemplateListApi: return_value=created, ) as create_external_knowledge_api, ): - resp, status = method(api, session, "tenant-1", current_user) + resp, status = method( + api, ExternalKnowledgeApiPayload.model_validate(payload), session, "tenant-1", current_user + ) assert status == 201 assert resp == _external_api_dict("api-created") @@ -239,7 +247,7 @@ class TestExternalApiTemplateListApi: patch.object(ExternalDatasetService, "validate_api_list"), ): with pytest.raises(Forbidden): - method(api, MagicMock(), "tenant-1", current_user) + method(api, ExternalKnowledgeApiPayload.model_validate(payload), MagicMock(), "tenant-1", current_user) def test_post_duplicate_name(self, app: Flask, current_user: Account): api = ExternalApiTemplateListApi() @@ -258,7 +266,7 @@ class TestExternalApiTemplateListApi: ), ): with pytest.raises(DatasetNameDuplicateError): - method(api, MagicMock(), "tenant-1", current_user) + method(api, ExternalKnowledgeApiPayload.model_validate(payload), MagicMock(), "tenant-1", current_user) class TestExternalApiTemplateApi: @@ -325,7 +333,14 @@ class TestExternalApiTemplateApi: return_value=updated, ) as update_external_knowledge_api, ): - resp, status = method(api, session, "tenant-1", current_user, "api-updated") + resp, status = method( + api, + ExternalKnowledgeApiPayload.model_validate(payload), + session, + "tenant-1", + current_user, + "api-updated", + ) assert status == 200 assert resp == _external_api_dict("api-updated") @@ -405,7 +420,9 @@ class TestExternalDatasetCreateApi: ) as dataset_response_source, ): session = MagicMock() - resp, status = method(api, session, "tenant-1", current_user) + resp, status = method( + api, ExternalDatasetCreatePayload.model_validate(payload), session, "tenant-1", current_user + ) assert status == 201 assert resp == _expected_dataset_detail_payload() @@ -433,7 +450,7 @@ class TestExternalDatasetCreateApi: patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload), ): with pytest.raises(Forbidden): - method(api, MagicMock(), "tenant-1", current_user) + method(api, ExternalDatasetCreatePayload.model_validate(payload), MagicMock(), "tenant-1", current_user) class TestExternalKnowledgeHitTestingApi: @@ -450,7 +467,7 @@ class TestExternalKnowledgeHitTestingApi: ), ): with pytest.raises(NotFound): - method(api, MagicMock(), current_user, "dataset-id") + method(api, ExternalHitTestingPayload(query="test"), MagicMock(), current_user, "dataset-id") def test_hit_testing_success(self, app: Flask, current_user: Account): api = ExternalKnowledgeHitTestingApi() @@ -496,7 +513,7 @@ class TestExternalKnowledgeHitTestingApi: ) as external_retrieve, patch("controllers.console.datasets.external.dump_response", side_effect=lambda _model, value: value), ): - resp = method(api, session, current_user, "dataset-id") + resp = method(api, ExternalHitTestingPayload.model_validate(payload), session, current_user, "dataset-id") assert resp == retrieve_response check_dataset_permission.assert_called_once_with(dataset, current_user, session) @@ -549,7 +566,7 @@ class TestBedrockRetrievalApi: return_value=retrieval_response, ) as knowledge_retrieval, ): - resp, status = method() + resp, status = method(api, BedrockRetrievalPayload.model_validate(payload)) assert status == 200 assert resp == retrieval_response @@ -576,7 +593,7 @@ class TestExternalApiTemplateListApiAdvanced: ), ): with pytest.raises(DatasetNameDuplicateError): - method(api, MagicMock(), "tenant-1", current_user) + method(api, ExternalKnowledgeApiPayload.model_validate(payload), MagicMock(), "tenant-1", current_user) def test_get_with_pagination(self, app: Flask): api = ExternalApiTemplateListApi() @@ -591,7 +608,7 @@ class TestExternalApiTemplateListApiAdvanced: return_value=(templates, 25), ) as get_external_knowledge_apis, ): - resp, status = method(api, MagicMock(), "tenant-1") + resp, status = method(api, ExternalApiTemplateListQuery(page=2, limit=3), MagicMock(), "tenant-1") assert status == 200 assert resp == { @@ -621,7 +638,7 @@ class TestExternalDatasetCreateApiAdvanced: with app.test_request_context("/", json=payload), patch.object(type(console_ns), "payload", payload): with pytest.raises(Forbidden): - method(api, MagicMock(), "tenant-1", current_user) + method(api, ExternalDatasetCreatePayload.model_validate(payload), MagicMock(), "tenant-1", current_user) class TestExternalKnowledgeHitTestingApiAdvanced: @@ -644,7 +661,7 @@ class TestExternalKnowledgeHitTestingApiAdvanced: ), ): with pytest.raises(NotFound): - method(api, MagicMock(), current_user, "ds-1") + method(api, ExternalHitTestingPayload.model_validate(payload), MagicMock(), current_user, "ds-1") def test_hit_testing_with_custom_retrieval_model(self, app: Flask, current_user: Account): api = ExternalKnowledgeHitTestingApi() @@ -682,7 +699,7 @@ class TestExternalKnowledgeHitTestingApiAdvanced: }, ) as external_retrieve, ): - resp = method(api, session, current_user, "ds-1") + resp = method(api, ExternalHitTestingPayload.model_validate(payload), session, current_user, "ds-1") assert resp == { "query": {"content": "test query"}, @@ -727,4 +744,4 @@ class TestBedrockRetrievalApiAdvanced: ), ): with pytest.raises(ValueError): - method() + method(api, BedrockRetrievalPayload.model_validate(payload)) diff --git a/api/tests/unit_tests/controllers/console/datasets/test_metadata.py b/api/tests/unit_tests/controllers/console/datasets/test_metadata.py index 00f45d488a4..bbc9bd6ecff 100644 --- a/api/tests/unit_tests/controllers/console/datasets/test_metadata.py +++ b/api/tests/unit_tests/controllers/console/datasets/test_metadata.py @@ -7,6 +7,7 @@ from flask import Flask from pytest_mock import MockerFixture from werkzeug.exceptions import NotFound +from controllers.common.controller_schemas import MetadataUpdatePayload from controllers.console import console_ns from controllers.console.datasets.metadata import ( DatasetMetadataApi, @@ -24,6 +25,7 @@ from services.metadata_service import MetadataService @pytest.fixture def app(): app = Flask("test_dataset_metadata") + app.config["TESTING"] = True return app @@ -76,7 +78,9 @@ class TestDatasetMetadataCreateApi: MetadataService, "create_metadata", return_value={"id": "m1", "type": "string", "name": "author"} ), ): - result, status = method(api, MagicMock(), "tenant-1", current_user, dataset_id) + result, status = method( + api, MetadataArgs(type="string", name="author"), MagicMock(), "tenant-1", current_user, dataset_id + ) assert status == 201 assert result["type"] == "string" assert result["name"] == "author" @@ -92,7 +96,9 @@ class TestDatasetMetadataCreateApi: patch.object(DatasetService, "get_dataset", return_value=None), ): with pytest.raises(NotFound, match="Dataset not found"): - method(api, MagicMock(), "tenant-1", current_user, dataset_id) + method( + api, MetadataArgs(type="string", name="author"), MagicMock(), "tenant-1", current_user, dataset_id + ) class TestDatasetMetadataGetApi: @@ -140,7 +146,15 @@ class TestDatasetMetadataApi: return_value={"id": "m1", "type": "string", "name": "updated-name"}, ), ): - result, status = method(api, MagicMock(), "tenant-1", current_user, dataset_id, metadata_id) + result, status = method( + api, + MetadataUpdatePayload(name="updated-name"), + MagicMock(), + "tenant-1", + current_user, + dataset_id, + metadata_id, + ) assert status == 200 assert result["type"] == "string" assert result["name"] == "updated-name" @@ -201,9 +215,14 @@ class TestDocumentMetadataEditApi: patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload), patch.object(DatasetService, "get_dataset", return_value=dataset), patch.object(DatasetService, "check_dataset_permission"), - patch.object(MetadataOperationData, "model_validate", return_value=MagicMock()), patch.object(MetadataService, "update_documents_metadata"), ): - result, status = method(api, MagicMock(), current_user, dataset_id) + result, status = method( + api, + MetadataOperationData(operation_data=[{"document_id": "doc-1", "metadata_list": []}]), + MagicMock(), + current_user, + dataset_id, + ) assert status == 204 assert result == ""