diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 249c2932f65..006e3fd55ba 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -57,13 +57,10 @@ # Backend - RAG (Retrieval Augmented Generation) /api/core/rag/ @JohnJyong +/api/services/knowledge/ @JohnJyong /api/services/rag_pipeline/ @JohnJyong -/api/services/dataset_service.py @JohnJyong /api/services/knowledge_service.py @JohnJyong -/api/services/external_knowledge_service.py @JohnJyong /api/services/hit_testing_service.py @JohnJyong -/api/services/metadata_service.py @JohnJyong -/api/services/vector_service.py @JohnJyong /api/services/entities/knowledge_entities/ @JohnJyong /api/services/entities/external_knowledge_entities/ @JohnJyong /api/controllers/console/datasets/ @JohnJyong diff --git a/api/.importlinter b/api/.importlinter index badb1a6e95a..ff8d7d6fe5d 100644 --- a/api/.importlinter +++ b/api/.importlinter @@ -1,12 +1,14 @@ [importlinter] root_packages = core + core.rag.extractor.entity constants context configs controllers extensions factories + fields libs machinery models @@ -16,7 +18,7 @@ root_packages = include_external_packages = True [importlinter:contract:backend-layers] -name = Backend modules follow the transport, application, domain, and shared-library layers +name = Backend modules follow the transport, application, execution-kernel, and shared-library layers type = layers layers = controllers @@ -27,7 +29,7 @@ layers = # error makes this list shrink whenever a legacy dependency is removed. unmatched_ignore_imports_alerting = error ignore_imports = - core.app.app_config.easy_ui_based_app.dataset.manager -> services.dataset_service + core.app.app_config.easy_ui_based_app.dataset.manager -> services.knowledge.dataset_service core.app.apps.advanced_chat.app_generator -> controllers.console.app.workflow core.app.apps.advanced_chat.app_generator -> services.conversation_service core.app.apps.advanced_chat.app_generator -> services.errors.conversation @@ -48,45 +50,44 @@ ignore_imports = core.app.apps.message_based_app_generator -> services.errors.conversation core.app.apps.message_based_app_generator -> services.errors.message core.app.apps.pipeline.pipeline_generator -> services - core.app.apps.pipeline.pipeline_generator -> services.dataset_service - core.app.apps.pipeline.pipeline_generator -> services.datasource_provider_service + core.app.apps.pipeline.pipeline_generator -> services.knowledge.dataset_service + core.app.apps.pipeline.pipeline_generator -> services.data_source.provider_service core.app.apps.pipeline.pipeline_generator -> services.feature_service core.app.apps.pipeline.pipeline_generator -> services.workflow_draft_variable_service - core.app.apps.pipeline.pipeline_runner -> services.dataset_ref_service core.app.apps.workflow.app_generator -> controllers.console.app.workflow core.app.apps.workflow.app_generator -> services.snippet_generate_service core.app.apps.workflow.app_generator -> services.workflow_draft_variable_service core.app.apps.workflow_app_runner -> tasks.mail_human_input_delivery_task core.app.features.annotation_reply.annotation_reply -> services.annotation_service - core.app.features.annotation_reply.annotation_reply -> services.dataset_service + core.app.features.annotation_reply.annotation_reply -> services.knowledge.dataset_service core.app.llm.quota -> services.credit_pool_service core.app.task_pipeline.message_cycle_manager -> services.account_service core.app.task_pipeline.message_cycle_manager -> services.annotation_service core.app.workflow.layers.persistence -> services.workflow.inspector_events - core.datasource.datasource_manager -> services.datasource_provider_service core.helper.credential_utils -> services.enterprise.plugin_manager_service core.helper.credential_utils -> services.system_feature_service - core.indexing_runner -> services.vector_space_admission_service core.mcp.auth_client -> services - core.provider_manager -> services.credential_permission_service core.provider_manager -> services.credit_pool_service core.provider_manager -> services.feature_service - core.rag.datasource.retrieval_service -> services.external_knowledge_service + core.rag.datasource.retrieval_service -> services.knowledge.external.service core.rag.index_processor.index_processor -> services.vector_space_admission_service core.rag.index_processor.index_processor_base -> services.file_service core.rag.index_processor.processor.paragraph_index_processor -> services.account_service - core.rag.index_processor.processor.paragraph_index_processor -> services.summary_index_service + core.rag.index_processor.processor.paragraph_index_processor -> services.knowledge.summaries.adapters core.rag.index_processor.processor.parent_child_index_processor -> services.account_service - core.rag.index_processor.processor.parent_child_index_processor -> services.summary_index_service - core.rag.index_processor.processor.qa_index_processor -> services.summary_index_service - core.rag.retrieval.dataset_retrieval -> services.external_knowledge_service + core.rag.index_processor.processor.parent_child_index_processor -> services.knowledge.summaries.adapters + core.rag.index_processor.processor.qa_index_processor -> services.knowledge.summaries.adapters + core.rag.retrieval.dataset_retrieval -> services.knowledge.external.service + core.rag.retrieval.dataset_retrieval -> services.knowledge.retrieval.attachments core.rag.retrieval.dataset_retrieval -> services.feature_service - core.rag.summary_index.summary_index -> services.summary_index_service + core.rag.summary_index.summary_index -> services.knowledge.summaries.adapters core.rag.summary_index.summary_index -> tasks.generate_summary_index_task core.repositories.sqlalchemy_workflow_node_execution_repository -> services.file_service core.repositories.sqlalchemy_workflow_node_execution_repository -> services.variable_truncator core.tools.tool_manager -> services core.workflow.node_runtime -> services + # Node construction delegates datasource dependency assembly to the composition root. + core.workflow.node_factory -> extensions.application_services.data_sources core.workflow.nodes.agent.message_transformer -> services core.workflow.nodes.agent_v2.agent_node -> services.agent.prompt_mentions core.workflow.nodes.agent_v2.agent_node -> services.agent.workspace_service @@ -99,7 +100,7 @@ ignore_imports = core.workflow.nodes.agent_v2.workspace_retirement_layer -> tasks.collect_agent_resources_task libs.email_i18n -> services.entities.feature_entities libs.email_i18n -> services.system_feature_service - libs.external_api -> core + libs.external_api -> core.plugin libs.external_api -> core.errors.error libs.external_api -> extensions.ext_logging libs.flask_utils -> models @@ -168,6 +169,141 @@ forbidden_modules = sqlalchemy werkzeug +[importlinter:contract:dataset-application-boundary] +name = Dataset application services and contracts are framework and persistence neutral +type = forbidden +source_modules = + core.rag.extractor.entity.extract_setting + services.knowledge.entities.segments + services.knowledge.entities.datasets + services.knowledge.entities.document_creation + services.data_source.binding_application_service + services.data_source.credential_gateway + services.data_source.notion_import_application_service + services.data_source.entities.credential + services.data_source.entities.notion_import + services.knowledge.entities.indexing_estimate + services.knowledge.dataset_access + services.knowledge.datasets.application + services.knowledge.datasets.retrieval + services.knowledge.document_sync + services.knowledge.documents.application + services.knowledge.external.application + services.knowledge.indexing.errors + services.knowledge.indexing.estimate + services.knowledge.indexing.execution + services.knowledge.resource_scope + services.knowledge.segments.application + services.knowledge.metadata.application + services.knowledge.summaries.application +forbidden_modules = + fields + configs + controllers + extensions + flask + models + repositories + sqlalchemy + tasks + werkzeug + services.knowledge.dataset_service + services.data_source.provider_service + services.knowledge.indexing.adapters + services.knowledge.segments.adapters + services.knowledge.segments.indexing + core.datasource.datasource_manager + core.rag.extractor.notion_extractor + +[importlinter:contract:knowledge-repository-boundary] +name = Knowledge persistence adapters only depend on their consumer-owned application contracts +type = forbidden +source_modules = + repositories.knowledge +forbidden_modules = + controllers + extensions + flask + services + tasks + werkzeug +allow_indirect_imports = True +ignore_imports = + repositories.knowledge.dataset_api_key_repository -> services.auth.api_key_contracts + repositories.knowledge.dataset_api_key_repository -> services.knowledge.api_key_service + repositories.knowledge.dataset_repository -> services.knowledge.segments.application + repositories.knowledge.dataset_repository -> services.knowledge.dataset_access + repositories.knowledge.dataset_repository -> services.knowledge.indexing.estimate + repositories.knowledge.dataset_repository -> services.knowledge.resource_scope + repositories.knowledge.document_repository -> services.knowledge.document_sync + repositories.knowledge.document_repository -> services.knowledge.indexing.estimate + repositories.knowledge.document_repository -> services.knowledge.indexing.errors + repositories.knowledge.document_repository -> services.knowledge.indexing.execution + repositories.knowledge.document_repository -> services.knowledge.resource_scope + repositories.knowledge.segment_repository -> services.knowledge.resource_scope + repositories.knowledge.segment_repository -> services.knowledge.indexing.execution + repositories.knowledge.segment_repository -> services.knowledge.segments.application + repositories.knowledge.segment_repository -> services.knowledge.entities.segments + repositories.knowledge.metadata_repository -> services.errors.metadata + repositories.knowledge.metadata_repository -> services.knowledge.metadata.application + repositories.knowledge.metadata_repository -> services.knowledge.resource_scope + repositories.knowledge.summary_repository -> services.knowledge.summaries.application + +[importlinter:contract:datasource-credential-repository-boundary] +name = Data source credential persistence does not depend on SQLAlchemy service helpers +type = forbidden +source_modules = + repositories.credentials.query_repository + repositories.data_source.credential_repository +forbidden_modules = + services.data_source.provider_service +allow_indirect_imports = True + +[importlinter:contract:datasource-provider-persistence-boundary] +name = Datasource provider service consumes injected persistence contracts +type = forbidden +source_modules = + services.data_source.provider_service +forbidden_modules = + repositories + extensions.application_services + extensions.ext_application_services +allow_indirect_imports = True + +[importlinter:contract:credential-permission-query-policy-boundary] +name = Credential query contracts do not expose persistence implementations +type = forbidden +source_modules = + services.credentials.query +forbidden_modules = + repositories + sqlalchemy + models + extensions + flask +allow_indirect_imports = True + +[importlinter:contract:credential-query-implementation-boundary] +name = Credential query implementations are wired only by composition or other repositories +type = forbidden +source_modules = + controllers + services + core +forbidden_modules = + repositories.credentials.query_repository +allow_indirect_imports = True + +[importlinter:contract:pipeline-document-lookup-boundary] +name = Pipeline core depends on its document store port +type = forbidden +source_modules = + core.app.apps.pipeline.pipeline_generator + core.app.apps.pipeline.pipeline_runner +forbidden_modules = + repositories +allow_indirect_imports = True + [importlinter:contract:init-validation-service-boundary] name = Initialization validation application service is framework and persistence neutral type = forbidden @@ -479,7 +615,7 @@ forbidden_modules = name = Data source OAuth application service is framework and persistence neutral type = forbidden source_modules = - services.data_source_oauth_service + services.data_source.oauth_service forbidden_modules = configs controllers @@ -800,6 +936,46 @@ forbidden_modules = tasks werkzeug +[importlinter:contract:knowledge-segment-presentation-boundary] +name = Knowledge operations do not serialize HTTP responses +type = forbidden +allow_indirect_imports = True +source_modules = + services.knowledge.segments + services.knowledge.documents + services.knowledge.datasets + services.knowledge.external +forbidden_modules = + controllers + fields + +[importlinter:contract:knowledge-response-persistence-boundary] +name = Knowledge response schemas do not load database state +type = forbidden +allow_indirect_imports = True +source_modules = + fields.dataset_fields + fields.document_fields + fields.segment_fields +forbidden_modules = + sqlalchemy + models + repositories + services.knowledge + +[importlinter:contract:dataset-model-infrastructure-boundary] +name = Dataset models do not execute infrastructure or application operations +type = forbidden +allow_indirect_imports = True +source_modules = + models.dataset +forbidden_modules = + repositories + services + extensions + core.tools + models.engine + [importlinter:contract:app-mcp-server-service-boundary] name = App MCP server application service is framework and persistence neutral type = forbidden diff --git a/api/AGENTS.md b/api/AGENTS.md index 160ae644a28..86f2ef34664 100644 --- a/api/AGENTS.md +++ b/api/AGENTS.md @@ -15,7 +15,8 @@ Run direct Python commands through `uv run --project api`. Docker-backed integra ## Architecture And Boundaries -- Keep transport parsing and serialization in controllers, orchestration in services, and domain policy in `core/` or its domain owner. Keep `libs/` business-agnostic and reuse existing owners before adding abstractions. +- Treat `api/core/` as a migration-only area: do not add files or move implementations into it. Move responsibilities out; limit changes to existing callers to the adaptations needed for that migration. +- Keep transport parsing and serialization in controllers, orchestration in services, and domain policy in its domain owner. New or extracted implementations must stay outside `core/`. Keep `libs/` business-agnostic and reuse existing owners before adding abstractions. - Before changing controller schemas, generated API contracts, or `SystemFeatureModel`, read `controllers/API_SCHEMA_GUIDE.md`. Treat `/system-features` as a minimal unauthenticated bootstrap allowlist, not a general configuration registry. - Scope tenant-owned reads and writes by the complete owner chain, and propagate `tenant_id` across every affected layer. Reconstruct trusted internal references from validated database state after payload or async boundaries. - Keep write transactions explicit and bounded. Do not perform external I/O inside an open transaction unless a documented consistency contract requires it. diff --git a/api/commands/vector.py b/api/commands/vector.py index 437ea40f529..d0555332e5e 100644 --- a/api/commands/vector.py +++ b/api/commands/vector.py @@ -19,6 +19,7 @@ from models.dataset import Dataset, DatasetCollectionBinding, DatasetMetadata, D from models.dataset import Document as DatasetDocument from models.enums import DatasetMetadataType, IndexingStatus, SegmentStatus from models.model import App, AppAnnotationSetting, MessageAnnotation +from repositories.knowledge.dataset_read_repository import get_segment_child_chunks @click.command("vdb-migrate", help="Migrate vector db.") @@ -278,7 +279,7 @@ def migrate_knowledge_vector_database(): }, ) if dataset_document.doc_form == IndexStructureType.PARENT_CHILD_INDEX: - child_chunks = segment.get_child_chunks(session=db_session) + child_chunks = get_segment_child_chunks(segment, session=db_session) if child_chunks: child_documents = [] for child_chunk in child_chunks: diff --git a/api/controllers/common/rbac/__init__.py b/api/controllers/common/rbac/__init__.py index 6f227d5034d..7ed48814fc5 100644 --- a/api/controllers/common/rbac/__init__.py +++ b/api/controllers/common/rbac/__init__.py @@ -4,6 +4,7 @@ from .checks import RBAC_CHECKS_ATTR, RBACCheck, enforce_rbac_checks from .locators import ( AgentBehindApp, AgentId, + DatasetByDocument, DatasetByPipeline, DatasetId, PlainApp, @@ -16,6 +17,7 @@ __all__ = [ "RBAC_CHECKS_ATTR", "AgentBehindApp", "AgentId", + "DatasetByDocument", "DatasetByPipeline", "DatasetId", "PlainApp", diff --git a/api/controllers/common/rbac/locators.py b/api/controllers/common/rbac/locators.py index 16790059045..1fde75f517a 100644 --- a/api/controllers/common/rbac/locators.py +++ b/api/controllers/common/rbac/locators.py @@ -15,6 +15,7 @@ if TYPE_CHECKING: __all__ = [ "AgentBehindApp", "AgentId", + "DatasetByDocument", "DatasetByPipeline", "DatasetId", "PlainApp", @@ -133,6 +134,18 @@ class DatasetId(_ParamLocator): return RBACResourceService.get_dataset_maintainer(db.session, tenant_id, identity.id) +class DatasetByDocument(DatasetId): + default_param = "document_id" + + @override + def locate(self, tenant_id: str, path_args: Mapping[str, object]) -> ResourceIdentity | None: + document_id = _required(path_args, self.param) + dataset_id = RBACResourceService.get_dataset_id_by_document(db.session, tenant_id, document_id) + if dataset_id is None: + raise NotFound("Document not found") + return ResourceIdentity(self.scope, dataset_id) + + class DatasetByPipeline(DatasetId): default_param = "pipeline_id" diff --git a/api/controllers/console/app/app.py b/api/controllers/console/app/app.py index 58019cbdda1..e1db9164dca 100644 --- a/api/controllers/console/app/app.py +++ b/api/controllers/console/app/app.py @@ -79,7 +79,8 @@ from services.entities.app_entities import ( UpdateAppParams, ) from services.entities.dsl_entities import ImportStatus -from services.entities.knowledge_entities.knowledge_entities import ( +from services.errors.account import NoPermissionError +from services.knowledge.entities.knowledge_entities import ( DataSource, InfoList, NotionIcon, @@ -91,7 +92,6 @@ from services.entities.knowledge_entities.knowledge_entities import ( WeightModel, WeightVectorSetting, ) -from services.errors.account import NoPermissionError register_enum_models(console_ns, IconType) diff --git a/api/controllers/console/auth/data_source_bearer_auth.py b/api/controllers/console/auth/data_source_bearer_auth.py index f42aca2997e..279748d7720 100644 --- a/api/controllers/console/auth/data_source_bearer_auth.py +++ b/api/controllers/console/auth/data_source_bearer_auth.py @@ -27,7 +27,7 @@ from services.auth.errors import ( InvalidDataSourceApiKeyAuthCredentialsError, UnsupportedDataSourceApiKeyAuthProviderError, ) -from services.entities.data_source_api_key_auth_entities import ( +from services.data_source.entities.api_key_auth import ( DataSourceApiKeyAuthBindingCreate, DataSourceApiKeyAuthCredentials, ) @@ -80,7 +80,7 @@ class ApiKeyAuthDataSource(Resource): @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[ApiKeyAuthDataSourceListResponse.__name__]) @console_account_admission() def get(self, request_context: RequestContext): - bindings = application_services().data_source_api_key_auth.list_bindings(request_context) + bindings = application_services().data_sources.api_key_auth.list_bindings(request_context) return dump_response( ApiKeyAuthDataSourceListResponse, { @@ -120,7 +120,7 @@ class ApiKeyAuthDataSourceBinding(Resource): ), ) try: - application_services().data_source_api_key_auth.create_binding(request_context, command) + application_services().data_sources.api_key_auth.create_binding(request_context, command) except UnsupportedDataSourceApiKeyAuthProviderError as exc: raise DataSourceApiKeyAuthProviderNotSupportedError() from exc except InvalidDataSourceApiKeyAuthCredentialsError as exc: @@ -140,5 +140,5 @@ class ApiKeyAuthDataSourceBindingDelete(Resource): rbac_checks=[RBACCheck(RBACPermission.CREDENTIAL_MANAGE, Workspace())], ) def delete(self, request_context: RequestContext, binding_id: UUID): - application_services().data_source_api_key_auth.delete_binding(request_context, str(binding_id)) + application_services().data_sources.api_key_auth.delete_binding(request_context, str(binding_id)) return "", HTTPStatus.NO_CONTENT diff --git a/api/controllers/console/auth/data_source_oauth.py b/api/controllers/console/auth/data_source_oauth.py index 8c4147ea7fa..0fce33ff038 100644 --- a/api/controllers/console/auth/data_source_oauth.py +++ b/api/controllers/console/auth/data_source_oauth.py @@ -24,13 +24,13 @@ from fields.base import ResponseModel from libs.helper import dump_response from machinery.context import RequestContext from models.account import TenantAccountRole -from services.data_source_oauth_service import ( +from services.data_source.entities.oauth import DataSourceOAuthCallback +from services.data_source.oauth_service import ( DataSourceOAuthConfigurationError, DataSourceOAuthError, InvalidDataSourceOAuthCodeError, InvalidDataSourceOAuthProviderError, ) -from services.entities.data_source_oauth_entities import DataSourceOAuthCallback from .. import console_ns @@ -116,7 +116,7 @@ class OAuthDataSource(Resource): ) def get(self, request_context: RequestContext, provider: str): try: - service = application_services().resolve_data_source_oauth(provider) + service = application_services().data_sources.resolve_oauth(provider) authorization = service.start_authorization(request_context) except InvalidDataSourceOAuthProviderError: return _invalid_provider_response() @@ -141,7 +141,7 @@ class OAuthDataSourceCallback(Resource): def get(self, provider: str): query = query_params_from_request(OAuthDataSourceCallbackQuery) try: - service = application_services().resolve_data_source_oauth(provider) + service = application_services().data_sources.resolve_oauth(provider) callback = service.complete_callback( code=query.code, error=query.error, @@ -172,7 +172,7 @@ class OAuthDataSourceBinding(Resource): if not query.code: return {"error": "Invalid code"}, HTTPStatus.BAD_REQUEST try: - service = application_services().resolve_data_source_oauth(provider) + service = application_services().data_sources.resolve_oauth(provider) service.bind(request_context, code=query.code) except InvalidDataSourceOAuthProviderError: return _invalid_provider_response() @@ -198,7 +198,7 @@ class OAuthDataSourceSync(Resource): @console_account_admission() def get(self, request_context: RequestContext, provider: str, binding_id: UUID): try: - service = application_services().resolve_data_source_oauth(provider) + service = application_services().data_sources.resolve_oauth(provider) service.sync( request_context, binding_id=str(binding_id), diff --git a/api/controllers/console/datasets/data_source.py b/api/controllers/console/datasets/data_source.py index e4ad4b3b819..2272ed89894 100644 --- a/api/controllers/console/datasets/data_source.py +++ b/api/controllers/console/datasets/data_source.py @@ -1,53 +1,59 @@ -import json -from collections.abc import Generator from datetime import datetime -from typing import Any, Literal, cast +from typing import Never from uuid import UUID -from flask import request from flask_restx import Resource from pydantic import BaseModel, Field, field_serializer -from sqlalchemy import select -from sqlalchemy.orm import Session -from werkzeug.exceptions import NotFound +from configs import dify_config +from controllers.common.errors import InvalidArgumentError, NotFoundError from controllers.common.fields import SimpleResultResponse, TextContentResponse from controllers.common.rbac import DatasetId, RBACCheck, Workspace from controllers.common.schema import query_params_from_model, register_response_schema_models, register_schema_models -from controllers.common.session import with_session -from core.datasource.entities.datasource_entities import DatasourceProviderType, OnlineDocumentPagesMessage -from core.datasource.online_document.online_document_plugin import OnlineDocumentDatasourcePlugin -from core.entities.knowledge_entities import IndexingEstimate -from core.indexing_runner import IndexingRunner -from core.rag.extractor.entity.datasource_type import DatasourceType -from core.rag.extractor.entity.extract_setting import ExtractSetting, NotionInfo -from core.rag.extractor.notion_extractor import NotionExtractor -from extensions.ext_database import db -from fields.base import ResponseModel -from libs.datetime_utils import naive_utc_now -from libs.helper import dump_response, to_timestamp -from libs.login import login_required -from models import Account, DataSourceOauthBinding, Document -from services.dataset_service import DatasetService, DocumentService -from services.datasource_provider_service import DatasourceProviderService -from tasks.document_indexing_sync_task import document_indexing_sync_task - -from .. import console_ns -from ..wraps import ( - RBACPermission, - account_initialization_required, - is_admin_or_owner_required, - model_validate, - rbac_permission_required, - setup_required, - with_current_tenant_id, - with_current_user, +from controllers.console import console_ns +from controllers.console.app.error import ProviderNotInitializeError +from controllers.console.datasets.error import ( + DatasetAccessDeniedRequestError, + IndexingEstimateError, + InvalidActionError, ) +from controllers.console.flask_admission import console_account_admission +from controllers.console.wraps import model_validate +from core.entities.knowledge_entities import IndexingEstimate +from core.rbac import RBACPermission +from extensions.ext_application_services import application_services +from fields.base import ResponseModel +from fields.dataset_fields import NotionEstimateWorkspacePayload +from libs.helper import dump_response, to_timestamp +from machinery.context import RequestContext +from models.account import TenantAccountRole +from services.data_source.binding_application_service import ( + BindingAction, + DataSourceBindingNotFoundError, + DataSourceBindingStateError, +) +from services.data_source.entities.notion_import import NotionPageType +from services.data_source.notion_import_application_service import ( + DatasetIsNotNotionSourceError, + NotionImportCredentialUnavailableError, +) +from services.knowledge.dataset_access import DatasetAccessDeniedError, DatasetNotFoundError +from services.knowledge.document_sync import SyncDocumentNotFoundError, SyncDocumentSourceError +from services.knowledge.entities.indexing_estimate import NewSourcesEstimateCommand, NotionEstimateSource +from services.knowledge.indexing.estimate import ( + EstimateSourceNotFoundError, + IndexingEstimateCredentialUnavailableError, + IndexingEstimateExecutionError, + IndexingEstimateProviderUnavailableError, +) + +_ADMIN_OR_OWNER_ROLES = frozenset({TenantAccountRole.ADMIN, TenantAccountRole.OWNER}) +_DATA_SOURCE_OAUTH_BASE_PATH = "/console/api/oauth/data-source" class NotionEstimatePayload(BaseModel): - notion_info_list: list[dict[str, Any]] - process_rule: dict[str, Any] + notion_info_list: list[NotionEstimateWorkspacePayload] = Field(min_length=1) + process_rule: dict[str, object] doc_form: str = Field(default="text_model") doc_language: str = Field(default="English") @@ -72,7 +78,7 @@ class DataSourceIntegratePageResponse(ResponseModel): page_id: str page_icon: DataSourceIntegrateIconResponse | None parent_id: str - type: str + type: NotionPageType class DataSourceIntegrateWorkspaceResponse(ResponseModel): @@ -106,7 +112,7 @@ class NotionIntegratePageResponse(ResponseModel): page_id: str page_icon: DataSourceIntegrateIconResponse | None parent_id: str | None - type: str + page_type: NotionPageType = Field(alias="type") is_bound: bool @@ -132,325 +138,204 @@ register_response_schema_models( ) -@console_ns.route( - "/data-source/integrates", - "/data-source/integrates//", -) -class DataSourceApi(Resource): - @setup_required - @login_required - @is_admin_or_owner_required - @rbac_permission_required(RBACCheck(RBACPermission.CREDENTIAL_MANAGE, Workspace())) - @account_initialization_required +def _raise_dataset_access_error(error: DatasetNotFoundError | DatasetAccessDeniedError) -> Never: + if isinstance(error, DatasetNotFoundError): + raise NotFoundError(description="Dataset not found.") from None + raise DatasetAccessDeniedRequestError(description=str(error)) from None + + +@console_ns.route("/data-source/integrates") +class DataSourceIntegrationListApi(Resource): @console_ns.response(200, "Success", console_ns.models[DataSourceIntegrateListResponse.__name__]) - @with_current_tenant_id - def get(self, current_tenant_id: str) -> tuple[dict[str, Any], int]: - # get workspace data source integrates - data_source_integrates = db.session.scalars( - select(DataSourceOauthBinding).where( - DataSourceOauthBinding.tenant_id == current_tenant_id, - DataSourceOauthBinding.disabled == False, - ) - ).all() + @console_account_admission( + allowed_roles=_ADMIN_OR_OWNER_ROLES, + rbac_checks=(RBACCheck(RBACPermission.CREDENTIAL_MANAGE, Workspace()),), + ) + def get(self, request_context: RequestContext) -> tuple[dict[str, object], int]: + bindings = application_services().data_sources.bindings.list_integrations(request_context) + base_url = dify_config.CONSOLE_API_URL.rstrip("/") + data = [ + { + "id": binding.id, + "provider": binding.provider, + "created_at": binding.created_at, + "is_bound": True, + "disabled": binding.disabled, + "source_info": binding.source_info, + "link": f"{base_url}{_DATA_SOURCE_OAUTH_BASE_PATH}/{binding.provider}", + } + for binding in bindings + ] + return dump_response(DataSourceIntegrateListResponse, {"data": data}), 200 - base_url = request.url_root.rstrip("/") - data_source_oauth_base_path = "/console/api/oauth/data-source" - providers = ["notion"] - integrate_data: list[dict[str, bool | datetime | dict[str, Any] | str | None]] = [] - for provider in providers: - # existing_integrate = next((ai for ai in data_source_integrates if ai.provider == provider), None) - existing_integrates = filter(lambda item: item.provider == provider, data_source_integrates) - if existing_integrates: - for existing_integrate in list(existing_integrates): - integrate_data.append( - { - "id": existing_integrate.id, - "provider": provider, - "created_at": existing_integrate.created_at, - "is_bound": True, - "disabled": existing_integrate.disabled, - "source_info": existing_integrate.source_info, - "link": f"{base_url}{data_source_oauth_base_path}/{provider}", - } - ) - else: - integrate_data.append( - { - "id": None, - "provider": provider, - "created_at": None, - "source_info": None, - "is_bound": False, - "disabled": None, - "link": f"{base_url}{data_source_oauth_base_path}/{provider}", - } - ) - return dump_response(DataSourceIntegrateListResponse, {"data": integrate_data}), 200 - - @setup_required - @login_required - @is_admin_or_owner_required - @rbac_permission_required(RBACCheck(RBACPermission.CREDENTIAL_MANAGE, Workspace())) - @account_initialization_required +@console_ns.route("/data-source/integrates//") +class DataSourceIntegrationApi(Resource): @console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__]) - @with_current_tenant_id - @with_session + @console_account_admission( + allowed_roles=_ADMIN_OR_OWNER_ROLES, + rbac_checks=(RBACCheck(RBACPermission.CREDENTIAL_MANAGE, Workspace()),), + ) def patch( - self, session: Session, current_tenant_id: str, binding_id: UUID, action: Literal["enable", "disable"] - ) -> tuple[dict[str, str], int]: - binding_id_str = str(binding_id) - data_source_binding = session.scalar( - select(DataSourceOauthBinding).where( - DataSourceOauthBinding.id == binding_id_str, DataSourceOauthBinding.tenant_id == current_tenant_id + self, + request_context: RequestContext, + binding_id: UUID, + action: str, + ) -> tuple[dict[str, object], int]: + try: + validated_action = BindingAction(action) + except ValueError as error: + raise InvalidActionError() from error + try: + application_services().data_sources.bindings.change_state( + request_context, + str(binding_id), + validated_action, ) - ) - if data_source_binding is None: - raise NotFound("Data source binding not found.") - # enable binding - match action: - case "enable": - if data_source_binding.disabled: - data_source_binding.disabled = False - data_source_binding.updated_at = naive_utc_now() - else: - raise ValueError("Data source is not disabled.") - # disable binding - case "disable": - if not data_source_binding.disabled: - data_source_binding.disabled = True - data_source_binding.updated_at = naive_utc_now() - else: - raise ValueError("Data source is disabled.") - return {"result": "success"}, 200 + except DataSourceBindingNotFoundError as error: + raise NotFoundError(description="Data source binding not found.") from error + except DataSourceBindingStateError as error: + raise InvalidArgumentError(description=str(error)) from error + return dump_response(SimpleResultResponse, {"result": "success"}), 200 @console_ns.route("/notion/pre-import/pages") class DataSourceNotionListApi(Resource): - @setup_required - @login_required - @account_initialization_required @console_ns.doc(params=query_params_from_model(DataSourceNotionListQuery)) @console_ns.response(200, "Success", console_ns.models[NotionIntegrateInfoListResponse.__name__]) - @with_current_user - @with_current_tenant_id - @with_session(write=False) + @console_account_admission() @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=req_data.credential_id, - provider="notion_datasource", - plugin_id="langgenius/notion_datasource", - ) - if not credential: - raise NotFound("Credential not found.") - exist_page_ids = [] - # import notion in the exist dataset - 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": - raise ValueError("Dataset is not notion type.") - - documents = session.scalars( - select(Document).where( - Document.dataset_id == req_data.dataset_id, - Document.tenant_id == current_tenant_id, - Document.data_source_type == "notion_import", - Document.enabled.is_(True), - ) - ).all() - if documents: - for document in documents: - data_source_info = json.loads(document.data_source_info) - exist_page_ids.append(data_source_info["notion_page_id"]) - # get all authorized pages - from core.datasource.datasource_manager import DatasourceManager - - datasource_runtime = DatasourceManager.get_datasource_runtime( - provider_id="langgenius/notion_datasource/notion_datasource", - datasource_name="notion_datasource", - tenant_id=current_tenant_id, - datasource_type=DatasourceProviderType.ONLINE_DOCUMENT, - ) - datasource_provider_service = DatasourceProviderService() - if credential: - datasource_runtime.runtime.credentials = credential - datasource_runtime = cast(OnlineDocumentDatasourcePlugin, datasource_runtime) - online_document_result: Generator[OnlineDocumentPagesMessage, None, None] = ( - datasource_runtime.get_online_document_pages( - user_id=current_user.id, - datasource_parameters={}, - provider_type=datasource_runtime.datasource_provider_type(), - ) - ) + request_context: RequestContext, + ) -> tuple[dict[str, object], int]: try: - pages = [] - workspace_info = {} - for message in online_document_result: - result = message.result - for info in result: - workspace_info = { - "workspace_id": info.workspace_id, - "workspace_name": info.workspace_name, - "workspace_icon": info.workspace_icon, - } - for page in info.pages: - page_info = { - "page_id": page.page_id, - "page_name": page.page_name, - "type": page.type, - "parent_id": page.parent_id, - "is_bound": page.page_id in exist_page_ids, - "page_icon": page.page_icon, - } - pages.append(page_info) - except Exception as e: - raise e - notion_info = [{**workspace_info, "pages": pages}] if workspace_info else [] - return dump_response(NotionIntegrateInfoListResponse, {"notion_info": notion_info}), 200 + result = application_services().data_sources.notion_imports.list_pages( + request_context, + credential_id=req_data.credential_id, + dataset_id=req_data.dataset_id, + ) + except NotionImportCredentialUnavailableError as error: + raise NotFoundError(description="Credential not found.") from error + except (DatasetNotFoundError, DatasetAccessDeniedError) as error: + _raise_dataset_access_error(error) + except DatasetIsNotNotionSourceError as error: + raise InvalidArgumentError(description="Dataset is not notion type.") from error + return dump_response(NotionIntegrateInfoListResponse, {"notion_info": result.workspaces}), 200 @console_ns.route("/notion/pages///preview") class DataSourceNotionPreviewApi(Resource): """Preview one authorized Notion page through the datasource credential.""" - @setup_required - @login_required - @account_initialization_required @console_ns.doc(params=query_params_from_model(DataSourceNotionPreviewQuery)) @console_ns.response(200, "Success", console_ns.models[TextContentResponse.__name__]) - @with_current_tenant_id + @console_account_admission() @model_validate(DataSourceNotionPreviewQuery) def get( self, req_data: DataSourceNotionPreviewQuery, - current_tenant_id: str, + request_context: RequestContext, 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=req_data.credential_id, - provider="notion_datasource", - plugin_id="langgenius/notion_datasource", - ) - - page_id_str = str(page_id) - - extractor = NotionExtractor( - notion_workspace_id="", - notion_obj_id=page_id_str, - notion_page_type=page_type, - notion_access_token=credential.get("integration_secret"), - tenant_id=current_tenant_id, - ) - - text_docs = extractor.extract() - return {"content": "\n".join([doc.page_content for doc in text_docs])}, 200 + ) -> tuple[dict[str, object], int]: + try: + validated_page_type = NotionPageType(page_type) + except ValueError as error: + raise InvalidArgumentError(description="Invalid Notion page type.") from error + try: + content = application_services().data_sources.notion_imports.preview_page( + request_context, + credential_id=req_data.credential_id, + page_id=str(page_id), + page_type=validated_page_type, + ) + except NotionImportCredentialUnavailableError as error: + raise NotFoundError(description="Credential not found.") from error + return dump_response(TextContentResponse, {"content": content}), 200 @console_ns.route("/datasets/notion-indexing-estimate") class DataSourceNotionIndexingEstimateApi(Resource): """Estimate indexing work for selected Notion pages.""" - @setup_required - @login_required - @account_initialization_required @console_ns.expect(console_ns.models[NotionEstimatePayload.__name__]) @console_ns.response(200, "Success", console_ns.models[IndexingEstimate.__name__]) - @with_current_tenant_id - @with_session + @console_account_admission() @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 = req_data.notion_info_list - extract_settings = [] - for notion_info in notion_info_list: - workspace_id = notion_info["workspace_id"] - credential_id = notion_info.get("credential_id") - for page in notion_info["pages"]: - extract_setting = ExtractSetting( - datasource_type=DatasourceType.NOTION, - notion_info=NotionInfo.model_validate( - { - "credential_id": credential_id, - "notion_workspace_id": workspace_id, - "notion_obj_id": page["page_id"], - "notion_page_type": page["type"], - "tenant_id": current_tenant_id, - } - ), - document_model=args["doc_form"], + request_context: RequestContext, + ) -> tuple[dict[str, object], int]: + command = NewSourcesEstimateCommand( + sources=tuple( + NotionEstimateSource( + workspace_id=workspace.workspace_id, + page_id=page.page_id, + page_type=page.page_type, + credential_id=workspace.credential_id, ) - extract_settings.append(extract_setting) - indexing_runner = IndexingRunner() - response = indexing_runner.indexing_estimate( - tenant_id=current_tenant_id, - extract_settings=extract_settings, - tmp_processing_rule=args["process_rule"], - doc_form=args["doc_form"], - doc_language=args["doc_language"], - session=session, + for workspace in req_data.notion_info_list + for page in workspace.pages + ), + process_rule=req_data.process_rule, + doc_form=req_data.doc_form, + doc_language=req_data.doc_language, ) + try: + response = application_services().knowledge.indexing_estimates.estimate_new_sources( + request_context, + command, + ) + except IndexingEstimateCredentialUnavailableError as error: + raise NotFoundError(description="Credential not found.") from error + except EstimateSourceNotFoundError as error: + raise NotFoundError(description=str(error)) from error + except IndexingEstimateProviderUnavailableError as error: + raise ProviderNotInitializeError(str(error)) from error + except IndexingEstimateExecutionError as error: + raise IndexingEstimateError(str(error)) from error return dump_response(IndexingEstimate, response), 200 @console_ns.route("/datasets//notion/sync") class DataSourceNotionDatasetSyncApi(Resource): - @setup_required - @login_required - @account_initialization_required @console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__]) - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId())) - @with_session(write=False) - def get(self, session: Session, dataset_id: UUID) -> tuple[dict[str, str], int]: - dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, session) - if dataset is None: - raise NotFound("Dataset not found.") - - documents = DocumentService.get_document_by_dataset_id(dataset_id_str, session) - for document in documents: - document_indexing_sync_task.delay(dataset_id_str, document.id) - return {"result": "success"}, 200 + @console_account_admission( + rbac_checks=(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()),), + ) + def get(self, request_context: RequestContext, dataset_id: UUID) -> tuple[dict[str, object], int]: + try: + application_services().knowledge.document_sync.sync_dataset(request_context, str(dataset_id)) + except (DatasetNotFoundError, DatasetAccessDeniedError) as error: + _raise_dataset_access_error(error) + return dump_response(SimpleResultResponse, {"result": "success"}), 200 @console_ns.route("/datasets//documents//notion/sync") class DataSourceNotionDocumentSyncApi(Resource): - @setup_required - @login_required - @account_initialization_required @console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__]) - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId())) - @with_session(write=False) - def get(self, session: Session, dataset_id: UUID, document_id: UUID) -> tuple[dict[str, str], int]: - dataset_id_str = str(dataset_id) - document_id_str = str(document_id) - dataset = DatasetService.get_dataset(dataset_id_str, session) - if dataset is None: - raise NotFound("Dataset not found.") - - document = DocumentService.get_document(dataset_id_str, document_id_str, session=session) - if document is None: - raise NotFound("Document not found.") - document_indexing_sync_task.delay(dataset_id_str, document_id_str) - return {"result": "success"}, 200 + @console_account_admission( + rbac_checks=(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()),), + ) + def get( + self, + request_context: RequestContext, + dataset_id: UUID, + document_id: UUID, + ) -> tuple[dict[str, object], int]: + try: + application_services().knowledge.document_sync.sync_document( + request_context, + str(dataset_id), + str(document_id), + ) + except (DatasetNotFoundError, DatasetAccessDeniedError) as error: + _raise_dataset_access_error(error) + except SyncDocumentNotFoundError as error: + raise NotFoundError(description="Document not found.") from error + except SyncDocumentSourceError as error: + raise InvalidArgumentError(description=str(error)) from error + return dump_response(SimpleResultResponse, {"result": "success"}), 200 diff --git a/api/controllers/console/datasets/datasets.py b/api/controllers/console/datasets/datasets.py index 24a3dd81d50..fa4f19fd10e 100644 --- a/api/controllers/console/datasets/datasets.py +++ b/api/controllers/console/datasets/datasets.py @@ -1,92 +1,89 @@ -from dataclasses import dataclass from datetime import datetime from http import HTTPStatus -from typing import Any +from typing import Any, Never from uuid import UUID from flask import request from flask_restx import Resource -from pydantic import BaseModel, Field, field_validator, model_validator -from sqlalchemy import select -from sqlalchemy.orm import Session +from pydantic import BaseModel, Field, TypeAdapter, field_validator, model_validator from werkzeug.exceptions import Forbidden, NotFound -import services -from configs import dify_config +from controllers.common.errors import InvalidArgumentError, NotFoundError from controllers.common.fields import ApiBaseUrlResponse, SimpleResultResponse, UsageCheckResponse from controllers.common.rbac import DatasetId, RBACCheck, Workspace, enforce_rbac_checks from controllers.common.schema import query_params_from_model, register_response_schema_models, register_schema_models -from controllers.common.session import with_session from controllers.console import console_ns from controllers.console.apikey import API_KEY_DELETE_ROLES, api_key_errors from controllers.console.app.error import ProviderNotInitializeError -from controllers.console.datasets.error import DatasetInUseError, DatasetNameDuplicateError, IndexingEstimateError +from controllers.console.datasets.error import ( + DatasetAccessDeniedRequestError, + DatasetInUseError, + DatasetNameDuplicateError, + IndexingEstimateError, +) from controllers.console.flask_admission import console_account_admission from controllers.console.wraps import ( RBACPermission, - account_initialization_required, cloud_edition_billing_rate_limit_check, - enterprise_license_required, model_validate, - rbac_permission_required, - setup_required, - with_current_tenant_id, - with_current_user, ) from core.entities.knowledge_entities import IndexingEstimate -from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError -from core.indexing_runner import IndexingRunner -from core.plugin.impl.model_runtime_factory import create_plugin_provider_manager -from core.rag.datasource.vdb.vector_type import VectorType -from core.rag.extractor.entity.datasource_type import DatasourceType -from core.rag.extractor.entity.extract_setting import ExtractSetting, NotionInfo, WebsiteInfo -from core.rag.index_processor.constant.index_type import IndexTechniqueType -from core.rag.retrieval.retrieval_methods import RetrievalMethod +from core.rag.index_processor.constant.index_type import IndexStructureType from extensions.ext_application_services import application_services from fields.api_key_fields import ApiKeyItem, ApiKeyList, build_masked_api_key_list from fields.base import ResponseModel from fields.dataset_fields import ( DatasetDetailResponse, - build_dataset_detail_prefetch, - dataset_detail_response_source, + NotionEstimateWorkspacePayload, ) -from graphon.model_runtime.entities.model_entities import ModelType from libs.helper import build_icon_url, dump_response, to_timestamp -from libs.login import login_required -from libs.pagination import clamp_pagination -from libs.url_utils import normalize_api_base_url from machinery.context import RequestContext -from models import Account, Dataset, Document, UploadFile -from models.dataset import DatasetPermission, DatasetPermissionEnum, DatasetQuery -from models.provider_ids import ModelProviderID -from services.dataset_ref_service import DatasetRefService -from services.dataset_service import DatasetPermissionService, DatasetService, DocumentService -from services.enterprise import rbac_service as enterprise_rbac_service +from models.account import TenantAccountRole +from models.enums import PermissionEnum as DatasetPermissionEnum +from services.errors.dataset import DatasetInUseError as DatasetInUseFailure +from services.errors.dataset import DatasetNameDuplicateError as DatasetNameDuplicateFailure +from services.knowledge.dataset_access import DatasetAccessDeniedError, DatasetNotFoundError +from services.knowledge.datasets.application import DatasetListFilter +from services.knowledge.entities.indexing_estimate import ( + NewEstimateSource, + NewSourcesEstimateCommand, + NotionEstimateSource, + UploadFileEstimateSource, + WebsiteEstimateSource, +) +from services.knowledge.entities.knowledge_entities import FileInfo, WebsiteInfo +from services.knowledge.indexing.estimate import ( + EstimateSourceNotFoundError, + IndexingEstimateCredentialUnavailableError, + IndexingEstimateExecutionError, + IndexingEstimateProviderUnavailableError, + UnsupportedEstimateSourceError, +) + +_DATASET_EDIT_ROLES = frozenset( + {TenantAccountRole.OWNER, TenantAccountRole.ADMIN, TenantAccountRole.EDITOR, TenantAccountRole.DATASET_OPERATOR} +) + + +def _raise_dataset_error(error: Exception) -> Never: + if isinstance(error, DatasetNotFoundError): + raise NotFound("Dataset not found.") from error + if isinstance(error, DatasetAccessDeniedError): + raise Forbidden(str(error)) from error + if isinstance(error, DatasetNameDuplicateFailure): + raise DatasetNameDuplicateError() from error + if isinstance(error, DatasetInUseFailure): + raise DatasetInUseError() from error + raise error + register_response_schema_models(console_ns, ApiBaseUrlResponse, SimpleResultResponse, UsageCheckResponse) -DATASET_LIST_PERMISSION_KEYS = frozenset({"dataset.preview", "dataset.acl.preview", "dataset.full_access"}) - - -def _has_dataset_list_permission(permission_keys: list[str]) -> bool: - return any(permission_key in DATASET_LIST_PERMISSION_KEYS for permission_key in permission_keys) - - -def _get_accessible_dataset(dataset_id: UUID, tenant_id: str, current_user: Account, session: Session) -> Dataset: - dataset = DatasetService.get_dataset_for_tenant(str(dataset_id), tenant_id, session=session) - if dataset is None: - raise NotFound("Dataset not found.") - try: - DatasetService.check_dataset_permission(dataset, current_user, session) - except services.errors.account.NoPermissionError as e: - raise Forbidden(str(e)) - return dataset - def _validate_indexing_technique(value: str | None) -> str | None: if value is None: return value - if value not in Dataset.INDEXING_TECHNIQUE_LIST: + if value not in {"high_quality", "economy"}: raise ValueError("Invalid indexing technique.") return value @@ -94,7 +91,7 @@ def _validate_indexing_technique(value: str | None) -> str | None: def _validate_doc_form(value: str | None) -> str | None: if value is None: return value - if value not in Dataset.DOC_FORM_LIST: + if value not in {member.value for member in IndexStructureType}: raise ValueError("Invalid doc_form.") return value @@ -116,7 +113,7 @@ class DatasetCreatePayload(BaseModel): @field_validator("provider") @classmethod def validate_provider(cls, value: str) -> str: - if value not in Dataset.PROVIDER_LIST: + if value not in {"vendor", "external"}: raise ValueError("Invalid provider.") return value @@ -168,6 +165,44 @@ class IndexingEstimatePayload(BaseModel): return result +_NOTION_SELECTIONS = TypeAdapter(list[NotionEstimateWorkspacePayload]) + + +def _new_estimate_sources(info_list: dict[str, Any]) -> tuple[NewEstimateSource, ...]: + match info_list.get("data_source_type"): + case "upload_file": + files = FileInfo.model_validate(info_list.get("file_info_list")) + return tuple(UploadFileEstimateSource(file_id=file_id) for file_id in dict.fromkeys(files.file_ids)) + case "notion_import": + workspaces = _NOTION_SELECTIONS.validate_python(info_list.get("notion_info_list")) + return tuple( + NotionEstimateSource( + workspace_id=workspace.workspace_id, + credential_id=workspace.credential_id, + page_id=page.page_id, + page_type=page.page_type, + ) + for workspace in workspaces + for page in workspace.pages + ) + case "website_crawl": + values = info_list.get("website_info_list") + if not isinstance(values, dict): + raise ValueError("Website info list is required") + website = WebsiteInfo.model_validate({"only_main_content": False, **values}) + return tuple( + WebsiteEstimateSource( + provider=website.provider, + job_id=website.job_id, + url=url, + only_main_content=website.only_main_content, + ) + for url in website.urls + ) + case _: + raise ValueError("Data source type not support") + + class DatasetApiKeyCreatePayload(BaseModel): # Knowledge bases to scope the key to. Absent/empty => the key can access every # dataset in the tenant (default). Declared so the generated client can send it. @@ -229,21 +264,6 @@ class DatasetQueryDetailResponse(ResponseModel): return to_timestamp(value) -@dataclass(frozen=True) -class _DatasetQueryResponseSource: - """Expose query content through the request's database session.""" - - query: DatasetQuery - session: Session - - @property - def queries(self) -> list[dict[str, Any]]: - return self.query.get_queries(session=self.session) - - def __getattr__(self, name: str) -> Any: - return getattr(self.query, name) # guard-ignore: no-new-getattr -- delegates model fields - - class DatasetQueryListResponse(ResponseModel): data: list[DatasetQueryDetailResponse] has_more: bool @@ -353,84 +373,6 @@ register_response_schema_models( ) -def _get_retrieval_methods_by_vector_type(vector_type: str | None, is_mock: bool = False) -> dict[str, list[str]]: - """ - Get supported retrieval methods based on vector database type. - - Args: - vector_type: Vector database type, can be None - is_mock: Whether this is a Mock API, affects MILVUS handling - - Returns: - Dictionary containing supported retrieval methods - - Raises: - ValueError: If vector_type is None or unsupported - """ - if vector_type is None: - raise ValueError("Vector store type is not configured.") - - # Define vector database types that only support semantic search - semantic_only_types = { - VectorType.RELYT, - VectorType.CHROMA, - VectorType.PGVECTO_RS, - VectorType.VIKINGDB, - VectorType.UPSTASH, - } - - # Define vector database types that support all retrieval methods - full_search_types = { - VectorType.QDRANT, - VectorType.WEAVIATE, - VectorType.OPENSEARCH, - VectorType.ANALYTICDB, - VectorType.MYSCALE, - VectorType.ORACLE, - VectorType.ELASTICSEARCH, - VectorType.ELASTICSEARCH_JA, - VectorType.PGVECTOR, - VectorType.VASTBASE, - VectorType.TIDB_ON_QDRANT, - VectorType.LINDORM, - VectorType.COUCHBASE, - VectorType.OPENGAUSS, - VectorType.OCEANBASE, - VectorType.SEEKDB, - VectorType.TABLESTORE, - VectorType.HUAWEI_CLOUD, - VectorType.TENCENT, - VectorType.MATRIXONE, - VectorType.CLICKZETTA, - VectorType.BAIDU, - VectorType.ALIBABACLOUD_MYSQL, - VectorType.IRIS, - VectorType.HOLOGRES, - } - - semantic_methods = {"retrieval_method": [RetrievalMethod.SEMANTIC_SEARCH.value]} - full_methods = { - "retrieval_method": [ - RetrievalMethod.SEMANTIC_SEARCH.value, - RetrievalMethod.FULL_TEXT_SEARCH.value, - RetrievalMethod.HYBRID_SEARCH.value, - ] - } - - if vector_type == VectorType.MILVUS: - return semantic_methods if is_mock else full_methods - - if vector_type == VectorType.TIDB_VECTOR: - return full_methods if dify_config.TIDB_VECTOR_ENABLE_FULLTEXT_SEARCH else semantic_methods - - if vector_type in semantic_only_types: - return semantic_methods - elif vector_type in full_search_types: - return full_methods - else: - raise ValueError(f"Unsupported vector db type {vector_type}.") - - @console_ns.route("/datasets") class DatasetListApi(Resource): @console_ns.doc("get_datasets") @@ -439,208 +381,39 @@ class DatasetListApi(Resource): @console_ns.response( HTTPStatus.OK, "Datasets retrieved successfully", console_ns.models[DatasetListResponse.__name__] ) - @setup_required - @login_required - @account_initialization_required - @enterprise_license_required - @with_current_user - @with_current_tenant_id - @with_session(write=False) - def get(self, session: Session, current_tenant_id: str, current_user: Account): - # Convert query parameters to dict, handling list parameters correctly + @console_account_admission(require_valid_enterprise_license=True) + def get(self, request_context: RequestContext): query_params: dict[str, str | list[str]] = dict(request.args.to_dict()) - # Handle ids and tag_ids as lists (Flask request.args.getlist returns list even for single value) - if "ids" in request.args: - query_params["ids"] = request.args.getlist("ids") - if "tag_ids" in request.args: - query_params["tag_ids"] = request.args.getlist("tag_ids") + for key in ("ids", "tag_ids"): + if key in request.args: + query_params[key] = request.args.getlist(key) query = ConsoleDatasetListQuery.model_validate(query_params) - - permissions = enterprise_rbac_service.RBACService.MyPermissions.get( - str(current_tenant_id), - current_user.id, - session=session, + result = application_services().knowledge.datasets.list_datasets( + request_context, DatasetListFilter(**query.model_dump()) ) - - accessible_dataset_ids: list[str] | None = None - include_own_datasets = False - if dify_config.RBAC_ENABLED: - whitelist_scope = enterprise_rbac_service.RBACService.DatasetAccess.whitelist_resources( - str(current_tenant_id), - current_user.id, - ) - has_default_readonly = _has_dataset_list_permission( - permissions.dataset.default_permission_keys - ) or _has_dataset_list_permission(permissions.workspace.permission_keys) - permission_dataset_ids: set[str] | None = None - if not has_default_readonly: - permission_dataset_ids = { - override.resource_id - for override in permissions.dataset.overrides - if _has_dataset_list_permission(override.permission_keys) - } - if getattr(whitelist_scope, "unrestricted", False): - filtered_dataset_ids = permission_dataset_ids - include_own_datasets = "dataset.create_and_management" in permissions.workspace.permission_keys - else: - # A restricted dataset whitelist is the highest-priority visibility gate: - # default readonly, per-dataset permission overrides, and own-dataset - # management must not expose datasets outside this set. - filtered_dataset_ids = set(whitelist_scope.resource_ids) - if filtered_dataset_ids is not None: - accessible_dataset_ids = sorted(filtered_dataset_ids) - - effective_page, effective_limit = clamp_pagination(query.page, query.limit, 100) - if query.ids: - datasets, total = DatasetService.get_datasets_by_ids( - query.ids, - current_tenant_id, - user=current_user, - accessible_dataset_ids=accessible_dataset_ids, - include_own_datasets=include_own_datasets, - session=session, - ) - # This branch resolves the ids it was handed in a single page - # (`per_page=len(ids)`), so there is never a next one to ask for. - has_more = False - else: - datasets, total = DatasetService.get_datasets( - effective_page, - effective_limit, - session, - current_tenant_id, - current_user, - query.keyword, - query.tag_ids, - query.include_all, - tags=application_services().tags, - accessible_dataset_ids=accessible_dataset_ids, - include_own_datasets=include_own_datasets, - ) - has_more = effective_page * effective_limit < total - - permission_keys_map = {} - if datasets: - dataset_ids = [str(dataset.id) for dataset in datasets] - permission_keys_map = permissions.dataset.permission_keys_by_resource_ids(dataset_ids) - - # check embedding setting - provider_manager = create_plugin_provider_manager(tenant_id=current_tenant_id) - configurations = provider_manager.get_configurations(tenant_id=current_tenant_id) - - embedding_models = configurations.get_models(model_type=ModelType.TEXT_EMBEDDING, only_active=True) - - model_names = [] - for embedding_model in embedding_models: - model_names.append(f"{embedding_model.model}:{embedding_model.provider.provider}") - - prefetch = build_dataset_detail_prefetch(datasets, session=session) - data = [ - dump_response( - DatasetDetailResponse, dataset_detail_response_source(dataset, session=session, prefetch=prefetch) - ) - for dataset in datasets - ] - dataset_ids = [item["id"] for item in data if item.get("permission") == "partial_members"] - partial_members_map: dict[str, list[str]] = {} - if dataset_ids: - partial_member_rows = session.execute( - select(DatasetPermission.dataset_id, DatasetPermission.account_id).where( - DatasetPermission.dataset_id.in_(dataset_ids) - ) - ).all() - - for dataset_id, account_id in partial_member_rows: - partial_members_map.setdefault(dataset_id, []).append(account_id) - - for item in data: - # convert embedding_model_provider to plugin standard format - if item["indexing_technique"] == IndexTechniqueType.HIGH_QUALITY and item["embedding_model_provider"]: - item["embedding_model_provider"] = str(ModelProviderID(item["embedding_model_provider"])) - item_model = f"{item['embedding_model']}:{item['embedding_model_provider']}" - if item_model in model_names: - item["embedding_available"] = True - else: - item["embedding_available"] = False - else: - item["embedding_available"] = True - - if item.get("permission") == "partial_members": - item.update({"partial_member_list": partial_members_map.get(item["id"], [])}) - else: - item.update({"partial_member_list": []}) - item["permission_keys"] = permission_keys_map.get(str(item["id"]), []) - - response = { - "data": data, - "has_more": has_more, - "limit": effective_limit, - "total": total, - "page": effective_page, - } - return dump_response(DatasetListResponse, response), HTTPStatus.OK + return dump_response(DatasetListResponse, result), HTTPStatus.OK @console_ns.doc("create_dataset") @console_ns.doc(description="Create a new dataset") @console_ns.expect(console_ns.models[DatasetCreatePayload.__name__]) - @console_ns.response(201, "Dataset created successfully", console_ns.models[DatasetDetailResponse.__name__]) + @console_ns.response( + 201, "Dataset created successfully", console_ns.models[DatasetDetailWithPartialMembersResponse.__name__] + ) @console_ns.response(400, "Invalid request parameters") - @setup_required - @login_required - @account_initialization_required - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_CREATE_AND_MANAGEMENT, Workspace())) + @console_account_admission( + allowed_roles=_DATASET_EDIT_ROLES, + rbac_checks=(RBACCheck(RBACPermission.DATASET_CREATE_AND_MANAGEMENT, Workspace()),), + ) @cloud_edition_billing_rate_limit_check("knowledge") - @with_current_user - @with_current_tenant_id - @with_session @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() - - if dify_config.RBAC_ENABLED: - permission = DatasetPermissionEnum.ALL_TEAM - else: - permission = req_data.permission or DatasetPermissionEnum.ONLY_ME - + def post(self, req_data: DatasetCreatePayload, request_context: RequestContext): try: - dataset = DatasetService.create_empty_dataset( - session=session, - tenant_id=current_tenant_id, - name=req_data.name, - description=req_data.description, - indexing_technique=req_data.indexing_technique, - account=current_user, - permission=permission, - provider=req_data.provider, - external_knowledge_api_id=req_data.external_knowledge_api_id, - external_knowledge_id=req_data.external_knowledge_id, + result = application_services().knowledge.datasets.create_dataset( + request_context, values=req_data.model_dump() ) - except services.errors.dataset.DatasetNameDuplicateError: - raise DatasetNameDuplicateError() - - permission_keys_map = enterprise_rbac_service.RBACService.DatasetPermissions.batch_get( - current_tenant_id, - current_user.id, - [dataset.id], - session=session, - ) - - item = DatasetDetailWithPartialMembersResponse.model_validate( - dataset_detail_response_source(dataset, session=session), from_attributes=True - ).model_dump(mode="json") - item["permission_keys"] = permission_keys_map.get(dataset.id, []) - - if dify_config.RBAC_ENABLED: - enterprise_rbac_service.RBACService.DatasetAccess.replace_whitelist( - current_tenant_id, - current_user.id, - dataset.id, - enterprise_rbac_service.ReplaceMemberBindings(automatic_include_workspace_members=False), - ) - - return item, 201 + except Exception as error: + _raise_dataset_error(error) + return dump_response(DatasetDetailWithPartialMembersResponse, result), 201 @console_ns.route("/datasets/") @@ -655,59 +428,13 @@ class DatasetApi(Resource): ) @console_ns.response(404, "Dataset not found") @console_ns.response(403, "Permission denied") - @setup_required - @login_required - @account_initialization_required - @with_current_user - @with_current_tenant_id - @with_session(write=False) - def get(self, 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.") - + @console_account_admission() + def get(self, request_context: RequestContext, dataset_id: UUID): try: - DatasetService.check_dataset_permission(dataset, current_user, session) - except services.errors.account.NoPermissionError as e: - raise Forbidden(str(e)) - permissions = enterprise_rbac_service.RBACService.MyPermissions.get( - current_tenant_id, - current_user.id, - dataset_id=dataset_id_str, - session=session, - ) - permission_keys_map = permissions.dataset.permission_keys_by_resource_ids([dataset_id_str]) - data = dump_response(DatasetDetailResponse, dataset_detail_response_source(dataset, session=session)) - data["permission_keys"] = permission_keys_map.get(dataset_id_str, []) - if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY: - if dataset.embedding_model_provider: - provider_id = ModelProviderID(dataset.embedding_model_provider) - data["embedding_model_provider"] = str(provider_id) - if data.get("permission") == "partial_members": - part_users_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str, session) - data.update({"partial_member_list": part_users_list}) - - # check embedding setting - provider_manager = create_plugin_provider_manager(tenant_id=current_tenant_id) - configurations = provider_manager.get_configurations(tenant_id=current_tenant_id) - - embedding_models = configurations.get_models(model_type=ModelType.TEXT_EMBEDDING, only_active=True) - - model_names = [] - for embedding_model in embedding_models: - model_names.append(f"{embedding_model.model}:{embedding_model.provider.provider}") - - if data["indexing_technique"] == IndexTechniqueType.HIGH_QUALITY: - item_model = f"{data['embedding_model']}:{data['embedding_model_provider']}" - if item_model in model_names: - data["embedding_available"] = True - else: - data["embedding_available"] = False - else: - data["embedding_available"] = True - - return dump_response(DatasetDetailWithPartialMembersResponse, data), 200 + result = application_services().knowledge.datasets.get_dataset(request_context, dataset_id=str(dataset_id)) + except Exception as error: + _raise_dataset_error(error) + return dump_response(DatasetDetailWithPartialMembersResponse, result), 200 @console_ns.doc("update_dataset") @console_ns.doc(description="Update dataset details") @@ -719,95 +446,31 @@ class DatasetApi(Resource): ) @console_ns.response(404, "Dataset not found") @console_ns.response(403, "Permission denied") - @setup_required - @login_required - @account_initialization_required + @console_account_admission( + allowed_roles=_DATASET_EDIT_ROLES, rbac_checks=(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()),) + ) @cloud_edition_billing_rate_limit_check("knowledge") - @with_current_user - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId())) - @with_session @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.") - - # check embedding model setting - if ( - 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, req_data.embedding_model_provider, req_data.embedding_model - ) - 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, req_data.permission, req_data.partial_member_list, session=session - ) - - dataset = DatasetService.update_dataset(dataset_id_str, payload_data, current_user, session=session) - - if dataset is None: - raise NotFound("Dataset not found.") - - permission_keys_map = enterprise_rbac_service.RBACService.DatasetPermissions.batch_get( - current_tenant_id, - current_user.id, - [dataset_id_str], - session=session, - ) - result_data = dump_response(DatasetDetailResponse, dataset_detail_response_source(dataset, session=session)) - result_data["permission_keys"] = permission_keys_map.get(dataset_id_str, []) - tenant_id = current_tenant_id - - 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, req_data.partial_member_list, session - ) - # clear partial member list when permission is only_me or all_team_members - 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) - result_data.update({"partial_member_list": partial_member_list}) - - return dump_response(DatasetDetailWithPartialMembersResponse, result_data), 200 - - @setup_required - @login_required - @account_initialization_required - @cloud_edition_billing_rate_limit_check("knowledge") - @console_ns.response(204, "Dataset deleted successfully") - @with_current_user - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_DELETE, DatasetId())) - @with_session - def delete(self, session: Session, current_user: Account, dataset_id: UUID): - dataset_id_str = str(dataset_id) - - if not (current_user.has_edit_permission or current_user.is_dataset_operator): - raise Forbidden() - + def patch(self, req_data: DatasetUpdatePayload, request_context: RequestContext, dataset_id: UUID): try: - if DatasetService.delete_dataset(dataset_id_str, current_user, session): - DatasetPermissionService.clear_partial_member_list(dataset_id_str, session) - return "", 204 - else: - raise NotFound("Dataset not found.") - except services.errors.dataset.DatasetInUseError: - raise DatasetInUseError() + result = application_services().knowledge.datasets.update_dataset( + request_context, dataset_id=str(dataset_id), values=req_data.model_dump(exclude_unset=True) + ) + except Exception as error: + _raise_dataset_error(error) + return dump_response(DatasetDetailWithPartialMembersResponse, result), 200 + + @console_ns.response(204, "Dataset deleted successfully") + @console_account_admission( + allowed_roles=_DATASET_EDIT_ROLES, rbac_checks=(RBACCheck(RBACPermission.DATASET_DELETE, DatasetId()),) + ) + @cloud_edition_billing_rate_limit_check("knowledge") + def delete(self, request_context: RequestContext, dataset_id: UUID): + try: + application_services().knowledge.datasets.delete_dataset(request_context, dataset_id=str(dataset_id)) + except Exception as error: + _raise_dataset_error(error) + return "", 204 @console_ns.route("/datasets//use-check") @@ -820,17 +483,13 @@ class DatasetUseCheckApi(Resource): "Dataset use status retrieved successfully", console_ns.models[UsageCheckResponse.__name__], ) - @setup_required - @login_required - @account_initialization_required - @with_current_user - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_READONLY, DatasetId())) - @with_session(write=False) - def get(self, session: Session, current_tenant_id: str, current_user: Account, dataset_id: UUID): - dataset = _get_accessible_dataset(dataset_id, current_tenant_id, current_user, session) - dataset_is_using = DatasetService.dataset_use_check(DatasetRefService.create_dataset_ref(dataset), session) - return UsageCheckResponse(is_using=dataset_is_using).model_dump(mode="json"), 200 + @console_account_admission(rbac_checks=(RBACCheck(RBACPermission.DATASET_READONLY, DatasetId()),)) + def get(self, request_context: RequestContext, dataset_id: UUID): + try: + result = application_services().knowledge.datasets.is_in_use(request_context, dataset_id=str(dataset_id)) + except Exception as error: + _raise_dataset_error(error) + return dump_response(UsageCheckResponse, {"is_using": result}), 200 @console_ns.route("/datasets//queries") @@ -843,39 +502,18 @@ class DatasetQueryApi(Resource): "Query history retrieved successfully", console_ns.models[DatasetQueryListResponse.__name__], ) - @setup_required - @login_required - @account_initialization_required - @with_current_user - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_READONLY, DatasetId())) - @with_session(write=False) - def get(self, 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.") - + @console_account_admission(rbac_checks=(RBACCheck(RBACPermission.DATASET_READONLY, DatasetId()),)) + def get(self, request_context: RequestContext, dataset_id: UUID): try: - DatasetService.check_dataset_permission(dataset, current_user, session) - except services.errors.account.NoPermissionError as e: - raise Forbidden(str(e)) - - page = request.args.get("page", default=1, type=int) - limit = request.args.get("limit", default=20, type=int) - - effective_page, effective_limit = clamp_pagination(page, limit, 100) - dataset_queries, total = DatasetService.get_dataset_queries( - dataset_id=dataset.id, page=effective_page, per_page=effective_limit, session=session - ) - - response = { - "data": [_DatasetQueryResponseSource(query=query, session=session) for query in dataset_queries], - "has_more": effective_page * effective_limit < total, - "limit": effective_limit, - "total": total, - "page": effective_page, - } - return dump_response(DatasetQueryListResponse, response), 200 + result = application_services().knowledge.datasets.queries( + request_context, + dataset_id=str(dataset_id), + page=request.args.get("page", default=1, type=int), + limit=request.args.get("limit", default=20, type=int), + ) + except Exception as error: + _raise_dataset_error(error) + return dump_response(DatasetQueryListResponse, result), 200 @console_ns.route("/datasets/indexing-estimate") @@ -887,115 +525,49 @@ class DatasetIndexingEstimateApi(Resource): "Indexing estimate calculated successfully", console_ns.models[IndexingEstimateResponse.__name__], ) - @setup_required - @login_required - @account_initialization_required @console_ns.expect(console_ns.models[IndexingEstimatePayload.__name__]) - @with_current_user - @with_current_tenant_id - @with_session + @console_account_admission() @model_validate(IndexingEstimatePayload) - def post( - self, - req_data: IndexingEstimatePayload, - session: Session, - current_tenant_id: str, - current_user: Account, - ): - args = req_data.model_dump() - dataset_id = req_data.dataset_id - if dataset_id: + def post(self, req_data: IndexingEstimatePayload, request_context: RequestContext): + if req_data.dataset_id: checks = [RBACCheck(RBACPermission.DATASET_USE, DatasetId())] - path_args = {"dataset_id": dataset_id} + path_args = {"dataset_id": req_data.dataset_id} else: checks = [RBACCheck(RBACPermission.DATASET_CREATE_AND_MANAGEMENT, Workspace())] path_args = None enforce_rbac_checks( - tenant_id=current_tenant_id, - account_id=current_user.id, + tenant_id=request_context.active_workspace_id, + account_id=request_context.account_id, checks=checks, path_args=path_args, ) - # validate args - DocumentService.estimate_args_validate(args) - extract_settings = [] - match args["info_list"]["data_source_type"]: - case "upload_file": - file_ids = args["info_list"]["file_info_list"]["file_ids"] - file_details = session.scalars( - select(UploadFile).where(UploadFile.tenant_id == current_tenant_id, UploadFile.id.in_(file_ids)) - ).all() - if not file_details: - raise NotFound("File not found.") - - if file_details: - for file_detail in file_details: - extract_setting = ExtractSetting( - datasource_type=DatasourceType.FILE, - upload_file=file_detail, - document_model=args["doc_form"], - ) - extract_settings.append(extract_setting) - case "notion_import": - notion_info_list = args["info_list"]["notion_info_list"] - for notion_info in notion_info_list: - workspace_id = notion_info["workspace_id"] - credential_id = notion_info.get("credential_id") - for page in notion_info["pages"]: - extract_setting = ExtractSetting( - datasource_type=DatasourceType.NOTION, - notion_info=NotionInfo.model_validate( - { - "credential_id": credential_id, - "notion_workspace_id": workspace_id, - "notion_obj_id": page["page_id"], - "notion_page_type": page["type"], - "tenant_id": current_tenant_id, - } - ), - document_model=args["doc_form"], - ) - extract_settings.append(extract_setting) - case "website_crawl": - website_info_list = args["info_list"]["website_info_list"] - for url in website_info_list["urls"]: - extract_setting = ExtractSetting( - datasource_type=DatasourceType.WEBSITE, - website_info=WebsiteInfo.model_validate( - { - "provider": website_info_list["provider"], - "job_id": website_info_list["job_id"], - "url": url, - "tenant_id": current_tenant_id, - "mode": "crawl", - "only_main_content": website_info_list["only_main_content"], - } - ), - document_model=args["doc_form"], - ) - extract_settings.append(extract_setting) - case _: - raise ValueError("Data source type not support") - indexing_runner = IndexingRunner() + command = NewSourcesEstimateCommand( + sources=_new_estimate_sources(req_data.info_list), + process_rule=req_data.process_rule, + doc_form=req_data.doc_form, + doc_language=req_data.doc_language, + dataset_id=req_data.dataset_id, + indexing_technique=req_data.indexing_technique, + ) try: - response = indexing_runner.indexing_estimate( - tenant_id=current_tenant_id, - extract_settings=extract_settings, - tmp_processing_rule=args["process_rule"], - doc_form=args["doc_form"], - doc_language=args["doc_language"], - dataset_id=args["dataset_id"], - indexing_technique=args["indexing_technique"], - session=session, + response = application_services().knowledge.indexing_estimates.estimate_new_sources( + request_context, + command, ) - except LLMBadRequestError: - raise ProviderNotInitializeError( - "No Embedding Model available. Please configure a valid provider in the Settings -> Model Provider." - ) - except ProviderTokenNotInitError as ex: - raise ProviderNotInitializeError(ex.description) - except Exception as e: - raise IndexingEstimateError(str(e)) + except ( + IndexingEstimateCredentialUnavailableError, + EstimateSourceNotFoundError, + DatasetNotFoundError, + ) as error: + raise NotFoundError(description=str(error)) from error + except DatasetAccessDeniedError as error: + raise DatasetAccessDeniedRequestError(description=str(error)) from error + except UnsupportedEstimateSourceError as error: + raise InvalidArgumentError(description=str(error)) from error + except IndexingEstimateProviderUnavailableError as error: + raise ProviderNotInitializeError(str(error)) from error + except IndexingEstimateExecutionError as error: + raise IndexingEstimateError(str(error)) from error return ( IndexingEstimateResponse( @@ -1020,30 +592,13 @@ class DatasetRelatedAppListApi(Resource): "Related apps retrieved successfully", console_ns.models[RelatedAppListResponse.__name__], ) - @setup_required - @login_required - @account_initialization_required - @with_current_user - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_READONLY, DatasetId())) - @with_session(write=False) - def get(self, 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.") - + @console_account_admission(rbac_checks=(RBACCheck(RBACPermission.DATASET_READONLY, DatasetId()),)) + def get(self, request_context: RequestContext, dataset_id: UUID): try: - DatasetService.check_dataset_permission(dataset, current_user, session) - except services.errors.account.NoPermissionError as e: - raise Forbidden(str(e)) - - app_dataset_joins = DatasetService.get_related_apps(dataset.id, session) - - related_apps = application_services().apps.queries.related_apps( - dataset.tenant_id, [join.app_id for join in app_dataset_joins] - ) - - return dump_response(RelatedAppListResponse, {"data": related_apps, "total": len(related_apps)}), HTTPStatus.OK + result = application_services().knowledge.datasets.related_apps(request_context, dataset_id=str(dataset_id)) + except Exception as error: + _raise_dataset_error(error) + return dump_response(RelatedAppListResponse, result), HTTPStatus.OK @console_ns.route("/datasets//indexing-status") @@ -1056,43 +611,15 @@ class DatasetIndexingStatusApi(Resource): "Indexing status retrieved successfully", console_ns.models[DocumentStatusListResponse.__name__], ) - @setup_required - @login_required - @account_initialization_required - @with_current_user - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_READONLY, DatasetId())) - @with_session(write=False) - def get(self, session: Session, current_tenant_id: str, current_user: Account, dataset_id: UUID): - dataset = _get_accessible_dataset(dataset_id, current_tenant_id, current_user, session) - - documents = session.scalars( - select(Document).where(Document.dataset_id == dataset.id, Document.tenant_id == dataset.tenant_id) - ).all() - segment_counts = DocumentService.get_document_segment_counts( - documents, - session=session, - ) - documents_status = [] - for document in documents: - completed_segments, total_segments = segment_counts.get(str(document.id), (0, 0)) - # Create a dictionary with document attributes and additional fields - document_dict = { - "id": document.id, - "indexing_status": document.indexing_status, - "processing_started_at": document.processing_started_at, - "parsing_completed_at": document.parsing_completed_at, - "cleaning_completed_at": document.cleaning_completed_at, - "splitting_completed_at": document.splitting_completed_at, - "completed_at": document.completed_at, - "paused_at": document.paused_at, - "error": document.error, - "stopped_at": document.stopped_at, - "completed_segments": completed_segments, - "total_segments": total_segments, - } - documents_status.append(document_dict) - return dump_response(DocumentStatusListResponse, {"data": documents_status}), 200 + @console_account_admission(rbac_checks=(RBACCheck(RBACPermission.DATASET_READONLY, DatasetId()),)) + def get(self, request_context: RequestContext, dataset_id: UUID): + try: + result = application_services().knowledge.datasets.indexing_status( + request_context, dataset_id=str(dataset_id) + ) + except Exception as error: + _raise_dataset_error(error) + return dump_response(DocumentStatusListResponse, result), 200 @console_ns.route("/datasets/api-keys") @@ -1102,7 +629,7 @@ class DatasetApiKeyApi(Resource): @console_ns.response(200, "API keys retrieved successfully", console_ns.models[ApiKeyList.__name__]) @console_account_admission( allowed_roles=API_KEY_DELETE_ROLES, - rbac_checks=[RBACCheck(RBACPermission.DATASET_API_KEY_MANAGE, Workspace())], + rbac_checks=(RBACCheck(RBACPermission.DATASET_API_KEY_MANAGE, Workspace()),), ) def get(self, request_context: RequestContext): with api_key_errors(): @@ -1114,16 +641,16 @@ class DatasetApiKeyApi(Resource): @console_ns.response(400, "Maximum keys exceeded") @console_account_admission( allowed_roles=API_KEY_DELETE_ROLES, - rbac_checks=[RBACCheck(RBACPermission.DATASET_API_KEY_MANAGE, Workspace())], + rbac_checks=(RBACCheck(RBACPermission.DATASET_API_KEY_MANAGE, Workspace()),), ) def post(self, request_context: RequestContext): # Preserve the optional scope and validation response of the workspace route. payload = request.get_json(silent=True) or {} - raw_dataset_ids = payload.get("dataset_ids") or [] - if not isinstance(raw_dataset_ids, list) or any(not isinstance(item, str) for item in raw_dataset_ids): + dataset_ids = payload.get("dataset_ids") or [] + if not isinstance(dataset_ids, list) or any(not isinstance(item, str) for item in dataset_ids): console_ns.abort(400, message="dataset_ids must be a list of strings.") with api_key_errors(): - key = application_services().dataset_api_keys.create_workspace_key(request_context, tuple(raw_dataset_ids)) + key = application_services().dataset_api_keys.create_workspace_key(request_context, tuple(dataset_ids)) return dump_response(ApiKeyItem, key), 200 @@ -1135,7 +662,7 @@ class DatasetApiDeleteApi(Resource): @console_ns.response(204, "API key deleted successfully") @console_account_admission( allowed_roles=API_KEY_DELETE_ROLES, - rbac_checks=[RBACCheck(RBACPermission.DATASET_API_KEY_MANAGE, Workspace())], + rbac_checks=(RBACCheck(RBACPermission.DATASET_API_KEY_MANAGE, Workspace()),), ) def delete(self, request_context: RequestContext, api_key_id: UUID): with api_key_errors(): @@ -1145,21 +672,17 @@ class DatasetApiDeleteApi(Resource): @console_ns.route("/datasets//api-keys/") class DatasetEnableApiApi(Resource): - @setup_required - @login_required - @account_initialization_required @console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__]) - @with_current_user - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId())) - @with_session - def post(self, session: Session, current_tenant_id: str, current_user: Account, dataset_id: UUID, status: str): - dataset = _get_accessible_dataset(dataset_id, current_tenant_id, current_user, session) - if not current_user.is_dataset_editor: - raise Forbidden() - - DatasetService.update_dataset_api_status(dataset, status == "enable", current_user, session) - + @console_account_admission( + allowed_roles=_DATASET_EDIT_ROLES, rbac_checks=(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()),) + ) + def post(self, request_context: RequestContext, dataset_id: UUID, status: str): + try: + application_services().knowledge.datasets.set_api_enabled( + request_context, dataset_id=str(dataset_id), status=status + ) + except Exception as error: + _raise_dataset_error(error) return SimpleResultResponse(result="success").model_dump(mode="json"), 200 @@ -1168,12 +691,12 @@ class DatasetApiBaseUrlApi(Resource): @console_ns.doc("get_dataset_api_base_info") @console_ns.doc(description="Get dataset API base information") @console_ns.response(200, "API base info retrieved successfully", console_ns.models[ApiBaseUrlResponse.__name__]) - @setup_required - @login_required - @account_initialization_required - def get(self): - base = dify_config.SERVICE_API_URL or request.host_url.rstrip("/") - return ApiBaseUrlResponse(api_base_url=normalize_api_base_url(base)).model_dump(mode="json") + @console_account_admission() + def get(self, request_context: RequestContext): + base = application_services().knowledge.datasets.api_base_url( + request_context, request_base_url=request.host_url.rstrip("/") + ) + return ApiBaseUrlResponse(api_base_url=base).model_dump(mode="json") @console_ns.route("/datasets/retrieval-setting") @@ -1183,15 +706,10 @@ class DatasetRetrievalSettingApi(Resource): @console_ns.response( 200, "Retrieval settings retrieved successfully", console_ns.models[RetrievalSettingResponse.__name__] ) - @setup_required - @login_required - @account_initialization_required - def get(self): - vector_type = dify_config.VECTOR_STORE - return dump_response( - RetrievalSettingResponse, - _get_retrieval_methods_by_vector_type(vector_type, is_mock=False), - ) + @console_account_admission() + def get(self, request_context: RequestContext): + result = application_services().knowledge.datasets.retrieval_settings(request_context) + return dump_response(RetrievalSettingResponse, result) @console_ns.route("/datasets/retrieval-setting/") @@ -1204,14 +722,12 @@ class DatasetRetrievalSettingMockApi(Resource): "Mock retrieval settings retrieved successfully", console_ns.models[RetrievalSettingResponse.__name__], ) - @setup_required - @login_required - @account_initialization_required - def get(self, vector_type: str): - return dump_response( - RetrievalSettingResponse, - _get_retrieval_methods_by_vector_type(vector_type, is_mock=True), + @console_account_admission() + def get(self, request_context: RequestContext, vector_type: str): + result = application_services().knowledge.datasets.retrieval_settings( + request_context, vector_type=vector_type, is_mock=True ) + return dump_response(RetrievalSettingResponse, result) @console_ns.route("/datasets//error-docs") @@ -1221,20 +737,15 @@ class DatasetErrorDocs(Resource): @console_ns.doc(params={"dataset_id": "Dataset ID"}) @console_ns.response(200, "Error documents retrieved successfully", console_ns.models[ErrorDocsResponse.__name__]) @console_ns.response(404, "Dataset not found") - @setup_required - @login_required - @account_initialization_required - @with_current_user - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_READONLY, DatasetId())) - @with_session(write=False) - def get(self, session: Session, current_tenant_id: str, current_user: Account, dataset_id: UUID): - dataset = _get_accessible_dataset(dataset_id, current_tenant_id, current_user, session) - results = DocumentService.get_error_documents_by_dataset_ref( - DatasetRefService.create_dataset_ref(dataset), session - ) - - return dump_response(ErrorDocsResponse, {"data": results, "total": len(results)}), 200 + @console_account_admission(rbac_checks=(RBACCheck(RBACPermission.DATASET_READONLY, DatasetId()),)) + def get(self, request_context: RequestContext, dataset_id: UUID): + try: + result = application_services().knowledge.datasets.error_documents( + request_context, dataset_id=str(dataset_id) + ) + except Exception as error: + _raise_dataset_error(error) + return dump_response(ErrorDocsResponse, result), 200 @console_ns.route("/datasets//permission-part-users") @@ -1249,25 +760,15 @@ class DatasetPermissionUserListApi(Resource): ) @console_ns.response(404, "Dataset not found") @console_ns.response(403, "Permission denied") - @setup_required - @login_required - @account_initialization_required - @with_current_user - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_READONLY, DatasetId())) - @with_session(write=False) - def get(self, 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.") + @console_account_admission(rbac_checks=(RBACCheck(RBACPermission.DATASET_READONLY, DatasetId()),)) + def get(self, request_context: RequestContext, dataset_id: UUID): try: - DatasetService.check_dataset_permission(dataset, current_user, session) - except services.errors.account.NoPermissionError as e: - raise Forbidden(str(e)) - - partial_members_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str, session) - - return dump_response(PartialMemberListResponse, {"data": partial_members_list}), 200 + result = application_services().knowledge.datasets.partial_members( + request_context, dataset_id=str(dataset_id) + ) + except Exception as error: + _raise_dataset_error(error) + return dump_response(PartialMemberListResponse, {"data": result}), 200 @console_ns.route("/datasets//auto-disable-logs") @@ -1281,16 +782,12 @@ class DatasetAutoDisableLogApi(Resource): console_ns.models[AutoDisableLogsResponse.__name__], ) @console_ns.response(404, "Dataset not found") - @setup_required - @login_required - @account_initialization_required - @with_current_user - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_READONLY, DatasetId())) - @with_session(write=False) - def get(self, session: Session, current_tenant_id: str, current_user: Account, dataset_id: UUID): - dataset = _get_accessible_dataset(dataset_id, current_tenant_id, current_user, session) - auto_disable_logs = DatasetService.get_dataset_auto_disable_logs( - DatasetRefService.create_dataset_ref(dataset), session - ) - return dump_response(AutoDisableLogsResponse, auto_disable_logs), 200 + @console_account_admission(rbac_checks=(RBACCheck(RBACPermission.DATASET_READONLY, DatasetId()),)) + def get(self, request_context: RequestContext, dataset_id: UUID): + try: + result = application_services().knowledge.datasets.auto_disable_logs( + request_context, dataset_id=str(dataset_id) + ) + except Exception as error: + _raise_dataset_error(error) + return dump_response(AutoDisableLogsResponse, result), 200 diff --git a/api/controllers/console/datasets/datasets_document.py b/api/controllers/console/datasets/datasets_document.py index 72abbbbd264..76e6c21a727 100644 --- a/api/controllers/console/datasets/datasets_document.py +++ b/api/controllers/console/datasets/datasets_document.py @@ -1,95 +1,106 @@ -import json -import logging -from argparse import ArgumentTypeError -from collections.abc import Mapping, Sequence from contextlib import ExitStack from datetime import datetime -from typing import Any, Literal, cast +from typing import Any, Literal, Never from uuid import UUID -import sqlalchemy as sa from flask import request, send_file from flask_restx import Resource from pydantic import BaseModel, Field, JsonValue, field_validator -from sqlalchemy import asc, desc, func, select -from sqlalchemy.orm import Session -from werkzeug.exceptions import Forbidden, NotFound +from werkzeug.exceptions import BadRequest, Forbidden, NotFound -import services -from configs import dify_config from controllers.common.controller_schemas import DocumentBatchDownloadZipPayload +from controllers.common.errors import InvalidArgumentError, NotFoundError from controllers.common.fields import SimpleResultMessageResponse, SimpleResultResponse, UrlResponse -from controllers.common.rbac import DatasetId, RBACCheck, enforce_rbac_checks +from controllers.common.rbac import DatasetByDocument, DatasetId, RBACCheck, enforce_rbac_checks from controllers.common.schema import register_response_schema_models, register_schema_models -from controllers.common.session import with_session from controllers.console import console_ns -from controllers.console.wraps import RBACPermission, model_validate, rbac_permission_required -from core.entities.knowledge_entities import IndexingEstimate -from core.errors.error import ( - LLMBadRequestError, - ModelCurrentlyNotSupportError, - ProviderTokenNotInitError, - QuotaExceededError, -) -from core.indexing_runner import IndexingRunner -from core.model_manager import ModelManager -from core.plugin.impl.exc import PluginDaemonClientSideError -from core.rag.entities import Rule -from core.rag.extractor.entity.datasource_type import DatasourceType -from core.rag.extractor.entity.extract_setting import ExtractSetting, NotionInfo, WebsiteInfo -from core.rag.index_processor.constant.index_type import IndexTechniqueType -from fields.base import ResponseModel -from fields.document_fields import ( - DocumentMetadataResponse, - DocumentResponse, - DocumentStatusListResponse, - DocumentStatusResponse, - DocumentWithSession, - document_response, - document_responses, - normalize_enum, -) -from graphon.model_runtime.entities.model_entities import ModelType -from graphon.model_runtime.errors.invoke import InvokeAuthorizationError -from libs.datetime_utils import naive_utc_now -from libs.helper import dump_response, to_timestamp -from libs.login import login_required -from libs.pagination import paginate_query -from models import Account, Document, DocumentSegment, UploadFile -from models.dataset import DatasetPermissionEnum, DocumentPipelineExecutionLog -from models.enums import IndexingStatus, ProcessRuleMode, SegmentStatus -from services.dataset_ref_service import DatasetRefService -from services.dataset_service import DatasetService, DocumentService -from services.enterprise import rbac_service as enterprise_rbac_service -from services.entities.knowledge_entities.knowledge_entities import KnowledgeConfig, ProcessRule, RetrievalModel -from services.file_service import FileService -from services.vector_space_admission_service import get_vector_space_admission_error_fields -from tasks.generate_summary_index_task import generate_summary_index_task - -from ..app.error import ( +from controllers.console.app.error import ( ProviderModelCurrentlyNotSupportError, ProviderNotInitializeError, ProviderQuotaExceededError, ) -from ..datasets.error import ( +from controllers.console.datasets.error import ( ArchivedDocumentImmutableError, + DatasetAccessDeniedRequestError, DocumentAlreadyFinishedError, DocumentIndexingError, IndexingEstimateError, InvalidActionError, InvalidMetadataError, ) -from ..wraps import ( - account_initialization_required, +from controllers.console.flask_admission import console_account_admission +from controllers.console.wraps import ( + RBACPermission, check_knowledge_rate_limit, cloud_edition_billing_rate_limit_check, cloud_edition_billing_resource_check, - setup_required, - with_current_tenant_id, - with_current_user, + model_validate, +) +from core.entities.knowledge_entities import IndexingEstimate +from core.rag.entities import Rule +from extensions.ext_application_services import application_services +from fields.base import ResponseModel +from fields.document_fields import ( + DocumentMetadataResponse, + DocumentResponse, + DocumentStatusListResponse, + DocumentStatusResponse, + normalize_enum, +) +from libs.helper import dump_response, to_timestamp +from machinery.context import RequestContext +from models.account import TenantAccountRole +from models.enums import ProcessRuleMode +from services.knowledge.dataset_access import DatasetAccessDeniedError, DatasetNotFoundError +from services.knowledge.documents.application import ( + DocumentArchivedError, + DocumentIndexingStateError, + DocumentInvalidActionError, + DocumentListFilter, + DocumentNotFoundError, + DocumentProviderError, +) +from services.knowledge.entities.knowledge_entities import KnowledgeConfig, ProcessRule, RetrievalModel +from services.knowledge.indexing.estimate import ( + EstimateDocumentAlreadyFinishedError, + EstimateDocumentNotFoundError, + EstimateSourceNotFoundError, + IndexingEstimateCredentialUnavailableError, + IndexingEstimateExecutionError, + IndexingEstimateProviderUnavailableError, + UnsupportedEstimateSourceError, ) -logger = logging.getLogger(__name__) +_DATASET_EDIT_ROLES = frozenset( + { + TenantAccountRole.OWNER, + TenantAccountRole.ADMIN, + TenantAccountRole.EDITOR, + TenantAccountRole.DATASET_OPERATOR, + } +) + + +def _raise_document_error(error: Exception) -> Never: + if isinstance(error, DatasetNotFoundError): + raise NotFound("Dataset not found.") from error + if isinstance(error, DocumentNotFoundError): + raise NotFound(str(error)) from error + if isinstance(error, DatasetAccessDeniedError): + raise Forbidden(str(error)) from error + if isinstance(error, DocumentArchivedError): + raise ArchivedDocumentImmutableError() from error + if isinstance(error, DocumentIndexingStateError): + raise DocumentIndexingError(str(error)) from error + if isinstance(error, DocumentInvalidActionError): + raise InvalidActionError(str(error)) from error + if isinstance(error, DocumentProviderError): + if error.kind == "quota": + raise ProviderQuotaExceededError() from error + if error.kind == "unsupported": + raise ProviderModelCurrentlyNotSupportError() from error + raise ProviderNotInitializeError(str(error)) from error + raise error class DatasetResponse(ResponseModel): @@ -119,24 +130,6 @@ class DocumentWithSegmentsResponse(DocumentResponse): total_segments: int | None = Field(default=None, exclude_if=lambda value: value is None) -class DocumentWithSegmentsSession(DocumentWithSession): - @property - def process_rule_dict(self) -> Any: - process_rule = self.document.get_dataset_process_rule(session=self.session) - return process_rule.to_dict() if process_rule else None - - -def document_with_segments_responses( - documents: Sequence[Document], *, session: Session -) -> list[DocumentWithSegmentsResponse]: - return [ - DocumentWithSegmentsResponse.model_validate( - DocumentWithSegmentsSession(document=document, session=session), from_attributes=True - ) - for document in documents - ] - - class DatasetAndDocumentResponse(ResponseModel): dataset: DatasetResponse documents: list[DocumentResponse] @@ -183,6 +176,20 @@ class IndexingEstimateResponse(IndexingEstimate): currency: str +def _serialize_indexing_estimate(estimate: IndexingEstimate) -> tuple[dict[str, object], int]: + return ( + IndexingEstimateResponse( + tokens=0, + total_price=0, + currency="USD", + total_segments=estimate.total_segments, + preview=estimate.preview, + qa_preview=estimate.qa_preview, + ).model_dump(mode="json", exclude_none=True), + 200, + ) + + class DocumentDetailResponse(ResponseModel): id: str position: int | None = None @@ -294,96 +301,29 @@ register_response_schema_models( ) -class DocumentResource(Resource): - def get_document( - self, session: Session, dataset_id: str, document_id: str, current_user: Account, current_tenant_id: str - ) -> Document: - dataset = DatasetService.get_dataset_for_tenant(dataset_id, current_tenant_id, session=session) - if not dataset: - raise NotFound("Dataset not found.") - - try: - DatasetService.check_dataset_permission(dataset, current_user, session) - except services.errors.account.NoPermissionError as e: - raise Forbidden(str(e)) - - dataset_ref = DatasetRefService.create_dataset_ref(dataset) - document_ref = DatasetRefService.create_document_ref_from_id(dataset_ref, document_id) - document = DatasetRefService.get_document_by_ref(document_ref, session=session) - - if not document: - raise NotFound("Document not found.") - - return document - - def get_batch_documents( - self, session: Session, dataset_id: str, batch: str, current_user: Account - ) -> Sequence[Document]: - dataset = DatasetService.get_dataset(dataset_id, session) - if not dataset: - raise NotFound("Dataset not found.") - - try: - DatasetService.check_dataset_permission(dataset, current_user, session) - except services.errors.account.NoPermissionError as e: - raise Forbidden(str(e)) - - documents = DocumentService.get_batch_documents(dataset_id, batch, session) - - if not documents: - raise NotFound("Documents not found.") - - return documents - - @console_ns.route("/datasets/process-rule") class GetProcessRuleApi(Resource): @console_ns.doc("get_process_rule") @console_ns.doc(description="Get dataset document processing rules") @console_ns.doc(params={"document_id": "Document ID (optional)"}) @console_ns.response(200, "Process rules retrieved successfully", console_ns.models[ProcessRuleResponse.__name__]) - @setup_required - @login_required - @account_initialization_required - @with_current_user - @with_session(write=False) - def get(self, session: Session, current_user: Account): - req_data = request.args - - document_id = req_data.get("document_id") - - # get default rules - mode = DocumentService.DEFAULT_RULES["mode"] - rules = DocumentService.DEFAULT_RULES["rules"] - limits = DocumentService.DEFAULT_RULES["limits"] + @console_account_admission() + def get(self, request_context: RequestContext): + document_id = request.args.get("document_id") if document_id: - document = DocumentService.get_document_by_id(document_id, session) - if document is None: - raise NotFound("Document not found.") - - dataset = DatasetService.get_dataset(document.dataset_id, session) - - if not dataset: - raise NotFound("Dataset not found.") - enforce_rbac_checks( - tenant_id=dataset.tenant_id, - account_id=current_user.id, - checks=[RBACCheck(RBACPermission.DATASET_READONLY, DatasetId())], - path_args={"dataset_id": dataset.id}, + tenant_id=request_context.active_workspace_id, + account_id=request_context.account_id, + checks=[RBACCheck(RBACPermission.DATASET_READONLY, DatasetByDocument())], + path_args={"document_id": document_id}, ) - - try: - DatasetService.check_dataset_permission(dataset, current_user, session) - except services.errors.account.NoPermissionError as e: - raise Forbidden(str(e)) - - dataset_process_rule = dataset.get_latest_process_rule(session=session) - if dataset_process_rule: - mode = dataset_process_rule.mode - rules = dataset_process_rule.rules_dict - - return dump_response(ProcessRuleResponse, {"mode": mode, "rules": rules, "limits": limits}) + try: + result = application_services().knowledge.documents.get_process_rule( + request_context, document_id=document_id + ) + except Exception as error: + _raise_document_error(error) + return dump_response(ProcessRuleResponse, result) @console_ns.route("/datasets//documents") @@ -406,224 +346,56 @@ class DatasetDocumentListApi(Resource): "Documents retrieved successfully", console_ns.models[DocumentWithSegmentsListResponse.__name__], ) - @setup_required - @login_required - @account_initialization_required - @with_current_user - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_READONLY, DatasetId())) - @with_session(write=False) - def get(self, session: Session, current_tenant_id: str, current_user: Account, dataset_id: UUID): - dataset_id_str = str(dataset_id) - raw_args = request.args.to_dict() - param = DocumentDatasetListParam.model_validate(raw_args) - page = param.page - limit = param.limit - search = param.search - sort = param.sort_by - status = param.status - # "yes", "true", "t", "y", "1" convert to True, while others convert to False. + @console_account_admission(rbac_checks=(RBACCheck(RBACPermission.DATASET_READONLY, DatasetId()),)) + def get(self, request_context: RequestContext, dataset_id: UUID): + args = DocumentDatasetListParam.model_validate(request.args.to_dict()) try: - fetch_val = param.fetch_val - if isinstance(fetch_val, bool): - fetch = fetch_val - else: - if fetch_val.lower() in ("yes", "true", "t", "y", "1"): - fetch = True - elif fetch_val.lower() in ("no", "false", "f", "n", "0"): - fetch = False - else: - raise ArgumentTypeError( - f"Truthy value expected: got {fetch_val} but expected one of yes/no, true/false, t/f, y/n, 1/0 " - f"(case insensitive)." - ) - except (ArgumentTypeError, ValueError, Exception): - fetch = False - dataset = DatasetService.get_dataset(dataset_id_str, session) - if not dataset: - raise NotFound("Dataset not found.") + result = application_services().knowledge.documents.list_documents( + request_context, + dataset_id=str(dataset_id), + query=DocumentListFilter( + page=args.page, + limit=args.limit, + search=args.search or "", + sort=args.sort_by, + status=args.status, + fetch=args.fetch_val.lower() in {"yes", "true", "t", "y", "1"}, + ), + ) + except Exception as error: + _raise_document_error(error) + return dump_response(DocumentWithSegmentsListResponse, result) - try: - DatasetService.check_dataset_permission(dataset, current_user, session) - except services.errors.account.NoPermissionError as e: - raise Forbidden(str(e)) - - query = select(Document).where(Document.dataset_id == dataset_id_str, Document.tenant_id == current_tenant_id) - - if status: - query = DocumentService.apply_display_status_filter(query, status) - - if search: - search = f"%{search}%" - query = query.where(Document.name.like(search)) - - if sort.startswith("-"): - sort_logic = desc - sort = sort[1:] - else: - sort_logic = asc - - match sort: - case "hit_count": - sub_query = ( - sa.select( - DocumentSegment.document_id, sa.func.sum(DocumentSegment.hit_count).label("total_hit_count") - ) - .where(DocumentSegment.dataset_id == dataset_id_str) - .group_by(DocumentSegment.document_id) - .subquery() - ) - - query = query.outerjoin(sub_query, sub_query.c.document_id == Document.id).order_by( - sort_logic(sa.func.coalesce(sub_query.c.total_hit_count, 0)), - sort_logic(Document.position), - ) - case "created_at": - query = query.order_by( - sort_logic(Document.created_at), - sort_logic(Document.position), - ) - case _: - query = query.order_by( - desc(Document.created_at), - desc(Document.position), - ) - - paginated_documents = paginate_query(query, session=session, page=page, per_page=limit, max_per_page=100) - documents = paginated_documents.items - - DocumentService.enrich_documents_with_summary_index_status( - documents=documents, - dataset=dataset, - tenant_id=current_tenant_id, - session=session, - ) - - if fetch: - for document in documents: - completed_segments = ( - session.scalar( - select(func.count(DocumentSegment.id)).where( - DocumentSegment.completed_at.isnot(None), - DocumentSegment.document_id == str(document.id), - DocumentSegment.status != SegmentStatus.RE_SEGMENT, - ) - ) - or 0 - ) - total_segments = ( - session.scalar( - select(func.count(DocumentSegment.id)).where( - DocumentSegment.document_id == str(document.id), - DocumentSegment.status != SegmentStatus.RE_SEGMENT, - ) - ) - or 0 - ) - document.completed_segments = completed_segments - document.total_segments = total_segments - response = { - "data": document_with_segments_responses(documents, session=session), - # The result object already knows: it was built from the page the query - # ran with, where `len(documents) == limit` is only ever a guess that a - # full page means another one follows. - "has_more": paginated_documents.has_next, - "limit": paginated_documents.per_page, - "total": paginated_documents.total, - "page": paginated_documents.page, - } - - return dump_response(DocumentWithSegmentsListResponse, response) - - @setup_required - @login_required - @account_initialization_required - @cloud_edition_billing_resource_check("vector_space") - @cloud_edition_billing_rate_limit_check("knowledge") @console_ns.expect(console_ns.models[KnowledgeConfig.__name__]) @console_ns.response(200, "Documents created successfully", console_ns.models[DatasetAndDocumentResponse.__name__]) - @with_current_user - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_USE, DatasetId())) - @with_session - def post(self, session: Session, current_user: Account, dataset_id: UUID): - dataset_id_str = str(dataset_id) - - dataset = DatasetService.get_dataset(dataset_id_str, session) - - if not dataset: - raise NotFound("Dataset not found.") - - # The role of the current user in the ta table must be admin, owner, or editor - if not current_user.is_dataset_editor: - raise Forbidden() - + @console_account_admission( + allowed_roles=_DATASET_EDIT_ROLES, rbac_checks=(RBACCheck(RBACPermission.DATASET_USE, DatasetId()),) + ) + @cloud_edition_billing_resource_check("vector_space") + @cloud_edition_billing_rate_limit_check("knowledge") + def post(self, request_context: RequestContext, dataset_id: UUID): + config = KnowledgeConfig.model_validate(console_ns.payload or {}) try: - DatasetService.check_dataset_permission(dataset, current_user, session) - except services.errors.account.NoPermissionError as e: - raise Forbidden(str(e)) - - knowledge_config = KnowledgeConfig.model_validate(console_ns.payload or {}) - - if not dataset.indexing_technique and not knowledge_config.indexing_technique: - raise ValueError("indexing_technique is required.") - - # validate args - DocumentService.document_create_args_validate(knowledge_config) - - try: - documents, batch = DocumentService.save_document_with_dataset_id( - dataset, knowledge_config, current_user, session=session + result = application_services().knowledge.documents.create_documents( + request_context, dataset_id=str(dataset_id), settings=config.model_dump() ) - dataset = DatasetService.get_dataset(dataset_id_str, session) + except Exception as error: + _raise_document_error(error) + return dump_response(DatasetAndDocumentResponse, result) - except ProviderTokenNotInitError as ex: - raise ProviderNotInitializeError(ex.description) - except QuotaExceededError: - raise ProviderQuotaExceededError() - except ModelCurrentlyNotSupportError: - raise ProviderModelCurrentlyNotSupportError() - - return dump_response( - DatasetAndDocumentResponse, - {"dataset": dataset, "documents": document_responses(documents, session=session), "batch": batch}, - ) - - @setup_required - @login_required - @account_initialization_required @console_ns.response(204, "Documents deleted successfully") - @with_current_user - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_DELETE_FILE, DatasetId())) - @with_session - def delete( - self, - session: Session, - current_tenant_id: str, - current_user: Account, - dataset_id: UUID, - ): - dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset_for_tenant(dataset_id_str, current_tenant_id, session=session) - if dataset is None: - raise NotFound("Dataset not found.") - - if not current_user.is_dataset_editor: - raise Forbidden() - - try: - DatasetService.check_dataset_permission(dataset, current_user, session) - except services.errors.account.NoPermissionError as e: - raise Forbidden(str(e)) - + @console_account_admission( + allowed_roles=_DATASET_EDIT_ROLES, + rbac_checks=(RBACCheck(RBACPermission.DATASET_DELETE_FILE, DatasetId()),), + ) + def delete(self, request_context: RequestContext, dataset_id: UUID): check_knowledge_rate_limit() try: - document_ids = request.args.getlist("document_id") - dataset_ref = DatasetRefService.create_dataset_ref(dataset) - DocumentService.delete_documents(dataset_ref, document_ids, dataset.get_doc_form(session=session), session) - except services.errors.document.DocumentIndexingError: - raise DocumentIndexingError("Cannot delete document during indexing.") - + application_services().knowledge.documents.delete_documents( + request_context, dataset_id=str(dataset_id), document_ids=request.args.getlist("document_id") + ) + except Exception as error: + _raise_document_error(error) return "", 204 @@ -636,87 +408,22 @@ class DatasetInitApi(Resource): 200, "Dataset initialized successfully", console_ns.models[DatasetAndDocumentResponse.__name__] ) @console_ns.response(400, "Invalid request parameters") - @setup_required - @login_required - @account_initialization_required + @console_account_admission(allowed_roles=_DATASET_EDIT_ROLES) @cloud_edition_billing_resource_check("vector_space") @cloud_edition_billing_rate_limit_check("knowledge") - @with_current_user - @with_current_tenant_id - @with_session - def post(self, session: Session, current_tenant_id: str, current_user: Account): - # The role of the current user in the ta table must be admin, owner, dataset_operator, or editor - if not current_user.is_dataset_editor: - raise Forbidden() - - knowledge_config = KnowledgeConfig.model_validate(console_ns.payload or {}) - if knowledge_config.indexing_technique == IndexTechniqueType.HIGH_QUALITY: - if knowledge_config.embedding_model is None or knowledge_config.embedding_model_provider is None: - raise ValueError("embedding model and embedding model provider are required for high quality indexing.") - try: - model_manager = ModelManager.for_tenant(tenant_id=current_tenant_id) - model_manager.get_model_instance( - tenant_id=current_tenant_id, - provider=knowledge_config.embedding_model_provider, - model_type=ModelType.TEXT_EMBEDDING, - model=knowledge_config.embedding_model, - ) - is_multimodal = DatasetService.check_is_multimodal_model( - current_tenant_id, knowledge_config.embedding_model_provider, knowledge_config.embedding_model - ) - knowledge_config.is_multimodal = is_multimodal # pyrefly: ignore[bad-assignment] - except InvokeAuthorizationError: - raise ProviderNotInitializeError( - "No Embedding Model available. Please configure a valid provider in the Settings -> Model Provider." - ) - except ProviderTokenNotInitError as ex: - raise ProviderNotInitializeError(ex.description) - - # validate args - DocumentService.document_create_args_validate(knowledge_config) - + def post(self, request_context: RequestContext): + config = KnowledgeConfig.model_validate(console_ns.payload or {}) try: - dataset, documents, batch = DocumentService.save_document_without_dataset_id( - tenant_id=current_tenant_id, - knowledge_config=knowledge_config, - account=current_user, - session=session, + result = application_services().knowledge.documents.initialize_dataset( + request_context, settings=config.model_dump() ) - except ProviderTokenNotInitError as ex: - raise ProviderNotInitializeError(ex.description) - except QuotaExceededError: - raise ProviderQuotaExceededError() - except ModelCurrentlyNotSupportError: - raise ProviderModelCurrentlyNotSupportError() - - if dify_config.RBAC_ENABLED: - dataset.permission = DatasetPermissionEnum.ALL_TEAM - else: - dataset.permission = DatasetPermissionEnum.ONLY_ME - session.flush() - - if dify_config.RBAC_ENABLED: - enterprise_rbac_service.RBACService.DatasetAccess.replace_whitelist( - current_tenant_id, - current_user.id, - dataset.id, - enterprise_rbac_service.ReplaceMemberBindings(automatic_include_workspace_members=False), - ) - enterprise_rbac_service.try_sync_creator_access_policy_member_bindings( - current_tenant_id, - current_user.id, - enterprise_rbac_service.RBACResourceType.DATASET, - dataset.id, - ) - - return dump_response( - DatasetAndDocumentResponse, - {"dataset": dataset, "documents": document_responses(documents, session=session), "batch": batch}, - ) + except Exception as error: + _raise_document_error(error) + return dump_response(DatasetAndDocumentResponse, result) @console_ns.route("/datasets//documents//indexing-estimate") -class DocumentIndexingEstimateApi(DocumentResource): +class DocumentIndexingEstimateApi(Resource): @console_ns.doc("estimate_document_indexing") @console_ns.doc(description="Estimate document indexing cost") @console_ns.doc(params={"dataset_id": "Dataset ID", "document_id": "Document ID"}) @@ -727,256 +434,91 @@ class DocumentIndexingEstimateApi(DocumentResource): ) @console_ns.response(404, "Document not found") @console_ns.response(400, "Document already finished") - @setup_required - @login_required - @account_initialization_required - @with_current_user - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_USE, DatasetId())) - @with_session - def get(self, session: Session, current_tenant_id: str, current_user: Account, dataset_id: UUID, document_id: UUID): - dataset_id_str = str(dataset_id) - document_id_str = str(document_id) - document = self.get_document(session, dataset_id_str, document_id_str, current_user, current_tenant_id) - - if document.indexing_status in {IndexingStatus.COMPLETED, IndexingStatus.ERROR}: - raise DocumentAlreadyFinishedError() - - data_process_rule = document.get_dataset_process_rule(session=session) - data_process_rule_dict: Mapping[str, Any] = data_process_rule.to_dict() if data_process_rule else {} - - if document.data_source_type == "upload_file": - data_source_info = document.data_source_info_dict - if data_source_info and "upload_file_id" in data_source_info: - file_id = data_source_info["upload_file_id"] - - file = session.scalar( - select(UploadFile) - .where(UploadFile.tenant_id == document.tenant_id, UploadFile.id == file_id) - .limit(1) - ) - - # raise error if file not found - if not file: - raise NotFound("File not found.") - - extract_setting = ExtractSetting( - datasource_type=DatasourceType.FILE, upload_file=file, document_model=document.doc_form - ) - - indexing_runner = IndexingRunner() - - try: - estimate_response = indexing_runner.indexing_estimate( - tenant_id=current_tenant_id, - extract_settings=[extract_setting], - tmp_processing_rule=data_process_rule_dict, - doc_form=document.doc_form, - doc_language="English", - dataset_id=dataset_id_str, - session=session, - ) - return ( - # TODO: why using zero here? the same for the below endpoint - IndexingEstimateResponse( - tokens=0, - total_price=0, - currency="USD", - total_segments=estimate_response.total_segments, - preview=estimate_response.preview, - qa_preview=estimate_response.qa_preview, - ).model_dump(mode="json", exclude_none=True), - 200, - ) - except LLMBadRequestError: - raise ProviderNotInitializeError( - "No Embedding Model available. Please configure a valid provider " - "in the Settings -> Model Provider." - ) - except ProviderTokenNotInitError as ex: - raise ProviderNotInitializeError(ex.description) - except PluginDaemonClientSideError as ex: - raise ProviderNotInitializeError(ex.description) - except Exception as e: - raise IndexingEstimateError(str(e)) - - return ( - IndexingEstimateResponse( - tokens=0, - total_price=0, - currency="USD", - total_segments=0, - preview=[], - ).model_dump(mode="json", exclude_none=True), - 200, - ) + @console_account_admission( + rbac_checks=(RBACCheck(RBACPermission.DATASET_USE, DatasetId()),), + ) + def get(self, request_context: RequestContext, dataset_id: UUID, document_id: UUID): + try: + estimate = application_services().knowledge.indexing_estimates.estimate_document( + request_context, + dataset_id=str(dataset_id), + document_id=str(document_id), + ) + except ( + DatasetNotFoundError, + EstimateDocumentNotFoundError, + EstimateSourceNotFoundError, + IndexingEstimateCredentialUnavailableError, + ) as error: + raise NotFoundError(description=str(error)) from error + except DatasetAccessDeniedError as error: + raise DatasetAccessDeniedRequestError(description=str(error)) from error + except EstimateDocumentAlreadyFinishedError as error: + raise DocumentAlreadyFinishedError() from error + except UnsupportedEstimateSourceError as error: + raise InvalidArgumentError(description=str(error)) from error + except IndexingEstimateProviderUnavailableError as error: + raise ProviderNotInitializeError(str(error)) from error + except IndexingEstimateExecutionError as error: + raise IndexingEstimateError(str(error)) from error + return _serialize_indexing_estimate(estimate) @console_ns.route("/datasets//batch//indexing-estimate") -class DocumentBatchIndexingEstimateApi(DocumentResource): +class DocumentBatchIndexingEstimateApi(Resource): @console_ns.response( 200, "Indexing estimate calculated successfully", console_ns.models[IndexingEstimateResponse.__name__], ) - @setup_required - @login_required - @account_initialization_required - @with_current_user - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_USE, DatasetId())) - @with_session - def get(self, session: Session, current_tenant_id: str, current_user: Account, dataset_id: UUID, batch: str): - dataset_id_str = str(dataset_id) - documents = self.get_batch_documents(session, dataset_id_str, batch, current_user) - if not documents: - return ( - IndexingEstimateResponse( - tokens=0, - total_price=0, - currency="USD", - total_segments=0, - preview=[], - ).model_dump(mode="json", exclude_none=True), - 200, + @console_account_admission( + rbac_checks=(RBACCheck(RBACPermission.DATASET_USE, DatasetId()),), + ) + def get(self, request_context: RequestContext, dataset_id: UUID, batch: str): + try: + estimate = application_services().knowledge.indexing_estimates.estimate_batch( + request_context, + dataset_id=str(dataset_id), + batch=batch, ) - data_process_rule = documents[0].get_dataset_process_rule(session=session) - data_process_rule_dict: Mapping[str, Any] = data_process_rule.to_dict() if data_process_rule else {} - extract_settings = [] - for document in documents: - if document.indexing_status in {IndexingStatus.COMPLETED, IndexingStatus.ERROR}: - raise DocumentAlreadyFinishedError() - data_source_info = document.data_source_info_dict - match document.data_source_type: - case "upload_file": - if not data_source_info: - continue - file_id = data_source_info["upload_file_id"] - file_detail = session.scalar( - select(UploadFile) - .where(UploadFile.tenant_id == current_tenant_id, UploadFile.id == file_id) - .limit(1) - ) - - if file_detail is None: - raise NotFound("File not found.") - - extract_setting = ExtractSetting( - datasource_type=DatasourceType.FILE, upload_file=file_detail, document_model=document.doc_form - ) - extract_settings.append(extract_setting) - case "notion_import": - if not data_source_info: - continue - extract_setting = ExtractSetting( - datasource_type=DatasourceType.NOTION, - notion_info=NotionInfo.model_validate( - { - "credential_id": data_source_info.get("credential_id"), - "notion_workspace_id": data_source_info["notion_workspace_id"], - "notion_obj_id": data_source_info["notion_page_id"], - "notion_page_type": data_source_info["type"], - "tenant_id": current_tenant_id, - } - ), - document_model=document.doc_form, - ) - extract_settings.append(extract_setting) - case "website_crawl": - if not data_source_info: - continue - extract_setting = ExtractSetting( - datasource_type=DatasourceType.WEBSITE, - website_info=WebsiteInfo.model_validate( - { - "provider": data_source_info["provider"], - "job_id": data_source_info["job_id"], - "url": data_source_info["url"], - "tenant_id": current_tenant_id, - "mode": data_source_info["mode"], - "only_main_content": data_source_info["only_main_content"], - } - ), - document_model=document.doc_form, - ) - extract_settings.append(extract_setting) - - case _: - raise ValueError("Data source type not support") - indexing_runner = IndexingRunner() - try: - response = indexing_runner.indexing_estimate( - tenant_id=current_tenant_id, - extract_settings=extract_settings, - tmp_processing_rule=data_process_rule_dict, - doc_form=document.doc_form, - doc_language="English", - dataset_id=dataset_id_str, - session=session, - ) - return ( - IndexingEstimateResponse( - tokens=0, - total_price=0, - currency="USD", - total_segments=response.total_segments, - preview=response.preview, - qa_preview=response.qa_preview, - ).model_dump(mode="json", exclude_none=True), - 200, - ) - except LLMBadRequestError: - raise ProviderNotInitializeError( - "No Embedding Model available. Please configure a valid provider in the Settings -> Model Provider." - ) - except ProviderTokenNotInitError as ex: - raise ProviderNotInitializeError(ex.description) - except PluginDaemonClientSideError as ex: - raise ProviderNotInitializeError(ex.description) - except Exception as e: - raise IndexingEstimateError(str(e)) + except ( + DatasetNotFoundError, + EstimateDocumentNotFoundError, + EstimateSourceNotFoundError, + IndexingEstimateCredentialUnavailableError, + ) as error: + raise NotFoundError(description=str(error)) from error + except DatasetAccessDeniedError as error: + raise DatasetAccessDeniedRequestError(description=str(error)) from error + except EstimateDocumentAlreadyFinishedError as error: + raise DocumentAlreadyFinishedError() from error + except UnsupportedEstimateSourceError as error: + raise InvalidArgumentError(description=str(error)) from error + except IndexingEstimateProviderUnavailableError as error: + raise ProviderNotInitializeError(str(error)) from error + except IndexingEstimateExecutionError as error: + raise IndexingEstimateError(str(error)) from error + return _serialize_indexing_estimate(estimate) @console_ns.route("/datasets//batch//indexing-status") -class DocumentBatchIndexingStatusApi(DocumentResource): +class DocumentBatchIndexingStatusApi(Resource): @console_ns.response( 200, "Indexing status retrieved successfully", console_ns.models[DocumentStatusListResponse.__name__] ) - @setup_required - @login_required - @account_initialization_required - @with_current_user - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_READONLY, DatasetId())) - @with_session(write=False) - def get(self, session: Session, current_user: Account, dataset_id: UUID, batch: str): - dataset_id_str = str(dataset_id) - documents = self.get_batch_documents(session, dataset_id_str, batch, current_user) - segment_counts = DocumentService.get_document_segment_counts(documents, session=session) - documents_status = [] - for document in documents: - completed_segments, total_segments = segment_counts.get(str(document.id), (0, 0)) - # Create a dictionary with document attributes and additional fields - document_dict = { - "id": document.id, - "indexing_status": IndexingStatus.PAUSED if document.is_paused else document.indexing_status, - "processing_started_at": document.processing_started_at, - "parsing_completed_at": document.parsing_completed_at, - "cleaning_completed_at": document.cleaning_completed_at, - "splitting_completed_at": document.splitting_completed_at, - "completed_at": document.completed_at, - "paused_at": document.paused_at, - "error": document.error, - **get_vector_space_admission_error_fields(document.error), - "stopped_at": document.stopped_at, - "completed_segments": completed_segments, - "total_segments": total_segments, - } - documents_status.append(document_dict) - return dump_response(DocumentStatusListResponse, {"data": documents_status}) + @console_account_admission(rbac_checks=(RBACCheck(RBACPermission.DATASET_READONLY, DatasetId()),)) + def get(self, request_context: RequestContext, dataset_id: UUID, batch: str): + try: + result = application_services().knowledge.documents.get_batch_indexing_status( + request_context, dataset_id=str(dataset_id), batch=batch + ) + except Exception as error: + _raise_document_error(error) + return dump_response(DocumentStatusListResponse, result) @console_ns.route("/datasets//documents//indexing-status") -class DocumentIndexingStatusApi(DocumentResource): +class DocumentIndexingStatusApi(Resource): @console_ns.doc("get_document_indexing_status") @console_ns.doc(description="Get document indexing status") @console_ns.doc(params={"dataset_id": "Dataset ID", "document_id": "Document ID"}) @@ -984,59 +526,19 @@ class DocumentIndexingStatusApi(DocumentResource): 200, "Indexing status retrieved successfully", console_ns.models[DocumentStatusResponse.__name__] ) @console_ns.response(404, "Document not found") - @setup_required - @login_required - @account_initialization_required - @with_current_user - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_READONLY, DatasetId())) - @with_session(write=False) - def get(self, session: Session, current_tenant_id: str, current_user: Account, dataset_id: UUID, document_id: UUID): - dataset_id_str = str(dataset_id) - document_id_str = str(document_id) - document = self.get_document(session, dataset_id_str, document_id_str, current_user, current_tenant_id) - - completed_segments = ( - session.scalar( - select(func.count(DocumentSegment.id)).where( - DocumentSegment.completed_at.isnot(None), - DocumentSegment.document_id == document_id_str, - DocumentSegment.status != SegmentStatus.RE_SEGMENT, - ) + @console_account_admission(rbac_checks=(RBACCheck(RBACPermission.DATASET_READONLY, DatasetId()),)) + def get(self, request_context: RequestContext, dataset_id: UUID, document_id: UUID): + try: + result = application_services().knowledge.documents.get_indexing_status( + request_context, dataset_id=str(dataset_id), document_id=str(document_id) ) - or 0 - ) - total_segments = ( - session.scalar( - select(func.count(DocumentSegment.id)).where( - DocumentSegment.document_id == document_id_str, - DocumentSegment.status != SegmentStatus.RE_SEGMENT, - ) - ) - or 0 - ) - - # Create a dictionary with document attributes and additional fields - document_dict = { - "id": document.id, - "indexing_status": IndexingStatus.PAUSED if document.is_paused else document.indexing_status, - "processing_started_at": document.processing_started_at, - "parsing_completed_at": document.parsing_completed_at, - "cleaning_completed_at": document.cleaning_completed_at, - "splitting_completed_at": document.splitting_completed_at, - "completed_at": document.completed_at, - "paused_at": document.paused_at, - "error": document.error, - **get_vector_space_admission_error_fields(document.error), - "stopped_at": document.stopped_at, - "completed_segments": completed_segments, - "total_segments": total_segments, - } - return dump_response(DocumentStatusResponse, document_dict) + except Exception as error: + _raise_document_error(error) + return dump_response(DocumentStatusResponse, result) @console_ns.route("/datasets//documents/") -class DocumentApi(DocumentResource): +class DocumentApi(Resource): METADATA_CHOICES = {"all", "only", "without"} @console_ns.doc("get_document") @@ -1050,179 +552,103 @@ class DocumentApi(DocumentResource): ) @console_ns.response(200, "Document retrieved successfully", console_ns.models[DocumentDetailResponse.__name__]) @console_ns.response(404, "Document not found") - @setup_required - @login_required - @account_initialization_required - @with_current_user - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_READONLY, DatasetId())) - @with_session(write=False) - def get(self, session: Session, current_tenant_id: str, current_user: Account, dataset_id: UUID, document_id: UUID): - dataset_id_str = str(dataset_id) - document_id_str = str(document_id) - document = self.get_document(session, dataset_id_str, document_id_str, current_user, current_tenant_id) - + @console_account_admission(rbac_checks=(RBACCheck(RBACPermission.DATASET_READONLY, DatasetId()),)) + def get(self, request_context: RequestContext, dataset_id: UUID, document_id: UUID): metadata = request.args.get("metadata", "all") if metadata not in self.METADATA_CHOICES: raise InvalidMetadataError(f"Invalid metadata value: {metadata}") + try: + result = application_services().knowledge.documents.get_document( + request_context, + dataset_id=str(dataset_id), + document_id=str(document_id), + metadata_only=metadata == "only", + ) + except Exception as error: + _raise_document_error(error) metadata_fields = {"doc_type", "doc_metadata"} - if metadata == "only": - return ( - dump_response( - DocumentDetailResponse, - { - "id": document.id, - "doc_type": document.doc_type, - "doc_metadata": document.get_doc_metadata_details(session=session), - }, - include={"id", *metadata_fields}, - exclude_unset=True, - ), - 200, - ) + return dump_response( + DocumentDetailResponse, + result, + include={"id", *metadata_fields} if metadata == "only" else None, + exclude=metadata_fields if metadata == "without" else None, + exclude_unset=True, + ), 200 - dataset_process_rules = DatasetService.get_process_rules(dataset_id_str, session) - document_process_rule = document.get_dataset_process_rule(session=session) - document_process_rules: Mapping[str, Any] = document_process_rule.to_dict() if document_process_rule else {} - segment_count = document.get_segment_count(session=session) - response_payload = { - "id": document.id, - "position": document.position, - "data_source_type": document.data_source_type, - "data_source_info": document.data_source_info_dict, - "data_source_detail_dict": document.get_data_source_detail_dict(session=session), - "dataset_process_rule_id": document.dataset_process_rule_id, - "dataset_process_rule": dataset_process_rules, - "document_process_rule": document_process_rules, - "name": document.name, - "created_from": document.created_from, - "created_by": document.created_by, - "created_at": int(document.created_at.timestamp()), - "tokens": document.tokens, - "indexing_status": document.indexing_status, - "completed_at": int(document.completed_at.timestamp()) if document.completed_at else None, - "updated_at": int(document.updated_at.timestamp()) if document.updated_at else None, - "indexing_latency": document.indexing_latency, - "error": document.error, - "enabled": document.enabled, - "disabled_at": int(document.disabled_at.timestamp()) if document.disabled_at else None, - "disabled_by": document.disabled_by, - "archived": document.archived, - "doc_type": document.doc_type, - "doc_metadata": document.get_doc_metadata_details(session=session), - "segment_count": segment_count, - "average_segment_length": (document.word_count or 0) // segment_count if segment_count else 0, - "hit_count": document.get_hit_count(session=session), - "display_status": document.display_status, - "doc_form": document.doc_form, - "doc_language": document.doc_language, - "need_summary": document.need_summary if document.need_summary is not None else False, - } - exclude = metadata_fields if metadata == "without" else None - return dump_response(DocumentDetailResponse, response_payload, exclude=exclude, exclude_unset=True), 200 - - @setup_required - @login_required - @account_initialization_required - @cloud_edition_billing_rate_limit_check("knowledge") @console_ns.response(204, "Document deleted successfully") - @with_current_user - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_DELETE_FILE, DatasetId())) - @with_session - def delete( - self, session: Session, current_tenant_id: str, current_user: Account, dataset_id: UUID, document_id: UUID - ): - dataset_id_str = str(dataset_id) - document_id_str = str(document_id) - dataset = DatasetService.get_dataset(dataset_id_str, session) - if dataset is None: - raise NotFound("Dataset not found.") - # check user's model setting - DatasetService.check_dataset_model_setting(dataset) - - document = self.get_document(session, dataset_id_str, document_id_str, current_user, current_tenant_id) - + @console_account_admission( + allowed_roles=_DATASET_EDIT_ROLES, + rbac_checks=(RBACCheck(RBACPermission.DATASET_DELETE_FILE, DatasetId()),), + ) + @cloud_edition_billing_rate_limit_check("knowledge") + def delete(self, request_context: RequestContext, dataset_id: UUID, document_id: UUID): try: - DocumentService.delete_document(document, session) - except services.errors.document.DocumentIndexingError: - raise DocumentIndexingError("Cannot delete document during indexing.") - + application_services().knowledge.documents.delete_document( + request_context, dataset_id=str(dataset_id), document_id=str(document_id) + ) + except Exception as error: + _raise_document_error(error) return "", 204 @console_ns.route("/datasets//documents//download") -class DocumentDownloadApi(DocumentResource): +class DocumentDownloadApi(Resource): """Return a signed download URL for a dataset document's original uploaded file.""" @console_ns.doc("get_dataset_document_download_url") @console_ns.doc(description="Get a signed download URL for a dataset document's original uploaded file") @console_ns.response(200, "Download URL generated successfully", console_ns.models[UrlResponse.__name__]) - @setup_required - @login_required - @account_initialization_required + @console_account_admission(rbac_checks=(RBACCheck(RBACPermission.DATASET_DOCUMENT_DOWNLOAD, DatasetId()),)) @cloud_edition_billing_rate_limit_check("knowledge") - @with_current_user - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_DOCUMENT_DOWNLOAD, DatasetId())) - @with_session(write=False) - def get( - self, session: Session, current_tenant_id: str, current_user: Account, dataset_id: UUID, document_id: UUID - ) -> dict[str, Any]: - # Reuse the shared permission/tenant checks implemented in DocumentResource. - document = self.get_document(session, str(dataset_id), str(document_id), current_user, current_tenant_id) - return UrlResponse(url=DocumentService.get_document_download_url(document, session)).model_dump(mode="json") + def get(self, request_context: RequestContext, dataset_id: UUID, document_id: UUID): + try: + result = application_services().knowledge.documents.get_download_url( + request_context, dataset_id=str(dataset_id), document_id=str(document_id) + ) + except Exception as error: + _raise_document_error(error) + return UrlResponse(url=result).model_dump(mode="json") @console_ns.route("/datasets//documents/download-zip") -class DocumentBatchDownloadZipApi(DocumentResource): +class DocumentBatchDownloadZipApi(Resource): """Download multiple uploaded-file documents as a single ZIP (avoids browser multi-download limits).""" @console_ns.doc("download_dataset_documents_as_zip") @console_ns.doc(description="Download selected dataset documents as a single ZIP archive (upload-file only)") @console_ns.response(200, "ZIP archive downloaded successfully") - @setup_required - @login_required - @account_initialization_required - @cloud_edition_billing_rate_limit_check("knowledge") @console_ns.expect(console_ns.models[DocumentBatchDownloadZipPayload.__name__]) - @with_current_user - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_DOCUMENT_DOWNLOAD, DatasetId())) - @with_session(write=False) - def post(self, session: Session, current_tenant_id: str, current_user: Account, dataset_id: UUID): + @console_account_admission( + allowed_roles=_DATASET_EDIT_ROLES, + rbac_checks=(RBACCheck(RBACPermission.DATASET_DOCUMENT_DOWNLOAD, DatasetId()),), + ) + @cloud_edition_billing_rate_limit_check("knowledge") + def post(self, request_context: RequestContext, dataset_id: UUID): """Stream a ZIP archive containing the requested uploaded documents.""" - # Parse and validate request payload. payload = DocumentBatchDownloadZipPayload.model_validate(console_ns.payload or {}) - - dataset_id_str = str(dataset_id) - document_ids: list[str] = [str(document_id) for document_id in payload.document_ids] - upload_files, download_name = DocumentService.prepare_document_batch_download_zip( - dataset_id=dataset_id_str, - document_ids=document_ids, - tenant_id=current_tenant_id, - current_user=current_user, - session=session, - ) - - # Delegate ZIP packing to FileService, but keep Flask response+cleanup in the route. - with ExitStack() as stack: - zip_path = stack.enter_context(FileService.build_upload_files_zip_tempfile(upload_files=upload_files)) - response = send_file( - zip_path, - mimetype="application/zip", - as_attachment=True, - download_name=download_name, - ) - cleanup = stack.pop_all() - response.call_on_close(cleanup.close) + try: + with ExitStack() as stack: + archive = stack.enter_context( + application_services().knowledge.documents.build_download_zip( + request_context, + dataset_id=str(dataset_id), + document_ids=[str(value) for value in payload.document_ids], + ) + ) + response = send_file( + archive.path, mimetype="application/zip", as_attachment=True, download_name=archive.filename + ) + cleanup = stack.pop_all() + response.call_on_close(cleanup.close) + except Exception as error: + _raise_document_error(error) # response-contract:ignore binary ZIP download response return response @console_ns.route("/datasets//documents//processing/") -class DocumentProcessingApi(DocumentResource): +class DocumentProcessingApi(Resource): @console_ns.doc("update_document_processing") @console_ns.doc(description="Update document processing status (pause/resume)") @console_ns.doc( @@ -1235,53 +661,24 @@ class DocumentProcessingApi(DocumentResource): ) @console_ns.response(404, "Document not found") @console_ns.response(400, "Invalid action") - @setup_required - @login_required - @account_initialization_required + @console_account_admission( + allowed_roles=_DATASET_EDIT_ROLES, rbac_checks=(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()),) + ) @cloud_edition_billing_rate_limit_check("knowledge") - @with_current_user - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId())) - @with_session def patch( - self, - session: Session, - current_tenant_id: str, - current_user: Account, - dataset_id: UUID, - document_id: UUID, - action: Literal["pause", "resume"], + self, request_context: RequestContext, dataset_id: UUID, document_id: UUID, action: Literal["pause", "resume"] ): - dataset_id_str = str(dataset_id) - document_id_str = str(document_id) - document = self.get_document(session, dataset_id_str, document_id_str, current_user, current_tenant_id) - - # The role of the current user in the ta table must be admin, owner, dataset_operator, or editor - if not current_user.is_dataset_editor: - raise Forbidden() - - match action: - case "pause": - if document.indexing_status != IndexingStatus.INDEXING: - raise InvalidActionError("Document not in indexing state.") - - document.paused_by = current_user.id - document.paused_at = naive_utc_now() - document.is_paused = True - - case "resume": - if document.indexing_status not in {IndexingStatus.PAUSED, IndexingStatus.ERROR}: - raise InvalidActionError("Document not in paused or error state.") - - document.paused_by = None - document.paused_at = None - document.is_paused = False - + try: + application_services().knowledge.documents.update_processing( + request_context, dataset_id=str(dataset_id), document_id=str(document_id), action=action + ) + except Exception as error: + _raise_document_error(error) return SimpleResultResponse(result="success").model_dump(mode="json"), 200 @console_ns.route("/datasets//documents//metadata") -class DocumentMetadataApi(DocumentResource): +class DocumentMetadataApi(Resource): @console_ns.doc("update_document_metadata") @console_ns.doc(description="Update document metadata") @console_ns.doc(params={"dataset_id": "Dataset ID", "document_id": "Document ID"}) @@ -1293,368 +690,167 @@ class DocumentMetadataApi(DocumentResource): ) @console_ns.response(404, "Document not found") @console_ns.response(403, "Permission denied") - @setup_required - @login_required - @account_initialization_required - @with_current_user - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId())) - @with_session + @console_account_admission( + allowed_roles=_DATASET_EDIT_ROLES, rbac_checks=(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()),) + ) @model_validate(DocumentMetadataUpdatePayload) def put( self, req_data: DocumentMetadataUpdatePayload, - session: Session, - current_tenant_id: str, - current_user: Account, + request_context: RequestContext, dataset_id: UUID, document_id: UUID, ): - dataset_id_str = str(dataset_id) - document_id_str = str(document_id) - document = self.get_document(session, dataset_id_str, document_id_str, current_user, current_tenant_id) - - doc_type = req_data.doc_type - doc_metadata = req_data.doc_metadata - - # The role of the current user in the ta table must be admin, owner, dataset_operator, or editor - if not current_user.is_dataset_editor: - raise Forbidden() - - if doc_type is None or doc_metadata is None: - raise ValueError("Both doc_type and doc_metadata must be provided.") - - if doc_type not in DocumentService.DOCUMENT_METADATA_SCHEMA: - raise ValueError("Invalid doc_type.") - - if not isinstance(doc_metadata, dict): - raise ValueError("doc_metadata must be a dictionary.") - metadata_schema: dict[str, Any] = cast(dict[str, Any], DocumentService.DOCUMENT_METADATA_SCHEMA[doc_type]) - - document.doc_metadata = {} - if doc_type == "others": - document.doc_metadata = doc_metadata - else: - for key, value_type in metadata_schema.items(): - value = doc_metadata.get(key) - if value is not None and isinstance(value, value_type): - document.doc_metadata[key] = value - - document.doc_type = doc_type - document.updated_at = naive_utc_now() - + try: + application_services().knowledge.documents.update_metadata( + request_context, + dataset_id=str(dataset_id), + document_id=str(document_id), + doc_type=req_data.doc_type, + doc_metadata=req_data.doc_metadata, + ) + except Exception as error: + _raise_document_error(error) return SimpleResultMessageResponse(result="success", message="Document metadata updated.").model_dump( mode="json" ), 200 @console_ns.route("/datasets//documents/status//batch") -class DocumentStatusApi(DocumentResource): - @setup_required - @login_required - @account_initialization_required +class DocumentStatusApi(Resource): + @console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__]) + @console_account_admission( + allowed_roles=_DATASET_EDIT_ROLES, rbac_checks=(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()),) + ) @cloud_edition_billing_resource_check("vector_space") @cloud_edition_billing_rate_limit_check("knowledge") - @console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__]) - @with_current_user - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId())) - @with_session def patch( self, - session: Session, - current_user: Account, + request_context: RequestContext, dataset_id: UUID, action: Literal["enable", "disable", "archive", "un_archive"], ): - dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, session) - if dataset is None: - raise NotFound("Dataset not found.") - - # The role of the current user in the ta table must be admin, owner, or editor - if not current_user.is_dataset_editor: - raise Forbidden() - - # check user's model setting - DatasetService.check_dataset_model_setting(dataset) - - # check user's permission - DatasetService.check_dataset_permission(dataset, current_user, session) - - document_ids = request.args.getlist("document_id") - try: - DocumentService.batch_update_document_status(dataset, document_ids, action, current_user, session) - except services.errors.document.DocumentIndexingError as e: - raise InvalidActionError(str(e)) - except ValueError as e: - raise InvalidActionError(str(e)) - except NotFound as e: - raise NotFound(str(e)) - + application_services().knowledge.documents.change_status( + request_context, + dataset_id=str(dataset_id), + document_ids=request.args.getlist("document_id"), + action=action, + ) + except Exception as error: + _raise_document_error(error) return SimpleResultResponse(result="success").model_dump(mode="json"), 200 @console_ns.route("/datasets//documents//processing/pause") -class DocumentPauseApi(DocumentResource): - @setup_required - @login_required - @account_initialization_required +class DocumentPauseApi(Resource): @console_ns.response(204, "Document paused successfully") - @with_current_user - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId())) - @with_session - def patch( - self, - session: Session, - current_tenant_id: str, - current_user: Account, - dataset_id: UUID, - document_id: UUID, - ): + @console_account_admission( + allowed_roles=_DATASET_EDIT_ROLES, rbac_checks=(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()),) + ) + def patch(self, request_context: RequestContext, dataset_id: UUID, document_id: UUID): """pause document.""" - dataset_id_str = str(dataset_id) - document_id_str = str(document_id) - - document = self.get_document(session, dataset_id_str, document_id_str, current_user, current_tenant_id) - if not current_user.is_dataset_editor: - raise Forbidden() - - # 403 if document is archived - if DocumentService.check_archived(document): - raise ArchivedDocumentImmutableError() - check_knowledge_rate_limit() try: - # pause document - DocumentService.pause_document(document, session) - except services.errors.document.DocumentIndexingError: - raise DocumentIndexingError("Cannot pause completed document.") - + application_services().knowledge.documents.pause_document( + request_context, dataset_id=str(dataset_id), document_id=str(document_id) + ) + except Exception as error: + _raise_document_error(error) return "", 204 @console_ns.route("/datasets//documents//processing/resume") -class DocumentRecoverApi(DocumentResource): - @setup_required - @login_required - @account_initialization_required +class DocumentRecoverApi(Resource): @console_ns.response(204, "Document resumed successfully") - @with_current_user - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId())) - @with_session - def patch( - self, - session: Session, - current_tenant_id: str, - current_user: Account, - dataset_id: UUID, - document_id: UUID, - ): + @console_account_admission( + allowed_roles=_DATASET_EDIT_ROLES, rbac_checks=(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()),) + ) + def patch(self, request_context: RequestContext, dataset_id: UUID, document_id: UUID): """recover document.""" - dataset_id_str = str(dataset_id) - document_id_str = str(document_id) - - document = self.get_document(session, dataset_id_str, document_id_str, current_user, current_tenant_id) - if not current_user.is_dataset_editor: - raise Forbidden() - - # 403 if document is archived - if DocumentService.check_archived(document): - raise ArchivedDocumentImmutableError() check_knowledge_rate_limit() try: - # pause document - DocumentService.recover_document(document, session) - except services.errors.document.DocumentIndexingError: - raise DocumentIndexingError("Document is not in paused status.") - + application_services().knowledge.documents.recover_document( + request_context, dataset_id=str(dataset_id), document_id=str(document_id) + ) + except Exception as error: + _raise_document_error(error) return "", 204 @console_ns.route("/datasets//retry") -class DocumentRetryApi(DocumentResource): - @setup_required - @login_required - @account_initialization_required +class DocumentRetryApi(Resource): @console_ns.expect(console_ns.models[DocumentRetryPayload.__name__]) @console_ns.response(204, "Documents retry started successfully") - @with_current_user - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId())) - @with_session + @console_account_admission( + allowed_roles=_DATASET_EDIT_ROLES, rbac_checks=(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()),) + ) @model_validate(DocumentRetryPayload) - def post( - self, - req_data: DocumentRetryPayload, - session: Session, - current_tenant_id: str, - current_user: Account, - dataset_id: UUID, - ): + def post(self, req_data: DocumentRetryPayload, request_context: RequestContext, dataset_id: UUID): """retry document.""" - dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset_for_tenant(dataset_id_str, current_tenant_id, session=session) - if not dataset: - raise NotFound("Dataset not found.") - - if not current_user.is_dataset_editor: - raise Forbidden() - - try: - DatasetService.check_dataset_permission(dataset, current_user, session) - except services.errors.account.NoPermissionError as e: - raise Forbidden(str(e)) - - documents = DocumentService.get_documents_by_ids( - DatasetRefService.create_dataset_ref(dataset), req_data.document_ids, session - ) - documents_by_id = {document.id: document for document in documents} - retry_documents = [] - for document_id in req_data.document_ids: - try: - document = documents_by_id.get(document_id) - - # 404 if document not found - if document is None: - raise NotFound("Document Not Exists.") - - # 403 if document is archived - if DocumentService.check_archived(document): - raise ArchivedDocumentImmutableError() - - # 400 if document is completed - if document.indexing_status == IndexingStatus.COMPLETED: - raise DocumentAlreadyFinishedError() - retry_documents.append(document) - except Exception: - logger.exception("Failed to retry document, document id: %s", document_id) - continue - # retry document check_knowledge_rate_limit() - DocumentService.retry_document(dataset_id_str, retry_documents, session) - + try: + application_services().knowledge.documents.retry_documents( + request_context, dataset_id=str(dataset_id), document_ids=req_data.document_ids + ) + except Exception as error: + _raise_document_error(error) return "", 204 @console_ns.route("/datasets//documents//rename") -class DocumentRenameApi(DocumentResource): - @setup_required - @login_required - @account_initialization_required +class DocumentRenameApi(Resource): @console_ns.response(200, "Document renamed successfully", console_ns.models[DocumentResponse.__name__]) @console_ns.expect(console_ns.models[DocumentRenamePayload.__name__]) - @with_current_user - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId())) - @with_session + @console_account_admission( + allowed_roles=_DATASET_EDIT_ROLES, rbac_checks=(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()),) + ) @model_validate(DocumentRenamePayload) def post( - self, - req_data: DocumentRenamePayload, - session: Session, - current_user: Account, - dataset_id: UUID, - document_id: UUID, + self, req_data: DocumentRenamePayload, request_context: RequestContext, dataset_id: UUID, document_id: UUID ): - # The role of the current user in the ta table must be admin, owner, editor, or dataset_operator - if not current_user.is_dataset_editor: - raise Forbidden() - dataset = DatasetService.get_dataset(dataset_id, session) - if not dataset: - raise NotFound("Dataset not found.") - DatasetService.check_dataset_operator_permission(current_user, dataset, session=session) - try: - document = DocumentService.rename_document(str(dataset_id), str(document_id), req_data.name, session) - except services.errors.document.DocumentIndexingError: - raise DocumentIndexingError("Cannot delete document during indexing.") - - return dump_response(DocumentResponse, document_response(document, session=session)) + result = application_services().knowledge.documents.rename_document( + request_context, dataset_id=str(dataset_id), document_id=str(document_id), name=req_data.name + ) + except Exception as error: + _raise_document_error(error) + return dump_response(DocumentResponse, result) @console_ns.route("/datasets//documents//website-sync") -class WebsiteDocumentSyncApi(DocumentResource): - @setup_required - @login_required - @account_initialization_required +class WebsiteDocumentSyncApi(Resource): @console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__]) - @with_current_user - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId())) - @with_session - def get( - self, - session: Session, - current_tenant_id: str, - current_user: Account, - dataset_id: UUID, - document_id: UUID, - ): + @console_account_admission( + allowed_roles=_DATASET_EDIT_ROLES, rbac_checks=(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()),) + ) + def get(self, request_context: RequestContext, dataset_id: UUID, document_id: UUID): """sync website document.""" - dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset_for_tenant(dataset_id_str, current_tenant_id, session=session) - if not dataset: - raise NotFound("Dataset not found.") - if not current_user.is_dataset_editor: - raise Forbidden() - document_id_str = str(document_id) - document = self.get_document(session, dataset.id, document_id_str, current_user, current_tenant_id) - if document.data_source_type != "website_crawl": - raise ValueError("Document is not a website document.") - # 403 if document is archived - if DocumentService.check_archived(document): - raise ArchivedDocumentImmutableError() - # sync document - DocumentService.sync_website_document(dataset, document, session) - + try: + application_services().knowledge.documents.sync_website( + request_context, dataset_id=str(dataset_id), document_id=str(document_id) + ) + except Exception as error: + _raise_document_error(error) return SimpleResultResponse(result="success").model_dump(mode="json"), 200 @console_ns.route("/datasets//documents//pipeline-execution-log") -class DocumentPipelineExecutionLogApi(DocumentResource): +class DocumentPipelineExecutionLogApi(Resource): @console_ns.response( 200, "Pipeline execution log retrieved successfully", console_ns.models[DocumentPipelineExecutionLogResponse.__name__], ) - @setup_required - @login_required - @account_initialization_required - @with_current_user - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_READONLY, DatasetId())) - @with_session(write=False) - def get( - self, - session: Session, - current_tenant_id: str, - current_user: Account, - dataset_id: UUID, - document_id: UUID, - ): - dataset_id_str = str(dataset_id) - document_id_str = str(document_id) - - document = self.get_document(session, dataset_id_str, document_id_str, current_user, current_tenant_id) - log = session.scalar( - select(DocumentPipelineExecutionLog) - .where(DocumentPipelineExecutionLog.document_id == document.id) - .order_by(DocumentPipelineExecutionLog.created_at.desc()) - .limit(1) - ) - if not log: - return DocumentPipelineExecutionLogResponse().model_dump(mode="json"), 200 - return dump_response( - DocumentPipelineExecutionLogResponse, - { - "datasource_info": json.loads(log.datasource_info), - "datasource_type": log.datasource_type, - "input_data": log.input_data, - "datasource_node_id": log.datasource_node_id, - }, - ), 200 + @console_account_admission(rbac_checks=(RBACCheck(RBACPermission.DATASET_READONLY, DatasetId()),)) + def get(self, request_context: RequestContext, dataset_id: UUID, document_id: UUID): + try: + result = application_services().knowledge.documents.get_execution_log( + request_context, dataset_id=str(dataset_id), document_id=str(document_id) + ) + except Exception as error: + _raise_document_error(error) + return dump_response(DocumentPipelineExecutionLogResponse, result), 200 @console_ns.route("/datasets//documents/generate-summary") @@ -1671,15 +867,12 @@ class DocumentGenerateSummaryApi(Resource): @console_ns.response(400, "Invalid request or dataset configuration") @console_ns.response(403, "Permission denied") @console_ns.response(404, "Dataset not found") - @setup_required - @login_required - @account_initialization_required + @console_account_admission( + allowed_roles=_DATASET_EDIT_ROLES, rbac_checks=(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()),) + ) @cloud_edition_billing_rate_limit_check("knowledge") - @with_current_user - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId())) - @with_session @model_validate(GenerateSummaryPayload) - def post(self, req_data: GenerateSummaryPayload, session: Session, current_user: Account, dataset_id: UUID): + def post(self, req_data: GenerateSummaryPayload, request_context: RequestContext, dataset_id: UUID): """ Generate summary index for specified documents. @@ -1687,83 +880,19 @@ class DocumentGenerateSummaryApi(Resource): (indexing_technique must be 'high_quality' and summary_index_setting.enable must be true), then asynchronously generates summary indexes for the provided documents. """ - dataset_id_str = str(dataset_id) - - # Get dataset - dataset = DatasetService.get_dataset(dataset_id_str, session) - if not dataset: - raise NotFound("Dataset not found.") - - # Check permissions - if not current_user.is_dataset_editor: - raise Forbidden() - - try: - DatasetService.check_dataset_permission(dataset, current_user, session) - except services.errors.account.NoPermissionError as e: - raise Forbidden(str(e)) - - document_list = req_data.document_list - - if not document_list: - from werkzeug.exceptions import BadRequest - + if not req_data.document_list: raise BadRequest("document_list cannot be empty.") - - # Check if dataset configuration supports summary generation - if dataset.indexing_technique != IndexTechniqueType.HIGH_QUALITY: - raise ValueError( - f"Summary generation is only available for 'high_quality' indexing technique. " - f"Current indexing technique: {dataset.indexing_technique}" + try: + application_services().knowledge.documents.generate_summary( + request_context, dataset_id=str(dataset_id), document_ids=req_data.document_list ) - - summary_index_setting = dataset.summary_index_setting - if not summary_index_setting or not summary_index_setting.get("enable"): - raise ValueError("Summary index is not enabled for this dataset. Please enable it in the dataset settings.") - - # Verify all documents exist and belong to the dataset - documents = DocumentService.get_documents_by_ids( - DatasetRefService.create_dataset_ref(dataset), document_list, session - ) - - if len(documents) != len(document_list): - found_ids = {doc.id for doc in documents} - missing_ids = set(document_list) - found_ids - raise NotFound(f"Some documents not found: {list(missing_ids)}") - - # Update need_summary to True for documents that don't have it set - # This handles the case where documents were created when summary_index_setting was disabled - documents_to_update = [doc for doc in documents if not doc.need_summary and doc.doc_form != "qa_model"] - - if documents_to_update: - document_ids_to_update = [str(doc.id) for doc in documents_to_update] - DocumentService.update_documents_need_summary( - dataset_id=dataset_id_str, - document_ids=document_ids_to_update, - session=session, - need_summary=True, - ) - - # Dispatch async tasks for each document - for document in documents: - # Skip qa_model documents as they don't generate summaries - if document.doc_form == "qa_model": - logger.info("Skipping summary generation for qa_model document %s", document.id) - continue - - # Dispatch async task - generate_summary_index_task.delay(dataset_id_str, document.id) - logger.info( - "Dispatched summary generation task for document %s in dataset %s", - document.id, - dataset_id_str, - ) - + except Exception as error: + _raise_document_error(error) return SimpleResultResponse(result="success").model_dump(mode="json"), 200 @console_ns.route("/datasets//documents//summary-status") -class DocumentSummaryStatusApi(DocumentResource): +class DocumentSummaryStatusApi(Resource): @console_ns.doc("get_document_summary_status") @console_ns.doc(description="Get summary index generation status for a document") @console_ns.doc(params={"dataset_id": "Dataset ID", "document_id": "Document ID"}) @@ -1773,13 +902,8 @@ class DocumentSummaryStatusApi(DocumentResource): console_ns.models[DocumentSummaryStatusResponse.__name__], ) @console_ns.response(404, "Document not found") - @setup_required - @login_required - @account_initialization_required - @with_current_user - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_READONLY, DatasetId())) - @with_session(write=False) - def get(self, session: Session, current_user: Account, dataset_id: UUID, document_id: UUID): + @console_account_admission(rbac_checks=(RBACCheck(RBACPermission.DATASET_READONLY, DatasetId()),)) + def get(self, request_context: RequestContext, dataset_id: UUID, document_id: UUID): """ Get summary index generation status for a document. @@ -1793,27 +917,10 @@ class DocumentSummaryStatusApi(DocumentResource): - timeout: Number of summaries that timed out - summaries: List of summary records with status and content preview """ - dataset_id_str = str(dataset_id) - document_id_str = str(document_id) - - # Get dataset - dataset = DatasetService.get_dataset(dataset_id_str, session) - if not dataset: - raise NotFound("Dataset not found.") - - # Check permissions try: - DatasetService.check_dataset_permission(dataset, current_user, session) - except services.errors.account.NoPermissionError as e: - raise Forbidden(str(e)) - - # Get summary status detail from service - from services.summary_index_service import SummaryIndexService - - result = SummaryIndexService.get_document_summary_status_detail( - document_id=document_id_str, - dataset_id=dataset_id_str, - session=session, - ) - + result = application_services().knowledge.documents.get_summary_status( + request_context, dataset_id=str(dataset_id), document_id=str(document_id) + ) + except Exception as error: + _raise_document_error(error) return dump_response(DocumentSummaryStatusResponse, result), 200 diff --git a/api/controllers/console/datasets/datasets_segments.py b/api/controllers/console/datasets/datasets_segments.py index fdaa47c13f5..44e35cf0fbe 100644 --- a/api/controllers/console/datasets/datasets_segments.py +++ b/api/controllers/console/datasets/datasets_segments.py @@ -1,18 +1,10 @@ -import uuid -from typing import Literal -from typing import cast as type_cast +from typing import Literal, Never from uuid import UUID -from flask import request from flask_restx import Resource from pydantic import BaseModel, Field -from sqlalchemy import String, case, cast, func, literal, or_, select -from sqlalchemy.dialects.postgresql import JSONB -from sqlalchemy.orm import Session from werkzeug.exceptions import Forbidden, NotFound -import services -from configs import dify_config from controllers.common.controller_schemas import ChildChunkCreatePayload, ChildChunkUpdatePayload from controllers.common.fields import SimpleResultResponse from controllers.common.rbac import DatasetId, RBACCheck @@ -22,7 +14,6 @@ from controllers.common.schema import ( register_response_schema_models, register_schema_models, ) -from controllers.common.session import with_session from controllers.console import console_ns from controllers.console.app.error import ProviderNotInitializeError from controllers.console.datasets.error import ( @@ -30,22 +21,15 @@ from controllers.console.datasets.error import ( ChildChunkIndexingError, InvalidActionError, ) +from controllers.console.flask_admission import console_account_admission from controllers.console.wraps import ( RBACPermission, - account_initialization_required, 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, - with_current_user, ) -from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError -from core.model_manager import ModelManager -from core.rag.index_processor.constant.index_type import IndexTechniqueType -from extensions.ext_redis import redis_client +from extensions.ext_application_services import application_services from fields.base import ResponseModel from fields.segment_fields import ( ChildChunkDetailResponse, @@ -53,23 +37,39 @@ from fields.segment_fields import ( ChildChunkResponse, SegmentDetailResponse, SegmentResponse, - segment_response_with_summary, - segment_responses_with_summaries, ) -from graphon.model_runtime.entities.model_entities import ModelType -from libs.helper import dump_response, escape_like_pattern -from libs.login import login_required -from libs.pagination import paginate_query -from models import Account -from models.dataset import Dataset, Document, DocumentSegment -from models.model import UploadFile -from services.dataset_ref_service import DatasetRefService, SegmentRef -from services.dataset_service import DatasetService, DocumentService, SegmentService -from services.entities.knowledge_entities.knowledge_entities import ChildChunkUpdateArgs, SegmentUpdateArgs -from services.errors.chunk import ChildChunkDeleteIndexError as ChildChunkDeleteIndexServiceError -from services.errors.chunk import ChildChunkIndexingError as ChildChunkIndexingServiceError -from services.summary_index_service import SummaryIndexService -from tasks.batch_create_segment_to_index_task import batch_create_segment_to_index_task +from libs.helper import dump_response +from machinery.context import RequestContext +from models.account import TenantAccountRole +from services.knowledge.entities.segments import ChildChunkUpdateArgs +from services.knowledge.segments.application import ( + ChildChunkDeleteIndexApplicationError, + ChildChunkIndexingApplicationError, + ChildChunkListFilter, + ChildChunkNotFoundError, + SegmentBatchImportDispatchError, + SegmentBatchImportNotFoundError, + SegmentDatasetModelUnavailableError, + SegmentDatasetNotFoundError, + SegmentDocumentIndexingError, + SegmentDocumentNotFoundError, + SegmentEmbeddingModelUnavailableError, + SegmentInvalidFileTypeError, + SegmentListFilter, + SegmentNotFoundError, + SegmentPermissionDeniedError, + SegmentStatusUpdateError, + SegmentUploadFileNotFoundError, +) + +_DATASET_EDIT_ROLES = frozenset( + { + TenantAccountRole.OWNER, + TenantAccountRole.ADMIN, + TenantAccountRole.EDITOR, + TenantAccountRole.DATASET_OPERATOR, + } +) class SegmentListQuery(BaseModel): @@ -166,19 +166,32 @@ register_response_schema_models( ) -def _get_segment_for_document( - session: Session, dataset: Dataset, document: Document, segment_id: str -) -> tuple[SegmentRef, DocumentSegment]: - dataset_ref = DatasetRefService.create_dataset_ref(dataset) - document_ref = DatasetRefService.create_document_ref(dataset_ref, document) - if document_ref is None: - raise NotFound("Document not found.") - - segment_ref = DatasetRefService.create_segment_ref(document_ref, segment_id) - segment = SegmentService.get_segment_by_ref(segment_ref, session=session) - if not segment: - raise NotFound("Segment not found.") - return segment_ref, segment +def _raise_segment_error(error: Exception) -> Never: + if isinstance(error, SegmentDatasetNotFoundError): + raise NotFound("Dataset not found.") from error + if isinstance(error, SegmentDocumentNotFoundError): + raise NotFound("Document not found.") from error + if isinstance(error, SegmentNotFoundError): + raise NotFound("Segment not found.") from error + if isinstance(error, ChildChunkNotFoundError): + raise NotFound("Child chunk not found.") from error + if isinstance(error, SegmentUploadFileNotFoundError): + raise NotFound("UploadFile not found.") from error + if isinstance(error, SegmentPermissionDeniedError): + raise Forbidden(str(error)) from error + if isinstance(error, SegmentEmbeddingModelUnavailableError): + raise ProviderNotInitializeError(str(error)) from error + if isinstance(error, SegmentDatasetModelUnavailableError): + raise ValueError(str(error)) from error + if isinstance(error, (SegmentDocumentIndexingError, SegmentStatusUpdateError)): + raise InvalidActionError(str(error)) from error + if isinstance(error, SegmentInvalidFileTypeError): + raise ValueError(str(error)) from error + if isinstance(error, ChildChunkIndexingApplicationError): + raise ChildChunkIndexingError(str(error)) from error + if isinstance(error, ChildChunkDeleteIndexApplicationError): + raise ChildChunkDeleteIndexError(str(error)) from error + raise error @console_ns.route("/datasets//documents//segments") @@ -186,147 +199,60 @@ class DatasetDocumentSegmentListApi(Resource): @console_ns.doc(params=SegmentDocParams.DATASET_DOCUMENT) @console_ns.doc(params=query_params_from_model(SegmentListQuery)) @console_ns.response(200, "Segments retrieved successfully", console_ns.models[ConsoleSegmentListResponse.__name__]) - @setup_required - @login_required - @account_initialization_required - @with_current_user - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_READONLY, DatasetId())) - @with_session(write=False) - def get(self, session: Session, current_tenant_id: str, current_user: Account, dataset_id: UUID, document_id: UUID): - dataset_id_str = str(dataset_id) - document_id_str = str(document_id) - dataset = DatasetService.get_dataset(dataset_id_str, session) - if not dataset: - raise NotFound("Dataset not found.") - - try: - DatasetService.check_dataset_permission(dataset, current_user, session) - except services.errors.account.NoPermissionError as e: - raise Forbidden(str(e)) - - document = DocumentService.get_document(dataset_id_str, document_id_str, session=session) - - if not document: - raise NotFound("Document not found.") - + @console_account_admission( + rbac_checks=(RBACCheck(RBACPermission.DATASET_READONLY, DatasetId()),), + ) + def get(self, request_context: RequestContext, dataset_id: UUID, document_id: UUID): args = query_params_from_request(SegmentListQuery, list_fields=("status",)) - - page = args.page - limit = min(args.limit, 100) - status_list = args.status - hit_count_gte = args.hit_count_gte - keyword = args.keyword - - query = ( - select(DocumentSegment) - .where( - DocumentSegment.document_id == document_id_str, - DocumentSegment.tenant_id == current_tenant_id, + try: + result = application_services().knowledge.segments.list_segments( + request_context, + dataset_id=str(dataset_id), + document_id=str(document_id), + query=SegmentListFilter( + page=args.page, + limit=args.limit, + statuses=tuple(args.status), + hit_count_gte=args.hit_count_gte, + enabled=args.enabled, + keyword=args.keyword, + ), ) - .order_by(DocumentSegment.position.asc()) - ) - - if status_list: - query = query.where(DocumentSegment.status.in_(status_list)) - - if hit_count_gte is not None: - query = query.where(DocumentSegment.hit_count >= hit_count_gte) - - if keyword: - # Escape special characters in keyword to prevent SQL injection via LIKE wildcards - escaped_keyword = escape_like_pattern(keyword) - # Search in both content and keywords fields - # Use database-specific methods for JSON array search - if dify_config.SQLALCHEMY_DATABASE_URI_SCHEME == "postgresql": - # PostgreSQL: Use jsonb_array_elements_text to properly handle Unicode/Chinese text - # Feed the set-returning function a JSON array in every row. Filtering in - # the subquery is not enough because PostgreSQL can still evaluate the - # SRF on scalar JSON before applying the predicate. - keywords_jsonb = cast(DocumentSegment.keywords, JSONB) - keywords_array = case( - (func.jsonb_typeof(keywords_jsonb) == "array", keywords_jsonb), - else_=cast(literal("[]"), JSONB), - ) - keywords_condition = func.array_to_string( - func.array( - select(func.jsonb_array_elements_text(keywords_array)) - .correlate(DocumentSegment) - .scalar_subquery() - ), - ",", - ).ilike(f"%{escaped_keyword}%", escape="\\") - else: - # MySQL: Cast JSON to string for pattern matching - # MySQL stores Chinese text directly in JSON without Unicode escaping - keywords_condition = cast(DocumentSegment.keywords, String).ilike(f"%{escaped_keyword}%", escape="\\") - - query = query.where( - or_( - DocumentSegment.content.ilike(f"%{escaped_keyword}%", escape="\\"), - keywords_condition, - ) - ) - - if args.enabled.lower() != "all": - if args.enabled.lower() == "true": - query = query.where(DocumentSegment.enabled == True) - elif args.enabled.lower() == "false": - query = query.where(DocumentSegment.enabled == False) - - segments = paginate_query(query, session=session, page=page, per_page=limit, max_per_page=100) - - segment_list = list(segments.items) - segment_ids = [segment.id for segment in segment_list] - summaries: dict[str, str | None] = {} - if segment_ids: - summary_records = SummaryIndexService.get_segments_summaries( - segment_ids=segment_ids, dataset_id=dataset_id_str, session=session - ) - summaries = {chunk_id: summary.summary_content for chunk_id, summary in summary_records.items()} - + except Exception as error: + _raise_segment_error(error) response = { - "data": segment_responses_with_summaries(segment_list, summaries, session=session), - "limit": limit, - "total": segments.total, - "total_pages": segments.pages, - "page": page, + "data": result.items, + "limit": result.limit, + "total": result.total, + "total_pages": result.total_pages, + "page": result.page, } return dump_response(ConsoleSegmentListResponse, response), 200 - @setup_required - @login_required - @account_initialization_required + @console_account_admission( + allowed_roles=_DATASET_EDIT_ROLES, + rbac_checks=(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()),), + ) @cloud_edition_billing_rate_limit_check("knowledge") @console_ns.doc(params=SegmentDocParams.DATASET_DOCUMENT) @console_ns.doc(params=query_params_from_model(SegmentIdListQuery)) @console_ns.response(204, "Segments deleted successfully") - @with_current_user - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId())) - @with_session - def delete(self, session: Session, current_user: Account, dataset_id: UUID, document_id: UUID): - # check dataset - dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, session) - if not dataset: - raise NotFound("Dataset not found.") - # check user's model setting - DatasetService.check_dataset_model_setting(dataset) - # check document - document_id_str = str(document_id) - document = DocumentService.get_document(dataset_id_str, document_id_str, session=session) - if not document: - raise NotFound("Document not found.") - segment_ids = request.args.getlist("segment_id") - - # The role of the current user in the ta table must be admin, owner, dataset_operator, or editor - if not current_user.is_dataset_editor: - raise Forbidden() + def delete( + self, + request_context: RequestContext, + dataset_id: UUID, + document_id: UUID, + ): + args = query_params_from_request(SegmentIdListQuery, list_fields=("segment_id",)) try: - DatasetService.check_dataset_permission(dataset, current_user, session) - except services.errors.account.NoPermissionError as e: - raise Forbidden(str(e)) - SegmentService.delete_segments(segment_ids, document, dataset, session) + application_services().knowledge.segments.delete_segments( + request_context, + dataset_id=str(dataset_id), + document_id=str(document_id), + segment_ids=args.segment_id, + ) + except Exception as error: + _raise_segment_error(error) return "", 204 @@ -334,141 +260,68 @@ class DatasetDocumentSegmentListApi(Resource): class DatasetDocumentSegmentApi(Resource): @console_ns.doc(params=SegmentDocParams.DATASET_DOCUMENT_ACTION) @console_ns.doc(params=query_params_from_model(SegmentIdListQuery)) - @setup_required - @login_required - @account_initialization_required + @console_account_admission( + allowed_roles=_DATASET_EDIT_ROLES, + rbac_checks=(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()),), + ) @cloud_edition_billing_resource_check("vector_space") @cloud_edition_billing_rate_limit_check("knowledge") @console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__]) - @with_current_user - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId())) - @with_session def patch( self, - session: Session, - current_tenant_id: str, - current_user: Account, + request_context: RequestContext, dataset_id: UUID, document_id: UUID, action: Literal["enable", "disable"], ): - dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, session) - if not dataset: - raise NotFound("Dataset not found.") - document_id_str = str(document_id) - document = DocumentService.get_document(dataset_id_str, document_id_str, session=session) - if not document: - raise NotFound("Document not found.") - # check user's model setting - DatasetService.check_dataset_model_setting(dataset) - # The role of the current user in the ta table must be admin, owner, dataset_operator, or editor - if not current_user.is_dataset_editor: - raise Forbidden() - + if action not in ("enable", "disable"): + raise InvalidActionError() + args = query_params_from_request(SegmentIdListQuery, list_fields=("segment_id",)) try: - DatasetService.check_dataset_permission(dataset, current_user, session) - except services.errors.account.NoPermissionError as e: - raise Forbidden(str(e)) - if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY: - # check embedding model setting - try: - model_manager = ModelManager.for_tenant(tenant_id=current_tenant_id) - model_manager.get_model_instance( - tenant_id=current_tenant_id, - provider=dataset.embedding_model_provider, - model_type=ModelType.TEXT_EMBEDDING, - model=dataset.embedding_model, - ) - except LLMBadRequestError: - raise ProviderNotInitializeError( - "No Embedding Model available. Please configure a valid provider in the Settings -> Model Provider." - ) - except ProviderTokenNotInitError as ex: - raise ProviderNotInitializeError(ex.description) - segment_ids = request.args.getlist("segment_id") - - document_indexing_cache_key = f"document_{document.id}_indexing" - cache_result = redis_client.get(document_indexing_cache_key) - if cache_result is not None: - raise InvalidActionError("Document is being indexed, please try again later") - try: - SegmentService.update_segments_status(segment_ids, action, dataset, document, session) - except Exception as e: - raise InvalidActionError(str(e)) + application_services().knowledge.segments.change_segment_status( + request_context, + dataset_id=str(dataset_id), + document_id=str(document_id), + segment_ids=args.segment_id, + action=action, + ) + except Exception as error: + _raise_segment_error(error) return SimpleResultResponse(result="success").model_dump(mode="json"), 200 @console_ns.route("/datasets//documents//segment") class DatasetDocumentSegmentAddApi(Resource): @console_ns.doc(params=SegmentDocParams.DATASET_DOCUMENT) - @setup_required - @login_required - @account_initialization_required + @console_account_admission( + allowed_roles=_DATASET_EDIT_ROLES, + rbac_checks=(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()),), + ) @cloud_edition_billing_resource_check("vector_space") @cloud_edition_billing_knowledge_limit_check("add_segment") @cloud_edition_billing_rate_limit_check("knowledge") @console_ns.expect(console_ns.models[SegmentCreatePayload.__name__]) @console_ns.response(200, "Segment created successfully", console_ns.models[SegmentDetailResponse.__name__]) - @with_current_user - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId())) - @with_session @model_validate(SegmentCreatePayload) def post( self, req_data: SegmentCreatePayload, - session: Session, - current_tenant_id: str, - current_user: Account, + request_context: RequestContext, dataset_id: UUID, document_id: UUID, ): - # check dataset - dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, session) - if not dataset: - raise NotFound("Dataset not found.") - # check document - document_id_str = str(document_id) - document = DocumentService.get_document(dataset_id_str, document_id_str, session=session) - if not document: - raise NotFound("Document not found.") - if not current_user.is_dataset_editor: - raise Forbidden() - # check embedding model setting - if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY: - try: - model_manager = ModelManager.for_tenant(tenant_id=current_tenant_id) - model_manager.get_model_instance( - tenant_id=current_tenant_id, - provider=dataset.embedding_model_provider, - model_type=ModelType.TEXT_EMBEDDING, - model=dataset.embedding_model, - ) - except LLMBadRequestError: - raise ProviderNotInitializeError( - "No Embedding Model available. Please configure a valid provider in the Settings -> Model Provider." - ) - except ProviderTokenNotInitError as ex: - raise ProviderNotInitializeError(ex.description) try: - DatasetService.check_dataset_permission(dataset, current_user, session) - except services.errors.account.NoPermissionError as e: - raise Forbidden(str(e)) - # validate args - 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( - segment_id=segment.id, dataset_id=dataset_id_str, session=session - ) + result = application_services().knowledge.segments.create_segment( + request_context, + dataset_id=str(dataset_id), + document_id=str(document_id), + values=req_data.model_dump(exclude_none=True), + ) + except Exception as error: + _raise_segment_error(error) response = { - "data": segment_response_with_summary( - segment, summary.summary_content if summary else None, session=session - ), - "doc_form": document.doc_form, + "data": result.data, + "doc_form": result.doc_form, } return dump_response(SegmentDetailResponse, response), 200 @@ -476,338 +329,173 @@ class DatasetDocumentSegmentAddApi(Resource): @console_ns.route("/datasets//documents//segments/") class DatasetDocumentSegmentUpdateApi(Resource): @console_ns.doc(params=SegmentDocParams.DATASET_DOCUMENT_SEGMENT) - @setup_required - @login_required - @account_initialization_required + @console_account_admission( + allowed_roles=_DATASET_EDIT_ROLES, + rbac_checks=(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()),), + ) @cloud_edition_billing_resource_check("vector_space") @cloud_edition_billing_rate_limit_check("knowledge") @console_ns.expect(console_ns.models[SegmentUpdatePayload.__name__]) @console_ns.response(200, "Segment updated successfully", console_ns.models[SegmentDetailResponse.__name__]) - @with_current_user - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId())) - @with_session @model_validate(SegmentUpdatePayload) def patch( self, req_data: SegmentUpdatePayload, - session: Session, - current_tenant_id: str, - current_user: Account, + request_context: RequestContext, dataset_id: UUID, document_id: UUID, segment_id: UUID, ): - # check dataset - dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, session) - if not dataset: - raise NotFound("Dataset not found.") - # check user's model setting - DatasetService.check_dataset_model_setting(dataset) - # check document - document_id_str = str(document_id) - document = DocumentService.get_document(dataset_id_str, document_id_str, session=session) - if not document: - raise NotFound("Document not found.") - # The role of the current user in the ta table must be admin, owner, dataset_operator, or editor - if not current_user.is_dataset_editor: - raise Forbidden() try: - DatasetService.check_dataset_permission(dataset, current_user, session) - except services.errors.account.NoPermissionError as e: - raise Forbidden(str(e)) - if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY: - # check embedding model setting - try: - model_manager = ModelManager.for_tenant(tenant_id=current_tenant_id) - model_manager.get_model_instance( - tenant_id=current_tenant_id, - provider=dataset.embedding_model_provider, - model_type=ModelType.TEXT_EMBEDDING, - model=dataset.embedding_model, - ) - except LLMBadRequestError: - raise ProviderNotInitializeError( - "No Embedding Model available. Please configure a valid provider in the Settings -> Model Provider." - ) - except ProviderTokenNotInitError as ex: - raise ProviderNotInitializeError(ex.description) - segment_id_str = str(segment_id) - _, segment = _get_segment_for_document(session, dataset, document, segment_id_str) - # validate args - 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(req_data.model_dump(exclude_none=True)), - segment, - document, - dataset, - session, - ) - summary = SummaryIndexService.get_segment_summary( - segment_id=segment.id, dataset_id=dataset_id_str, session=session - ) + result = application_services().knowledge.segments.update_segment( + request_context, + dataset_id=str(dataset_id), + document_id=str(document_id), + segment_id=str(segment_id), + values=req_data.model_dump(exclude_none=True), + ) + except Exception as error: + _raise_segment_error(error) response = { - "data": segment_response_with_summary( - segment, summary.summary_content if summary else None, session=session - ), - "doc_form": document.doc_form, + "data": result.data, + "doc_form": result.doc_form, } return dump_response(SegmentDetailResponse, response), 200 - @setup_required - @login_required - @account_initialization_required + @console_account_admission( + allowed_roles=_DATASET_EDIT_ROLES, + rbac_checks=(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()),), + ) @cloud_edition_billing_rate_limit_check("knowledge") @console_ns.doc(params=SegmentDocParams.DATASET_DOCUMENT_SEGMENT) @console_ns.response(204, "Segment deleted successfully") - @with_current_user - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId())) - @with_session def delete( self, - session: Session, - current_tenant_id: str, - current_user: Account, + request_context: RequestContext, dataset_id: UUID, document_id: UUID, segment_id: UUID, ): - # check dataset - dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, session) - if not dataset: - raise NotFound("Dataset not found.") - # check user's model setting - DatasetService.check_dataset_model_setting(dataset) - # check document - document_id_str = str(document_id) - document = DocumentService.get_document(dataset_id_str, document_id_str, session=session) - if not document: - raise NotFound("Document not found.") - # The role of the current user in the ta table must be admin, owner, dataset_operator, or editor - if not current_user.is_dataset_editor: - raise Forbidden() try: - DatasetService.check_dataset_permission(dataset, current_user, session) - except services.errors.account.NoPermissionError as e: - raise Forbidden(str(e)) - segment_id_str = str(segment_id) - _, segment = _get_segment_for_document(session, dataset, document, segment_id_str) - SegmentService.delete_segment(segment, document, dataset, session) + application_services().knowledge.segments.delete_segment( + request_context, + dataset_id=str(dataset_id), + document_id=str(document_id), + segment_id=str(segment_id), + ) + except Exception as error: + _raise_segment_error(error) return "", 204 -@console_ns.route( - "/datasets//documents//segments/batch_import", - "/datasets/batch_import_status/", -) +@console_ns.route("/datasets//documents//segments/batch_import") class DatasetDocumentSegmentBatchImportApi(Resource): @console_ns.response(200, "Batch import started", console_ns.models[SegmentBatchImportStatusResponse.__name__]) - @setup_required - @login_required - @account_initialization_required + @console_account_admission( + rbac_checks=(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()),), + ) @cloud_edition_billing_resource_check("vector_space") @cloud_edition_billing_knowledge_limit_check("add_segment") @cloud_edition_billing_rate_limit_check("knowledge") @console_ns.expect(console_ns.models[BatchImportPayload.__name__]) - @with_current_user - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId())) - @with_session @model_validate(BatchImportPayload) def post( self, req_data: BatchImportPayload, - session: Session, - current_tenant_id: str, - current_user: Account, + request_context: RequestContext, dataset_id: UUID, document_id: UUID, ): - # check dataset - dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset_for_tenant(dataset_id_str, current_tenant_id, session=session) - if not dataset: - raise NotFound("Dataset not found.") - - # check document - document_id_str = str(document_id) - document_ref = DatasetRefService.create_document_ref_from_id( - DatasetRefService.create_dataset_ref(dataset), document_id_str - ) - document = DatasetRefService.get_document_by_ref(document_ref, session=session) - if not document: - raise NotFound("Document not found.") - - 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: - raise NotFound("UploadFile not found.") - - # check file type - if not upload_file.name or not upload_file.name.lower().endswith(".csv"): - raise ValueError("Invalid file type. Only CSV files are allowed") - try: - # async job - job_id = str(uuid.uuid4()) - indexing_cache_key = f"segment_batch_import_{job_id}" - # send batch add segments task - redis_client.setnx(indexing_cache_key, "waiting") - batch_create_segment_to_index_task.delay( - job_id, - upload_file_id, - dataset_id_str, - document_id_str, - current_tenant_id, - current_user.id, + result = application_services().knowledge.segments.start_batch_import( + request_context, + dataset_id=str(dataset_id), + document_id=str(document_id), + upload_file_id=req_data.upload_file_id, ) - except Exception as e: - return {"error": str(e)}, 500 - return dump_response(SegmentBatchImportStatusResponse, {"job_id": job_id, "job_status": "waiting"}), 200 + except SegmentBatchImportDispatchError as error: + return {"error": str(error)}, 500 + except Exception as error: + _raise_segment_error(error) + return dump_response(SegmentBatchImportStatusResponse, result), 200 + +@console_ns.route("/datasets/batch_import_status/") +class DatasetDocumentSegmentBatchImportStatusApi(Resource): @console_ns.response(200, "Batch import status", console_ns.models[SegmentBatchImportStatusResponse.__name__]) - @setup_required - @login_required - @account_initialization_required - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_READONLY, DatasetId())) - def get(self, job_id=None, dataset_id: UUID | None = None, document_id: UUID | None = None): - if job_id is None: - raise NotFound("The job does not exist.") - job_id = str(job_id) - indexing_cache_key = f"segment_batch_import_{job_id}" - cache_result = redis_client.get(indexing_cache_key) - if cache_result is None: - raise ValueError("The job does not exist.") - - response = {"job_id": job_id, "job_status": cache_result.decode()} - return dump_response(SegmentBatchImportStatusResponse, response), 200 + @console_account_admission() + def get(self, _request_context: RequestContext, job_id: UUID): + try: + result = application_services().knowledge.segments.get_batch_import_status(str(job_id)) + except SegmentBatchImportNotFoundError as error: + raise ValueError(str(error)) from error + return dump_response(SegmentBatchImportStatusResponse, result), 200 @console_ns.route("/datasets//documents//segments//child_chunks") class ChildChunkAddApi(Resource): @console_ns.doc(params=SegmentDocParams.DATASET_DOCUMENT_PARENT_SEGMENT) - @setup_required - @login_required - @account_initialization_required + @console_account_admission( + allowed_roles=_DATASET_EDIT_ROLES, + rbac_checks=(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()),), + ) @cloud_edition_billing_resource_check("vector_space") @cloud_edition_billing_knowledge_limit_check("add_segment") @cloud_edition_billing_rate_limit_check("knowledge") @console_ns.expect(console_ns.models[ChildChunkCreatePayload.__name__]) @console_ns.response(200, "Child chunk created successfully", console_ns.models[ChildChunkDetailResponse.__name__]) - @with_current_user - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId())) - @with_session @model_validate(ChildChunkCreatePayload) def post( self, req_data: ChildChunkCreatePayload, - session: Session, - current_tenant_id: str, - current_user: Account, + request_context: RequestContext, dataset_id: UUID, document_id: UUID, segment_id: UUID, ): - # check dataset - dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, session) - if not dataset: - raise NotFound("Dataset not found.") - # check document - document_id_str = str(document_id) - document = DocumentService.get_document(dataset_id_str, document_id_str, session=session) - if not document: - raise NotFound("Document not found.") - if not current_user.is_dataset_editor: - raise Forbidden() try: - DatasetService.check_dataset_permission(dataset, current_user, session) - except services.errors.account.NoPermissionError as e: - raise Forbidden(str(e)) - # check embedding model setting - if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY: - try: - model_manager = ModelManager.for_tenant(tenant_id=current_tenant_id) - model_manager.get_model_instance( - tenant_id=current_tenant_id, - provider=dataset.embedding_model_provider, - model_type=ModelType.TEXT_EMBEDDING, - model=dataset.embedding_model, - ) - except LLMBadRequestError: - raise ProviderNotInitializeError( - "No Embedding Model available. Please configure a valid provider in the Settings -> Model Provider." - ) - except ProviderTokenNotInitError as ex: - raise ProviderNotInitializeError(ex.description) - segment_id_str = str(segment_id) - _, segment = _get_segment_for_document(session, dataset, document, segment_id_str) - # validate args - try: - child_chunk = SegmentService.create_child_chunk(req_data.content, segment, document, dataset, session) - except ChildChunkIndexingServiceError as e: - raise ChildChunkIndexingError(str(e)) + child_chunk = application_services().knowledge.segments.create_child_chunk( + request_context, + dataset_id=str(dataset_id), + document_id=str(document_id), + segment_id=str(segment_id), + content=req_data.content, + ) + except Exception as error: + _raise_segment_error(error) return dump_response(ChildChunkDetailResponse, {"data": child_chunk}), 200 @console_ns.doc(params=SegmentDocParams.DATASET_DOCUMENT_PARENT_SEGMENT) @console_ns.doc(params=query_params_from_model(ChildChunkListQuery)) @console_ns.response(200, "Child chunks retrieved successfully", console_ns.models[ChildChunkListResponse.__name__]) - @setup_required - @login_required - @account_initialization_required - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_READONLY, DatasetId())) - @with_session(write=False) - def get(self, session: Session, current_tenant_id: str, dataset_id: UUID, document_id: UUID, segment_id: UUID): - # check dataset - dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, session) - if not dataset: - raise NotFound("Dataset not found.") - # check user's model setting - DatasetService.check_dataset_model_setting(dataset) - # check document - document_id_str = str(document_id) - document = DocumentService.get_document(dataset_id_str, document_id_str, session=session) - if not document: - raise NotFound("Document not found.") - segment_id_str = str(segment_id) - _get_segment_for_document(session, dataset, document, segment_id_str) + @console_account_admission( + rbac_checks=(RBACCheck(RBACPermission.DATASET_READONLY, DatasetId()),), + ) + def get(self, request_context: RequestContext, dataset_id: UUID, document_id: UUID, segment_id: UUID): args = query_params_from_request(ChildChunkListQuery, use_defaults_for_malformed_ints=True) - - page = args.page - limit = min(args.limit, 100) - keyword = args.keyword - - child_chunks = SegmentService.get_child_chunks( - segment_id_str, - document_id_str, - dataset_id_str, - page, - limit, - keyword, - session=session, - ) + try: + result = application_services().knowledge.segments.list_child_chunks( + workspace_id=request_context.active_workspace_id, + dataset_id=str(dataset_id), + document_id=str(document_id), + segment_id=str(segment_id), + query=ChildChunkListFilter(page=args.page, limit=args.limit, keyword=args.keyword), + ) + except Exception as error: + _raise_segment_error(error) response = { - "data": child_chunks.items, - "total": child_chunks.total, - "total_pages": child_chunks.pages, - "page": page, - "limit": limit, + "data": result.items, + "total": result.total, + "total_pages": result.total_pages, + "page": result.page, + "limit": result.limit, } return dump_response(ChildChunkListResponse, response), 200 - @setup_required - @login_required - @account_initialization_required + @console_account_admission( + allowed_roles=_DATASET_EDIT_ROLES, + rbac_checks=(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()),), + ) @cloud_edition_billing_resource_check("vector_space") @cloud_edition_billing_rate_limit_check("knowledge") @console_ns.doc(params=SegmentDocParams.DATASET_DOCUMENT_PARENT_SEGMENT) @@ -817,47 +505,25 @@ class ChildChunkAddApi(Resource): console_ns.models[ChildChunkBatchUpdateResponse.__name__], ) @console_ns.expect(console_ns.models[ChildChunkBatchUpdatePayload.__name__]) - @with_current_user - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId())) - @with_session @model_validate(ChildChunkBatchUpdatePayload) def patch( self, req_data: ChildChunkBatchUpdatePayload, - session: Session, - current_tenant_id: str, - current_user: Account, + request_context: RequestContext, dataset_id: UUID, document_id: UUID, segment_id: UUID, ): - # check dataset - dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, session) - if not dataset: - raise NotFound("Dataset not found.") - # check user's model setting - DatasetService.check_dataset_model_setting(dataset) - # check document - document_id_str = str(document_id) - document = DocumentService.get_document(dataset_id_str, document_id_str, session=session) - if not document: - raise NotFound("Document not found.") - # The role of the current user in the ta table must be admin, owner, dataset_operator, or editor - if not current_user.is_dataset_editor: - raise Forbidden() try: - DatasetService.check_dataset_permission(dataset, current_user, session) - except services.errors.account.NoPermissionError as e: - raise Forbidden(str(e)) - segment_id_str = str(segment_id) - _, segment = _get_segment_for_document(session, dataset, document, segment_id_str) - # validate args - try: - child_chunks = SegmentService.update_child_chunks(req_data.chunks, segment, document, dataset, session) - except ChildChunkIndexingServiceError as e: - raise ChildChunkIndexingError(str(e)) + child_chunks = application_services().knowledge.segments.update_child_chunks( + request_context, + dataset_id=str(dataset_id), + document_id=str(document_id), + segment_id=str(segment_id), + chunks=req_data.chunks, + ) + except Exception as error: + _raise_segment_error(error) return dump_response(ChildChunkBatchUpdateResponse, {"data": child_chunks}), 200 @@ -865,111 +531,61 @@ class ChildChunkAddApi(Resource): "/datasets//documents//segments//child_chunks/" ) class ChildChunkUpdateApi(Resource): - @setup_required - @login_required - @account_initialization_required + @console_account_admission( + allowed_roles=_DATASET_EDIT_ROLES, + rbac_checks=(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()),), + ) @cloud_edition_billing_rate_limit_check("knowledge") @console_ns.doc(params=SegmentDocParams.DATASET_DOCUMENT_CHILD_CHUNK) @console_ns.response(204, "Child chunk deleted successfully") - @with_current_user - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId())) - @with_session def delete( self, - session: Session, - current_tenant_id: str, - current_user: Account, + request_context: RequestContext, dataset_id: UUID, document_id: UUID, segment_id: UUID, child_chunk_id: UUID, ): - # check dataset - dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, session) - if not dataset: - raise NotFound("Dataset not found.") - # check user's model setting - DatasetService.check_dataset_model_setting(dataset) - # check document - document_id_str = str(document_id) - document = DocumentService.get_document(dataset_id_str, document_id_str, session=session) - if not document: - raise NotFound("Document not found.") - # The role of the current user in the ta table must be admin, owner, dataset_operator, or editor - if not current_user.is_dataset_editor: - raise Forbidden() try: - DatasetService.check_dataset_permission(dataset, current_user, session) - except services.errors.account.NoPermissionError as e: - raise Forbidden(str(e)) - segment_id_str = str(segment_id) - segment_ref, _ = _get_segment_for_document(session, dataset, document, segment_id_str) - child_chunk_id_str = str(child_chunk_id) - child_chunk = SegmentService.get_child_chunk_by_segment_ref(child_chunk_id_str, segment_ref, session=session) - if not child_chunk: - raise NotFound("Child chunk not found.") - try: - SegmentService.delete_child_chunk(child_chunk, dataset, session) - except ChildChunkDeleteIndexServiceError as e: - raise ChildChunkDeleteIndexError(str(e)) + application_services().knowledge.segments.delete_child_chunk( + request_context, + dataset_id=str(dataset_id), + document_id=str(document_id), + segment_id=str(segment_id), + child_chunk_id=str(child_chunk_id), + ) + except Exception as error: + _raise_segment_error(error) return "", 204 - @setup_required - @login_required - @account_initialization_required + @console_account_admission( + allowed_roles=_DATASET_EDIT_ROLES, + rbac_checks=(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()),), + ) @cloud_edition_billing_resource_check("vector_space") @cloud_edition_billing_rate_limit_check("knowledge") @console_ns.doc(params=SegmentDocParams.DATASET_DOCUMENT_CHILD_CHUNK) @console_ns.expect(console_ns.models[ChildChunkUpdatePayload.__name__]) @console_ns.response(200, "Child chunk updated successfully", console_ns.models[ChildChunkDetailResponse.__name__]) - @with_current_user - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId())) - @with_session @model_validate(ChildChunkUpdatePayload) def patch( self, req_data: ChildChunkUpdatePayload, - session: Session, - current_tenant_id: str, - current_user: Account, + request_context: RequestContext, dataset_id: UUID, document_id: UUID, segment_id: UUID, child_chunk_id: UUID, ): - # check dataset - dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, session) - if not dataset: - raise NotFound("Dataset not found.") - # check user's model setting - DatasetService.check_dataset_model_setting(dataset) - # check document - document_id_str = str(document_id) - document = DocumentService.get_document(dataset_id_str, document_id_str, session=session) - if not document: - raise NotFound("Document not found.") - # The role of the current user in the ta table must be admin, owner, dataset_operator, or editor - if not current_user.is_dataset_editor: - raise Forbidden() try: - DatasetService.check_dataset_permission(dataset, current_user, session) - except services.errors.account.NoPermissionError as e: - raise Forbidden(str(e)) - segment_id_str = str(segment_id) - segment_ref, segment = _get_segment_for_document(session, dataset, document, segment_id_str) - child_chunk_id_str = str(child_chunk_id) - child_chunk = SegmentService.get_child_chunk_by_segment_ref(child_chunk_id_str, segment_ref, session=session) - if not child_chunk: - raise NotFound("Child chunk not found.") - # validate args - try: - child_chunk = SegmentService.update_child_chunk( - req_data.content, child_chunk, segment, document, dataset, session + child_chunk = application_services().knowledge.segments.update_child_chunk( + request_context, + dataset_id=str(dataset_id), + document_id=str(document_id), + segment_id=str(segment_id), + child_chunk_id=str(child_chunk_id), + content=req_data.content, ) - except ChildChunkIndexingServiceError as e: - raise ChildChunkIndexingError(str(e)) + except Exception as error: + _raise_segment_error(error) return dump_response(ChildChunkDetailResponse, {"data": child_chunk}), 200 diff --git a/api/controllers/console/datasets/error.py b/api/controllers/console/datasets/error.py index ac09ec16b27..924bd391784 100644 --- a/api/controllers/console/datasets/error.py +++ b/api/controllers/console/datasets/error.py @@ -25,6 +25,12 @@ class InvalidActionError(BaseHTTPException): code = 400 +class DatasetAccessDeniedRequestError(BaseHTTPException): + error_code = "dataset_access_denied" + description = "Dataset access denied." + code = 403 + + class DocumentAlreadyFinishedError(BaseHTTPException): error_code = "document_already_finished" description = "The document has been processed. Please refresh the page or go to the document details." diff --git a/api/controllers/console/datasets/external.py b/api/controllers/console/datasets/external.py index fd93100d9fc..1ceb31469d9 100644 --- a/api/controllers/console/datasets/external.py +++ b/api/controllers/console/datasets/external.py @@ -1,16 +1,13 @@ from __future__ import annotations -from dataclasses import dataclass from datetime import datetime -from typing import Any +from typing import Any, Never from uuid import UUID from flask_restx import Resource from pydantic import AliasChoices, BaseModel, Field, field_validator -from sqlalchemy.orm import Session from werkzeug.exceptions import Forbidden, InternalServerError, NotFound -import services from controllers.common.fields import UsageCountResponse from controllers.common.rbac import DatasetId, RBACCheck, Workspace from controllers.common.schema import ( @@ -18,31 +15,42 @@ from controllers.common.schema import ( register_response_schema_models, register_schema_models, ) -from controllers.common.session import with_session from controllers.console import console_ns from controllers.console.datasets.error import DatasetNameDuplicateError +from controllers.console.flask_admission import console_account_admission from controllers.console.wraps import ( RBACPermission, - account_initialization_required, - edit_permission_required, model_validate, - rbac_permission_required, - setup_required, - with_current_tenant_id, - with_current_user, ) +from extensions.ext_application_services import application_services from fields.base import ResponseModel -from fields.dataset_fields import DatasetDetailResponse, dataset_detail_response_source +from fields.dataset_fields import ( + DatasetDetailResponse, +) from libs.helper import dump_response -from libs.login import login_required -from libs.pagination import clamp_pagination -from models import Account -from models.dataset import ExternalKnowledgeApis -from services.dataset_service import DatasetService -from services.enterprise import rbac_service as enterprise_rbac_service +from machinery.context import RequestContext +from models.account import TenantAccountRole from services.entities.external_knowledge_entities.external_knowledge_entities import ExternalDatasetCreatePayload -from services.external_knowledge_service import ExternalDatasetService -from services.hit_testing_service import HitTestingService +from services.errors.dataset import DatasetNameDuplicateError as DatasetNameDuplicateFailure +from services.knowledge.dataset_access import DatasetAccessDeniedError, DatasetNotFoundError +from services.knowledge.external.application import ExternalHitTestingError, ExternalTemplateNotFoundError + +_DATASET_EDIT_ROLES = frozenset( + {TenantAccountRole.OWNER, TenantAccountRole.ADMIN, TenantAccountRole.EDITOR, TenantAccountRole.DATASET_OPERATOR} +) +_EXTERNAL_CONNECT_ROLES = frozenset({TenantAccountRole.OWNER, TenantAccountRole.ADMIN, TenantAccountRole.EDITOR}) + + +def _raise_external_error(error: Exception) -> Never: + if isinstance(error, (DatasetNotFoundError, ExternalTemplateNotFoundError)): + raise NotFound(str(error)) from error + if isinstance(error, DatasetAccessDeniedError): + raise Forbidden(str(error)) from error + if isinstance(error, DatasetNameDuplicateFailure): + raise DatasetNameDuplicateError() from error + if isinstance(error, ExternalHitTestingError): + raise InternalServerError(str(error)) from error + raise error class ExternalKnowledgeApiPayload(BaseModel): @@ -85,28 +93,6 @@ class ExternalKnowledgeApiResponse(ResponseModel): return value -@dataclass(frozen=True) -class ExternalKnowledgeApiResponseSource: - external_knowledge_api: ExternalKnowledgeApis - session: Session - - @property - def dataset_bindings(self) -> Any: - return self.external_knowledge_api.get_dataset_bindings(session=self.session) - - def __getattr__(self, name: str) -> Any: - return getattr(self.external_knowledge_api, name) # guard-ignore: no-new-getattr -- delegates model fields - - -def external_knowledge_api_response( - external_knowledge_api: ExternalKnowledgeApis, *, session: Session -) -> ExternalKnowledgeApiResponse: - return ExternalKnowledgeApiResponse.model_validate( - ExternalKnowledgeApiResponseSource(external_knowledge_api=external_knowledge_api, session=session), - from_attributes=True, - ) - - class ExternalKnowledgeApiListResponse(ResponseModel): data: list[ExternalKnowledgeApiResponse] has_more: bool @@ -161,25 +147,13 @@ class ExternalApiTemplateListApi(Resource): "External API templates retrieved successfully", console_ns.models[ExternalKnowledgeApiListResponse.__name__], ) - @setup_required - @login_required - @with_current_tenant_id - @account_initialization_required - @with_session(write=False) + @console_account_admission() @model_validate(ExternalApiTemplateListQuery) - def get(self, req_data: ExternalApiTemplateListQuery, session: Session, current_tenant_id: str): - - effective_page, effective_limit = clamp_pagination(req_data.page, req_data.limit, 100) - external_knowledge_apis, total = ExternalDatasetService.get_external_knowledge_apis( - effective_page, effective_limit, current_tenant_id, req_data.keyword, session=session + def get(self, req_data: ExternalApiTemplateListQuery, request_context: RequestContext): + result = application_services().knowledge.external.list_templates( + request_context, page=req_data.page, limit=req_data.limit, keyword=req_data.keyword ) - return ExternalKnowledgeApiListResponse( - data=[external_knowledge_api_response(item, session=session) for item in external_knowledge_apis], - has_more=effective_page * effective_limit < total, - limit=effective_limit, - total=total, - page=effective_page, - ).model_dump(mode="json"), 200 + return dump_response(ExternalKnowledgeApiListResponse, result), 200 @console_ns.doc("create_external_api_template") @console_ns.doc(description="Create external knowledge API template") @@ -190,48 +164,16 @@ class ExternalApiTemplateListApi(Resource): console_ns.models[ExternalKnowledgeApiResponse.__name__], ) @console_ns.response(403, "Permission denied") - @setup_required - @login_required - @account_initialization_required - @with_current_user - @with_current_tenant_id - @with_session + @console_account_admission(allowed_roles=_DATASET_EDIT_ROLES) @model_validate(ExternalKnowledgeApiPayload) - def post( - self, - req_data: ExternalKnowledgeApiPayload, - session: Session, - current_tenant_id: str, - current_user: Account, - ): - - 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: - raise Forbidden() - + def post(self, req_data: ExternalKnowledgeApiPayload, request_context: RequestContext): try: - external_knowledge_api = ExternalDatasetService.create_external_knowledge_api( - tenant_id=current_tenant_id, - user_id=current_user.id, - args=req_data.model_dump(), - session=session, + result = application_services().knowledge.external.create_template( + request_context, name=req_data.name, settings=req_data.settings ) - except services.errors.dataset.DatasetNameDuplicateError: - raise DatasetNameDuplicateError() - - # if dify_config.RBAC_ENABLED: - # enterprise_rbac_service.RBACService.DatasetAccess.replace_whitelist( - # current_tenant_id, - # current_user.id, - # external_knowledge_api.id, - # enterprise_rbac_service.ReplaceMemberBindings(automatic_include_workspace_members=True), - # ) - # initialize_created_app_rbac_access_task.delay( - # current_tenant_id, current_user.id, dataset_id=external_knowledge_api.id) - - return external_knowledge_api_response(external_knowledge_api, session=session).model_dump(mode="json"), 201 + except Exception as error: + _raise_external_error(error) + return dump_response(ExternalKnowledgeApiResponse, result), 201 @console_ns.route("/datasets/external-knowledge-api/") @@ -245,20 +187,15 @@ class ExternalApiTemplateApi(Resource): console_ns.models[ExternalKnowledgeApiResponse.__name__], ) @console_ns.response(404, "Template not found") - @setup_required - @login_required - @account_initialization_required - @with_current_tenant_id - @with_session - def get(self, session: Session, current_tenant_id: str, external_knowledge_api_id: UUID): - external_knowledge_api_id_str = str(external_knowledge_api_id) - external_knowledge_api = ExternalDatasetService.get_external_knowledge_api( - external_knowledge_api_id=external_knowledge_api_id_str, tenant_id=current_tenant_id, session=session - ) - if external_knowledge_api is None: - raise NotFound("API template not found.") - - return external_knowledge_api_response(external_knowledge_api, session=session).model_dump(mode="json"), 200 + @console_account_admission() + def get(self, request_context: RequestContext, external_knowledge_api_id: UUID): + try: + result = application_services().knowledge.external.get_template( + request_context, template_id=str(external_knowledge_api_id) + ) + except Exception as error: + _raise_external_error(error) + return dump_response(ExternalKnowledgeApiResponse, result), 200 @console_ns.doc("update_external_api_template") @console_ns.doc(description="Update external knowledge API template") @@ -270,51 +207,31 @@ class ExternalApiTemplateApi(Resource): console_ns.models[ExternalKnowledgeApiResponse.__name__], ) @console_ns.response(404, "Template not found") - @setup_required - @login_required - @account_initialization_required - @with_current_user - @with_current_tenant_id - @with_session + @console_account_admission() @model_validate(ExternalKnowledgeApiPayload) def patch( - self, - req_data: ExternalKnowledgeApiPayload, - session: Session, - current_tenant_id: str, - current_user: Account, - external_knowledge_api_id: UUID, + self, req_data: ExternalKnowledgeApiPayload, request_context: RequestContext, external_knowledge_api_id: UUID ): - external_knowledge_api_id_str = str(external_knowledge_api_id) + try: + result = application_services().knowledge.external.update_template( + request_context, + template_id=str(external_knowledge_api_id), + name=req_data.name, + settings=req_data.settings, + ) + except Exception as error: + _raise_external_error(error) + return dump_response(ExternalKnowledgeApiResponse, result), 200 - 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=req_data.model_dump(), - session=session, - ) - - return external_knowledge_api_response(external_knowledge_api, session=session).model_dump(mode="json"), 200 - - @setup_required - @login_required - @account_initialization_required @console_ns.response(204, "External knowledge API deleted successfully") - @with_current_user - @with_current_tenant_id - @with_session - def delete(self, session: Session, current_tenant_id: str, current_user: Account, external_knowledge_api_id: UUID): - external_knowledge_api_id_str = str(external_knowledge_api_id) - - if not (current_user.has_edit_permission or current_user.is_dataset_operator): - raise Forbidden() - - ExternalDatasetService.delete_external_knowledge_api( - current_tenant_id, external_knowledge_api_id_str, session=session - ) + @console_account_admission(allowed_roles=_DATASET_EDIT_ROLES) + def delete(self, request_context: RequestContext, external_knowledge_api_id: UUID): + try: + application_services().knowledge.external.delete_template( + request_context, template_id=str(external_knowledge_api_id) + ) + except Exception as error: + _raise_external_error(error) return "", 204 @@ -324,18 +241,12 @@ class ExternalApiUseCheckApi(Resource): @console_ns.doc(description="Check if external knowledge API is being used") @console_ns.doc(params={"external_knowledge_api_id": "External knowledge API ID"}) @console_ns.response(200, "Usage check completed successfully", console_ns.models[UsageCountResponse.__name__]) - @setup_required - @login_required - @account_initialization_required - @with_current_tenant_id - @with_session - def get(self, session: Session, current_tenant_id: str, external_knowledge_api_id: UUID): - external_knowledge_api_id_str = str(external_knowledge_api_id) - - external_knowledge_api_is_using, count = ExternalDatasetService.external_knowledge_api_use_check( - external_knowledge_api_id_str, current_tenant_id, session=session + @console_account_admission() + def get(self, request_context: RequestContext, external_knowledge_api_id: UUID): + is_using, count = application_services().knowledge.external.template_usage( + request_context, template_id=str(external_knowledge_api_id) ) - return UsageCountResponse(is_using=external_knowledge_api_is_using, count=count).model_dump(mode="json"), 200 + return dump_response(UsageCountResponse, {"is_using": is_using, "count": count}), 200 @console_ns.route("/datasets/external") @@ -348,50 +259,17 @@ class ExternalDatasetCreateApi(Resource): ) @console_ns.response(400, "Invalid parameters") @console_ns.response(403, "Permission denied") - @setup_required - @login_required - @account_initialization_required - @edit_permission_required - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_EXTERNAL_CONNECT, Workspace())) - @with_current_user - @with_current_tenant_id - @with_session + @console_account_admission( + allowed_roles=_EXTERNAL_CONNECT_ROLES, + rbac_checks=(RBACCheck(RBACPermission.DATASET_EXTERNAL_CONNECT, Workspace()),), + ) @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 - - # 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() - + def post(self, req_data: ExternalDatasetCreatePayload, request_context: RequestContext): try: - dataset = ExternalDatasetService.create_external_dataset( - tenant_id=current_tenant_id, - user_id=current_user.id, - args=req_data, - session=session, - ) - except services.errors.dataset.DatasetNameDuplicateError: - raise DatasetNameDuplicateError() - - dataset_id_str = str(dataset.id) - permission_keys_map = enterprise_rbac_service.RBACService.DatasetPermissions.batch_get( - str(current_tenant_id), - current_user.id, - [dataset_id_str], - session=session, - ) - data = DatasetDetailResponse.model_validate( - dataset_detail_response_source(dataset, session=session) - ).model_dump(mode="json") - data["permission_keys"] = permission_keys_map.get(dataset_id_str, []) - return data, 201 + result = application_services().knowledge.external.create_dataset(request_context, payload=req_data) + except Exception as error: + _raise_external_error(error) + return dump_response(DatasetDetailResponse, result), 201 @console_ns.route("/datasets//external-hit-testing") @@ -407,36 +285,17 @@ class ExternalKnowledgeHitTestingApi(Resource): ) @console_ns.response(404, "Dataset not found") @console_ns.response(400, "Invalid parameters") - @setup_required - @login_required - @account_initialization_required - @with_current_user - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_PIPELINE_TEST, DatasetId())) - @with_session + @console_account_admission(rbac_checks=(RBACCheck(RBACPermission.DATASET_PIPELINE_TEST, DatasetId()),)) @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: - raise NotFound("Dataset not found.") - + def post(self, req_data: ExternalHitTestingPayload, request_context: RequestContext, dataset_id: UUID): try: - DatasetService.check_dataset_permission(dataset, current_user, session) - except services.errors.account.NoPermissionError as e: - raise Forbidden(str(e)) - - HitTestingService.hit_testing_args_check(req_data.model_dump()) - - try: - response = HitTestingService.external_retrieve( - session=session, - dataset=dataset, + result = application_services().knowledge.external.hit_testing( + request_context, + dataset_id=str(dataset_id), query=req_data.query, - account=current_user, - external_retrieval_model=req_data.external_retrieval_model, - metadata_filtering_conditions=req_data.metadata_filtering_conditions, + retrieval_model=req_data.external_retrieval_model, + metadata_filters=req_data.metadata_filtering_conditions, ) - - return dump_response(ExternalHitTestingResponse, response) - except Exception as e: - raise InternalServerError(str(e)) + except Exception as error: + _raise_external_error(error) + return dump_response(ExternalHitTestingResponse, result) diff --git a/api/controllers/console/datasets/hit_testing_base.py b/api/controllers/console/datasets/hit_testing_base.py index e211507d5ab..0e0a3114d98 100644 --- a/api/controllers/console/datasets/hit_testing_base.py +++ b/api/controllers/console/datasets/hit_testing_base.py @@ -23,9 +23,9 @@ from graphon.model_runtime.errors.invoke import InvokeError from libs.login import resolve_account_fallback from models.account import Account from models.dataset import Dataset -from services.dataset_service import DatasetService -from services.entities.knowledge_entities.knowledge_entities import ExternalRetrievalModel, RetrievalModel from services.hit_testing_service import HitTestingService +from services.knowledge.dataset_service import DatasetService +from services.knowledge.entities.knowledge_entities import ExternalRetrievalModel, RetrievalModel logger = logging.getLogger(__name__) diff --git a/api/controllers/console/datasets/metadata.py b/api/controllers/console/datasets/metadata.py index 4aa6a9073e6..1b9255e87b8 100644 --- a/api/controllers/console/datasets/metadata.py +++ b/api/controllers/console/datasets/metadata.py @@ -1,43 +1,47 @@ -from typing import Literal +from collections.abc import Callable +from functools import wraps from uuid import UUID from flask_restx import Resource -from sqlalchemy.orm import Session from werkzeug.exceptions import Forbidden, NotFound -import services from controllers.common.controller_schemas import MetadataUpdatePayload from controllers.common.rbac import DatasetId, RBACCheck from controllers.common.schema import register_response_schema_models, register_schema_models -from controllers.common.session import with_session from controllers.console import console_ns -from controllers.console.wraps import ( - RBACPermission, - account_initialization_required, - enterprise_license_required, - model_validate, - rbac_permission_required, - setup_required, - with_current_tenant_id, - with_current_user, -) +from controllers.console.flask_admission import console_account_admission +from controllers.console.wraps import model_validate +from core.rbac import RBACPermission +from extensions.ext_application_services import application_services from fields.dataset_fields import ( DatasetMetadataBuiltInFieldsResponse, DatasetMetadataListResponse, DatasetMetadataResponse, ) from libs.helper import dump_response -from libs.login import login_required -from models.account import Account -from services.dataset_service import DatasetService -from services.entities.knowledge_entities.knowledge_entities import ( +from machinery.context import RequestContext +from services.errors.metadata import MetadataResourceNotFoundError +from services.knowledge.dataset_access import DatasetAccessDeniedError, DatasetNotFoundError +from services.knowledge.entities.knowledge_entities import ( DocumentMetadataOperation, MetadataArgs, MetadataDetail, MetadataOperationData, ) -from services.errors.metadata import MetadataResourceNotFoundError -from services.metadata_service import MetadataService + + +def metadata_errors[**P, T](method: Callable[P, T]) -> Callable[P, T]: + @wraps(method) + def wrapped(*args: P.args, **kwargs: P.kwargs) -> T: + try: + return method(*args, **kwargs) + except (DatasetNotFoundError, MetadataResourceNotFoundError) as error: + raise NotFound(str(error)) from error + except DatasetAccessDeniedError as error: + raise Forbidden(str(error)) from error + + return wrapped + register_schema_models( console_ns, MetadataArgs, MetadataOperationData, MetadataUpdatePayload, DocumentMetadataOperation, MetadataDetail @@ -52,199 +56,111 @@ register_response_schema_models( @console_ns.route("/datasets//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__]) - @with_current_user - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId())) - @with_session + @console_account_admission( + rbac_checks=(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()),), + ) + @metadata_errors @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: - raise NotFound("Dataset not found.") - DatasetService.check_dataset_permission(dataset, current_user, session) - - metadata = MetadataService.create_metadata( - dataset_id_str, req_data, current_user, current_tenant_id, session=session - ) + def post(self, req_data: MetadataArgs, request_context: RequestContext, dataset_id: UUID): + service = application_services().knowledge.metadata + ref = service.require_dataset(request_context, str(dataset_id)) + metadata = service.create_metadata(ref, req_data, actor_id=request_context.account_id) 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__] ) - @with_current_user - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_READONLY, DatasetId())) - @with_session(write=False) - def get(self, session: Session, current_tenant_id: str, current_user: Account, dataset_id: UUID): - dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset_for_tenant(dataset_id_str, current_tenant_id, session=session) - if dataset is None: - raise NotFound("Dataset not found.") - try: - DatasetService.check_dataset_permission(dataset, current_user, session) - except services.errors.account.NoPermissionError as e: - raise Forbidden(str(e)) - metadata = MetadataService.get_dataset_metadatas(dataset, session) - return dump_response(DatasetMetadataListResponse, metadata), 200 + @console_account_admission( + rbac_checks=(RBACCheck(RBACPermission.DATASET_READONLY, DatasetId()),), + ) + @metadata_errors + def get(self, request_context: RequestContext, dataset_id: UUID): + service = application_services().knowledge.metadata + ref = service.require_dataset(request_context, str(dataset_id)) + return dump_response(DatasetMetadataListResponse, service.get_dataset_metadatas(ref)), 200 @console_ns.route("/datasets//metadata/") 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__]) - @with_current_user - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId())) - @with_session + @console_account_admission( + rbac_checks=(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()),), + ) + @metadata_errors @model_validate(MetadataUpdatePayload) def patch( - self, - req_data: MetadataUpdatePayload, - session: Session, - current_tenant_id: str, - current_user: Account, - dataset_id: UUID, - metadata_id: UUID, + self, req_data: MetadataUpdatePayload, request_context: RequestContext, dataset_id: UUID, metadata_id: UUID ): - name = req_data.name - - dataset_id_str = str(dataset_id) - metadata_id_str = str(metadata_id) - dataset = DatasetService.get_dataset_for_tenant(dataset_id_str, current_tenant_id, session=session) - if dataset is None: - raise NotFound("Dataset not found.") - DatasetService.check_dataset_permission(dataset, current_user, session) - - metadata = MetadataService.update_metadata_name(dataset, metadata_id_str, name, current_user, session=session) + service = application_services().knowledge.metadata + ref = service.require_dataset(request_context, str(dataset_id)) + metadata = service.update_metadata_name( + ref, str(metadata_id), req_data.name, actor_id=request_context.account_id + ) return dump_response(DatasetMetadataResponse, metadata), 200 - @setup_required - @login_required - @account_initialization_required - @enterprise_license_required @console_ns.response(204, "Metadata deleted successfully") - @with_current_user - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId())) - @with_session - def delete( - self, - session: Session, - current_tenant_id: str, - current_user: Account, - dataset_id: UUID, - metadata_id: UUID, - ): - dataset_id_str = str(dataset_id) - metadata_id_str = str(metadata_id) - dataset = DatasetService.get_dataset_for_tenant(dataset_id_str, current_tenant_id, session=session) - if dataset is None: - raise NotFound("Dataset not found.") - DatasetService.check_dataset_permission(dataset, current_user, session) - - MetadataService.delete_metadata(dataset, metadata_id_str, session) - # Frontend callers only await success and invalidate metadata caches; no response body is consumed. + @console_account_admission( + rbac_checks=(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()),), + ) + @metadata_errors + def delete(self, request_context: RequestContext, dataset_id: UUID, metadata_id: UUID): + service = application_services().knowledge.metadata + ref = service.require_dataset(request_context, str(dataset_id)) + service.delete_metadata(ref, str(metadata_id)) 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_account_admission() + @metadata_errors + def get(self, request_context: RequestContext): + fields = application_services().knowledge.metadata.get_built_in_fields() + return dump_response(DatasetMetadataBuiltInFieldsResponse, {"fields": fields}), 200 @console_ns.route("/datasets//metadata/built-in/") class DatasetMetadataBuiltInFieldActionApi(Resource): - @setup_required - @login_required - @account_initialization_required - @enterprise_license_required @console_ns.response(204, "Action completed successfully") - @with_current_user - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId())) - @with_session - def post(self, session: Session, current_user: Account, dataset_id: UUID, action: Literal["enable", "disable"]): - 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) - - match action: - case "enable": - MetadataService.enable_built_in_field(dataset, session) - case "disable": - MetadataService.disable_built_in_field(dataset, session) - # Frontend callers only await success and invalidate metadata caches; no response body is consumed. + @console_account_admission( + rbac_checks=(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()),), + ) + @metadata_errors + def post(self, request_context: RequestContext, dataset_id: UUID, action: str): + service = application_services().knowledge.metadata + ref = service.require_dataset(request_context, str(dataset_id)) + if action == "enable": + service.enable_built_in_field(ref) + elif action == "disable": + service.disable_built_in_field(ref) + else: + raise ValueError("Invalid action.") return "", 204 @console_ns.route("/datasets//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", ) @console_ns.response(404, "Dataset, document, or metadata not found") - @with_current_user - @with_current_tenant_id - @rbac_permission_required(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId())) - @with_session + @console_account_admission( + rbac_checks=(RBACCheck(RBACPermission.DATASET_EDIT, DatasetId()),), + ) + @metadata_errors @model_validate(MetadataOperationData) - def post( - self, - req_data: MetadataOperationData, - session: Session, - current_tenant_id: str, - current_user: Account, - dataset_id: UUID, - ): - dataset = DatasetService.get_dataset_for_tenant(str(dataset_id), current_tenant_id, session=session) - if dataset is None: - raise NotFound("Dataset not found.") - DatasetService.check_dataset_permission(dataset, current_user, session) - - try: - MetadataService.update_documents_metadata(dataset, req_data, current_user, session=session) - except MetadataResourceNotFoundError as exc: - raise NotFound(str(exc)) from exc - - # Frontend callers only await success and invalidate caches; no response body is consumed. + def post(self, req_data: MetadataOperationData, request_context: RequestContext, dataset_id: UUID): + service = application_services().knowledge.metadata + ref = service.require_dataset(request_context, str(dataset_id)) + service.update_documents_metadata(ref, req_data, actor_id=request_context.account_id) return "", 204 diff --git a/api/controllers/console/datasets/rag_pipeline/datasource_auth.py b/api/controllers/console/datasets/rag_pipeline/datasource_auth.py index 672a8a3d83a..ab20f673d32 100644 --- a/api/controllers/console/datasets/rag_pipeline/datasource_auth.py +++ b/api/controllers/console/datasets/rag_pipeline/datasource_auth.py @@ -24,6 +24,7 @@ from core.entities.provider_entities import ProviderConfig from core.plugin.entities.plugin_daemon import PluginOAuthAuthorizationUrlResponse from core.plugin.impl.oauth import OAuthHandler from core.tools.entities.common_entities import I18nObject +from extensions.ext_application_services import application_services from extensions.ext_database import db from fields.base import ResponseModel from graphon.model_runtime.errors.validate import CredentialsValidateFailedError @@ -32,7 +33,6 @@ from libs.login import login_required from models import Account from models.enums import PermissionEnum from models.provider_ids import DatasourceProviderID -from services.datasource_provider_service import DatasourceProviderService from services.plugin.oauth_service import OAuthProxyService @@ -176,7 +176,7 @@ class DatasourcePluginOAuthAuthorizationUrl(Resource): datasource_provider_id = DatasourceProviderID(provider_id) provider_name = datasource_provider_id.provider_name plugin_id = datasource_provider_id.plugin_id - oauth_config = DatasourceProviderService().get_oauth_client( + oauth_config = application_services().data_sources.providers.get_oauth_client( tenant_id=tenant_id, datasource_provider_id=datasource_provider_id, ) @@ -246,7 +246,7 @@ class DatasourceOAuthCallback(Resource): tenant_id: str = context["tenant_id"] datasource_provider_id = DatasourceProviderID(provider_id) plugin_id = datasource_provider_id.plugin_id - datasource_provider_service = DatasourceProviderService() + datasource_provider_service = application_services().data_sources.providers oauth_client_params = datasource_provider_service.get_oauth_client( tenant_id=tenant_id, datasource_provider_id=datasource_provider_id, @@ -315,7 +315,7 @@ class DatasourceAuth(Resource): @model_validate(DatasourceCredentialPayload) def post(self, req_data: DatasourceCredentialPayload, current_tenant_id: str, provider_id: str): datasource_provider_id = DatasourceProviderID(provider_id) - datasource_provider_service = DatasourceProviderService() + datasource_provider_service = application_services().data_sources.providers try: datasource_provider_service.add_datasource_api_key_provider( @@ -342,14 +342,14 @@ class DatasourceAuth(Resource): @with_current_tenant_id def get(self, current_tenant_id: str, user: Account, provider_id: str): datasource_provider_id = DatasourceProviderID(provider_id) - datasource_provider_service = DatasourceProviderService() + datasource_provider_service = application_services().data_sources.providers datasources = datasource_provider_service.list_datasource_credentials( tenant_id=current_tenant_id, provider=datasource_provider_id.provider_name, plugin_id=datasource_provider_id.plugin_id, user=user, - session=db.session(), + credential_query=application_services().credential_queries, ) return dump_response(DatasourceCredentialListResponse, {"result": datasources}), 200 @@ -370,7 +370,7 @@ class DatasourceAuthDeleteApi(Resource): plugin_id = datasource_provider_id.plugin_id provider_name = datasource_provider_id.provider_name - datasource_provider_service = DatasourceProviderService() + datasource_provider_service = application_services().data_sources.providers datasource_provider_service.remove_datasource_credentials( tenant_id=current_tenant_id, auth_id=req_data.credential_id, @@ -397,7 +397,7 @@ class DatasourceAuthUpdateApi(Resource): def post(self, req_data: DatasourceCredentialUpdatePayload, current_tenant_id: str, provider_id: str): datasource_provider_id = DatasourceProviderID(provider_id) - datasource_provider_service = DatasourceProviderService() + datasource_provider_service = application_services().data_sources.providers datasource_provider_service.update_datasource_credentials( tenant_id=current_tenant_id, auth_id=req_data.credential_id, @@ -422,9 +422,12 @@ class DatasourceAuthListApi(Resource): @with_current_user @with_current_tenant_id def get(self, current_tenant_id: str, user: Account): - datasource_provider_service = DatasourceProviderService() + datasource_provider_service = application_services().data_sources.providers datasources = datasource_provider_service.get_all_datasource_credentials( - tenant_id=current_tenant_id, session=db.session(), user=user + tenant_id=current_tenant_id, + session=db.session(), + user=user, + credential_query=application_services().credential_queries, ) return dump_response(DatasourceProviderAuthListResponse, {"result": datasources}), 200 @@ -442,9 +445,12 @@ class DatasourceHardCodeAuthListApi(Resource): @with_current_user @with_current_tenant_id def get(self, current_tenant_id: str, user: Account): - datasource_provider_service = DatasourceProviderService() + datasource_provider_service = application_services().data_sources.providers datasources = datasource_provider_service.get_hard_code_datasource_credentials( - tenant_id=current_tenant_id, session=db.session(), user=user + tenant_id=current_tenant_id, + session=db.session(), + user=user, + credential_query=application_services().credential_queries, ) return dump_response(DatasourceProviderAuthListResponse, {"result": datasources}), 200 @@ -464,7 +470,7 @@ class DatasourceAuthOauthCustomClient(Resource): @model_validate(DatasourceCustomClientPayload) def post(self, req_data: DatasourceCustomClientPayload, current_tenant_id: str, provider_id: str): datasource_provider_id = DatasourceProviderID(provider_id) - datasource_provider_service = DatasourceProviderService() + datasource_provider_service = application_services().data_sources.providers datasource_provider_service.setup_oauth_custom_client_params( tenant_id=current_tenant_id, datasource_provider_id=datasource_provider_id, @@ -482,7 +488,7 @@ class DatasourceAuthOauthCustomClient(Resource): @with_current_tenant_id def delete(self, current_tenant_id: str, provider_id: str): datasource_provider_id = DatasourceProviderID(provider_id) - datasource_provider_service = DatasourceProviderService() + datasource_provider_service = application_services().data_sources.providers datasource_provider_service.remove_oauth_custom_client_params( tenant_id=current_tenant_id, datasource_provider_id=datasource_provider_id, @@ -503,7 +509,7 @@ class DatasourceAuthDefaultApi(Resource): @model_validate(DatasourceDefaultPayload) def post(self, req_data: DatasourceDefaultPayload, current_tenant_id: str, provider_id: str): datasource_provider_id = DatasourceProviderID(provider_id) - datasource_provider_service = DatasourceProviderService() + datasource_provider_service = application_services().data_sources.providers datasource_provider_service.set_default_datasource_provider( tenant_id=current_tenant_id, datasource_provider_id=datasource_provider_id, @@ -525,7 +531,7 @@ class DatasourceUpdateProviderNameApi(Resource): @model_validate(DatasourceUpdateNamePayload) def post(self, req_data: DatasourceUpdateNamePayload, current_tenant_id: str, provider_id: str): datasource_provider_id = DatasourceProviderID(provider_id) - datasource_provider_service = DatasourceProviderService() + datasource_provider_service = application_services().data_sources.providers datasource_provider_service.update_datasource_provider_name( tenant_id=current_tenant_id, datasource_provider_id=datasource_provider_id, diff --git a/api/controllers/console/datasets/rag_pipeline/datasource_content_preview.py b/api/controllers/console/datasets/rag_pipeline/datasource_content_preview.py index 07ae94d3924..af4993d9d51 100644 --- a/api/controllers/console/datasets/rag_pipeline/datasource_content_preview.py +++ b/api/controllers/console/datasets/rag_pipeline/datasource_content_preview.py @@ -9,6 +9,7 @@ from controllers.common.schema import register_schema_models from controllers.console import console_ns from controllers.console.datasets.wraps import get_rag_pipeline from controllers.console.wraps import account_initialization_required, model_validate, setup_required, with_current_user +from extensions.ext_application_services import application_services from extensions.ext_database import db from libs.login import login_required from models import Account @@ -51,5 +52,6 @@ class DataSourceContentPreviewApi(Resource): datasource_type=datasource_type, is_published=True, credential_id=req_data.credential_id, + datasource_providers=application_services().data_sources.providers, ) return preview_content, 200 diff --git a/api/controllers/console/datasets/rag_pipeline/rag_pipeline.py b/api/controllers/console/datasets/rag_pipeline/rag_pipeline.py index a3d49f149b2..a69f2fa37fb 100644 --- a/api/controllers/console/datasets/rag_pipeline/rag_pipeline.py +++ b/api/controllers/console/datasets/rag_pipeline/rag_pipeline.py @@ -35,10 +35,11 @@ from libs.helper import dump_response from libs.login import login_required from models.account import Account from models.dataset import Pipeline -from services.dataset_service import DatasetService +from repositories.knowledge.dataset_read_repository import get_pipeline_dataset from services.entities.knowledge_entities.rag_pipeline_entities import IconInfo, PipelineTemplateInfoEntity from services.errors.account import NoPermissionError from services.errors.rag_pipeline import RagPipelineResourceNotFoundError +from services.knowledge.dataset_service import DatasetService from services.rag_pipeline.rag_pipeline import RagPipelineService logger: logging.Logger = logging.getLogger(__name__) @@ -232,7 +233,7 @@ class PublishCustomizedPipelineTemplateApi(Resource): pipeline: Pipeline, ) -> tuple[str, int]: session = db.session() - dataset = pipeline.retrieve_dataset(session=session) + dataset = get_pipeline_dataset(pipeline, session=session) if dataset is None: raise NotFound("Dataset not found") diff --git a/api/controllers/console/datasets/rag_pipeline/rag_pipeline_datasets.py b/api/controllers/console/datasets/rag_pipeline/rag_pipeline_datasets.py index 73ca05aba37..b6d1f6ff35a 100644 --- a/api/controllers/console/datasets/rag_pipeline/rag_pipeline_datasets.py +++ b/api/controllers/console/datasets/rag_pipeline/rag_pipeline_datasets.py @@ -17,14 +17,17 @@ from controllers.console.wraps import ( with_current_user, ) from extensions.ext_database import db -from fields.dataset_fields import DatasetDetailResponse, dataset_detail_response_source +from fields.dataset_fields import ( + DatasetDetailResponse, +) from libs.helper import dump_response from libs.login import login_required from models import Account from models.dataset import DatasetPermissionEnum -from services.dataset_service import DatasetPermissionService, DatasetService from services.enterprise import rbac_service as enterprise_rbac_service from services.entities.knowledge_entities.rag_pipeline_entities import IconInfo, RagPipelineDatasetCreateEntity +from services.knowledge.dataset_read_service import load_dataset_detail +from services.knowledge.dataset_service import DatasetPermissionService, DatasetService from services.rag_pipeline.rag_pipeline_dsl_service import RagPipelineDslService @@ -136,4 +139,4 @@ class CreateEmptyRagPipelineDatasetApi(Resource): ), session=session, ) - return dump_response(DatasetDetailResponse, dataset_detail_response_source(dataset, session=session)), 201 + return dump_response(DatasetDetailResponse, load_dataset_detail(dataset, session=session)), 201 diff --git a/api/controllers/console/datasets/rag_pipeline/rag_pipeline_workflow.py b/api/controllers/console/datasets/rag_pipeline/rag_pipeline_workflow.py index 86f9ea3534e..52898b07c34 100644 --- a/api/controllers/console/datasets/rag_pipeline/rag_pipeline_workflow.py +++ b/api/controllers/console/datasets/rag_pipeline/rag_pipeline_workflow.py @@ -46,6 +46,7 @@ from core.app.apps.base_app_queue_manager import AppQueueManager from core.app.apps.pipeline.pipeline_generator import PipelineGenerator from core.app.entities.app_invoke_entities import InvokeFrom from core.plugin.entities.plugin_daemon import PluginDatasourceProviderEntity +from extensions.ext_application_services import application_services from extensions.ext_database import db from factories import variable_factory from fields.base import ResponseModel @@ -66,10 +67,10 @@ from models.dataset import Pipeline from models.model import EndUser from models.workflow import Workflow from services.agent.retirement_service import WorkflowAgentRetirementService -from services.dataset_service import DatasetService from services.errors.app import IsDraftWorkflowError, WorkflowHashNotEqualError, WorkflowNotFoundError from services.errors.llm import InvokeRateLimitError from services.errors.rag_pipeline import RagPipelineResourceNotFoundError +from services.knowledge.dataset_service import DatasetService from services.rag_pipeline.pipeline_generate_service import PipelineGenerateService from services.rag_pipeline.rag_pipeline import RagPipelineService from services.rag_pipeline.rag_pipeline_manage_service import RagPipelineManageService @@ -303,7 +304,13 @@ class RagPipelineDraftRunIterationNodeApi(Resource): try: response = PipelineGenerateService.generate_single_iteration( - pipeline=pipeline, user=current_user, node_id=node_id, args=args, session=db.session(), streaming=True + generator=application_services().knowledge.pipeline_generator, + pipeline=pipeline, + user=current_user, + node_id=node_id, + args=args, + session=db.session(), + streaming=True, ) return helper.compact_generate_response(response) @@ -338,7 +345,13 @@ class RagPipelineDraftRunLoopNodeApi(Resource): try: response = PipelineGenerateService.generate_single_loop( - pipeline=pipeline, user=current_user, node_id=node_id, args=args, session=db.session(), streaming=True + generator=application_services().knowledge.pipeline_generator, + pipeline=pipeline, + user=current_user, + node_id=node_id, + args=args, + session=db.session(), + streaming=True, ) return helper.compact_generate_response(response) @@ -374,6 +387,7 @@ class DraftRagPipelineRunApi(Resource): try: response = PipelineGenerateService.generate( + generator=application_services().knowledge.pipeline_generator, session=session, pipeline=pipeline, user=current_user, @@ -409,6 +423,7 @@ class PublishedRagPipelineRunApi(Resource): try: response = PipelineGenerateService.generate( + generator=application_services().knowledge.pipeline_generator, session=session, pipeline=pipeline, user=current_user, @@ -450,6 +465,7 @@ class RagPipelinePublishedDatasourceNodeRunApi(Resource): datasource_type=req_data.datasource_type, is_published=False, credential_id=req_data.credential_id, + datasource_providers=application_services().data_sources.providers, ) ) ) @@ -483,6 +499,7 @@ class RagPipelineDraftDatasourceNodeRunApi(Resource): datasource_type=req_data.datasource_type, is_published=False, credential_id=req_data.credential_id, + datasource_providers=application_services().data_sources.providers, ) ) ) @@ -1016,7 +1033,9 @@ class DatasourceListApi(Resource): @account_initialization_required @with_current_tenant_id def get(self, current_tenant_id: str): - providers = RagPipelineManageService.list_rag_pipeline_datasources(current_tenant_id) + providers = RagPipelineManageService.list_rag_pipeline_datasources( + current_tenant_id, datasource_providers=application_services().data_sources.providers + ) return RagPipelineDatasourceListResponse.model_validate(providers, from_attributes=True).model_dump(mode="json") diff --git a/api/controllers/console/datasets/website.py b/api/controllers/console/datasets/website.py index 9411d128195..4a9e9af0912 100644 --- a/api/controllers/console/datasets/website.py +++ b/api/controllers/console/datasets/website.py @@ -7,8 +7,9 @@ from controllers.common.schema import query_params_from_model, register_response from controllers.console import console_ns from controllers.console.datasets.error import WebsiteCrawlError from controllers.console.wraps import account_initialization_required, model_validate, setup_required +from extensions.ext_application_services import application_services from libs.login import login_required -from services.website_service import WebsiteCrawlApiRequest, WebsiteCrawlStatusApiRequest, WebsiteService +from services.data_source.website_service import WebsiteCrawlApiRequest, WebsiteCrawlStatusApiRequest, WebsiteService class WebsiteCrawlPayload(BaseModel): @@ -49,7 +50,7 @@ class WebsiteCrawlApi(Resource): # Crawl URL using typed request try: - result = WebsiteService.crawl_url(api_request) + result = WebsiteService(providers=application_services().data_sources.providers).crawl_url(api_request) except Exception as e: raise WebsiteCrawlError(str(e)) return result, 200 @@ -77,7 +78,9 @@ class WebsiteCrawlStatusApi(Resource): # Get crawl status using typed request try: - result = WebsiteService.get_crawl_status_typed(api_request) + result = WebsiteService(providers=application_services().data_sources.providers).get_crawl_status_typed( + api_request + ) except Exception as e: raise WebsiteCrawlError(str(e)) return result, 200 diff --git a/api/controllers/console/flask_admission.py b/api/controllers/console/flask_admission.py index cb923fb4606..7d492c04ee2 100644 --- a/api/controllers/console/flask_admission.py +++ b/api/controllers/console/flask_admission.py @@ -71,7 +71,7 @@ def console_account_admission[T, **P, R]( ) -> Callable[Concatenate[T, P], R | Response]: admitted_view = cloud_edition_billing_resource_check(billing_resource)(view) if billing_resource else view - @wraps(view, updated=()) + @wraps(view) def inject_request_context(self: T, /, *args: P.args, **kwargs: P.kwargs) -> R: account_with_tenant = current_account_with_tenant() account = account_with_tenant.account diff --git a/api/controllers/console/workspace/model_providers.py b/api/controllers/console/workspace/model_providers.py index dd67b9ec131..23f467ef8be 100644 --- a/api/controllers/console/workspace/model_providers.py +++ b/api/controllers/console/workspace/model_providers.py @@ -33,7 +33,7 @@ from services.entities.model_provider_entities import ( ModelProviderSummaryResponse, ProviderResponse, ) -from services.model_provider_service import ModelProviderService +from services.model_provider.service import ModelProviderService from services.workspace_service import WorkspaceService diff --git a/api/controllers/console/workspace/models.py b/api/controllers/console/workspace/models.py index 85488610667..759887ba5f9 100644 --- a/api/controllers/console/workspace/models.py +++ b/api/controllers/console/workspace/models.py @@ -25,6 +25,7 @@ from controllers.console.wraps import ( ) from core.entities.model_entities import DefaultModelSetting from core.entities.provider_entities import CredentialConfiguration +from extensions.ext_application_services import application_services from extensions.ext_database import db from fields.base import ResponseModel from graphon.model_runtime.entities.model_entities import ModelType, ParameterRule @@ -38,7 +39,7 @@ from services.entities.model_provider_entities import ( ProviderWithModelsResponse, ) from services.model_load_balancing_service import ModelLoadBalancingService -from services.model_provider_service import ModelProviderService +from services.model_provider.service import ModelProviderService logger = logging.getLogger(__name__) @@ -376,6 +377,7 @@ class ModelProviderModelCredentialApi(Resource): tenant_id=tenant_id, provider=provider, user=user, + credential_query=application_services().credential_queries, ) else: available_credentials = model_provider_service.get_provider_model_available_credentials( diff --git a/api/controllers/console/workspace/tool_providers.py b/api/controllers/console/workspace/tool_providers.py index 34b89d5e767..01b1e661289 100644 --- a/api/controllers/console/workspace/tool_providers.py +++ b/api/controllers/console/workspace/tool_providers.py @@ -62,6 +62,7 @@ from core.tools.entities.tool_entities import ( WorkflowToolParameterConfiguration, ) from enums import DeploymentEdition +from extensions.ext_application_services import application_services from extensions.ext_database import db from fields.base import ResponseModel from libs.helper import alphanumeric, dump_response, uuid_value @@ -664,9 +665,9 @@ class ToolBuiltinProviderGetCredentialsApi(Resource): BuiltinToolManageService.get_builtin_tool_provider_credentials( tenant_id=tenant_id, provider_name=provider, - session=db.session(), user=user, include_credential_ids=query.include_credential_ids or None, + credential_query=application_services().credential_queries, ), ) @@ -1363,9 +1364,9 @@ class ToolBuiltinProviderGetCredentialInfoApi(Resource): BuiltinToolManageService.get_builtin_tool_provider_credential_info( tenant_id=tenant_id, provider=provider, - session=db.session(), user=user, include_credential_ids=query.include_credential_ids or None, + credential_query=application_services().credential_queries, ), ) diff --git a/api/controllers/console/workspace/trigger_providers.py b/api/controllers/console/workspace/trigger_providers.py index 8d15c3a99b4..c9d6074aa44 100644 --- a/api/controllers/console/workspace/trigger_providers.py +++ b/api/controllers/console/workspace/trigger_providers.py @@ -22,6 +22,7 @@ from core.trigger.entities.api_entities import ( ) from core.trigger.entities.entities import RequestLog, SubscriptionBuilderUpdater from core.trigger.trigger_manager import TriggerManager +from extensions.ext_application_services import application_services from extensions.ext_database import db from fields.base import ResponseModel from libs.helper import dump_response @@ -213,6 +214,7 @@ class TriggerSubscriptionListApi(Resource): tenant_id=tenant_id, provider_id=TriggerProviderID(provider), user=user, + credential_query=application_services().credential_queries, ), ) except ValueError as e: diff --git a/api/controllers/service_api/dataset/dataset.py b/api/controllers/service_api/dataset/dataset.py index 00c8d18894d..9c6a2d6750d 100644 --- a/api/controllers/service_api/dataset/dataset.py +++ b/api/controllers/service_api/dataset/dataset.py @@ -37,11 +37,6 @@ from core.plugin.impl.model_runtime_factory import create_plugin_provider_manage from core.rag.index_processor.constant.index_type import IndexTechniqueType from extensions.ext_application_services import application_services from fields.base import ResponseModel -from fields.dataset_fields import ( - DatasetDetailPrefetch, - build_dataset_detail_prefetch, - dataset_detail_response_source, -) from fields.dataset_fields import DatasetDetailResponse as BaseDatasetDetailResponse from graphon.model_runtime.entities.model_entities import ModelType from libs.helper import dump_response @@ -51,9 +46,11 @@ from models.account import Account from models.dataset import DatasetPermissionEnum from models.enums import TagType from models.provider_ids import ModelProviderID -from services.dataset_service import DatasetPermissionService, DatasetService, DocumentService from services.enterprise import rbac_service as enterprise_rbac_service -from services.entities.knowledge_entities.knowledge_entities import ( +from services.knowledge.dataset_read_service import load_dataset_detail, load_dataset_details +from services.knowledge.dataset_service import DatasetPermissionService, DatasetService, DocumentService +from services.knowledge.entities.datasets import DatasetDetailRecord +from services.knowledge.entities.knowledge_entities import ( ExternalRetrievalModel, KnowledgeProvider, RetrievalModel, @@ -104,12 +101,8 @@ _SERVICE_DATASET_DETAIL_EXCLUDE = {"permission_keys"} _SERVICE_DATASET_LIST_EXCLUDE = {"data": {"__all__": _SERVICE_DATASET_DETAIL_EXCLUDE}} -def _dump_service_dataset_detail( - dataset: Any, *, session: Session, prefetch: DatasetDetailPrefetch | None = None -) -> dict[str, Any]: - return DatasetDetailResponse.model_validate( - dataset_detail_response_source(dataset, session=session, prefetch=prefetch), from_attributes=True - ).model_dump( +def _dump_service_dataset_detail(detail: DatasetDetailRecord) -> dict[str, Any]: + return DatasetDetailResponse.model_validate(detail).model_dump( mode="json", exclude=_SERVICE_DATASET_DETAIL_EXCLUDE, ) @@ -470,8 +463,7 @@ class DatasetListApi(DatasetApiResource): for embedding_model in embedding_models: model_names.append(f"{embedding_model.model}:{embedding_model.provider.provider}") - prefetch = build_dataset_detail_prefetch(datasets, session=session) - data = [_dump_service_dataset_detail(dataset, session=session, prefetch=prefetch) for dataset in datasets] + data = [_dump_service_dataset_detail(detail) for detail in load_dataset_details(datasets, session=session)] for item in data: if item["indexing_technique"] == IndexTechniqueType.HIGH_QUALITY and item["embedding_model_provider"]: item["embedding_model_provider"] = str(ModelProviderID(item["embedding_model_provider"])) @@ -571,7 +563,7 @@ class DatasetListApi(DatasetApiResource): enterprise_rbac_service.ReplaceMemberBindings(automatic_include_workspace_members=False), ) - return _dump_service_dataset_detail(dataset, session=session), 200 + return _dump_service_dataset_detail(load_dataset_detail(dataset, session=session)), 200 @service_api_ns.route("/datasets/") @@ -617,7 +609,7 @@ class DatasetApi(DatasetApiResource): DatasetService.check_dataset_permission(dataset, current_user, session) except services.errors.account.NoPermissionError as e: raise Forbidden(str(e)) - data = _dump_service_dataset_detail(dataset, session=session) + data = _dump_service_dataset_detail(load_dataset_detail(dataset, session=session)) # check embedding setting assert isinstance(current_user, Account) cid = current_user.current_tenant_id @@ -734,7 +726,7 @@ class DatasetApi(DatasetApiResource): if dataset is None: raise NotFound("Dataset not found.") - result_data = _dump_service_dataset_detail(dataset, session=session) + result_data = _dump_service_dataset_detail(load_dataset_detail(dataset, session=session)) assert isinstance(current_user, Account) tenant_id = current_user.current_tenant_id diff --git a/api/controllers/service_api/dataset/document.py b/api/controllers/service_api/dataset/document.py index 50e5ebb3308..72870918ced 100644 --- a/api/controllers/service_api/dataset/document.py +++ b/api/controllers/service_api/dataset/document.py @@ -69,25 +69,36 @@ from fields.document_fields import ( DocumentMetadataResponse, DocumentResponse, DocumentStatusListResponse, - document_response, - document_responses, normalize_enum, ) from libs.helper import dump_response from libs.login import current_user from libs.pagination import clamp_pagination, paginate_query from models.dataset import Dataset, Document -from services.dataset_service import DatasetService, DocumentService -from services.entities.knowledge_entities.knowledge_entities import ( +from repositories.knowledge.dataset_read_repository import ( + get_dataset_creator, + get_document_hit_count, + get_document_process_rule, + get_document_segment_count, + get_latest_dataset_process_rule, +) +from services.feature_service import FeatureService +from services.file_service import FileService +from services.knowledge.dataset_read_service import ( + get_document_metadata_details, + get_document_source_detail, + load_document_detail, + load_document_details, +) +from services.knowledge.dataset_service import DatasetService, DocumentService +from services.knowledge.entities.knowledge_entities import ( DocForm, IndexingTechnique, KnowledgeConfig, ProcessRule, RetrievalModel, ) -from services.feature_service import FeatureService -from services.file_service import FileService -from services.summary_index_service import SummaryIndexService +from services.knowledge.summaries.adapters import SummaryIndexAdapter class DocumentTextCreatePayload(BaseModel): @@ -304,7 +315,7 @@ class DocumentAndBatchResponse(ResponseModel): def _document_and_batch_response(document: Document, batch: str, *, session: Session) -> dict[str, Any]: return dump_response( DocumentAndBatchResponse, - {"document": document_response(document, session=session), "batch": batch}, + {"document": load_document_detail(document, session=session), "batch": batch}, ) @@ -445,7 +456,7 @@ def _create_document_by_text(session: Session, tenant_id: str, dataset_id: UUID) dataset=dataset, knowledge_config=knowledge_config, account=current_user, - dataset_process_rule=dataset.get_latest_process_rule(session=session) + dataset_process_rule=get_latest_dataset_process_rule(dataset, session=session) if "process_rule" not in args else None, created_from="api", @@ -509,7 +520,7 @@ def _update_document_by_text( dataset=dataset, knowledge_config=knowledge_config, account=current_user, - dataset_process_rule=dataset.get_latest_process_rule(session=session) + dataset_process_rule=get_latest_dataset_process_rule(dataset, session=session) if "process_rule" not in args else None, created_from="api", @@ -822,7 +833,9 @@ class DocumentAddByFileApi(DatasetApiResource): knowledge_config = KnowledgeConfig.model_validate(args) DocumentService.document_create_args_validate(knowledge_config) - dataset_process_rule = dataset.get_latest_process_rule(session=session) if "process_rule" not in args else None + dataset_process_rule = ( + get_latest_dataset_process_rule(dataset, session=session) if "process_rule" not in args else None + ) if not knowledge_config.original_document_id and not dataset_process_rule and not knowledge_config.process_rule: raise ValueError("process_rule is required.") @@ -830,7 +843,7 @@ class DocumentAddByFileApi(DatasetApiResource): documents, batch = DocumentService.save_document_with_dataset_id( dataset=dataset, knowledge_config=knowledge_config, - account=dataset.get_created_by_account(session=session), + account=get_dataset_creator(dataset, session=session), dataset_process_rule=dataset_process_rule, created_from="api", session=session, @@ -909,8 +922,8 @@ def _update_document_by_file( documents, _ = DocumentService.save_document_with_dataset_id( dataset=dataset, knowledge_config=knowledge_config, - account=dataset.get_created_by_account(session=session), - dataset_process_rule=dataset.get_latest_process_rule(session=session) + account=get_dataset_creator(dataset, session=session), + dataset_process_rule=get_latest_dataset_process_rule(dataset, session=session) if "process_rule" not in args else None, created_from="api", @@ -1046,9 +1059,7 @@ class DocumentListApi(DatasetApiResource): ) response = { - "data": document_responses(documents, session=session), - # The result object already knows: it was built from the page the query - # ran with, while the requested values are only ever a request. + "data": load_document_details(documents, session=session), "has_more": paginated_documents.has_next, "limit": paginated_documents.per_page, "total": paginated_documents.total, @@ -1303,7 +1314,7 @@ class DocumentApi(DatasetApiResource): summary_index_status = None has_summary_index = dataset.summary_index_setting and dataset.summary_index_setting.get("enable") is True if has_summary_index and document.need_summary is True: - summary_index_status = SummaryIndexService.get_document_summary_index_status( + summary_index_status = SummaryIndexAdapter.get_document_summary_index_status( document_id=document_id_str, dataset_id=dataset_id_str, tenant_id=tenant_id, @@ -1315,15 +1326,15 @@ class DocumentApi(DatasetApiResource): response = { "id": document.id, "doc_type": document.doc_type, - "doc_metadata": document.get_doc_metadata_details(session=session), + "doc_metadata": get_document_metadata_details(document, session=session), } elif metadata == "without": dataset_process_rules = DatasetService.get_process_rules(dataset_id_str, session) response_exclude = {"doc_type", "doc_metadata"} - document_process_rule = document.get_dataset_process_rule(session=session) + document_process_rule = get_document_process_rule(document, session=session) document_process_rules: Mapping[str, Any] = document_process_rule.to_dict() if document_process_rule else {} - data_source_info = document.get_data_source_detail_dict(session=session) - segment_count = document.get_segment_count(session=session) + data_source_info = get_document_source_detail(document, session=session) + segment_count = get_document_segment_count(document, session=session) response = { "id": document.id, "position": document.position, @@ -1348,7 +1359,7 @@ class DocumentApi(DatasetApiResource): "archived": document.archived, "segment_count": segment_count, "average_segment_length": (document.word_count or 0) // segment_count if segment_count else 0, - "hit_count": document.get_hit_count(session=session), + "hit_count": get_document_hit_count(document, session=session), "display_status": document.display_status, "doc_form": document.doc_form, "doc_language": document.doc_language, @@ -1357,10 +1368,10 @@ class DocumentApi(DatasetApiResource): } else: dataset_process_rules = DatasetService.get_process_rules(dataset_id_str, session) - document_process_rule = document.get_dataset_process_rule(session=session) + document_process_rule = get_document_process_rule(document, session=session) document_process_rules = document_process_rule.to_dict() if document_process_rule else {} - data_source_info = document.get_data_source_detail_dict(session=session) - segment_count = document.get_segment_count(session=session) + data_source_info = get_document_source_detail(document, session=session) + segment_count = get_document_segment_count(document, session=session) response = { "id": document.id, "position": document.position, @@ -1384,10 +1395,10 @@ class DocumentApi(DatasetApiResource): "disabled_by": document.disabled_by, "archived": document.archived, "doc_type": document.doc_type, - "doc_metadata": document.get_doc_metadata_details(session=session), + "doc_metadata": get_document_metadata_details(document, session=session), "segment_count": segment_count, "average_segment_length": (document.word_count or 0) // segment_count if segment_count else 0, - "hit_count": document.get_hit_count(session=session), + "hit_count": get_document_hit_count(document, session=session), "display_status": document.display_status, "doc_form": document.doc_form, "doc_language": document.doc_language, diff --git a/api/controllers/service_api/dataset/metadata.py b/api/controllers/service_api/dataset/metadata.py index 693ba8377b9..99c7c2458ec 100644 --- a/api/controllers/service_api/dataset/metadata.py +++ b/api/controllers/service_api/dataset/metadata.py @@ -1,4 +1,4 @@ -from typing import Literal, cast +from typing import Literal from uuid import UUID from flask_login import current_user @@ -11,6 +11,7 @@ from controllers.common.session import with_session from controllers.console.wraps import model_validate from controllers.service_api import service_api_ns from controllers.service_api.wraps import DatasetApiResource, cloud_edition_billing_rate_limit_check +from extensions.ext_application_services import application_services from fields.dataset_fields import ( DatasetMetadataActionResponse, DatasetMetadataBuiltInFieldsResponse, @@ -18,16 +19,15 @@ from fields.dataset_fields import ( DatasetMetadataResponse, ) from libs.helper import dump_response -from models import Account -from services.dataset_service import DatasetService -from services.entities.knowledge_entities.knowledge_entities import ( +from services.errors.metadata import MetadataResourceNotFoundError +from services.knowledge.dataset_service import DatasetService +from services.knowledge.entities.knowledge_entities import ( DocumentMetadataOperation, MetadataArgs, MetadataDetail, MetadataOperationData, ) -from services.errors.metadata import MetadataResourceNotFoundError -from services.metadata_service import MetadataService +from services.knowledge.resource_scope import DatasetRef BUILT_IN_METADATA_ACTION_PARAM = { "description": "`enable` to activate built-in metadata fields, `disable` to deactivate them.", @@ -87,12 +87,18 @@ class DatasetMetadataCreateServiceApi(DatasetApiResource): """Create metadata for a dataset.""" dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, session) + dataset = DatasetService.get_dataset_for_tenant(dataset_id_str, str(tenant_id), session=session) if dataset is None: raise NotFound("Dataset not found.") DatasetService.check_dataset_permission(dataset, current_user, session) - metadata = MetadataService.create_metadata(dataset_id_str, metadata_args, session=session) + session.commit() + try: + metadata = application_services().knowledge.metadata.create_metadata( + DatasetRef(str(tenant_id), dataset_id_str), metadata_args, actor_id=current_user.id + ) + except MetadataResourceNotFoundError as exc: + raise NotFound(str(exc)) from exc return dump_response(DatasetMetadataResponse, metadata), 201 @service_api_ns.doc( @@ -123,10 +129,16 @@ class DatasetMetadataCreateServiceApi(DatasetApiResource): def get(self, session: Session, tenant_id, dataset_id: UUID): """Get all metadata for a dataset.""" dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, session) + dataset = DatasetService.get_dataset_for_tenant(dataset_id_str, str(tenant_id), session=session) if dataset is None: raise NotFound("Dataset not found.") - metadata = MetadataService.get_dataset_metadatas(dataset, session) + session.commit() + try: + metadata = application_services().knowledge.metadata.get_dataset_metadatas( + DatasetRef(str(tenant_id), dataset_id_str) + ) + except MetadataResourceNotFoundError as exc: + raise NotFound(str(exc)) from exc return dump_response(DatasetMetadataListResponse, metadata), 200 @@ -168,9 +180,13 @@ class DatasetMetadataServiceApi(DatasetApiResource): raise NotFound("Dataset not found.") DatasetService.check_dataset_permission(dataset, current_user, session) - metadata = MetadataService.update_metadata_name( - dataset, metadata_id_str, payload.name, cast(Account, current_user), session=session - ) + session.commit() + try: + metadata = application_services().knowledge.metadata.update_metadata_name( + DatasetRef(str(tenant_id), dataset_id_str), metadata_id_str, payload.name, actor_id=current_user.id + ) + except MetadataResourceNotFoundError as exc: + raise NotFound(str(exc)) from exc return dump_response(DatasetMetadataResponse, metadata), 200 @service_api_ns.doc( @@ -206,7 +222,13 @@ class DatasetMetadataServiceApi(DatasetApiResource): raise NotFound("Dataset not found.") DatasetService.check_dataset_permission(dataset, current_user, session) - MetadataService.delete_metadata(dataset, metadata_id_str, session) + session.commit() + try: + application_services().knowledge.metadata.delete_metadata( + DatasetRef(str(tenant_id), dataset_id_str), metadata_id_str + ) + except MetadataResourceNotFoundError as exc: + raise NotFound(str(exc)) from exc return "", 204 @@ -239,7 +261,7 @@ class DatasetMetadataBuiltInFieldServiceApi(DatasetApiResource): ) def get(self, tenant_id, dataset_id: UUID): """Get all built-in metadata fields.""" - built_in_fields = MetadataService.get_built_in_fields() + built_in_fields = application_services().knowledge.metadata.get_built_in_fields() return dump_response(DatasetMetadataBuiltInFieldsResponse, {"fields": built_in_fields}), 200 @@ -271,16 +293,24 @@ class DatasetMetadataBuiltInFieldActionServiceApi(DatasetApiResource): def post(self, session: Session, tenant_id, dataset_id: UUID, action: Literal["enable", "disable"]): """Enable or disable built-in metadata field.""" dataset_id_str = str(dataset_id) - dataset = DatasetService.get_dataset(dataset_id_str, session) + dataset = DatasetService.get_dataset_for_tenant(dataset_id_str, str(tenant_id), session=session) if dataset is None: raise NotFound("Dataset not found.") DatasetService.check_dataset_permission(dataset, current_user, session) - match action: - case "enable": - MetadataService.enable_built_in_field(dataset, session) - case "disable": - MetadataService.disable_built_in_field(dataset, session) + session.commit() + try: + match action: + case "enable": + application_services().knowledge.metadata.enable_built_in_field( + DatasetRef(str(tenant_id), dataset_id_str) + ) + case "disable": + application_services().knowledge.metadata.disable_built_in_field( + DatasetRef(str(tenant_id), dataset_id_str) + ) + except MetadataResourceNotFoundError as exc: + raise NotFound(str(exc)) from exc return dump_response(DatasetMetadataActionResponse, {"result": "success"}), 200 @@ -324,9 +354,10 @@ class DocumentMetadataEditServiceApi(DatasetApiResource): raise NotFound("Dataset not found.") DatasetService.check_dataset_permission(dataset, current_user, session) + session.commit() try: - MetadataService.update_documents_metadata( - dataset, metadata_args, cast(Account, current_user), session=session + application_services().knowledge.metadata.update_documents_metadata( + DatasetRef(str(tenant_id), str(dataset_id)), metadata_args, actor_id=current_user.id ) except MetadataResourceNotFoundError as exc: raise NotFound(str(exc)) from exc diff --git a/api/controllers/service_api/dataset/rag_pipeline/rag_pipeline_workflow.py b/api/controllers/service_api/dataset/rag_pipeline/rag_pipeline_workflow.py index 245466367b7..64c5ab76006 100644 --- a/api/controllers/service_api/dataset/rag_pipeline/rag_pipeline_workflow.py +++ b/api/controllers/service_api/dataset/rag_pipeline/rag_pipeline_workflow.py @@ -32,6 +32,7 @@ from controllers.service_api.wraps import DatasetApiResource from core.app.apps.pipeline.pipeline_generator import PipelineGenerator from core.app.entities.app_invoke_entities import InvokeFrom from core.entities.knowledge_entities import PipelineDataset, PipelineDocument +from extensions.ext_application_services import application_services from fields.base import ResponseModel from libs import helper from libs.helper import dump_response @@ -183,7 +184,11 @@ class DatasourcePluginsApi(DatasetApiResource): rag_pipeline_service = RagPipelineService(db.session()) datasource_plugins: list[dict[Any, Any]] = rag_pipeline_service.get_datasource_plugins( - tenant_id=tenant_id, dataset_id=dataset_id_str, is_published=query.is_published + tenant_id=tenant_id, + dataset_id=dataset_id_str, + is_published=query.is_published, + credential_query=application_services().credential_queries, + datasource_providers=application_services().data_sources.providers, ) return dump_response(DatasourcePluginListResponse, datasource_plugins), 200 @@ -247,6 +252,7 @@ class DatasourceNodeRunApi(DatasetApiResource): datasource_type=datasource_node_run_api_entity.datasource_type, is_published=datasource_node_run_api_entity.is_published, credential_id=datasource_node_run_api_entity.credential_id, + datasource_providers=application_services().data_sources.providers, ) ) ) @@ -310,6 +316,7 @@ class PipelineRunApi(DatasetApiResource): pipeline = rag_pipeline_service.get_pipeline(tenant_id=tenant_id, dataset_id=dataset_id_str) try: response: dict[Any, Any] | Generator[str, Any, None] = PipelineGenerateService.generate( + generator=application_services().knowledge.pipeline_generator, session=session, pipeline=pipeline, user=current_user, diff --git a/api/controllers/service_api/dataset/segment.py b/api/controllers/service_api/dataset/segment.py index 921c59a0942..7d93859970c 100644 --- a/api/controllers/service_api/dataset/segment.py +++ b/api/controllers/service_api/dataset/segment.py @@ -27,26 +27,26 @@ from controllers.service_api.wraps import ( from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError from core.model_manager import ModelManager from core.rag.index_processor.constant.index_type import IndexTechniqueType +from extensions.ext_application_services import application_services from fields.base import ResponseModel from fields.segment_fields import ( ChildChunkDetailResponse, ChildChunkListResponse, SegmentDetailResponse, SegmentResponse, - segment_response_with_summary, - segment_responses_with_summaries, ) from graphon.model_runtime.entities.model_entities import ModelType from libs.helper import dump_response from libs.login import current_account_with_tenant from models.dataset import Dataset, Document, DocumentSegment -from services.dataset_ref_service import DatasetRefService, SegmentRef -from services.dataset_service import DatasetService, DocumentService, SegmentService -from services.entities.knowledge_entities.knowledge_entities import SegmentUpdateArgs from services.errors.chunk import ChildChunkDeleteIndexError, ChildChunkIndexingError from services.errors.chunk import ChildChunkDeleteIndexError as ChildChunkDeleteIndexServiceError from services.errors.chunk import ChildChunkIndexingError as ChildChunkIndexingServiceError -from services.summary_index_service import SummaryIndexService +from services.knowledge.dataset_read_service import load_segment_detail, load_segment_details +from services.knowledge.dataset_service import DatasetService, DocumentService, SegmentService +from services.knowledge.entities.segments import SegmentUpdateArgs +from services.knowledge.resource_scope import DatasetRef, SegmentRef +from services.knowledge.summaries.adapters import SummaryIndexAdapter class SegmentCreateItemPayload(BaseModel): @@ -133,12 +133,11 @@ register_response_schema_models( def _get_segment_for_document( session: Session, dataset: Dataset, document: Document, segment_id: str ) -> tuple[SegmentRef, DocumentSegment]: - dataset_ref = DatasetRefService.create_dataset_ref(dataset) - document_ref = DatasetRefService.create_document_ref(dataset_ref, document) - if document_ref is None: + if document.tenant_id != dataset.tenant_id or document.dataset_id != dataset.id: raise NotFound("Document not found.") - segment_ref = DatasetRefService.create_segment_ref(document_ref, segment_id) + dataset_ref = DatasetRef(tenant_id=dataset.tenant_id, dataset_id=dataset.id) + segment_ref = dataset_ref.document(document.id).segment(segment_id) segment = SegmentService.get_segment_by_ref(segment_ref, session=session) if not segment: raise NotFound("Segment not found.") @@ -230,17 +229,24 @@ class SegmentApi(DatasetApiResource): for args_item in segment_items: SegmentService.segment_create_args_validate(args_item, document) segments = cast( - list[DocumentSegment], SegmentService.multi_create_segment(segment_items, document, dataset, session) + list[DocumentSegment], + SegmentService.multi_create_segment( + segment_items, + document, + dataset, + session, + mutations=application_services().knowledge.segments.mutations, + ), ) segment_ids = [segment.id for segment in segments] summaries: dict[str, str | None] = {} if segment_ids: - summary_records = SummaryIndexService.get_segments_summaries( + summary_records = SummaryIndexAdapter.get_segments_summaries( segment_ids=segment_ids, dataset_id=dataset_id_str, session=session ) summaries = {chunk_id: record.summary_content for chunk_id, record in summary_records.items()} response = { - "data": segment_responses_with_summaries(segments, summaries, session=session), + "data": load_segment_details(segments, summaries, session=session), "doc_form": document.doc_form, } return dump_response(SegmentCreateListResponse, response), 200 @@ -322,13 +328,13 @@ class SegmentApi(DatasetApiResource): segment_ids = [segment.id for segment in segments] summaries: dict[str, str | None] = {} if segment_ids: - summary_records = SummaryIndexService.get_segments_summaries( + summary_records = SummaryIndexAdapter.get_segments_summaries( segment_ids=segment_ids, dataset_id=dataset_id_str, session=session ) summaries = {chunk_id: record.summary_content for chunk_id, record in summary_records.items()} response = { - "data": segment_responses_with_summaries(segments, summaries, session=session), + "data": load_segment_details(segments, summaries, session=session), "doc_form": document.doc_form, "total": total, "has_more": page * limit < total, @@ -380,7 +386,9 @@ class DatasetSegmentApi(DatasetApiResource): raise NotFound("Document not found.") segment_id_str = str(segment_id) _, segment = _get_segment_for_document(session, dataset, document, segment_id_str) - SegmentService.delete_segment(segment, document, dataset, session) + SegmentService.delete_segment( + segment, document, dataset, session, mutations=application_services().knowledge.segments.mutations + ) return "", 204 @service_api_ns.doc( @@ -417,7 +425,7 @@ class DatasetSegmentApi(DatasetApiResource): document_id: UUID, segment_id: UUID, ): - _, current_tenant_id = current_account_with_tenant() + current_account, current_tenant_id = current_account_with_tenant() dataset_id_str = str(dataset_id) # check dataset dataset = session.scalar( @@ -451,14 +459,20 @@ class DatasetSegmentApi(DatasetApiResource): segment_id_str = str(segment_id) _, segment = _get_segment_for_document(session, dataset, document, segment_id_str) - updated_segment = SegmentService.update_segment(payload.segment, segment, document, dataset, session) - summary = SummaryIndexService.get_segment_summary( + updated_segment = SegmentService.update_segment( + payload.segment, + segment, + document, + dataset, + session, + actor_id=current_account.id, + mutations=application_services().knowledge.segments.mutations, + ) + summary = SummaryIndexAdapter.get_segment_summary( segment_id=updated_segment.id, dataset_id=dataset_id_str, session=session ) response = { - "data": segment_response_with_summary( - updated_segment, summary.summary_content if summary else None, session=session - ), + "data": load_segment_detail(updated_segment, summary.summary_content if summary else None, session=session), "doc_form": document.doc_form, } return dump_response(SegmentDetailResponse, response), 200 @@ -510,13 +524,11 @@ class DatasetSegmentApi(DatasetApiResource): segment_id_str = str(segment_id) _, segment = _get_segment_for_document(session, dataset, document, segment_id_str) - summary = SummaryIndexService.get_segment_summary( + summary = SummaryIndexAdapter.get_segment_summary( segment_id=segment.id, dataset_id=dataset_id_str, session=session ) response = { - "data": segment_response_with_summary( - segment, summary.summary_content if summary else None, session=session - ), + "data": load_segment_detail(segment, summary.summary_content if summary else None, session=session), "doc_form": document.doc_form, } return dump_response(SegmentDetailResponse, response), 200 @@ -567,7 +579,7 @@ class ChildChunkApi(DatasetApiResource): document_id: UUID, segment_id: UUID, ): - _, current_tenant_id = current_account_with_tenant() + current_account, current_tenant_id = current_account_with_tenant() """Create child chunk.""" dataset_id_str = str(dataset_id) # check dataset @@ -604,7 +616,15 @@ class ChildChunkApi(DatasetApiResource): raise ProviderNotInitializeError(ex.description) try: - child_chunk = SegmentService.create_child_chunk(payload.content, segment, document, dataset, session) + child_chunk = SegmentService.create_child_chunk( + payload.content, + segment, + document, + dataset, + session, + actor_id=current_account.id, + mutations=application_services().knowledge.segments.mutations, + ) except ChildChunkIndexingServiceError as e: raise ChildChunkIndexingError(str(e)) @@ -744,7 +764,9 @@ class DatasetChildChunkApi(DatasetApiResource): raise NotFound("Child chunk not found.") try: - SegmentService.delete_child_chunk(child_chunk, dataset, session) + SegmentService.delete_child_chunk( + child_chunk, dataset, session, mutations=application_services().knowledge.segments.mutations + ) except ChildChunkDeleteIndexServiceError as e: raise ChildChunkDeleteIndexError(str(e)) @@ -790,7 +812,7 @@ class DatasetChildChunkApi(DatasetApiResource): segment_id: UUID, child_chunk_id: UUID, ): - current_account_with_tenant() + current_account, _ = current_account_with_tenant() """Update child chunk.""" dataset_id_str = str(dataset_id) # check dataset @@ -817,7 +839,14 @@ class DatasetChildChunkApi(DatasetApiResource): try: child_chunk = SegmentService.update_child_chunk( - payload.content, child_chunk, segment, document, dataset, session + payload.content, + child_chunk, + segment, + document, + dataset, + session, + actor_id=current_account.id, + mutations=application_services().knowledge.segments.mutations, ) except ChildChunkIndexingServiceError as e: raise ChildChunkIndexingError(str(e)) diff --git a/api/controllers/service_api/workspace/models.py b/api/controllers/service_api/workspace/models.py index c2ea1a72cf5..0606906fe40 100644 --- a/api/controllers/service_api/workspace/models.py +++ b/api/controllers/service_api/workspace/models.py @@ -7,7 +7,7 @@ from controllers.service_api.wraps import validate_dataset_token from fields.base import ResponseModel from graphon.model_runtime.utils.encoders import jsonable_encoder from services.entities.model_provider_entities import ProviderWithModelsResponse -from services.model_provider_service import ModelProviderService +from services.model_provider.service import ModelProviderService MODEL_TYPE_PARAM = { "description": "Type of model to retrieve.", diff --git a/api/core/app/app_config/easy_ui_based_app/dataset/manager.py b/api/core/app/app_config/easy_ui_based_app/dataset/manager.py index 0cba996ad6f..802952b3d16 100644 --- a/api/core/app/app_config/easy_ui_based_app/dataset/manager.py +++ b/api/core/app/app_config/easy_ui_based_app/dataset/manager.py @@ -12,7 +12,7 @@ from core.app.app_config.entities import ( from core.entities.agent_entities import PlanningStrategy from core.rag.data_post_processor.data_post_processor import RerankingModelDict, WeightsDict from models.model import AppMode, AppModelConfigDict -from services.dataset_service import DatasetService +from services.knowledge.dataset_service import DatasetService class DatasetConfigManager: diff --git a/api/core/app/apps/pipeline/pipeline_generator.py b/api/core/app/apps/pipeline/pipeline_generator.py index c3e1be5a95b..17f84ff313d 100644 --- a/api/core/app/apps/pipeline/pipeline_generator.py +++ b/api/core/app/apps/pipeline/pipeline_generator.py @@ -22,7 +22,7 @@ from core.app.apps.draft_variable_saver import DraftVariableSaverFactory from core.app.apps.exc import GenerateTaskStoppedError from core.app.apps.pipeline.pipeline_config_manager import PipelineConfigManager from core.app.apps.pipeline.pipeline_queue_manager import PipelineQueueManager -from core.app.apps.pipeline.pipeline_runner import PipelineRunner +from core.app.apps.pipeline.pipeline_runner import PipelineDocumentStore, PipelineRunner from core.app.apps.workflow.generate_response_converter import WorkflowAppGenerateResponseConverter from core.app.apps.workflow.generate_task_pipeline import WorkflowAppGenerateTaskPipeline from core.app.entities.app_invoke_entities import InvokeFrom, RagPipelineGenerateEntity @@ -52,7 +52,7 @@ from models import Account, EndUser, Workflow, WorkflowNodeExecutionTriggeredFro from models.dataset import Document, DocumentPipelineExecutionLog, Pipeline from models.enums import WorkflowRunTriggeredFrom from models.model import AppMode -from services.datasource_provider_service import DatasourceProviderService +from services.data_source.provider_service import DatasourceProviderService from services.rag_pipeline.rag_pipeline_task_proxy import RagPipelineTaskProxy from services.workflow_draft_variable_service import DraftVarLoader, WorkflowDraftVariableService @@ -60,6 +60,10 @@ logger = logging.getLogger(__name__) class PipelineGenerator(BaseAppGenerator): + def __init__(self, *, documents: PipelineDocumentStore, datasource_providers: DatasourceProviderService) -> None: + self._documents = documents + self._datasource_providers = datasource_providers + @overload def generate( self, @@ -124,7 +128,7 @@ class PipelineGenerator(BaseAppGenerator): ) -> Mapping[str, Any] | Generator[Mapping | str, None, None] | None: # Add null check for dataset - dataset = pipeline.retrieve_dataset(session) + dataset = self._documents.get_pipeline_dataset(pipeline, session=session) if not dataset: raise ValueError("Pipeline dataset is required") inputs: Mapping[str, Any] = args["inputs"] @@ -140,8 +144,8 @@ class PipelineGenerator(BaseAppGenerator): ) documents: list[Document] = [] if invoke_from == InvokeFrom.PUBLISHED_PIPELINE and not is_retry and not args.get("original_document_id"): - from services.dataset_service import DocumentService from services.feature_service import FeatureService + from services.knowledge.dataset_service import DocumentService features = FeatureService.get_features(pipeline.tenant_id) DocumentService.check_document_creation_limits(len(datasource_info_list), features) @@ -407,7 +411,7 @@ class PipelineGenerator(BaseAppGenerator): pipeline=pipeline, workflow=workflow, start_node_id=args.get("start_node_id", "shared") ) - dataset = pipeline.retrieve_dataset(session) + dataset = self._documents.get_pipeline_dataset(pipeline, session=session) if not dataset: raise ValueError("Pipeline dataset is required") @@ -504,7 +508,7 @@ class PipelineGenerator(BaseAppGenerator): if args.get("inputs") is None: raise ValueError("inputs is required") - dataset = pipeline.retrieve_dataset(session) + dataset = self._documents.get_pipeline_dataset(pipeline, session=session) if not dataset: raise ValueError("Pipeline dataset is required") @@ -635,6 +639,7 @@ class PipelineGenerator(BaseAppGenerator): system_user_id=system_user_id, workflow_execution_repository=workflow_execution_repository, workflow_node_execution_repository=workflow_node_execution_repository, + documents=self._documents, ) runner.run() @@ -779,8 +784,7 @@ class PipelineGenerator(BaseAppGenerator): tenant_id=pipeline.tenant_id, datasource_type=DatasourceProviderType(datasource_type), ) - datasource_provider_service = DatasourceProviderService() - credentials = datasource_provider_service.get_datasource_credentials( + credentials = self._datasource_providers.get_datasource_credentials( tenant_id=pipeline.tenant_id, provider=datasource_node_data.get("provider_name"), plugin_id=datasource_node_data.get("plugin_id"), diff --git a/api/core/app/apps/pipeline/pipeline_runner.py b/api/core/app/apps/pipeline/pipeline_runner.py index 651b670fcce..5a3f6142201 100644 --- a/api/core/app/apps/pipeline/pipeline_runner.py +++ b/api/core/app/apps/pipeline/pipeline_runner.py @@ -1,6 +1,6 @@ import logging import time -from typing import cast +from typing import Protocol, cast from sqlalchemy import select from sqlalchemy.orm import Session @@ -28,14 +28,23 @@ from graphon.graph_events import GraphEngineEvent, GraphRunFailedEvent from graphon.runtime import GraphRuntimeState, VariablePool from graphon.variable_loader import VariableLoader from graphon.variables.variables import RAGPipelineVariable, RAGPipelineVariableInput -from models.dataset import Pipeline +from models.dataset import Dataset, Pipeline from models.model import EndUser from models.workflow import Workflow -from services.dataset_ref_service import DatasetRefService, DocumentRef logger = logging.getLogger(__name__) +class PipelineDocumentStore(Protocol): + """Knowledge state required by a pipeline run, with explicit session ownership.""" + + def get_pipeline_dataset(self, pipeline: Pipeline, *, session: Session) -> Dataset | None: ... + + def exists(self, *, workspace_id: str, dataset_id: str, document_id: str) -> bool: ... + + def mark_failed(self, *, workspace_id: str, dataset_id: str, document_id: str, error: str) -> None: ... + + class PipelineRunner(WorkflowBasedAppRunner): """ Pipeline Application Runner @@ -50,6 +59,7 @@ class PipelineRunner(WorkflowBasedAppRunner): system_user_id: str, workflow_execution_repository: WorkflowExecutionRepository, workflow_node_execution_repository: WorkflowNodeExecutionRepository, + documents: PipelineDocumentStore, workflow_thread_pool_id: str | None = None, ) -> None: """ @@ -68,6 +78,7 @@ class PipelineRunner(WorkflowBasedAppRunner): self._sys_user_id = system_user_id self._workflow_execution_repository = workflow_execution_repository self._workflow_node_execution_repository = workflow_node_execution_repository + self._documents = documents def _get_app_id(self) -> str: return self.application_generate_entity.app_config.app_id @@ -98,7 +109,7 @@ class PipelineRunner(WorkflowBasedAppRunner): if not pipeline or pipeline.tenant_id != app_config.tenant_id: raise ValueError("Pipeline not found") - dataset = pipeline.retrieve_dataset(session) + dataset = self._documents.get_pipeline_dataset(pipeline, session=session) if ( not dataset or dataset.tenant_id != pipeline.tenant_id @@ -108,23 +119,20 @@ class PipelineRunner(WorkflowBasedAppRunner): document_id = self.application_generate_entity.document_id original_document_id = self.application_generate_entity.original_document_id - dataset_ref = DatasetRefService.create_dataset_ref(dataset) - document_ref = ( - DatasetRefService.create_document_ref_from_id( - dataset_ref, - document_id, - ) - if document_id - else None - ) - if document_ref and DatasetRefService.get_document_by_ref(document_ref, session=session) is None: + dataset_workspace_id = dataset.tenant_id + dataset_id = dataset.id + if document_id and not self._documents.exists( + workspace_id=dataset_workspace_id, + dataset_id=dataset_id, + document_id=document_id, + ): raise ValueError("Pipeline document not found") if original_document_id and original_document_id != document_id: - original_document_ref = DatasetRefService.create_document_ref_from_id( - dataset_ref, - original_document_id, - ) - if DatasetRefService.get_document_by_ref(original_document_ref, session=session) is None: + if not self._documents.exists( + workspace_id=dataset_workspace_id, + dataset_id=dataset_id, + document_id=original_document_id, + ): raise ValueError("Pipeline original document not found") workflow = self.get_workflow(session=session, pipeline=pipeline, workflow_id=app_config.workflow_id) @@ -237,7 +245,12 @@ class PipelineRunner(WorkflowBasedAppRunner): generator = workflow_entry.run() for event in generator: - self._update_document_status(event, document_ref) + self._update_document_status( + event, + workspace_id=dataset_workspace_id, + dataset_id=dataset_id, + document_id=document_id, + ) self._handle_event(workflow_entry, event) def get_workflow(self, session: Session, pipeline: Pipeline, workflow_id: str) -> Workflow | None: @@ -325,14 +338,21 @@ class PipelineRunner(WorkflowBasedAppRunner): return graph - def _update_document_status(self, event: GraphEngineEvent, document_ref: DocumentRef | None) -> None: + def _update_document_status( + self, + event: GraphEngineEvent, + *, + workspace_id: str, + dataset_id: str, + document_id: str | None, + ) -> None: """Set an owner-bound document to error after a failed graph run, if it exists.""" - if not isinstance(event, GraphRunFailedEvent) or document_ref is None: + if not isinstance(event, GraphRunFailedEvent) or document_id is None: return - with create_session() as session, session.begin(): - document = DatasetRefService.get_document_by_ref(document_ref, session=session) - if document: - document.indexing_status = "error" - document.error = event.error or "Unknown error" - session.add(document) + self._documents.mark_failed( + workspace_id=workspace_id, + dataset_id=dataset_id, + document_id=document_id, + error=event.error or "Unknown error", + ) diff --git a/api/core/app/features/annotation_reply/annotation_reply.py b/api/core/app/features/annotation_reply/annotation_reply.py index 049f7ac2d5a..454a2394942 100644 --- a/api/core/app/features/annotation_reply/annotation_reply.py +++ b/api/core/app/features/annotation_reply/annotation_reply.py @@ -10,7 +10,7 @@ from models.dataset import Dataset from models.enums import CollectionBindingType, ConversationFromSource from models.model import AnnotationReplyEnabledConfig, App, Message, MessageAnnotation, load_annotation_reply_config from services.annotation_service import AppAnnotationService -from services.dataset_service import DatasetCollectionBindingService +from services.knowledge.dataset_service import DatasetCollectionBindingService logger = logging.getLogger(__name__) diff --git a/api/core/datasource/datasource_manager.py b/api/core/datasource/datasource_manager.py index b96d98d3f9c..cdf563b7632 100644 --- a/api/core/datasource/datasource_manager.py +++ b/api/core/datasource/datasource_manager.py @@ -1,5 +1,5 @@ import logging -from collections.abc import Generator +from collections.abc import Generator, Mapping from threading import Lock from typing import Any, cast @@ -33,7 +33,6 @@ from graphon.file import File, FileTransferMethod, FileType, get_file_type_by_mi from graphon.node_events import NodeRunResult, StreamChunkEvent, StreamCompletedEvent from models.model import UploadFile from models.tools import ToolFile -from services.datasource_provider_service import DatasourceProviderService logger = logging.getLogger(__name__) _file_access_controller = DatabaseFileAccessController() @@ -147,9 +146,7 @@ class DatasourceManager: datasource_type: str, provider_id: str, tenant_id: str, - provider: str, - plugin_id: str, - credential_id: str, + credentials: Mapping[str, Any], datasource_param: DatasourceParameter | None = None, online_drive_request: OnlineDriveDownloadFileParam | None = None, ) -> Generator[DatasourceMessage, None, Any]: @@ -166,18 +163,10 @@ class DatasourceManager: datasource_type=ds_type, ) - dsp_service = DatasourceProviderService() - credentials = dsp_service.get_datasource_credentials( - tenant_id=tenant_id, - provider=provider, - plugin_id=plugin_id, - credential_id=credential_id, - ) - if ds_type == DatasourceProviderType.ONLINE_DOCUMENT: doc_runtime = cast(OnlineDocumentDatasourcePlugin, runtime) if credentials: - doc_runtime.runtime.credentials = credentials + doc_runtime.runtime.credentials = dict(credentials) if datasource_param is None: raise ValueError("datasource_param is required for ONLINE_DOCUMENT streaming") inner_gen: Generator[DatasourceMessage, None, None] = doc_runtime.get_online_document_page_content( @@ -192,7 +181,7 @@ class DatasourceManager: elif ds_type == DatasourceProviderType.ONLINE_DRIVE: drive_runtime = cast(OnlineDriveDatasourcePlugin, runtime) if credentials: - drive_runtime.runtime.credentials = credentials + drive_runtime.runtime.credentials = dict(credentials) if online_drive_request is None: raise ValueError("online_drive_request is required for ONLINE_DRIVE streaming") inner_gen = drive_runtime.online_drive_download_file( @@ -221,9 +210,7 @@ class DatasourceManager: datasource_type: str, provider_id: str, tenant_id: str, - provider: str, - plugin_id: str, - credential_id: str, + credentials: Mapping[str, Any], parameters_for_log: dict[str, Any], datasource_info: dict[str, Any], variable_pool: Any, @@ -238,9 +225,7 @@ class DatasourceManager: datasource_type=datasource_type, provider_id=provider_id, tenant_id=tenant_id, - provider=provider, - plugin_id=plugin_id, - credential_id=credential_id, + credentials=credentials, datasource_param=datasource_param, online_drive_request=online_drive_request, ) diff --git a/api/core/indexing_runner.py b/api/core/indexing_runner.py deleted file mode 100644 index 2163a5c7fbc..00000000000 --- a/api/core/indexing_runner.py +++ /dev/null @@ -1,938 +0,0 @@ -import concurrent.futures -import json -import logging -import re -import threading -import time -import uuid -from collections.abc import Mapping -from typing import Any - -from flask import Flask, current_app -from sqlalchemy import delete, func, select, update -from sqlalchemy.orm import Session -from sqlalchemy.orm.exc import ObjectDeletedError - -from configs import dify_config -from core.credit_usage import CreditUsageCreatedBy -from core.db.session_factory import session_factory -from core.entities.knowledge_entities import IndexingEstimate, PreviewDetail, QAPreviewDetail -from core.errors.error import ProviderTokenNotInitError -from core.model_context import with_credit_usage_created_by -from core.model_manager import ModelInstance, ModelManager -from core.rag.cleaner.clean_processor import CleanProcessor -from core.rag.datasource.keyword.keyword_factory import Keyword -from core.rag.docstore.dataset_docstore import DatasetDocumentStore -from core.rag.embedding.token_counter import calculate_segment_token_counts -from core.rag.extractor.entity.datasource_type import DatasourceType -from core.rag.extractor.entity.extract_setting import ExtractSetting, NotionInfo, WebsiteInfo -from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType -from core.rag.index_processor.index_processor_base import BaseIndexProcessor -from core.rag.index_processor.index_processor_factory import IndexProcessorFactory -from core.rag.models.document import ChildDocument, Document -from core.rag.splitter.fixed_text_splitter import ( - EnhanceRecursiveCharacterTextSplitter, - FixedRecursiveCharacterTextSplitter, -) -from core.rag.splitter.text_splitter import TextSplitter -from core.tools.utils.web_reader_tool import get_image_upload_file_ids -from enums import DeploymentEdition -from extensions.ext_redis import redis_client -from extensions.ext_storage import storage -from extensions.otel import propagate_context -from graphon.model_runtime.entities.model_entities import ModelType -from libs import helper -from libs.datetime_utils import naive_utc_now -from models import Account -from models.dataset import AutomaticRulesConfig, ChildChunk, Dataset, DatasetProcessRule, DocumentSegment -from models.dataset import Document as DatasetDocument -from models.enums import DataSourceType, IndexingStatus, ProcessRuleMode, SegmentStatus -from models.model import UploadFile -from services.vector_space_admission_service import VectorSpaceAdmissionService - -logger = logging.getLogger(__name__) - - -class IndexingRunner: - def __init__( - self, - *, - enforce_vector_space_admission: bool = False, - ): - self.storage = storage - self.enforce_vector_space_admission = enforce_vector_space_admission - - @staticmethod - def _get_model_manager(tenant_id: str) -> ModelManager: - return ModelManager.for_tenant(tenant_id=tenant_id) - - def _handle_indexing_error(self, document_id: str, error: Exception, session: Session) -> None: - """Handle indexing errors by updating document status.""" - logger.exception("consume document failed") - document = session.get(DatasetDocument, document_id) - if document: - document.indexing_status = IndexingStatus.ERROR - error_message = getattr(error, "description", str(error)) - document.error = str(error_message) - document.stopped_at = naive_utc_now() - session.flush() - - @with_credit_usage_created_by(CreditUsageCreatedBy.KNOWLEDGE_INDEXING) - def run(self, dataset_documents: list[DatasetDocument], session: Session): - """Run indexing with commits before slow transforms and parallel index workers. - - The phase commits keep document locks short and make newly created segments - visible to the worker sessions used for keyword and vector indexing. - """ - vector_space_admission = VectorSpaceAdmissionService() - for dataset_document in dataset_documents: - document_id = dataset_document.id - try: - # Re-query the document to ensure it's bound to the current session - requeried_document = session.get(DatasetDocument, document_id) - if not requeried_document: - logger.warning("Document not found, skipping document id: %s", document_id) - continue - - # get dataset - dataset = session.get(Dataset, requeried_document.dataset_id) - - if not dataset: - raise ValueError("no dataset found") - # get the process rule - stmt = select(DatasetProcessRule).where( - DatasetProcessRule.id == requeried_document.dataset_process_rule_id - ) - processing_rule = session.scalar(stmt) - if not processing_rule: - raise ValueError("no process rule found") - index_type = requeried_document.doc_form - index_processor = IndexProcessorFactory(index_type).init_index_processor() - # extract - text_docs = self._extract(index_processor, requeried_document, processing_rule.to_dict(), session) - session.commit() - - # transform - current_user = session.get(Account, requeried_document.created_by) - if not current_user: - raise ValueError("no current user found") - current_user.set_tenant_id_with_session(dataset.tenant_id, session=session) - documents = self._transform( - index_processor, - dataset, - text_docs, - requeried_document.doc_language, - processing_rule.to_dict(), - current_user=current_user, - session=session, - ) - if self.enforce_vector_space_admission: - vector_space_admission.ensure_document_can_be_indexed( - dataset=dataset, - document_id=requeried_document.id, - doc_form=requeried_document.doc_form, - documents=documents, - include_summaries=bool(requeried_document.need_summary), - session=session, - ) - token_counts = calculate_segment_token_counts(dataset=dataset, documents=documents) - total_tokens = sum(token_counts) - # save segment - self._load_segments( - session=session, - dataset=dataset, - dataset_document=requeried_document, - documents=documents, - token_counts=token_counts, - ) - session.commit() - - # load - self._load( - session=session, - dataset=dataset, - dataset_document=requeried_document, - documents=documents, - total_tokens=total_tokens, - ) - except DocumentIsPausedError: - raise DocumentIsPausedError(f"Document paused, document id: {document_id}") - except ProviderTokenNotInitError as e: - self._handle_indexing_error(document_id, e, session) - except ObjectDeletedError: - logger.warning("Document deleted, document id: %s", document_id) - except Exception as e: - self._handle_indexing_error(document_id, e, session) - - @with_credit_usage_created_by(CreditUsageCreatedBy.KNOWLEDGE_INDEXING) - def run_in_splitting_status(self, dataset_document: DatasetDocument, session: Session): - """Run the indexing process when the index_status is splitting.""" - document_id = dataset_document.id - try: - # Re-query the document to ensure it's bound to the current session - requeried_document = session.get(DatasetDocument, document_id) - if not requeried_document: - logger.warning("Document not found: %s", document_id) - return - - # get dataset - dataset = session.get(Dataset, requeried_document.dataset_id) - - if not dataset: - raise ValueError("no dataset found") - - # get exist document_segment list and delete - document_segments = session.scalars( - select(DocumentSegment).where( - DocumentSegment.dataset_id == dataset.id, - DocumentSegment.document_id == requeried_document.id, - ) - ).all() - - for document_segment in document_segments: - session.delete(document_segment) - if requeried_document.doc_form == IndexStructureType.PARENT_CHILD_INDEX: - # delete child chunks - session.execute(delete(ChildChunk).where(ChildChunk.segment_id == document_segment.id)) - session.commit() - # get the process rule - stmt = select(DatasetProcessRule).where(DatasetProcessRule.id == requeried_document.dataset_process_rule_id) - processing_rule = session.scalar(stmt) - if not processing_rule: - raise ValueError("no process rule found") - - index_type = requeried_document.doc_form - index_processor = IndexProcessorFactory(index_type).init_index_processor() - # extract - text_docs = self._extract(index_processor, requeried_document, processing_rule.to_dict(), session) - session.commit() - - # transform - current_user = session.get(Account, requeried_document.created_by) - if not current_user: - raise ValueError("no current user found") - current_user.set_tenant_id_with_session(dataset.tenant_id, session=session) - documents = self._transform( - index_processor, - dataset, - text_docs, - requeried_document.doc_language, - processing_rule.to_dict(), - current_user=current_user, - session=session, - ) - token_counts = calculate_segment_token_counts(dataset=dataset, documents=documents) - total_tokens = sum(token_counts) - # save segment - self._load_segments( - session=session, - dataset=dataset, - dataset_document=requeried_document, - documents=documents, - token_counts=token_counts, - ) - session.commit() - - # load - self._load( - session=session, - dataset=dataset, - dataset_document=requeried_document, - documents=documents, - total_tokens=total_tokens, - ) - except DocumentIsPausedError: - raise DocumentIsPausedError(f"Document paused, document id: {document_id}") - except ProviderTokenNotInitError as e: - self._handle_indexing_error(document_id, e, session) - except Exception as e: - self._handle_indexing_error(document_id, e, session) - - @with_credit_usage_created_by(CreditUsageCreatedBy.KNOWLEDGE_INDEXING) - def run_in_indexing_status(self, dataset_document: DatasetDocument, session: Session): - """Run the indexing process when the index_status is indexing.""" - document_id = dataset_document.id - try: - # Re-query the document to ensure it's bound to the current session - requeried_document = session.get(DatasetDocument, document_id) - if not requeried_document: - logger.warning("Document not found: %s", document_id) - return - - # get dataset - dataset = session.get(Dataset, requeried_document.dataset_id) - - if not dataset: - raise ValueError("no dataset found") - - # get existing document segments - document_segments = session.scalars( - select(DocumentSegment).where( - DocumentSegment.dataset_id == dataset.id, - DocumentSegment.document_id == requeried_document.id, - ) - ).all() - - documents = [] - if document_segments: - for document_segment in document_segments: - # transform segment to node - if document_segment.status != SegmentStatus.COMPLETED: - document = Document( - page_content=document_segment.content, - metadata={ - "doc_id": document_segment.index_node_id, - "doc_hash": document_segment.index_node_hash, - "document_id": document_segment.document_id, - "dataset_id": document_segment.dataset_id, - }, - ) - if requeried_document.doc_form == IndexStructureType.PARENT_CHILD_INDEX: - child_chunks = document_segment.get_child_chunks(session=session) - if child_chunks: - child_documents = [] - for child_chunk in child_chunks: - child_document = ChildDocument( - page_content=child_chunk.content, - metadata={ - "doc_id": child_chunk.index_node_id, - "doc_hash": child_chunk.index_node_hash, - "document_id": document_segment.document_id, - "dataset_id": document_segment.dataset_id, - }, - ) - child_documents.append(child_document) - document.children = child_documents - documents.append(document) - # Preserve the full document total even when only incomplete segments are re-indexed. - total_tokens = sum(document_segment.tokens for document_segment in document_segments) - # build index - self._load( - session=session, - dataset=dataset, - dataset_document=requeried_document, - documents=documents, - total_tokens=total_tokens, - ) - except DocumentIsPausedError: - raise DocumentIsPausedError(f"Document paused, document id: {document_id}") - except ProviderTokenNotInitError as e: - self._handle_indexing_error(document_id, e, session) - except Exception as e: - self._handle_indexing_error(document_id, e, session) - - @with_credit_usage_created_by(CreditUsageCreatedBy.KNOWLEDGE_INDEXING) - def indexing_estimate( - self, - tenant_id: str, - extract_settings: list[ExtractSetting], - tmp_processing_rule: Mapping[str, Any], - doc_form: str | None = None, - doc_language: str = "English", - dataset_id: str | None = None, - indexing_technique: str = IndexTechniqueType.ECONOMY, - *, - session: Session, - ) -> IndexingEstimate: - """ - Estimate the indexing for the document. - """ - # check document limit - if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD: - count = len(extract_settings) - batch_upload_limit = dify_config.BATCH_UPLOAD_LIMIT - if count > batch_upload_limit: - raise ValueError(f"You have reached the batch upload limit of {batch_upload_limit}.") - - embedding_model_instance = None - if dataset_id: - dataset = session.get(Dataset, dataset_id) - if not dataset: - raise ValueError("Dataset not found.") - if IndexTechniqueType.HIGH_QUALITY in {dataset.indexing_technique, indexing_technique}: - if dataset.embedding_model_provider: - embedding_model_instance = self._get_model_manager(tenant_id).get_model_instance( - tenant_id=tenant_id, - provider=dataset.embedding_model_provider, - model_type=ModelType.TEXT_EMBEDDING, - model=dataset.embedding_model, - ) - else: - embedding_model_instance = self._get_model_manager(tenant_id).get_default_model_instance( - tenant_id=tenant_id, - model_type=ModelType.TEXT_EMBEDDING, - ) - else: - if indexing_technique == IndexTechniqueType.HIGH_QUALITY: - embedding_model_instance = self._get_model_manager(tenant_id).get_default_model_instance( - tenant_id=tenant_id, - model_type=ModelType.TEXT_EMBEDDING, - ) - # keep separate, avoid union-list ambiguity - preview_texts: list[PreviewDetail] = [] - qa_preview_texts: list[QAPreviewDetail] = [] - - total_segments = 0 - # doc_form represents the segmentation method (general, parent-child, QA) - index_type = doc_form - index_processor = IndexProcessorFactory(index_type).init_index_processor() - # one extract_setting is one source document - for extract_setting in extract_settings: - # extract - processing_rule = { - "mode": tmp_processing_rule["mode"], - "rules": tmp_processing_rule.get("rules"), - } - # Extract document content - text_docs = index_processor.extract( - extract_setting, process_rule_mode=tmp_processing_rule["mode"], session=session - ) - # Cleaning and segmentation - documents = index_processor.transform( - text_docs, - current_user=None, - embedding_model_instance=embedding_model_instance, - process_rule=processing_rule, - tenant_id=tenant_id, - doc_language=doc_language, - preview=True, - session=session, - ) - total_segments += len(documents) - for document in documents: - preview_count = len(qa_preview_texts) if doc_form == "qa_model" else len(preview_texts) - if preview_count < 10: - if doc_form and doc_form == "qa_model": - qa_detail = QAPreviewDetail( - question=document.page_content, answer=document.metadata.get("answer") or "" - ) - qa_preview_texts.append(qa_detail) - else: - preview_detail = PreviewDetail(content=document.page_content) - if document.children: - preview_detail.child_chunks = [child.page_content for child in document.children] - preview_texts.append(preview_detail) - - # delete image files and related db records - image_upload_file_ids = get_image_upload_file_ids(document.page_content) - for upload_file_id in image_upload_file_ids: - stmt = select(UploadFile).where(UploadFile.id == upload_file_id) - image_file = session.scalar(stmt) - if image_file is None: - continue - try: - storage.delete(image_file.key) - except Exception: - logger.exception( - "Delete image_files failed while indexing_estimate, \ - image_upload_file_is: %s", - upload_file_id, - ) - session.delete(image_file) - - # Persist preview cleanup and release the caller transaction before - # summary workers query through their own sessions. - session.commit() - - if doc_form and doc_form == "qa_model": - return IndexingEstimate(total_segments=total_segments * 20, qa_preview=qa_preview_texts, preview=[]) - - # Generate summary preview - summary_index_setting = tmp_processing_rule.get("summary_index_setting") - if summary_index_setting and summary_index_setting.get("enable") and preview_texts: - preview_texts = index_processor.generate_summary_preview( - tenant_id, preview_texts, summary_index_setting, doc_language, session=session - ) - - return IndexingEstimate(total_segments=total_segments, preview=preview_texts) - - def _extract( - self, - index_processor: BaseIndexProcessor, - dataset_document: DatasetDocument, - process_rule: Mapping[str, Any], - session: Session, - ) -> list[Document]: - data_source_info = dataset_document.data_source_info_dict - text_docs = [] - match dataset_document.data_source_type: - case DataSourceType.UPLOAD_FILE: - if not data_source_info or "upload_file_id" not in data_source_info: - raise ValueError("no upload file found") - stmt = select(UploadFile).where(UploadFile.id == data_source_info["upload_file_id"]) - file_detail = session.scalars(stmt).one_or_none() - - if file_detail: - extract_setting = ExtractSetting( - datasource_type=DatasourceType.FILE, - upload_file=file_detail, - document_model=dataset_document.doc_form, - ) - text_docs = index_processor.extract( - extract_setting, process_rule_mode=process_rule["mode"], session=session - ) - case DataSourceType.NOTION_IMPORT: - if ( - not data_source_info - or "notion_workspace_id" not in data_source_info - or "notion_page_id" not in data_source_info - ): - raise ValueError("no notion import info found") - extract_setting = ExtractSetting( - datasource_type=DatasourceType.NOTION, - notion_info=NotionInfo.model_validate( - { - "credential_id": data_source_info.get("credential_id"), - "notion_workspace_id": data_source_info["notion_workspace_id"], - "notion_obj_id": data_source_info["notion_page_id"], - "notion_page_type": data_source_info["type"], - "document": dataset_document, - "tenant_id": dataset_document.tenant_id, - } - ), - document_model=dataset_document.doc_form, - ) - text_docs = index_processor.extract( - extract_setting, process_rule_mode=process_rule["mode"], session=session - ) - case DataSourceType.WEBSITE_CRAWL: - if ( - not data_source_info - or "provider" not in data_source_info - or "url" not in data_source_info - or "job_id" not in data_source_info - ): - raise ValueError("no website import info found") - extract_setting = ExtractSetting( - datasource_type=DatasourceType.WEBSITE, - website_info=WebsiteInfo.model_validate( - { - "provider": data_source_info["provider"], - "job_id": data_source_info["job_id"], - "tenant_id": dataset_document.tenant_id, - "url": data_source_info["url"], - "mode": data_source_info["mode"], - "only_main_content": data_source_info["only_main_content"], - } - ), - document_model=dataset_document.doc_form, - ) - text_docs = index_processor.extract( - extract_setting, process_rule_mode=process_rule["mode"], session=session - ) - case _: - return [] - # update document status to splitting - self._update_document_index_status( - session=session, - document_id=dataset_document.id, - after_indexing_status=IndexingStatus.SPLITTING, - extra_update_params={ - DatasetDocument.parsing_completed_at: naive_utc_now(), - }, - ) - - # replace doc id to document model id - for text_doc in text_docs: - if text_doc.metadata is not None: - text_doc.metadata["document_id"] = dataset_document.id - text_doc.metadata["dataset_id"] = dataset_document.dataset_id - - return text_docs - - @staticmethod - def filter_string(text: str): - text = re.sub(r"<\|", "<", text) - text = re.sub(r"\|>", ">", text) - text = re.sub(r"[\x00-\x08\x0B\x0C\x0E-\x1F\x7F]", "", text) - # Unicode U+FFFE - text = re.sub("\ufffe", "", text) - return text - - @staticmethod - def _get_splitter( - processing_rule_mode: str, - max_tokens: int, - chunk_overlap: int, - separator: str, - embedding_model_instance: ModelInstance | None, - ) -> TextSplitter: - """ - Get the NodeParser object according to the processing rule. - """ - character_splitter: TextSplitter - if processing_rule_mode in ["custom", "hierarchical"]: - # The user-defined segmentation rule - max_segmentation_tokens_length = dify_config.INDEXING_MAX_SEGMENTATION_TOKENS_LENGTH - if max_tokens < 50 or max_tokens > max_segmentation_tokens_length: - raise ValueError(f"Custom segment length should be between 50 and {max_segmentation_tokens_length}.") - - if separator: - separator = separator.replace("\\n", "\n") - - character_splitter = FixedRecursiveCharacterTextSplitter.from_encoder( - chunk_size=max_tokens, - chunk_overlap=chunk_overlap, - fixed_separator=separator, - separators=["\n\n", "。", ". ", " ", ""], - embedding_model_instance=embedding_model_instance, - ) - else: - # Automatic segmentation - automatic_rules: dict[str, Any] = dict(DatasetProcessRule.AUTOMATIC_RULES["segmentation"]) - character_splitter = EnhanceRecursiveCharacterTextSplitter.from_encoder( - chunk_size=automatic_rules["max_tokens"], - chunk_overlap=automatic_rules["chunk_overlap"], - separators=["\n\n", "。", ". ", " ", ""], - embedding_model_instance=embedding_model_instance, - ) - - return character_splitter - - def _split_to_documents_for_estimate( - self, text_docs: list[Document], splitter: TextSplitter, processing_rule: DatasetProcessRule - ) -> list[Document]: - """ - Split the text documents into nodes. - """ - all_documents: list[Document] = [] - for text_doc in text_docs: - # document clean - document_text = self._document_clean(text_doc.page_content, processing_rule) - text_doc.page_content = document_text - - # parse document to nodes - documents = splitter.split_documents([text_doc]) - - split_documents = [] - for document in documents: - if document.page_content is None or not document.page_content.strip(): - continue - if document.metadata is not None: - doc_id = str(uuid.uuid4()) - hash = helper.generate_text_hash(document.page_content) - document.metadata["doc_id"] = doc_id - document.metadata["doc_hash"] = hash - - split_documents.append(document) - - all_documents.extend(split_documents) - - return all_documents - - @staticmethod - def _document_clean(text: str, processing_rule: DatasetProcessRule) -> str: - """ - Clean the document text according to the processing rules. - """ - rules: AutomaticRulesConfig | dict[str, Any] - if processing_rule.mode == ProcessRuleMode.AUTOMATIC: - rules = DatasetProcessRule.AUTOMATIC_RULES - else: - rules = json.loads(processing_rule.rules) if processing_rule.rules else {} - document_text = CleanProcessor.clean(text, {"rules": rules}) - - return document_text - - @staticmethod - def format_split_text(text: str) -> list[QAPreviewDetail]: - regex = r"Q\d+:\s*(.*?)\s*A\d+:\s*([\s\S]*?)(?=Q\d+:|$)" - matches = re.findall(regex, text, re.UNICODE) - - return [QAPreviewDetail(question=q, answer=re.sub(r"\n\s*", "\n", a.strip())) for q, a in matches if q and a] - - def _load( - self, - session: Session, - dataset: Dataset, - dataset_document: DatasetDocument, - documents: list[Document], - total_tokens: int, - ) -> None: - """Build indexes and mark the document complete using the token total computed before hash sharding.""" - - # Build indexes using the existing hash-based worker groups. - indexing_start_at = time.perf_counter() - create_keyword_thread = None - if ( - dataset_document.doc_form != IndexStructureType.PARENT_CHILD_INDEX - and dataset.indexing_technique == IndexTechniqueType.ECONOMY - ): - # create keyword index - create_keyword_thread = threading.Thread( - target=propagate_context(self._process_keyword_index), - args=(current_app._get_current_object(), dataset.id, dataset_document.id, documents), # type: ignore - ) - create_keyword_thread.start() - - max_workers = 10 - if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY: - with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: - futures = [] - - # Distribute documents into multiple groups based on the hash values of page_content - # This is done to prevent multiple threads from processing the same document, - # Thereby avoiding potential database insertion deadlocks - document_groups: list[list[Document]] = [[] for _ in range(max_workers)] - for document in documents: - hash = helper.generate_text_hash(document.page_content) - group_index = int(hash, 16) % max_workers - document_groups[group_index].append(document) - for chunk_documents in document_groups: - if len(chunk_documents) == 0: - continue - futures.append( - executor.submit( - propagate_context(self._process_chunk), - current_app._get_current_object(), # type: ignore - dataset_document.doc_form, - chunk_documents, - dataset.id, - dataset_document.id, - ) - ) - - for future in futures: - future.result() - if ( - dataset_document.doc_form != IndexStructureType.PARENT_CHILD_INDEX - and dataset.indexing_technique == IndexTechniqueType.ECONOMY - and create_keyword_thread is not None - ): - create_keyword_thread.join() - indexing_end_at = time.perf_counter() - - # update document status to completed - self._update_document_index_status( - session=session, - document_id=dataset_document.id, - after_indexing_status=IndexingStatus.COMPLETED, - extra_update_params={ - DatasetDocument.tokens: total_tokens, - DatasetDocument.completed_at: naive_utc_now(), - DatasetDocument.indexing_latency: indexing_end_at - indexing_start_at, - DatasetDocument.error: None, - }, - ) - - @staticmethod - def _process_keyword_index(flask_app: Flask, dataset_id: str, document_id: str, documents: list[Document]): - with flask_app.app_context(): - with session_factory.create_session() as session: - dataset = session.get(Dataset, dataset_id) - if not dataset: - raise ValueError("no dataset found") - keyword = Keyword(dataset) - keyword.create(documents, session) - if dataset.indexing_technique != IndexTechniqueType.HIGH_QUALITY: - document_ids = [document.metadata["doc_id"] for document in documents] - session.execute( - update(DocumentSegment) - .where( - DocumentSegment.document_id == document_id, - DocumentSegment.dataset_id == dataset_id, - DocumentSegment.index_node_id.in_(document_ids), - DocumentSegment.status == SegmentStatus.INDEXING, - ) - .values( - status=SegmentStatus.COMPLETED, - enabled=True, - completed_at=naive_utc_now(), - ) - ) - session.commit() - - def _process_chunk( - self, - flask_app: Flask, - index_type: str, - chunk_documents: list[Document], - dataset_id: str, - dataset_document_id: str, - ) -> None: - with flask_app.app_context(): - with session_factory.create_session() as session: - dataset = session.get(Dataset, dataset_id) - if not dataset: - raise ValueError("no dataset found") - - dataset_document = session.get(DatasetDocument, dataset_document_id) - if not dataset_document: - raise ValueError("no document found") - - # check document is paused - self._check_document_paused_status(dataset_document.id) - - multimodal_documents = [] - for document in chunk_documents: - if document.attachments and dataset.is_multimodal: - multimodal_documents.extend(document.attachments) - - # load index - index_processor = IndexProcessorFactory(index_type).init_index_processor() - index_processor.load( - dataset, - chunk_documents, - multimodal_documents=multimodal_documents, - with_keywords=False, - session=session, - ) - - document_ids = [document.metadata["doc_id"] for document in chunk_documents] - session.execute( - update(DocumentSegment) - .where( - DocumentSegment.document_id == dataset_document.id, - DocumentSegment.dataset_id == dataset.id, - DocumentSegment.index_node_id.in_(document_ids), - DocumentSegment.status == SegmentStatus.INDEXING, - ) - .values( - status=SegmentStatus.COMPLETED, - enabled=True, - completed_at=naive_utc_now(), - ) - ) - - session.commit() - - @staticmethod - def _check_document_paused_status(document_id: str): - indexing_cache_key = f"document_{document_id}_is_paused" - result = redis_client.get(indexing_cache_key) - if result: - raise DocumentIsPausedError() - - @staticmethod - def _update_document_index_status( - document_id: str, - after_indexing_status: IndexingStatus, - extra_update_params: Mapping[Any, Any] | None = None, - *, - session: Session, - ): - """ - Update the document indexing status. - """ - count = ( - session.scalar( - select(func.count()) - .select_from(DatasetDocument) - .where(DatasetDocument.id == document_id, DatasetDocument.is_paused == True) - ) - or 0 - ) - if count > 0: - raise DocumentIsPausedError() - document = session.get(DatasetDocument, document_id) - if not document: - raise DocumentIsDeletedPausedError() - - update_params = {DatasetDocument.indexing_status: after_indexing_status} - - if extra_update_params: - update_params.update(extra_update_params) - session.execute(update(DatasetDocument).where(DatasetDocument.id == document_id).values(update_params)) # type: ignore - session.flush() - - @staticmethod - def _update_segments_by_document(dataset_document_id: str, update_params: Mapping[Any, Any], session: Session): - """ - Update the document segment by document id. - """ - session.execute( - update(DocumentSegment).where(DocumentSegment.document_id == dataset_document_id).values(update_params) - ) - session.flush() - - def _transform( - self, - index_processor: BaseIndexProcessor, - dataset: Dataset, - text_docs: list[Document], - doc_language: str, - process_rule: Mapping[str, Any], - current_user: Account | None = None, - *, - session: Session, - ) -> list[Document]: - # get embedding model instance - embedding_model_instance = None - if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY: - if dataset.embedding_model_provider: - embedding_model_instance = self._get_model_manager(dataset.tenant_id).get_model_instance( - tenant_id=dataset.tenant_id, - provider=dataset.embedding_model_provider, - model_type=ModelType.TEXT_EMBEDDING, - model=dataset.embedding_model, - ) - else: - embedding_model_instance = self._get_model_manager(dataset.tenant_id).get_default_model_instance( - tenant_id=dataset.tenant_id, - model_type=ModelType.TEXT_EMBEDDING, - ) - - documents = index_processor.transform( - text_docs, - current_user, - embedding_model_instance=embedding_model_instance, - process_rule=process_rule, - tenant_id=dataset.tenant_id, - doc_language=doc_language, - session=session, - ) - - return documents - - def _load_segments( - self, - session: Session, - dataset: Dataset, - dataset_document: DatasetDocument, - documents: list[Document], - token_counts: list[int], - ) -> None: - """Persist transformed documents and their precomputed token counts before indexing starts.""" - # save node to document segment - doc_store = DatasetDocumentStore( - dataset=dataset, user_id=dataset_document.created_by, document_id=dataset_document.id - ) - - # add document segments - doc_store.add_documents( - session=session, - docs=documents, - save_child=dataset_document.doc_form == IndexStructureType.PARENT_CHILD_INDEX, - token_counts=token_counts, - ) - - # update document status to indexing - cur_time = naive_utc_now() - self._update_document_index_status( - session=session, - document_id=dataset_document.id, - after_indexing_status=IndexingStatus.INDEXING, - extra_update_params={ - DatasetDocument.cleaning_completed_at: cur_time, - DatasetDocument.splitting_completed_at: cur_time, - DatasetDocument.word_count: sum(len(doc.page_content) for doc in documents), - }, - ) - - # update segment status to indexing - self._update_segments_by_document( - session=session, - dataset_document_id=dataset_document.id, - update_params={ - DocumentSegment.status: SegmentStatus.INDEXING, - DocumentSegment.indexing_at: naive_utc_now(), - }, - ) - - -class DocumentIsPausedError(Exception): - pass - - -class DocumentIsDeletedPausedError(Exception): - pass diff --git a/api/core/provider_manager.py b/api/core/provider_manager.py index 837beffc37c..b62c7747bb4 100644 --- a/api/core/provider_manager.py +++ b/api/core/provider_manager.py @@ -65,7 +65,6 @@ from services.feature_service import FeatureService if TYPE_CHECKING: from graphon.model_runtime.protocols.runtime import ModelRuntime - from models.account import Account logger = logging.getLogger(__name__) @@ -1162,50 +1161,6 @@ class ProviderManager: provider_names.append(str(model_provider_id)) return provider_names - @staticmethod - def get_provider_available_credentials( - tenant_id: str, - provider_name: str, - user: Account | None = None, - ) -> list[CredentialConfiguration]: - """ - Get provider all credentials, filtered by visibility. - - :param tenant_id: workspace id - :param provider_name: provider name - :param user: current user (id + admin flag drive the visibility filter) - :return: - """ - from models.credential_permission import CredentialType as CredPermType - from services.credential_permission_service import CredentialPermissionService - - with session_factory.create_session() as session: - stmt = ( - select(ProviderCredential) - .where( - ProviderCredential.tenant_id == tenant_id, - ProviderCredential.provider_name.in_(ProviderManager._get_provider_names(provider_name)), - ) - .order_by(ProviderCredential.created_at.desc()) - ) - - if user is not None: - stmt = CredentialPermissionService.apply_visibility_filter( - stmt, - model_id_column=ProviderCredential.id, - model_user_id_column=ProviderCredential.user_id, - model_visibility_column=ProviderCredential.visibility, - credential_type=CredPermType.PROVIDER_CREDENTIAL, - user=user, - ) - - available_credentials = session.scalars(stmt).all() - - return [ - CredentialConfiguration(credential_id=credential.id, credential_name=credential.credential_name) - for credential in available_credentials - ] - @staticmethod def get_provider_model_available_credentials( tenant_id: str, provider_name: str, model_name: str, model_type: str diff --git a/api/core/rag/datasource/keyword/jieba/jieba.py b/api/core/rag/datasource/keyword/jieba/jieba.py index 62e79d8c84d..eea88349388 100644 --- a/api/core/rag/datasource/keyword/jieba/jieba.py +++ b/api/core/rag/datasource/keyword/jieba/jieba.py @@ -1,7 +1,8 @@ from collections import defaultdict -from typing import Any, TypedDict, override +from collections.abc import Callable, Mapping, Sequence +from contextlib import AbstractContextManager +from typing import Any, override -import orjson from pydantic import BaseModel from sqlalchemy import select from sqlalchemy.orm import Session @@ -12,12 +13,15 @@ from core.rag.datasource.keyword.keyword_base import BaseKeyword from core.rag.models.document import Document from extensions.ext_redis import redis_client from extensions.ext_storage import storage -from models.dataset import Dataset, DatasetKeywordTable, DocumentSegment - - -class PreSegmentData(TypedDict): - segment: DocumentSegment - keywords: list[str] +from models.dataset import Dataset, DocumentSegment +from repositories.knowledge.dataset_read_repository import get_dataset_keyword_table +from repositories.knowledge.keyword_table_repository import ( + keyword_file_key, + load_dataset_keyword_table, + load_keyword_table, + persist_keyword_table, + save_keyword_table, +) class KeywordTableConfig(BaseModel): @@ -31,80 +35,114 @@ class Jieba(BaseKeyword): @override def create(self, texts: list[Document], session: Session, **kwargs: Any) -> BaseKeyword: - lock_name = f"keyword_indexing_lock_{self.dataset.id}" - with redis_client.lock(lock_name, timeout=600): - keyword_table_handler = JiebaKeywordTableHandler() - keyword_table = self._get_dataset_keyword_table(session=session) - keyword_number = self.dataset.keyword_number or self._config.max_keywords_per_chunk - - for text in texts: - keywords = keyword_table_handler.extract_keywords(text.page_content, keyword_number) - if text.metadata is not None: - self._update_segment_keywords(self.dataset.id, text.metadata["doc_id"], list(keywords), session) - keyword_table = self._add_text_to_keyword_table( - keyword_table or {}, text.metadata["doc_id"], list(keywords) - ) - - self._save_dataset_keyword_table(keyword_table, session) - - return self + self.add_texts(texts, session) + return self @override def add_texts(self, texts: list[Document], session: Session, **kwargs: Any): - lock_name = f"keyword_indexing_lock_{self.dataset.id}" - with redis_client.lock(lock_name, timeout=600): - keyword_table_handler = JiebaKeywordTableHandler() + self.update_texts( + texts, + read=lambda: self._read_keyword_record(session), + write=lambda storage_type, data, keywords: persist_keyword_table( + session, + tenant_id=self.dataset.tenant_id, + dataset_id=self.dataset.id, + storage_type=storage_type, + data=data, + keywords=keywords, + ), + lock=redis_client.lock(f"keyword_indexing_lock_{self.dataset.id}", timeout=600), + keywords_list=kwargs.get("keywords_list"), + ) - keyword_table = self._get_dataset_keyword_table(session=session) - keywords_list = kwargs.get("keywords_list") + def update_texts( + self, + texts: Sequence[Document], + *, + read: Callable[[], tuple[str, str | None]], + write: Callable[[str, str, Mapping[str, Sequence[str]]], None], + lock: AbstractContextManager[object], + keywords_list: Sequence[Sequence[str] | None] | None = None, + replace_existing: bool = False, + deleted_ids: Sequence[str] = (), + ) -> None: + """Mutate the shared index with explicit persistence and one dataset lock. + + Repository callbacks may close their sessions before returning, keeping + extraction and object-storage I/O outside database transactions. + """ + with lock: + storage_type, data = read() + payload = load_keyword_table( + storage=storage, + tenant_id=self.dataset.tenant_id, + dataset_id=self.dataset.id, + storage_type=storage_type, + data=data, + ) + table = dict(payload["__data__"]["table"] or {}) if payload else {} + removed = list(deleted_ids) + if replace_existing: + removed.extend(text.metadata["doc_id"] for text in texts) + if removed: + table = self._delete_ids_from_keyword_table(table, removed) + handler: JiebaKeywordTableHandler | None = None keyword_number = self.dataset.keyword_number or self._config.max_keywords_per_chunk - for i in range(len(texts)): - text = texts[i] - if keywords_list: - keywords = keywords_list[i] - if not keywords: - keywords = keyword_table_handler.extract_keywords(text.page_content, keyword_number) - else: - keywords = keyword_table_handler.extract_keywords(text.page_content, keyword_number) - if text.metadata is not None: - self._update_segment_keywords(self.dataset.id, text.metadata["doc_id"], list(keywords), session) - keyword_table = self._add_text_to_keyword_table( - keyword_table or {}, text.metadata["doc_id"], list(keywords) - ) + selected: dict[str, Sequence[str]] = {} + for position, text in enumerate(texts): + keywords = keywords_list[position] if keywords_list else None + if not keywords: + if handler is None: + handler = JiebaKeywordTableHandler() + keywords = list(handler.extract_keywords(text.page_content, keyword_number)) + node_id = text.metadata["doc_id"] + selected[node_id] = list(keywords) + table = self._add_text_to_keyword_table(table, node_id, list(keywords)) + encoded = save_keyword_table( + storage=storage, + tenant_id=self.dataset.tenant_id, + dataset_id=self.dataset.id, + storage_type=storage_type, + table=table, + ) + write(storage_type, encoded, selected) - self._save_dataset_keyword_table(keyword_table, session) + def _read_keyword_record(self, session: Session) -> tuple[str, str | None]: + row = get_dataset_keyword_table(self.dataset, session=session) + return (row.data_source_type, row.keyword_table) if row else (dify_config.KEYWORD_DATA_SOURCE_TYPE, None) @override def text_exists(self, id: str, *, session: Session) -> bool: - dataset_keyword_table = self.dataset.get_dataset_keyword_table(session=session) keyword_table = None - keyword_table_dict = ( - dataset_keyword_table.get_keyword_table_dict(session=session) if dataset_keyword_table else None - ) + keyword_table_dict = load_dataset_keyword_table(self.dataset, session=session, storage=storage) if keyword_table_dict: data: Any = keyword_table_dict["__data__"] keyword_table = dict(data["table"]) if keyword_table is None: return False - return id in set.union(*keyword_table.values()) + return any(id in node_ids for node_ids in keyword_table.values()) @override def delete_by_ids(self, ids: list[str], session: Session, **kwargs: Any): - lock_name = f"keyword_indexing_lock_{self.dataset.id}" - with redis_client.lock(lock_name, timeout=600): - keyword_table = self._get_dataset_keyword_table(session) - if keyword_table is not None: - keyword_table = self._delete_ids_from_keyword_table(keyword_table, ids) - - self._save_dataset_keyword_table(keyword_table, session) + self.update_texts( + [], + read=lambda: self._read_keyword_record(session), + write=lambda storage_type, data, keywords: persist_keyword_table( + session, + tenant_id=self.dataset.tenant_id, + dataset_id=self.dataset.id, + storage_type=storage_type, + data=data, + keywords=keywords, + ), + lock=redis_client.lock(f"keyword_indexing_lock_{self.dataset.id}", timeout=600), + deleted_ids=ids, + ) @override def search(self, query: str, *, session: Session, **kwargs: Any) -> list[Document]: - dataset_keyword_table = self.dataset.get_dataset_keyword_table(session=session) keyword_table = None - keyword_table_dict = ( - dataset_keyword_table.get_keyword_table_dict(session=session) if dataset_keyword_table else None - ) + keyword_table_dict = load_dataset_keyword_table(self.dataset, session=session, storage=storage) if keyword_table_dict: data: Any = keyword_table_dict["__data__"] keyword_table = dict(data["table"]) @@ -145,62 +183,14 @@ class Jieba(BaseKeyword): def delete(self, *, session: Session): lock_name = f"keyword_indexing_lock_{self.dataset.id}" with redis_client.lock(lock_name, timeout=600): - dataset_keyword_table = self.dataset.get_dataset_keyword_table(session=session) + dataset_keyword_table = get_dataset_keyword_table(self.dataset, session=session) if dataset_keyword_table: session.delete(dataset_keyword_table) session.commit() if dataset_keyword_table.data_source_type != "database": - file_key = "keyword_files/" + self.dataset.tenant_id + "/" + self.dataset.id + ".txt" + file_key = keyword_file_key(self.dataset.tenant_id, self.dataset.id) storage.delete(file_key) - def _save_dataset_keyword_table(self, keyword_table: dict[str, set[str]] | None, session: Session): - keyword_table_dict = { - "__type__": "keyword_table", - "__data__": {"index_id": self.dataset.id, "summary": None, "table": keyword_table}, - } - dataset_keyword_table = session.scalar( - select(DatasetKeywordTable).where(DatasetKeywordTable.dataset_id == self.dataset.id) - ) - keyword_data_source_type = dataset_keyword_table.data_source_type if dataset_keyword_table else "file" - if keyword_data_source_type == "database": - if dataset_keyword_table is None: - return - dataset_keyword_table.keyword_table = dumps_with_sets(keyword_table_dict) - session.flush() - else: - file_key = "keyword_files/" + self.dataset.tenant_id + "/" + self.dataset.id + ".txt" - if storage.exists(file_key): - storage.delete(file_key) - storage.save(file_key, dumps_with_sets(keyword_table_dict).encode("utf-8")) - - def _get_dataset_keyword_table(self, session: Session) -> dict[str, set[str]] | None: - dataset_keyword_table = session.scalar( - select(DatasetKeywordTable).where(DatasetKeywordTable.dataset_id == self.dataset.id) - ) - if dataset_keyword_table: - keyword_table_dict = dataset_keyword_table.get_keyword_table_dict(session=session) - if keyword_table_dict: - data: Any = keyword_table_dict["__data__"] - return dict(data["table"]) - else: - keyword_data_source_type = dify_config.KEYWORD_DATA_SOURCE_TYPE - dataset_keyword_table = DatasetKeywordTable( - dataset_id=self.dataset.id, - keyword_table="", - data_source_type=keyword_data_source_type, - ) - if keyword_data_source_type == "database": - dataset_keyword_table.keyword_table = dumps_with_sets( - { - "__type__": "keyword_table", - "__data__": {"index_id": self.dataset.id, "summary": None, "table": {}}, - } - ) - session.add(dataset_keyword_table) - session.flush() - - return {} - def _add_text_to_keyword_table( self, keyword_table: dict[str, set[str]], id: str, keywords: list[str] ) -> dict[str, set[str]]: @@ -245,58 +235,3 @@ class Jieba(BaseKeyword): ) return sorted_chunk_indices[:k] - - def _update_segment_keywords(self, dataset_id: str, node_id: str, keywords: list[str], session: Session): - stmt = select(DocumentSegment).where( - DocumentSegment.dataset_id == dataset_id, DocumentSegment.index_node_id == node_id - ) - document_segment = session.scalar(stmt) - if document_segment: - document_segment.keywords = keywords - session.add(document_segment) - session.flush() - - def create_segment_keywords(self, node_id: str, keywords: list[str], session: Session): - keyword_table = self._get_dataset_keyword_table(session) - self._update_segment_keywords(self.dataset.id, node_id, keywords, session) - keyword_table = self._add_text_to_keyword_table(keyword_table or {}, node_id, keywords) - self._save_dataset_keyword_table(keyword_table, session) - - def multi_create_segment_keywords(self, pre_segment_data_list: list[PreSegmentData], session: Session): - keyword_table_handler = JiebaKeywordTableHandler() - keyword_table = self._get_dataset_keyword_table(session) - for pre_segment_data in pre_segment_data_list: - segment = pre_segment_data["segment"] - if pre_segment_data["keywords"]: - segment.keywords = pre_segment_data["keywords"] - assert segment.index_node_id - keyword_table = self._add_text_to_keyword_table( - keyword_table or {}, segment.index_node_id, pre_segment_data["keywords"] - ) - else: - keyword_number = self.dataset.keyword_number or self._config.max_keywords_per_chunk - - keywords = keyword_table_handler.extract_keywords(segment.content, keyword_number) - segment.keywords = list(keywords) - assert segment.index_node_id - keyword_table = self._add_text_to_keyword_table( - keyword_table or {}, segment.index_node_id, list(keywords) - ) - self._save_dataset_keyword_table(keyword_table, session) - - def update_segment_keywords_index(self, node_id: str, keywords: list[str], session: Session): - keyword_table = self._get_dataset_keyword_table(session) - keyword_table = self._add_text_to_keyword_table(keyword_table or {}, node_id, keywords) - self._save_dataset_keyword_table(keyword_table, session) - - -def set_orjson_default(obj: Any): - """Default function for orjson serialization of set types""" - if isinstance(obj, set): - return list(obj) - raise TypeError(f"Object of type {type(obj).__name__} is not JSON serializable") - - -def dumps_with_sets(obj: Any) -> str: - """JSON dumps with set support using orjson""" - return orjson.dumps(obj, default=set_orjson_default).decode("utf-8") diff --git a/api/core/rag/datasource/retrieval_service.py b/api/core/rag/datasource/retrieval_service.py index 35a70103623..e97984d4291 100644 --- a/api/core/rag/datasource/retrieval_service.py +++ b/api/core/rag/datasource/retrieval_service.py @@ -5,7 +5,7 @@ from concurrent.futures import ThreadPoolExecutor from typing import Any, NotRequired, TypedDict from flask import Flask, current_app -from sqlalchemy import select +from sqlalchemy import select, tuple_ from sqlalchemy.orm import Session, load_only from configs import dify_config @@ -35,7 +35,7 @@ from models.dataset import ( ) from models.dataset import Document as DatasetDocument from models.model import UploadFile -from services.external_knowledge_service import ExternalDatasetService +from services.knowledge.external.service import ExternalDatasetService class SegmentAttachmentResult(TypedDict): @@ -486,7 +486,14 @@ class RetrievalService: for doc in session.scalars( select(DatasetDocument) .where(DatasetDocument.id.in_(document_ids)) - .options(load_only(DatasetDocument.id, DatasetDocument.doc_form, DatasetDocument.dataset_id)) + .options( + load_only( + DatasetDocument.id, + DatasetDocument.tenant_id, + DatasetDocument.doc_form, + DatasetDocument.dataset_id, + ) + ) ).all() } @@ -559,7 +566,14 @@ class RetrievalService: doc_segment_map: dict[str, list[str]] = {} segment_summary_map: dict[str, str] = {} # Map segment_id to summary content - attachments = cls.get_segment_attachment_infos(image_doc_ids, session) + dataset_refs = tuple( + {(document.tenant_id, document.dataset_id) for document in valid_dataset_documents.values()} + ) + attachments = cls.get_segment_attachment_infos( + image_doc_ids, + dataset_refs=dataset_refs, + session=session, + ) for attachment in attachments: segment_ids.append(attachment["segment_id"]) @@ -947,15 +961,30 @@ class RetrievalService: @classmethod def get_segment_attachment_infos( - cls, attachment_ids: list[str], session: Session + cls, + attachment_ids: list[str], + *, + dataset_refs: Sequence[tuple[str, str]], + session: Session, ) -> list[SegmentAttachmentInfoResult]: + if not attachment_ids or not dataset_refs: + return [] attachment_infos: list[SegmentAttachmentInfoResult] = [] granted_upload_file_ids: list[str] = [] - upload_files = session.scalars(select(UploadFile).where(UploadFile.id.in_(attachment_ids))).all() + tenant_ids = tuple({tenant_id for tenant_id, _ in dataset_refs}) + upload_files = session.scalars( + select(UploadFile).where( + UploadFile.id.in_(attachment_ids), + UploadFile.tenant_id.in_(tenant_ids), + ) + ).all() if upload_files: upload_file_ids = [upload_file.id for upload_file in upload_files] attachment_bindings = session.scalars( - select(SegmentAttachmentBinding).where(SegmentAttachmentBinding.attachment_id.in_(upload_file_ids)) + select(SegmentAttachmentBinding).where( + SegmentAttachmentBinding.attachment_id.in_(upload_file_ids), + tuple_(SegmentAttachmentBinding.tenant_id, SegmentAttachmentBinding.dataset_id).in_(dataset_refs), + ) ).all() attachment_binding_map = {binding.attachment_id: binding for binding in attachment_bindings} diff --git a/api/core/rag/datasource/vdb/vector_factory.py b/api/core/rag/datasource/vdb/vector_factory.py index 0806dfb2745..9bd9eed6245 100644 --- a/api/core/rag/datasource/vdb/vector_factory.py +++ b/api/core/rag/datasource/vdb/vector_factory.py @@ -2,6 +2,7 @@ import base64 import logging import time from abc import ABC, abstractmethod +from collections.abc import Mapping from typing import Any, override from sqlalchemy import select @@ -105,7 +106,14 @@ class _LazyEmbeddings(Embeddings): class Vector: - def __init__(self, dataset: Dataset, attributes: list | None = None, *, session: Session): + def __init__( + self, + dataset: Dataset, + attributes: list | None = None, + *, + session: Session | None, + vector_type: str | None = None, + ): if attributes is None: # `is_summary` and `original_chunk_id` are stored on summary vectors # by `SummaryIndexService` and read back by `RetrievalService` to @@ -132,7 +140,14 @@ class Vector: self._embeddings: Embeddings = _LazyEmbeddings(dataset) self._attributes = attributes self._session = session - self._vector_processor = self._init_vector(session=session) + if vector_type is not None: + self._vector_processor = self.get_vector_factory(vector_type)().init_vector( + dataset, self._attributes, self._embeddings + ) + else: + if session is None: + raise ValueError("A resolved vector type is required without a database session") + self._vector_processor = self._init_vector(session=session) @staticmethod def resolve_vector_type(dataset: Dataset, *, session: Session) -> str: @@ -192,7 +207,9 @@ class Vector: self._vector_processor.create(texts=batch, embeddings=batch_embeddings, **kwargs) logger.info("Embedding %s texts took %s s", len(texts), time.time() - start) - def create_multimodal(self, file_documents: list | None = None, **kwargs): + def create_multimodal( + self, file_documents: list | None = None, *, upload_files: Mapping[str, UploadFile] | None = None, **kwargs + ): if file_documents: start = time.time() logger.info("start embedding %s files %s", len(file_documents), start) @@ -205,9 +222,16 @@ class Vector: # Batch query all upload files to avoid N+1 queries attachment_ids = [doc.metadata["doc_id"] for doc in batch] - stmt = select(UploadFile).where(UploadFile.id.in_(attachment_ids)) - upload_files = self._session.scalars(stmt).all() - upload_file_map = {str(f.id): f for f in upload_files} + upload_file_map: Mapping[str, UploadFile] + if upload_files is None: + if self._session is None: + raise ValueError("Upload files must be loaded before indexing without a database session") + stmt = select(UploadFile).where( + UploadFile.id.in_(attachment_ids), UploadFile.tenant_id == self._dataset.tenant_id + ) + upload_file_map = {str(f.id): f for f in self._session.scalars(stmt).all()} + else: + upload_file_map = upload_files file_base64_list = [] real_batch = [] @@ -264,6 +288,8 @@ class Vector: return self._vector_processor.search_by_vector(query_vector, **kwargs) def search_by_file(self, file_id: str, **kwargs: Any) -> list[Document]: + if self._session is None: + raise ValueError("File search requires a database session") upload_file: UploadFile | None = self._session.get(UploadFile, file_id) if not upload_file: diff --git a/api/core/rag/docstore/dataset_docstore.py b/api/core/rag/docstore/dataset_docstore.py index 456df05daa4..e5f563fb170 100644 --- a/api/core/rag/docstore/dataset_docstore.py +++ b/api/core/rag/docstore/dataset_docstore.py @@ -235,8 +235,12 @@ class DatasetDocumentStore: def get_document_segment(self, doc_id: str, session: Session) -> DocumentSegment | None: stmt = select(DocumentSegment).where( - DocumentSegment.dataset_id == self._dataset.id, DocumentSegment.index_node_id == doc_id + DocumentSegment.tenant_id == self._dataset.tenant_id, + DocumentSegment.dataset_id == self._dataset.id, + DocumentSegment.index_node_id == doc_id, ) + if self._document_id is not None: + stmt = stmt.where(DocumentSegment.document_id == self._document_id) document_segment = session.scalar(stmt) return document_segment diff --git a/api/core/rag/entities/__init__.py b/api/core/rag/entities/__init__.py index 373b68894bf..1e87472b2c4 100644 --- a/api/core/rag/entities/__init__.py +++ b/api/core/rag/entities/__init__.py @@ -3,7 +3,13 @@ from core.rag.entities.context_entities import DocumentContext from core.rag.entities.event import DatasourceCompletedEvent, DatasourceErrorEvent, DatasourceProcessingEvent from core.rag.entities.index_entities import EconomySetting, EmbeddingSetting, IndexMethod from core.rag.entities.metadata_entities import Condition, MetadataFilteringCondition, SupportedComparisonOperator -from core.rag.entities.processing_entities import ParentMode, PreProcessingRule, Rule, Segmentation +from core.rag.entities.processing_entities import ( + ParentMode, + PreProcessingRule, + PreProcessingRuleKey, + Rule, + Segmentation, +) from core.rag.entities.retrieval_settings import ( KeywordSetting, RerankingModelConfig, @@ -24,6 +30,7 @@ __all__ = [ "MetadataFilteringCondition", "ParentMode", "PreProcessingRule", + "PreProcessingRuleKey", "RerankingModelConfig", "RetrievalSourceMetadata", "Rule", diff --git a/api/core/rag/entities/processing_entities.py b/api/core/rag/entities/processing_entities.py index 6b08bb858c8..3a568062f40 100644 --- a/api/core/rag/entities/processing_entities.py +++ b/api/core/rag/entities/processing_entities.py @@ -1,7 +1,7 @@ from enum import StrEnum -from typing import Annotated, Literal +from typing import Literal -from pydantic import AliasChoices, BaseModel, Field, WithJsonSchema +from pydantic import AliasChoices, BaseModel, Field class ParentMode(StrEnum): @@ -9,19 +9,11 @@ class ParentMode(StrEnum): PARAGRAPH = "paragraph" -PreProcessingRuleID = Annotated[ - str, - WithJsonSchema( - { - "enum": ["remove_stopwords", "remove_extra_spaces", "remove_urls_emails"], - "type": "string", - } - ), -] +type PreProcessingRuleKey = Literal["remove_stopwords", "remove_extra_spaces", "remove_urls_emails"] class PreProcessingRule(BaseModel): - id: PreProcessingRuleID = Field(description="Rule identifier.") + id: PreProcessingRuleKey = Field(description="Rule identifier.") enabled: bool = Field(description="Whether this preprocessing rule is enabled.") diff --git a/api/core/rag/extractor/entity/extract_setting.py b/api/core/rag/extractor/entity/extract_setting.py index 0c420340732..f31acb4ebed 100644 --- a/api/core/rag/extractor/entity/extract_setting.py +++ b/api/core/rag/extractor/entity/extract_setting.py @@ -1,7 +1,24 @@ -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, Field -from models.dataset import Document -from models.model import UploadFile + +class UploadFileExtractionInput(BaseModel): + """Detached values required to extract one tenant-owned upload.""" + + id: str + tenant_id: str + key: str + created_by: str + model_config = ConfigDict(from_attributes=True) + + +class StoredDocumentExtractionInput(BaseModel): + """Detached document values needed to record a source's last edited time.""" + + id: str + tenant_id: str + dataset_id: str + data_source_info_dict: dict[str, object] | None + model_config = ConfigDict(from_attributes=True) class NotionInfo(BaseModel): @@ -13,9 +30,9 @@ class NotionInfo(BaseModel): notion_workspace_id: str | None = "" notion_obj_id: str notion_page_type: str - document: Document | None = None + document: StoredDocumentExtractionInput | None = None tenant_id: str - model_config = ConfigDict(arbitrary_types_allowed=True) + notion_access_token: str | None = Field(default=None, exclude=True, repr=False) class WebsiteInfo(BaseModel): @@ -23,8 +40,6 @@ class WebsiteInfo(BaseModel): website import info. """ - model_config = ConfigDict(arbitrary_types_allowed=True) - provider: str job_id: str url: str @@ -39,8 +54,7 @@ class ExtractSetting(BaseModel): """ datasource_type: str - upload_file: UploadFile | None = None + upload_file: UploadFileExtractionInput | None = None notion_info: NotionInfo | None = None website_info: WebsiteInfo | None = None document_model: str | None = None - model_config = ConfigDict(arbitrary_types_allowed=True) diff --git a/api/core/rag/extractor/extract_processor.py b/api/core/rag/extractor/extract_processor.py index dd73b436089..a8d83d57d83 100644 --- a/api/core/rag/extractor/extract_processor.py +++ b/api/core/rag/extractor/extract_processor.py @@ -7,10 +7,11 @@ from urllib.parse import unquote from sqlalchemy.orm import Session from configs import dify_config +from core.db.session_factory import get_session_maker from core.file import remote_fetcher from core.rag.extractor.csv_extractor import CSVExtractor from core.rag.extractor.entity.datasource_type import DatasourceType -from core.rag.extractor.entity.extract_setting import ExtractSetting +from core.rag.extractor.entity.extract_setting import ExtractSetting, UploadFileExtractionInput from core.rag.extractor.excel_extractor import ExcelExtractor from core.rag.extractor.extractor_base import BaseExtractor from core.rag.extractor.firecrawl.firecrawl_web_extractor import FirecrawlWebExtractor @@ -59,7 +60,9 @@ class ExtractProcessor: cls, upload_file: UploadFile, return_text: bool = False, is_automatic: bool = False ) -> list[Document] | str: extract_setting = ExtractSetting( - datasource_type=DatasourceType.FILE, upload_file=upload_file, document_model="text_model" + datasource_type=DatasourceType.FILE, + upload_file=UploadFileExtractionInput.model_validate(upload_file), + document_model="text_model", ) if return_text: delimiter = "\n" @@ -223,17 +226,29 @@ class ExtractProcessor: extractor = TextExtractor(file_path, autodetect_encoding=True) return extractor.extract() elif extract_setting.datasource_type == DatasourceType.NOTION: + from extensions.application_services.data_sources import build_data_source_credentials + assert extract_setting.notion_info is not None, "notion_info is required" + notion_info = extract_setting.notion_info extractor = NotionExtractor( notion_workspace_id=extract_setting.notion_info.notion_workspace_id or "", notion_obj_id=extract_setting.notion_info.notion_obj_id, notion_page_type=extract_setting.notion_info.notion_page_type, document_model=extract_setting.notion_info.document, tenant_id=extract_setting.notion_info.tenant_id, + notion_access_token=extract_setting.notion_info.notion_access_token, credential_id=extract_setting.notion_info.credential_id, + notion_token_loader=lambda: build_data_source_credentials( + database_client=get_session_maker() + ).providers.get_stored_notion_access_token( + tenant_id=notion_info.tenant_id, credential_id=notion_info.credential_id + ), ) return extractor.extract() elif extract_setting.datasource_type == DatasourceType.WEBSITE: + from extensions.application_services.data_sources import build_website_service + + website_service = build_website_service(database_client=get_session_maker()) assert extract_setting.website_info is not None, "website_info is required" match extract_setting.website_info.provider: case "firecrawl": @@ -243,6 +258,7 @@ class ExtractProcessor: tenant_id=extract_setting.website_info.tenant_id, mode=extract_setting.website_info.mode, only_main_content=extract_setting.website_info.only_main_content, + website_service=website_service, ) return extractor.extract() case "watercrawl": @@ -252,6 +268,7 @@ class ExtractProcessor: tenant_id=extract_setting.website_info.tenant_id, mode=extract_setting.website_info.mode, only_main_content=extract_setting.website_info.only_main_content, + website_service=website_service, ) return extractor.extract() case "jinareader": @@ -261,6 +278,7 @@ class ExtractProcessor: tenant_id=extract_setting.website_info.tenant_id, mode=extract_setting.website_info.mode, only_main_content=extract_setting.website_info.only_main_content, + website_service=website_service, ) return extractor.extract() case _: diff --git a/api/core/rag/extractor/firecrawl/firecrawl_web_extractor.py b/api/core/rag/extractor/firecrawl/firecrawl_web_extractor.py index 740527a32d7..01ac232898d 100644 --- a/api/core/rag/extractor/firecrawl/firecrawl_web_extractor.py +++ b/api/core/rag/extractor/firecrawl/firecrawl_web_extractor.py @@ -2,7 +2,7 @@ from typing import override from core.rag.extractor.extractor_base import BaseExtractor from core.rag.models.document import Document -from services.website_service import WebsiteService +from services.data_source.website_service import WebsiteService class FirecrawlWebExtractor(BaseExtractor): @@ -24,8 +24,11 @@ class FirecrawlWebExtractor(BaseExtractor): tenant_id: str, mode: str = "crawl", only_main_content: bool = True, + *, + website_service: WebsiteService, ): """Initialize with url, api_key, base_url and mode.""" + self._website_service = website_service self._url = url self.job_id = job_id self.tenant_id = tenant_id @@ -37,7 +40,7 @@ class FirecrawlWebExtractor(BaseExtractor): """Extract content from the URL.""" documents = [] if self.mode == "crawl": - crawl_data = WebsiteService.get_crawl_url_data(self.job_id, "firecrawl", self._url, self.tenant_id) + crawl_data = self._website_service.get_crawl_url_data(self.job_id, "firecrawl", self._url, self.tenant_id) if crawl_data is None: return [] document = Document( @@ -50,7 +53,7 @@ class FirecrawlWebExtractor(BaseExtractor): ) documents.append(document) elif self.mode == "scrape": - scrape_data = WebsiteService.get_scrape_url_data( + scrape_data = self._website_service.get_scrape_url_data( "firecrawl", self._url, self.tenant_id, self.only_main_content ) diff --git a/api/core/rag/extractor/jina_reader_extractor.py b/api/core/rag/extractor/jina_reader_extractor.py index b184e8fac65..489ec1082f8 100644 --- a/api/core/rag/extractor/jina_reader_extractor.py +++ b/api/core/rag/extractor/jina_reader_extractor.py @@ -2,7 +2,7 @@ from typing import override from core.rag.extractor.extractor_base import BaseExtractor from core.rag.models.document import Document -from services.website_service import WebsiteService +from services.data_source.website_service import WebsiteService class JinaReaderWebExtractor(BaseExtractor): @@ -17,8 +17,11 @@ class JinaReaderWebExtractor(BaseExtractor): tenant_id: str, mode: str = "crawl", only_main_content: bool = False, + *, + website_service: WebsiteService, ): """Initialize with url, api_key, base_url and mode.""" + self._website_service = website_service self._url = url self.job_id = job_id self.tenant_id = tenant_id @@ -30,7 +33,7 @@ class JinaReaderWebExtractor(BaseExtractor): """Extract content from the URL.""" documents = [] if self.mode == "crawl": - crawl_data = WebsiteService.get_crawl_url_data(self.job_id, "jinareader", self._url, self.tenant_id) + crawl_data = self._website_service.get_crawl_url_data(self.job_id, "jinareader", self._url, self.tenant_id) if crawl_data is None: return [] document = Document( diff --git a/api/core/rag/extractor/notion_extractor.py b/api/core/rag/extractor/notion_extractor.py index 10f1bd984a6..680b049a9ea 100644 --- a/api/core/rag/extractor/notion_extractor.py +++ b/api/core/rag/extractor/notion_extractor.py @@ -1,17 +1,17 @@ import json import logging import operator +from collections.abc import Callable from typing import Any, cast, override import httpx from sqlalchemy import update -from configs import dify_config +from core.rag.extractor.entity.extract_setting import StoredDocumentExtractionInput from core.rag.extractor.extractor_base import BaseExtractor from core.rag.models.document import Document from extensions.ext_database import db from models.dataset import Document as DocumentModel -from services.datasource_provider_service import DatasourceProviderService logger = logging.getLogger(__name__) @@ -41,36 +41,21 @@ class NotionExtractor(BaseExtractor): notion_obj_id: str, notion_page_type: str, tenant_id: str, - document_model: DocumentModel | None = None, + document_model: StoredDocumentExtractionInput | None = None, notion_access_token: str | None = None, credential_id: str | None = None, + *, + notion_token_loader: Callable[[], str], ): self._notion_access_token = None self._document_model = document_model self._notion_workspace_id = notion_workspace_id self._notion_obj_id = notion_obj_id self._notion_page_type = notion_page_type - self._credential_id = credential_id if notion_access_token: self._notion_access_token = notion_access_token else: - try: - self._notion_access_token = self._get_access_token(tenant_id, self._credential_id) - except Exception as e: - logger.warning( - ( - "Failed to get Notion access token from datasource credentials: %s, " - "falling back to environment variable NOTION_INTEGRATION_TOKEN" - ), - e, - ) - integration_token = dify_config.NOTION_INTEGRATION_TOKEN - if integration_token is None: - raise ValueError( - "Must specify `integration_token` or set environment variable `NOTION_INTEGRATION_TOKEN`." - ) from e - - self._notion_access_token = integration_token + self._notion_access_token = notion_token_loader() @override def extract(self) -> list[Document]: @@ -344,7 +329,7 @@ class NotionExtractor(BaseExtractor): result_lines = "\n".join(result_lines_arr) return result_lines - def update_last_edited_time(self, document_model: DocumentModel | None): + def update_last_edited_time(self, document_model: StoredDocumentExtractionInput | None): if not document_model: return @@ -355,7 +340,11 @@ class NotionExtractor(BaseExtractor): db.session.execute( update(DocumentModel) - .where(DocumentModel.id == document_model.id) + .where( + DocumentModel.id == document_model.id, + DocumentModel.tenant_id == document_model.tenant_id, + DocumentModel.dataset_id == document_model.dataset_id, + ) .values({DocumentModel.data_source_info: json.dumps(data_source_info)}) ) db.session.commit() @@ -385,20 +374,3 @@ class NotionExtractor(BaseExtractor): data = res.json() return cast(str, data["last_edited_time"]) - - @classmethod - def _get_access_token(cls, tenant_id: str, credential_id: str | None) -> str: - # get credential from tenant_id and credential_id - if not credential_id: - raise Exception(f"No credential id found for tenant {tenant_id}") - datasource_provider_service = DatasourceProviderService() - credential = datasource_provider_service.get_datasource_credentials( - tenant_id=tenant_id, - credential_id=credential_id, - provider="notion_datasource", - plugin_id="langgenius/notion_datasource", - ) - if not credential: - raise Exception(f"No notion credential found for tenant {tenant_id} and credential {credential_id}") - - return cast(str, credential["integration_secret"]) diff --git a/api/core/rag/extractor/watercrawl/extractor.py b/api/core/rag/extractor/watercrawl/extractor.py index fdfe75b6770..0359a70491b 100644 --- a/api/core/rag/extractor/watercrawl/extractor.py +++ b/api/core/rag/extractor/watercrawl/extractor.py @@ -2,7 +2,7 @@ from typing import override from core.rag.extractor.extractor_base import BaseExtractor from core.rag.models.document import Document -from services.website_service import WebsiteService +from services.data_source.website_service import WebsiteService class WaterCrawlWebExtractor(BaseExtractor): @@ -25,8 +25,11 @@ class WaterCrawlWebExtractor(BaseExtractor): tenant_id: str, mode: str = "crawl", only_main_content: bool = True, + *, + website_service: WebsiteService, ): """Initialize with url, api_key, base_url and mode.""" + self._website_service = website_service self._url = url self.job_id = job_id self.tenant_id = tenant_id @@ -38,7 +41,7 @@ class WaterCrawlWebExtractor(BaseExtractor): """Extract content from the URL.""" documents = [] if self.mode == "crawl": - crawl_data = WebsiteService.get_crawl_url_data(self.job_id, "watercrawl", self._url, self.tenant_id) + crawl_data = self._website_service.get_crawl_url_data(self.job_id, "watercrawl", self._url, self.tenant_id) if crawl_data is None: return [] document = Document( @@ -51,7 +54,7 @@ class WaterCrawlWebExtractor(BaseExtractor): ) documents.append(document) elif self.mode == "scrape": - scrape_data = WebsiteService.get_scrape_url_data( + scrape_data = self._website_service.get_scrape_url_data( "watercrawl", self._url, self.tenant_id, self.only_main_content ) diff --git a/api/core/rag/index_processor/index_processor_base.py b/api/core/rag/index_processor/index_processor_base.py index 06bacd68669..a2fea9bde9e 100644 --- a/api/core/rag/index_processor/index_processor_base.py +++ b/api/core/rag/index_processor/index_processor_base.py @@ -147,7 +147,12 @@ class BaseIndexProcessor(ABC): return character_splitter def _get_content_files( - self, document: Document, current_user: Account | None = None, *, session: Session + self, + document: Document, + current_user: Account | None = None, + *, + tenant_id: str, + session: Session, ) -> list[AttachmentDocument]: """ Get the content files from the document. @@ -185,7 +190,12 @@ class BaseIndexProcessor(ABC): if match: if current_user: tool_file_id = match.group(1) - upload_file_id = self._download_tool_file(tool_file_id, current_user, session=session) + upload_file_id = self._download_tool_file( + tool_file_id, + current_user, + tenant_id=tenant_id, + session=session, + ) if upload_file_id: upload_file_id_list.append(upload_file_id) continue @@ -208,7 +218,12 @@ class BaseIndexProcessor(ABC): # Get unique IDs for database query unique_upload_file_ids = list(set(upload_file_id_list)) - upload_files = session.scalars(select(UploadFile).where(UploadFile.id.in_(unique_upload_file_ids))).all() + upload_files = session.scalars( + select(UploadFile).where( + UploadFile.tenant_id == tenant_id, + UploadFile.id.in_(unique_upload_file_ids), + ) + ).all() # Create a mapping from ID to UploadFile for quick lookup upload_file_map = {upload_file.id: upload_file for upload_file in upload_files} @@ -314,13 +329,22 @@ class BaseIndexProcessor(ABC): logging.warning("Unexpected error downloading image from %s", image_url, exc_info=True) return None - def _download_tool_file(self, tool_file_id: str, current_user: Account, *, session: Session) -> str | None: + def _download_tool_file( + self, + tool_file_id: str, + current_user: Account, + *, + tenant_id: str, + session: Session, + ) -> str | None: """ Download the tool file from the ID. """ from services.file_service import FileService - tool_file = session.get(ToolFile, tool_file_id) + tool_file = session.scalar( + select(ToolFile).where(ToolFile.id == tool_file_id, ToolFile.tenant_id == tenant_id).limit(1) + ) if not tool_file: return None blob = storage.load_once(tool_file.file_key) diff --git a/api/core/rag/index_processor/processor/paragraph_index_processor.py b/api/core/rag/index_processor/processor/paragraph_index_processor.py index d0b322db70c..2102fe966c6 100644 --- a/api/core/rag/index_processor/processor/paragraph_index_processor.py +++ b/api/core/rag/index_processor/processor/paragraph_index_processor.py @@ -3,6 +3,7 @@ import logging import re import uuid +from collections.abc import Callable from typing import Any, TypedDict, cast, override from sqlalchemy import select @@ -46,7 +47,7 @@ from models.account import Account from models.dataset import Dataset, DatasetProcessRule, DocumentSegment, SegmentAttachmentBinding from models.dataset import Document as DatasetDocument from services.account_service import AccountService -from services.summary_index_service import SummaryIndexService +from services.knowledge.summaries.adapters import SummaryIndexAdapter logger = logging.getLogger(__name__) @@ -113,7 +114,12 @@ class ParagraphIndexProcessor(BaseIndexProcessor): document_node.metadata["doc_id"] = doc_id document_node.metadata["doc_hash"] = hash multimodal_documents = ( - self._get_content_files(document_node, current_user, session=session) + self._get_content_files( + document_node, + current_user, + tenant_id=kwargs["tenant_id"], + session=session, + ) if document_node.metadata else None ) @@ -172,10 +178,10 @@ class ParagraphIndexProcessor(BaseIndexProcessor): ).all() segment_ids = [segment.id for segment in segments] if segment_ids: - SummaryIndexService.delete_summaries_for_segments(dataset, segment_ids, session=session) + SummaryIndexAdapter.delete_summaries_for_segments(dataset, segment_ids, session=session) else: # Delete all summaries for the dataset - SummaryIndexService.delete_summaries_for_segments(dataset, None, session=session) + SummaryIndexAdapter.delete_summaries_for_segments(dataset, None, session=session) if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY: vector = Vector(dataset, session=session) @@ -204,7 +210,7 @@ class ParagraphIndexProcessor(BaseIndexProcessor): "doc_hash": helper.generate_text_hash(content), } doc = Document(page_content=content, metadata=metadata) - attachments = self._get_content_files(doc, session=session) + attachments = self._get_content_files(doc, tenant_id=dataset.tenant_id, session=session) if attachments: doc.attachments = attachments all_multimodal_documents.extend(attachments) @@ -240,7 +246,12 @@ class ParagraphIndexProcessor(BaseIndexProcessor): account = AccountService.load_user(document.created_by, account_session) if not account: raise ValueError("Invalid account") - doc.attachments = self._get_content_files(doc, current_user=account, session=session) + doc.attachments = self._get_content_files( + doc, + current_user=account, + tenant_id=dataset.tenant_id, + session=session, + ) if doc.attachments: all_multimodal_documents.extend(doc.attachments) documents.append(doc) @@ -378,7 +389,6 @@ class ParagraphIndexProcessor(BaseIndexProcessor): return preview_texts @staticmethod - @with_credit_usage_created_by(CreditUsageCreatedBy.KNOWLEDGE_INDEXING) def generate_summary( tenant_id: str, text: str, @@ -404,6 +414,30 @@ class ParagraphIndexProcessor(BaseIndexProcessor): Returns: Tuple of (summary_content, llm_usage) where llm_usage is LLMUsage object """ + + def load_images() -> list[File]: + images = ( + ParagraphIndexProcessor._extract_images_from_segment_attachments(tenant_id, segment_id, session) + if segment_id + else [] + ) + return images or ParagraphIndexProcessor._extract_images_from_text(tenant_id, text, session) + + return ParagraphIndexProcessor.generate_summary_from_inputs( + tenant_id, text, summary_index_setting, document_language=document_language, image_loader=load_images + ) + + @staticmethod + @with_credit_usage_created_by(CreditUsageCreatedBy.KNOWLEDGE_INDEXING) + def generate_summary_from_inputs( + tenant_id: str, + text: str, + summary_index_setting: SummaryIndexSettingDict | None, + *, + document_language: str | None, + image_loader: Callable[[], list[File]], + ) -> tuple[str, LLMUsage]: + """Generate a summary after an explicit loader materializes any required images.""" if not summary_index_setting or not summary_index_setting.get("enable"): raise ValueError("summary_index_setting is required and must be enabled to generate summary.") @@ -447,15 +481,7 @@ class ParagraphIndexProcessor(BaseIndexProcessor): # Extract images if model supports vision image_files = [] if supports_vision: - # First, try to get images from SegmentAttachmentBinding (preferred method) - if segment_id: - image_files = ParagraphIndexProcessor._extract_images_from_segment_attachments( - tenant_id, segment_id, session - ) - - # If no images from attachments, fall back to extracting from text - if not image_files: - image_files = ParagraphIndexProcessor._extract_images_from_text(tenant_id, text, session) + image_files = image_loader() # Build prompt messages prompt_messages = [] @@ -603,6 +629,7 @@ class ParagraphIndexProcessor(BaseIndexProcessor): .where( SegmentAttachmentBinding.segment_id == segment_id, SegmentAttachmentBinding.tenant_id == tenant_id, + UploadFile.tenant_id == tenant_id, ) ).all() diff --git a/api/core/rag/index_processor/processor/parent_child_index_processor.py b/api/core/rag/index_processor/processor/parent_child_index_processor.py index b31875c2f20..0c98b401800 100644 --- a/api/core/rag/index_processor/processor/parent_child_index_processor.py +++ b/api/core/rag/index_processor/processor/parent_child_index_processor.py @@ -29,7 +29,7 @@ from models.dataset import ChildChunk, Dataset, DatasetProcessRule, DocumentSegm from models.dataset import Document as DatasetDocument from models.enums import ProcessRuleMode from services.account_service import AccountService -from services.summary_index_service import SummaryIndexService +from services.knowledge.summaries.adapters import SummaryIndexAdapter logger = logging.getLogger(__name__) @@ -99,11 +99,16 @@ class ParentChildIndexProcessor(BaseIndexProcessor): page_content = page_content if len(page_content) > 0: document_node.page_content = page_content - multimodel_documents = self._get_content_files(document_node, current_user, session=session) + multimodel_documents = self._get_content_files( + document_node, + current_user, + tenant_id=kwargs["tenant_id"], + session=session, + ) if multimodel_documents: document_node.attachments = multimodel_documents # parse document to child nodes - child_nodes = self._split_child_nodes( + child_nodes = self.split_child_nodes( document_node, rules, process_rule.get("mode"), kwargs.get("embedding_model_instance") ) document_node.children = child_nodes @@ -112,11 +117,15 @@ class ParentChildIndexProcessor(BaseIndexProcessor): elif rules.parent_mode == ParentMode.FULL_DOC: page_content = "\n".join([document.page_content for document in documents]) document = Document(page_content=page_content, metadata=documents[0].metadata) - multimodel_documents = self._get_content_files(document, session=session) + multimodel_documents = self._get_content_files( + document, + tenant_id=kwargs["tenant_id"], + session=session, + ) if multimodel_documents: document.attachments = multimodel_documents # parse document to child nodes - child_nodes = self._split_child_nodes( + child_nodes = self.split_child_nodes( document, rules, process_rule.get("mode"), kwargs.get("embedding_model_instance") ) if kwargs.get("preview"): @@ -176,10 +185,10 @@ class ParentChildIndexProcessor(BaseIndexProcessor): ).all() segment_ids = [segment.id for segment in segments] if segment_ids: - SummaryIndexService.delete_summaries_for_segments(dataset, segment_ids, session=session) + SummaryIndexAdapter.delete_summaries_for_segments(dataset, segment_ids, session=session) else: # Delete all summaries for the dataset - SummaryIndexService.delete_summaries_for_segments(dataset, None, session=session) + SummaryIndexAdapter.delete_summaries_for_segments(dataset, None, session=session) if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY: delete_child_chunks = kwargs.get("delete_child_chunks") or False @@ -227,7 +236,7 @@ class ParentChildIndexProcessor(BaseIndexProcessor): ) session.flush() - def _split_child_nodes( + def split_child_nodes( self, document_node: Document, rules: Rule, @@ -302,7 +311,12 @@ class ParentChildIndexProcessor(BaseIndexProcessor): account = AccountService.load_user(document.created_by, account_session) if not account: raise ValueError("Invalid account") - doc.attachments = self._get_content_files(doc, current_user=account, session=session) + doc.attachments = self._get_content_files( + doc, + current_user=account, + tenant_id=dataset.tenant_id, + session=session, + ) documents.append(doc) if documents: token_counts = calculate_segment_token_counts(dataset=dataset, documents=documents) diff --git a/api/core/rag/index_processor/processor/qa_index_processor.py b/api/core/rag/index_processor/processor/qa_index_processor.py index e1971da27fa..c39a83d4489 100644 --- a/api/core/rag/index_processor/processor/qa_index_processor.py +++ b/api/core/rag/index_processor/processor/qa_index_processor.py @@ -29,7 +29,7 @@ from libs import helper from models.account import Account from models.dataset import Dataset, DocumentSegment from models.dataset import Document as DatasetDocument -from services.summary_index_service import SummaryIndexService +from services.knowledge.summaries.adapters import SummaryIndexAdapter logger = logging.getLogger(__name__) @@ -180,10 +180,10 @@ class QAIndexProcessor(BaseIndexProcessor): ).all() segment_ids = [segment.id for segment in segments] if segment_ids: - SummaryIndexService.delete_summaries_for_segments(dataset, segment_ids, session=session) + SummaryIndexAdapter.delete_summaries_for_segments(dataset, segment_ids, session=session) else: # Delete all summaries for the dataset - SummaryIndexService.delete_summaries_for_segments(dataset, None, session=session) + SummaryIndexAdapter.delete_summaries_for_segments(dataset, None, session=session) vector = Vector(dataset, session=session) if node_ids: diff --git a/api/core/rag/retrieval/dataset_retrieval.py b/api/core/rag/retrieval/dataset_retrieval.py index 2caa266c37c..04002bc87d1 100644 --- a/api/core/rag/retrieval/dataset_retrieval.py +++ b/api/core/rag/retrieval/dataset_retrieval.py @@ -28,7 +28,6 @@ from core.app.entities.app_invoke_entities import ( ModelConfigWithCredentialsEntity, get_credit_usage_app_type, ) -from core.app.file_access import grant_retriever_segment_access, grant_upload_file_access from core.callback_handler.index_tool_callback_handler import DatasetIndexToolCallbackHandler from core.db.session_factory import session_factory from core.entities.agent_entities import PlanningStrategy @@ -83,7 +82,6 @@ from graphon.model_runtime.entities.model_entities import ModelFeature, ModelTyp from graphon.model_runtime.model_providers.base.large_language_model import LargeLanguageModel from libs.helper import parse_uuid_str_or_none from libs.json_in_md_parser import parse_and_check_json_markdown -from models import UploadFile from models.dataset import ( ChildChunk, Dataset, @@ -96,8 +94,11 @@ from models.dataset import ( from models.dataset import Document as DatasetDocument from models.dataset import Document as DocumentModel from models.enums import CreatorUserRole, DatasetQuerySource -from services.external_knowledge_service import ExternalDatasetService +from repositories.knowledge.dataset_read_repository import get_dataset_available_document_count +from repositories.knowledge.segment_read_adapter import sign_segment_content from services.feature_service import FeatureService +from services.knowledge.external.service import ExternalDatasetService +from services.knowledge.retrieval.attachments import authorize_retrieved_segment default_retrieval_model: DefaultRetrievalModelDict = { "search_method": RetrievalMethod.SEMANTIC_SEARCH, @@ -310,25 +311,36 @@ class DatasetRetrieval: retrieval_resource_list.append(source) # deal with dify documents if dify_documents: - with Session(bind=session.get_bind()) as format_session: - records = RetrievalService.format_retrieval_documents(format_session, dify_documents) - dataset_ids = [i.segment.dataset_id for i in records] - document_ids = [i.segment.document_id for i in records] + # Materialize authorized, signed response values before closing the owned session. + with Session(bind=session.get_bind()) as retrieval_session: + records = RetrievalService.format_retrieval_documents(retrieval_session, dify_documents) + dataset_ids = [i.segment.dataset_id for i in records] + document_ids = [i.segment.document_id for i in records] + datasets = retrieval_session.scalars(select(Dataset).where(Dataset.id.in_(dataset_ids))).all() + documents = retrieval_session.scalars( + select(DatasetDocument).where(DatasetDocument.id.in_(document_ids)) + ).all() - with session_factory.create_session() as session: - datasets = session.scalars(select(Dataset).where(Dataset.id.in_(dataset_ids))).all() - documents = session.scalars(select(DatasetDocument).where(DatasetDocument.id.in_(document_ids))).all() + dataset_map = {i.id: i for i in datasets} + document_map = {i.id: i for i in documents} - dataset_map = {i.id: i for i in datasets} - document_map = {i.id: i for i in documents} - - if records: for record in records: segment = record.segment + if ( + authorize_retrieved_segment( + segment, + tenant_id=request.tenant_id, + dataset_ids=available_datasets_ids, + session=retrieval_session, + ) + is None + ): + continue dataset = dataset_map.get(segment.dataset_id) document = document_map.get(segment.document_id) if dataset and document: + content = sign_segment_content(segment, session=retrieval_session) source = Source( metadata=SourceMetadata( source="knowledge", @@ -358,15 +370,14 @@ class DatasetRetrieval: ), title=document.name, files=list(record.files) if record.files else None, - content=segment.get_sign_content(), + content=content, ) if segment.answer: - source.content = f"question:{segment.get_sign_content()} \nanswer:{segment.answer}" + source.content = f"question:{content} \nanswer:{segment.answer}" if record.summary: source.summary = record.summary - grant_retriever_segment_access([str(segment.id)]) retrieval_resource_list.append(source) if retrieval_resource_list: @@ -533,13 +544,22 @@ class DatasetRetrieval: with Session(bind=session.get_bind()) as format_session: records = RetrievalService.format_retrieval_documents(format_session, dify_documents) if records: + authorized_records = [] for record in records: segment = record.segment + attachments = authorize_retrieved_segment( + segment, tenant_id=tenant_id, dataset_ids=available_datasets_ids, session=session + ) + if attachments is None: + continue + authorized_records.append(record) # Build content: if summary exists, add it before the segment content if segment.answer: - segment_content = f"question:{segment.get_sign_content()} answer:{segment.answer}" + segment_content = ( + f"question:{sign_segment_content(segment, session=session)} answer:{segment.answer}" + ) else: - segment_content = segment.get_sign_content() + segment_content = sign_segment_content(segment, session=session) # If summary exists, prepend it to the content if record.summary: @@ -554,18 +574,8 @@ class DatasetRetrieval: ) ) if vision_enabled: - attachments_with_bindings = session.execute( - select(SegmentAttachmentBinding, UploadFile) - .join(UploadFile, UploadFile.id == SegmentAttachmentBinding.attachment_id) - .where( - SegmentAttachmentBinding.segment_id == segment.id, - ) - ).all() - if attachments_with_bindings: - grant_upload_file_access( - str(upload_file.id) for _, upload_file in attachments_with_bindings - ) - for _, upload_file in attachments_with_bindings: + if attachments: + for upload_file in attachments: attachment_info = File( file_id=upload_file.id, filename=upload_file.name, @@ -583,8 +593,8 @@ class DatasetRetrieval: ) context_files.append(attachment_info) if show_retrieve_source: - dataset_ids = [record.segment.dataset_id for record in records] - document_ids = [record.segment.document_id for record in records] + dataset_ids = [record.segment.dataset_id for record in authorized_records] + document_ids = [record.segment.document_id for record in authorized_records] dataset_document_stmt = select(DatasetDocument).where( DatasetDocument.id.in_(document_ids), DatasetDocument.enabled == True, @@ -597,7 +607,7 @@ class DatasetRetrieval: datasets = session.execute(dataset_stmt).scalars().all() # type: ignore dataset_map = {i.id: i for i in datasets} document_map = {i.id: i for i in documents} - for record in records: + for record in authorized_records: segment = record.segment dataset_item = dataset_map.get(segment.dataset_id) document_item = document_map.get(segment.document_id) @@ -1333,7 +1343,7 @@ class DatasetRetrieval: if ( dataset and dataset.provider != "external" - and dataset.get_total_available_documents(session=session) == 0 + and get_dataset_available_document_count(dataset, session=session) == 0 ): continue diff --git a/api/core/rag/summary_index/summary_index.py b/api/core/rag/summary_index/summary_index.py index 89788c6ab53..81fe9a1be73 100644 --- a/api/core/rag/summary_index/summary_index.py +++ b/api/core/rag/summary_index/summary_index.py @@ -8,7 +8,7 @@ from core.rag.index_processor.constant.index_type import IndexTechniqueType from core.rag.index_processor.index_processor_base import SummaryIndexSettingDict from extensions.otel import propagate_context from models.dataset import Dataset, Document, DocumentSegment, DocumentSegmentSummary -from services.summary_index_service import SummaryIndexService +from services.knowledge.summaries.adapters import SummaryIndexAdapter from tasks.generate_summary_index_task import generate_summary_index_task logger = logging.getLogger(__name__) @@ -82,7 +82,7 @@ class SummaryIndex: if segment is None: return try: - SummaryIndexService.generate_and_vectorize_summary( + SummaryIndexAdapter.generate_and_vectorize_summary( segment, dataset, summary_index_setting, session=session ) except Exception: diff --git a/api/core/tools/builtin_tool/providers/audio/tools/asr.py b/api/core/tools/builtin_tool/providers/audio/tools/asr.py index c315764c3b3..44543667815 100644 --- a/api/core/tools/builtin_tool/providers/audio/tools/asr.py +++ b/api/core/tools/builtin_tool/providers/audio/tools/asr.py @@ -14,7 +14,7 @@ from core.tools.entities.tool_entities import ToolInvokeMessage, ToolParameter from graphon.file import FileType from graphon.file.file_manager import download from graphon.model_runtime.entities.model_entities import ModelType -from services.model_provider_service import ModelProviderService +from services.model_provider.service import ModelProviderService class ASRTool(BuiltinTool): diff --git a/api/core/tools/builtin_tool/providers/audio/tools/tts.py b/api/core/tools/builtin_tool/providers/audio/tools/tts.py index a8083e74e25..8350af1e040 100644 --- a/api/core/tools/builtin_tool/providers/audio/tools/tts.py +++ b/api/core/tools/builtin_tool/providers/audio/tools/tts.py @@ -12,7 +12,7 @@ from core.tools.builtin_tool.tool import BuiltinTool from core.tools.entities.common_entities import I18nObject from core.tools.entities.tool_entities import ToolInvokeMessage, ToolParameter from graphon.model_runtime.entities.model_entities import ModelPropertyKey, ModelType -from services.model_provider_service import ModelProviderService +from services.model_provider.service import ModelProviderService class TTSTool(BuiltinTool): diff --git a/api/core/tools/utils/dataset_retriever/dataset_multi_retriever_tool.py b/api/core/tools/utils/dataset_retriever/dataset_multi_retriever_tool.py index da7c10d5ac4..3f2d64e3434 100644 --- a/api/core/tools/utils/dataset_retriever/dataset_multi_retriever_tool.py +++ b/api/core/tools/utils/dataset_retriever/dataset_multi_retriever_tool.py @@ -18,6 +18,7 @@ from core.tools.utils.dataset_retriever.dataset_retriever_base_tool import Datas from extensions.ext_database import db from graphon.model_runtime.entities.model_entities import ModelType from models.dataset import Dataset, Document, DocumentSegment +from repositories.knowledge.segment_read_adapter import sign_segment_content default_retrieval_model: DefaultRetrievalModelDict = { "search_method": RetrievalMethod.SEMANTIC_SEARCH, @@ -105,9 +106,11 @@ class DatasetMultiRetrieverTool(DatasetRetrieverBaseTool): ) for segment in sorted_segments: if segment.answer: - document_context_list.append(f"question:{segment.get_sign_content()} answer:{segment.answer}") + document_context_list.append( + f"question:{sign_segment_content(segment, session=session)} answer:{segment.answer}" + ) else: - document_context_list.append(segment.get_sign_content()) + document_context_list.append(sign_segment_content(segment, session=session)) if self.return_resource: context_list: list[RetrievalSourceMetadata] = [] resource_number = 1 diff --git a/api/core/tools/utils/dataset_retriever/dataset_retriever_tool.py b/api/core/tools/utils/dataset_retriever/dataset_retriever_tool.py index 489170053a8..ff0a1c16930 100644 --- a/api/core/tools/utils/dataset_retriever/dataset_retriever_tool.py +++ b/api/core/tools/utils/dataset_retriever/dataset_retriever_tool.py @@ -14,7 +14,8 @@ from core.rag.retrieval.retrieval_methods import RetrievalMethod from core.tools.utils.dataset_retriever.dataset_retriever_base_tool import DatasetRetrieverBaseTool from models.dataset import Dataset from models.dataset import Document as DatasetDocument -from services.external_knowledge_service import ExternalDatasetService +from repositories.knowledge.segment_read_adapter import sign_segment_content +from services.knowledge.external.service import ExternalDatasetService default_retrieval_model: DefaultRetrievalModelDict = { "search_method": RetrievalMethod.SEMANTIC_SEARCH, @@ -175,9 +176,11 @@ class DatasetRetrieverTool(DatasetRetrieverBaseTool): segment = record.segment # Build content: if summary exists, add it before the segment content if segment.answer: - segment_content = f"question:{segment.get_sign_content()} answer:{segment.answer}" + segment_content = ( + f"question:{sign_segment_content(segment, session=session)} answer:{segment.answer}" + ) else: - segment_content = segment.get_sign_content() + segment_content = sign_segment_content(segment, session=session) # If summary exists, prepend it to the content if record.summary: diff --git a/api/core/workflow/node_factory.py b/api/core/workflow/node_factory.py index 1d868c58f9c..7f67415fb05 100644 --- a/api/core/workflow/node_factory.py +++ b/api/core/workflow/node_factory.py @@ -11,7 +11,7 @@ from configs import dify_config from core.app.entities.app_invoke_entities import DIFY_RUN_CONTEXT_KEY, DifyRunContext from core.app.llm.model_access import build_dify_model_access, fetch_model_config from core.credit_usage import created_by_from_app_type -from core.db.session_factory import session_factory +from core.db.session_factory import get_session_maker, session_factory from core.file import remote_fetcher from core.helper.code_executor.code_executor import ( CodeExecutionError, @@ -55,6 +55,7 @@ from core.workflow.nodes.human_input.callback import DifyHITLCallback from core.workflow.nodes.human_input.entities import HumanInputNodeData as DifyHumanInputNodeData from core.workflow.system_variables import SystemVariableKey, get_system_text, system_variable_selector from core.workflow.template_rendering import CodeExecutorJinja2TemplateRenderer +from extensions.application_services.data_sources import build_data_source_credentials from graphon.entities.base_node_data import BaseNodeData from graphon.entities.graph_config import NodeConfigDict, NodeConfigDictAdapter from graphon.enums import BuiltinNodeTypes, NodeType @@ -431,6 +432,11 @@ class DifyNodeFactory(NodeFactory): resolved_node_data = self._resolve_llm_model_reference(cast(LLMNodeData, resolved_node_data)) node: Node | None = None node_init_kwargs_factories: Mapping[NodeType, Callable[[], dict[str, object]]] = { + BuiltinNodeTypes.DATASOURCE: lambda: { + "datasource_credentials": build_data_source_credentials( + database_client=get_session_maker() + ).providers.get_datasource_credentials, + }, BuiltinNodeTypes.CODE: lambda: { "code_executor": self._code_executor, "code_limits": self._code_limits, diff --git a/api/core/workflow/nodes/datasource/datasource_node.py b/api/core/workflow/nodes/datasource/datasource_node.py index f2ded99db3e..f90c06537aa 100644 --- a/api/core/workflow/nodes/datasource/datasource_node.py +++ b/api/core/workflow/nodes/datasource/datasource_node.py @@ -1,4 +1,4 @@ -from collections.abc import Generator, Mapping, Sequence +from collections.abc import Callable, Generator, Mapping, Sequence from typing import TYPE_CHECKING, Any, override from core.app.entities.app_invoke_entities import DIFY_RUN_CONTEXT_KEY, DifyRunContext @@ -40,6 +40,7 @@ class DatasourceNode(Node[DatasourceNodeData]): *, graph_init_params: "GraphInitParams", graph_runtime_state: "GraphRuntimeState", + datasource_credentials: Callable[[str, str, str, str | None], dict[str, Any]], ) -> None: super().__init__( node_id=node_id, @@ -48,6 +49,7 @@ class DatasourceNode(Node[DatasourceNodeData]): graph_runtime_state=graph_runtime_state, ) self.datasource_manager = DatasourceManager + self._datasource_credentials = datasource_credentials @override def populate_start_event(self, event) -> None: @@ -117,9 +119,9 @@ class DatasourceNode(Node[DatasourceNodeData]): datasource_type=datasource_type.value, provider_id=provider_id, tenant_id=dify_ctx.tenant_id, - provider=node_data.provider_name, - plugin_id=node_data.plugin_id, - credential_id=credential_id, + credentials=self._datasource_credentials( + dify_ctx.tenant_id, node_data.provider_name, node_data.plugin_id, credential_id + ), parameters_for_log=parameters_for_log, datasource_info=datasource_info, variable_pool=variable_pool, diff --git a/api/core/workflow/nodes/datasource/protocols.py b/api/core/workflow/nodes/datasource/protocols.py deleted file mode 100644 index 776e2673172..00000000000 --- a/api/core/workflow/nodes/datasource/protocols.py +++ /dev/null @@ -1,35 +0,0 @@ -from collections.abc import Generator -from typing import Any, Protocol - -from graphon.file import File -from graphon.node_events import StreamChunkEvent, StreamCompletedEvent - -from .entities import DatasourceParameter, OnlineDriveDownloadFileParam - - -class DatasourceManagerProtocol(Protocol): - @classmethod - def get_icon_url(cls, provider_id: str, tenant_id: str, datasource_name: str, datasource_type: str) -> str: ... - - @classmethod - def stream_node_events( - cls, - *, - node_id: str, - user_id: str, - datasource_name: str, - datasource_type: str, - provider_id: str, - tenant_id: str, - provider: str, - plugin_id: str, - credential_id: str, - parameters_for_log: dict[str, Any], - datasource_info: dict[str, Any], - variable_pool: Any, - datasource_param: DatasourceParameter | None = None, - online_drive_request: OnlineDriveDownloadFileParam | None = None, - ) -> Generator[StreamChunkEvent | StreamCompletedEvent, None, None]: ... - - @classmethod - def get_upload_file_by_id(cls, file_id: str, tenant_id: str) -> File: ... diff --git a/api/events/event_handlers/clean_when_dataset_deleted.py b/api/events/event_handlers/clean_when_dataset_deleted.py index ce0cf66f8f1..209317b7f78 100644 --- a/api/events/event_handlers/clean_when_dataset_deleted.py +++ b/api/events/event_handlers/clean_when_dataset_deleted.py @@ -1,13 +1,14 @@ from events.dataset_event import dataset_was_deleted from extensions.ext_database import db from models import Dataset +from repositories.knowledge.dataset_read_repository import get_dataset_doc_form from tasks.clean_dataset_task import clean_dataset_task @dataset_was_deleted.connect def handle(sender: Dataset, **kwargs): dataset = sender - doc_form = dataset.get_doc_form(session=db.session()) + doc_form = get_dataset_doc_form(dataset, session=db.session()) if not doc_form or not dataset.indexing_technique: return clean_dataset_task.delay( diff --git a/api/events/event_handlers/create_document_index.py b/api/events/event_handlers/create_document_index.py index aba387dd187..427a340baa2 100644 --- a/api/events/event_handlers/create_document_index.py +++ b/api/events/event_handlers/create_document_index.py @@ -6,11 +6,13 @@ from sqlalchemy import select from werkzeug.exceptions import NotFound from core.db.session_factory import session_factory -from core.indexing_runner import DocumentIsPausedError, IndexingRunner from events.document_index_event import document_index_created from libs.datetime_utils import naive_utc_now from models.dataset import Document from models.enums import IndexingStatus +from services.knowledge.indexing.adapters.execution import build_document_indexing_service +from services.knowledge.indexing.errors import DocumentIsPausedError +from services.knowledge.resource_scope import DatasetRef logger = logging.getLogger(__name__) @@ -21,7 +23,9 @@ def handle(sender, **kwargs): document_ids = kwargs.get("document_ids", []) start_at = time.perf_counter() try: - indexing_runner = IndexingRunner(enforce_vector_space_admission=True) + indexing_service = build_document_indexing_service( + session_factory=session_factory.get_session_maker(), enforce_vector_space_admission=True + ) with session_factory.create_session() as session: documents = [] for document_id in document_ids: @@ -41,11 +45,10 @@ def handle(sender, **kwargs): document.processing_started_at = naive_utc_now() documents.append(document) session.add(document) + document_refs = [DatasetRef(doc.tenant_id, doc.dataset_id).document(doc.id) for doc in documents] # Persist the status transition before extraction and indexing begin. session.commit() - - indexing_runner.run(documents, session) - session.commit() + indexing_service.run(document_refs) end_at = time.perf_counter() logger.info(click.style(f"Processed dataset: {dataset_id} latency: {end_at - start_at}", fg="green")) except DocumentIsPausedError as ex: diff --git a/api/extensions/application_services/data_sources.py b/api/extensions/application_services/data_sources.py new file mode 100644 index 00000000000..0818c635e47 --- /dev/null +++ b/api/extensions/application_services/data_sources.py @@ -0,0 +1,124 @@ +"""Composition root for data-source services and shared credential gateways.""" + +from collections.abc import Mapping +from dataclasses import dataclass + +from sqlalchemy.orm import Session, sessionmaker + +from configs import dify_config +from core.datasource.datasource_manager import DatasourceManager +from core.helper.ssrf_proxy import ssrf_proxy +from core.plugin.impl.datasource import PluginDatasourceManager +from core.plugin.impl.oauth import OAuthHandler +from repositories.data_source.api_key_auth_repository import SQLAlchemyDataSourceApiKeyAuthBindingRepository +from repositories.data_source.credential_repository import SQLAlchemyDatasourceCredentialRepository +from repositories.data_source.oauth_binding_repository import SQLAlchemyDataSourceOAuthBindingRepository +from repositories.knowledge.dataset_repository import SQLAlchemyDatasetRepository +from repositories.knowledge.document_repository import SQLAlchemyDocumentRepository +from services.data_source.auth.api_key_gateways import ( + ProviderApiKeyAuthCredentialValidator, + TenantApiKeyAuthCredentialEncryptor, +) +from services.data_source.auth.api_key_service import DataSourceApiKeyAuthService +from services.data_source.binding_application_service import DataSourceBindingApplicationService +from services.data_source.credential_adapters import ( + OAuthDatasourceCredentialRefresher, + PluginDatasourceCredentialCodec, + PluginDatasourceOAuthClientResolver, +) +from services.data_source.credential_gateway import ( + ActorAwareDatasourceCredentialGateway, + TrustedStoredDatasourceCredentialGateway, +) +from services.data_source.notion_import_adapters import PluginNotionSourceGateway +from services.data_source.notion_import_application_service import NotionImportApplicationService +from services.data_source.notion_oauth_gateway import NotionDataSourceGateway +from services.data_source.oauth_service import DataSourceOAuthService, InvalidDataSourceOAuthProviderError +from services.data_source.provider_service import DatasourceProviderService +from services.data_source.website_service import WebsiteService +from services.knowledge.dataset_access import DatasetAccessService + + +@dataclass(frozen=True, slots=True) +class DataSourceServices: + providers: DatasourceProviderService + api_key_auth: DataSourceApiKeyAuthService + oauth: Mapping[str, DataSourceOAuthService] + bindings: DataSourceBindingApplicationService + notion_imports: NotionImportApplicationService + + def resolve_oauth(self, provider: str) -> DataSourceOAuthService: + service = self.oauth.get(provider) + if service is None: + raise InvalidDataSourceOAuthProviderError("Invalid provider") + return service + + +@dataclass(frozen=True, slots=True) +class DataSourceCredentials: + providers: DatasourceProviderService + actor: ActorAwareDatasourceCredentialGateway + stored: TrustedStoredDatasourceCredentialGateway + + +def build_data_source_credentials(*, database_client: sessionmaker[Session]) -> DataSourceCredentials: + credentials = SQLAlchemyDatasourceCredentialRepository(session_factory=database_client) + provider_manager = PluginDatasourceManager() + codec = PluginDatasourceCredentialCodec(provider_manager=provider_manager) + refresher = OAuthDatasourceCredentialRefresher( + oauth_clients=PluginDatasourceOAuthClientResolver(configs=credentials, provider_manager=provider_manager), + oauth_handler=OAuthHandler(), + ) + return DataSourceCredentials( + providers=DatasourceProviderService(credentials=credentials), + actor=ActorAwareDatasourceCredentialGateway(credentials=credentials, codec=codec, refresher=refresher), + stored=TrustedStoredDatasourceCredentialGateway(credentials=credentials, codec=codec, refresher=refresher), + ) + + +def build_data_source_services( + *, + database_client: sessionmaker[Session], + dataset_access: DatasetAccessService, + datasets: SQLAlchemyDatasetRepository, + documents: SQLAlchemyDocumentRepository, + actor_credentials: ActorAwareDatasourceCredentialGateway, + providers: DatasourceProviderService, +) -> DataSourceServices: + oauth_bindings = SQLAlchemyDataSourceOAuthBindingRepository(session_factory=database_client) + notion_gateway = NotionDataSourceGateway( + client_id=dify_config.NOTION_CLIENT_ID or "", + client_secret=dify_config.NOTION_CLIENT_SECRET or "", + redirect_uri=dify_config.CONSOLE_API_URL + "/console/api/oauth/data-source/callback/notion", + http_client=ssrf_proxy, + ) + return DataSourceServices( + providers=providers, + api_key_auth=DataSourceApiKeyAuthService( + bindings=SQLAlchemyDataSourceApiKeyAuthBindingRepository(session_factory=database_client), + validator=ProviderApiKeyAuthCredentialValidator(), + encryptor=TenantApiKeyAuthCredentialEncryptor(), + ), + oauth={ + "notion": DataSourceOAuthService( + provider_name="notion", + provider_gateway=notion_gateway, + bindings=oauth_bindings, + is_internal_provider=dify_config.NOTION_INTEGRATION_TYPE == "internal", + internal_access_token=dify_config.NOTION_INTERNAL_SECRET, + ) + }, + bindings=DataSourceBindingApplicationService(bindings=oauth_bindings), + notion_imports=NotionImportApplicationService( + dataset_access=dataset_access, + datasets=datasets, + documents=documents, + source=PluginNotionSourceGateway( + credentials=actor_credentials, runtime_loader=DatasourceManager.get_datasource_runtime + ), + ), + ) + + +def build_website_service(*, database_client: sessionmaker[Session]) -> WebsiteService: + return WebsiteService(providers=build_data_source_credentials(database_client=database_client).providers) diff --git a/api/extensions/application_services/datasets.py b/api/extensions/application_services/datasets.py new file mode 100644 index 00000000000..053282d6beb --- /dev/null +++ b/api/extensions/application_services/datasets.py @@ -0,0 +1,32 @@ +"""Shared dataset dependencies for knowledge and data-source use cases.""" + +from dataclasses import dataclass + +from sqlalchemy.orm import Session, sessionmaker + +from configs import dify_config +from repositories.knowledge.dataset_repository import SQLAlchemyDatasetRepository +from repositories.knowledge.document_repository import SQLAlchemyDocumentRepository +from services.knowledge.dataset_access import DatasetAccessService, WorkspaceRoleReader + + +@dataclass(frozen=True, slots=True) +class DatasetDependencies: + datasets: SQLAlchemyDatasetRepository + documents: SQLAlchemyDocumentRepository + access: DatasetAccessService + + +def build_dataset_dependencies( + *, database_client: sessionmaker[Session], workspace_roles: WorkspaceRoleReader +) -> DatasetDependencies: + datasets = SQLAlchemyDatasetRepository(session_factory=database_client) + return DatasetDependencies( + datasets=datasets, + documents=SQLAlchemyDocumentRepository(session_factory=database_client), + access=DatasetAccessService( + datasets=datasets, + workspace_roles=workspace_roles, + legacy_permissions_enabled=not dify_config.RBAC_ENABLED, + ), + ) diff --git a/api/extensions/application_services/knowledge.py b/api/extensions/application_services/knowledge.py index e3300814ae3..91733434af4 100644 --- a/api/extensions/application_services/knowledge.py +++ b/api/extensions/application_services/knowledge.py @@ -1,16 +1,174 @@ -"""Composition of knowledge base application services.""" +"""Composition root for dataset-controller application services.""" + +from dataclasses import dataclass +from uuid import uuid4 from sqlalchemy.orm import Session, sessionmaker from configs import dify_config +from core.app.apps.pipeline.pipeline_generator import PipelineGenerator +from core.rag.extractor.entity.datasource_type import DatasourceType +from libs.helper import generate_text_hash from repositories.knowledge.dataset_api_key_repository import DatasetApiKeyRepository +from repositories.knowledge.dataset_repository import SQLAlchemyDatasetRepository +from repositories.knowledge.document_repository import SQLAlchemyDocumentRepository +from repositories.knowledge.metadata_repository import SQLAlchemyMetadataRepository +from repositories.knowledge.segment_repository import SQLAlchemySegmentRepository +from repositories.knowledge.upload_file_repository import SQLAlchemyKnowledgeUploadRepository from services.api_token_service import ApiTokenCache +from services.app.query_service import AppQueryService +from services.data_source.credential_gateway import ( + ActorAwareDatasourceCredentialGateway, + TrustedStoredDatasourceCredentialGateway, +) +from services.data_source.provider_service import DatasourceProviderService from services.knowledge.api_key_service import DatasetApiKeyService +from services.knowledge.dataset_access import DatasetAccessService +from services.knowledge.dataset_service import DocumentService +from services.knowledge.datasets.adapters import SQLAlchemyDatasetOperations +from services.knowledge.datasets.application import DatasetApplicationService +from services.knowledge.document_sync import DocumentSyncApplicationService +from services.knowledge.document_sync_adapters import CeleryDocumentSyncDispatcher +from services.knowledge.documents.adapters import SQLAlchemyDocumentOperations +from services.knowledge.documents.application import DatasetDocumentApplicationService +from services.knowledge.external.adapters import SQLAlchemyExternalKnowledgeOperations +from services.knowledge.external.application import ExternalKnowledgeApplicationService +from services.knowledge.indexing.adapters.estimate import IndexingEstimateAdapter, SQLAlchemyProcessRuleReader +from services.knowledge.indexing.adapters.sources import ( + CompositeStoredSourceResolver, + FileSourceAdapter, + NotionSourceResolver, + WebsiteSourceAdapter, +) +from services.knowledge.indexing.estimate import IndexingEstimateApplicationService +from services.knowledge.metadata.application import MetadataService +from services.knowledge.segments.adapters import ( + CelerySegmentBatchImportDispatcher, + ModelManagerSegmentGuard, + RedisSegmentClient, + RedisSegmentIndexingState, +) +from services.knowledge.segments.application import DatasetSegmentApplicationService +from services.knowledge.segments.indexing import SegmentIndexingGateway +from services.tag_application_service import TagTargetQuery +from tasks.batch_create_segment_to_index_task import batch_create_segment_to_index_task +from tasks.delete_segment_from_index_task import delete_segment_from_index_task +from tasks.disable_segments_from_index_task import disable_segments_from_index_task +from tasks.document_indexing_sync_task import document_indexing_sync_task +from tasks.enable_segments_to_index_task import enable_segments_to_index_task -def build_dataset_api_key_service(*, database_client: sessionmaker[Session]) -> DatasetApiKeyService: +@dataclass(frozen=True, slots=True) +class KnowledgeServices: + metadata: MetadataService + datasets: DatasetApplicationService + external: ExternalKnowledgeApplicationService + documents: DatasetDocumentApplicationService + document_sync: DocumentSyncApplicationService + indexing_estimates: IndexingEstimateApplicationService + segments: DatasetSegmentApplicationService + pipeline_generator: PipelineGenerator + + +def build_dataset_api_key_service( + *, database_client: sessionmaker[Session], dataset_access: DatasetAccessService +) -> DatasetApiKeyService: return DatasetApiKeyService( keys=DatasetApiKeyRepository(session_factory=database_client), cache=ApiTokenCache, + access=dataset_access, rbac_enabled=lambda: dify_config.RBAC_ENABLED, ) + + +def build_knowledge_services( + *, + database_client: sessionmaker[Session], + dataset_access: DatasetAccessService, + datasets: SQLAlchemyDatasetRepository, + documents: SQLAlchemyDocumentRepository, + actor_credentials: ActorAwareDatasourceCredentialGateway, + stored_credentials: TrustedStoredDatasourceCredentialGateway, + providers: DatasourceProviderService, + redis: RedisSegmentClient, + tags: TagTargetQuery, + app_queries: AppQueryService, +) -> KnowledgeServices: + """Build the dataset-controller knowledge use cases.""" + + uploads = SQLAlchemyKnowledgeUploadRepository(session_factory=database_client) + notion_sources = NotionSourceResolver( + actor_credentials=actor_credentials, + stored_credentials=stored_credentials, + ) + file_sources = FileSourceAdapter(uploads=uploads) + website_sources = WebsiteSourceAdapter() + stored_sources = CompositeStoredSourceResolver( + adapters={ + DatasourceType.FILE.value: file_sources, + DatasourceType.NOTION.value: notion_sources, + DatasourceType.WEBSITE.value: website_sources, + } + ) + segments = SQLAlchemySegmentRepository(session_factory=database_client) + segment_index = SegmentIndexingGateway( + segments=segments, + uploads=uploads, + redis=redis, + new_session=database_client, + delete_task=delete_segment_from_index_task.delay, + enable_task=enable_segments_to_index_task.delay, + disable_task=disable_segments_from_index_task.delay, + ) + return KnowledgeServices( + metadata=MetadataService( + store=SQLAlchemyMetadataRepository(session_factory=database_client), dataset_access=dataset_access + ), + datasets=DatasetApplicationService( + dataset_access=dataset_access, + operations=SQLAlchemyDatasetOperations(session_factory=database_client, tags=tags, app_queries=app_queries), + rbac_enabled=dify_config.RBAC_ENABLED, + service_api_url=dify_config.SERVICE_API_URL, + vector_store=dify_config.VECTOR_STORE, + tidb_fulltext=dify_config.TIDB_VECTOR_ENABLE_FULLTEXT_SEARCH, + ), + external=ExternalKnowledgeApplicationService( + dataset_access=dataset_access, + operations=SQLAlchemyExternalKnowledgeOperations(session_factory=database_client), + ), + documents=DatasetDocumentApplicationService( + dataset_access=dataset_access, + operations=SQLAlchemyDocumentOperations(session_factory=database_client), + metadata_schema=DocumentService.DOCUMENT_METADATA_SCHEMA, + ), + pipeline_generator=PipelineGenerator(documents=documents, datasource_providers=providers), + document_sync=DocumentSyncApplicationService( + dataset_access=dataset_access, + documents=documents, + dispatcher=CeleryDocumentSyncDispatcher(delay=document_indexing_sync_task.delay), + ), + indexing_estimates=IndexingEstimateApplicationService( + dataset_access=dataset_access, + datasets=datasets, + documents=documents, + files=file_sources, + websites=website_sources, + stored_sources=stored_sources, + notion=notion_sources, + process_rules=SQLAlchemyProcessRuleReader(session_factory=database_client), + runner=IndexingEstimateAdapter(session_factory=database_client), + ), + segments=DatasetSegmentApplicationService( + dataset_access=dataset_access, + scopes=datasets, + store=segments, + index=segment_index, + limits=dify_config, + text_hash=generate_text_hash, + uploads=uploads, + model_guard=ModelManagerSegmentGuard(), + indexing_state=RedisSegmentIndexingState(redis), + batch_dispatcher=CelerySegmentBatchImportDispatcher(delay=batch_create_segment_to_index_task.delay), + job_id_factory=lambda: str(uuid4()), + ), + ) diff --git a/api/extensions/ext_application_services.py b/api/extensions/ext_application_services.py index 6c5c343cd0c..0aa35bf44b1 100644 --- a/api/extensions/ext_application_services.py +++ b/api/extensions/ext_application_services.py @@ -19,14 +19,23 @@ from configs import dify_config from constants.dsl_version import CURRENT_APP_DSL_VERSION from constants.languages import languages from core.db.session_factory import get_session_maker -from core.helper.ssrf_proxy import ssrf_proxy from core.schemas.schema_manager import SchemaManager from core.tools.tool_file_manager import ToolFileManager from enums import DeploymentEdition, WebAppAccessMode from extensions.application_services.agent import AgentAppServices, build_agent_app_services from extensions.application_services.app import AppServices, build_app_api_key_service, build_app_services +from extensions.application_services.data_sources import ( + DataSourceServices, + build_data_source_credentials, + build_data_source_services, +) +from extensions.application_services.datasets import build_dataset_dependencies from extensions.application_services.installed_app import InstalledAppServices, build_installed_app_services -from extensions.application_services.knowledge import build_dataset_api_key_service +from extensions.application_services.knowledge import ( + KnowledgeServices, + build_dataset_api_key_service, + build_knowledge_services, +) from extensions.application_services.trial_app import TrialAppServices, build_trial_app_services from extensions.ext_redis import RedisClientWrapper, redis_client from extensions.ext_storage import storage @@ -52,8 +61,7 @@ from repositories.app_definition_query_repository import AppDefinitionQueryRepos from repositories.app_preview_query_repository import AppPreviewQueryRepository from repositories.app_scoped_end_user_repository import AppScopedEndUserRepo from repositories.app_statistic_query_repository import AppStatisticQueryRepository -from repositories.data_source_api_key_auth_repository import SQLAlchemyDataSourceApiKeyAuthBindingRepository -from repositories.data_source_oauth_binding_repository import SQLAlchemyDataSourceOAuthBindingRepository +from repositories.credentials.query_repository import CredentialQueryRepository from repositories.explore_banner_query_repository import ExploreBannerQueryRepository from repositories.factory import DifyAPIRepositoryFactory from repositories.file_grant_repository import FileGrantRepository @@ -164,15 +172,10 @@ from services.app_statistic_query import AppStatisticQuery from services.app_task_service import AppTaskControlService from services.app_tracing_config_gateway import OpsTraceManagerGateway from services.app_tracing_config_service import AppTracingConfigService -from services.auth.data_source_api_key_auth_gateways import ( - ProviderApiKeyAuthCredentialValidator, - TenantApiKeyAuthCredentialEncryptor, -) -from services.auth.data_source_api_key_auth_service import DataSourceApiKeyAuthService from services.billing_portal_service import BillingPortalService from services.billing_service import BillingService from services.compliance_download_service import ComplianceDownloadService -from services.data_source_oauth_service import DataSourceOAuthService, InvalidDataSourceOAuthProviderError +from services.credentials.query import CredentialQuery from services.enterprise.enterprise_service import EnterpriseService from services.entities.file_grant_entities import FileGrantLimits from services.errors.enterprise import EnterpriseServiceError @@ -191,7 +194,6 @@ from services.message_suggested_questions_adapters import MessageSuggestedQuesti from services.message_suggested_questions_service import MessageSuggestedQuestions from services.notification_gateway import BillingNotificationGateway from services.notification_service import NotificationService -from services.notion_data_source_gateway import NotionDataSourceGateway from services.oauth_device_adapters import ( DifyConfigOAuthDeviceSettings, EnterpriseOAuthDeviceSSOGateway, @@ -337,6 +339,7 @@ class AppScopedEndUserServices: class ApplicationServices: agent_apps: AgentAppServices advanced_prompt_templates: AdvancedPromptTemplateService + credential_queries: CredentialQuery accounts: AccountServices app_api_keys: AppApiKeyService dataset_api_keys: DatasetApiKeyService @@ -351,8 +354,8 @@ class ApplicationServices: app_tracing_configs: AppTracingConfigService billing_portal: BillingPortalService compliance_downloads: ComplianceDownloadService - data_source_api_key_auth: DataSourceApiKeyAuthService - data_source_oauth: Mapping[str, DataSourceOAuthService] + data_sources: DataSourceServices + knowledge: KnowledgeServices app_scoped_end_users: AppScopedEndUserServices webapp_access: WebAppAccessQueryService web_app_runtime: WebAppRuntimeQueryService @@ -392,34 +395,6 @@ class ApplicationServices: tags: TagApplicationService workflow_statistics: WorkflowStatisticQueryService - def resolve_data_source_oauth(self, provider: str) -> DataSourceOAuthService: - service = self.data_source_oauth.get(provider) - if service is None: - raise InvalidDataSourceOAuthProviderError("Invalid provider") - return service - - -def _build_data_source_oauth_services( - *, - database_client: sessionmaker[Session], -) -> Mapping[str, DataSourceOAuthService]: - notion_data_source = NotionDataSourceGateway( - client_id=dify_config.NOTION_CLIENT_ID or "", - client_secret=dify_config.NOTION_CLIENT_SECRET or "", - redirect_uri=dify_config.CONSOLE_API_URL + "/console/api/oauth/data-source/callback/notion", - http_client=ssrf_proxy, - ) - bindings = SQLAlchemyDataSourceOAuthBindingRepository(session_factory=database_client) - return { - "notion": DataSourceOAuthService( - provider_name="notion", - provider_gateway=notion_data_source, - bindings=bindings, - is_internal_provider=dify_config.NOTION_INTEGRATION_TYPE == "internal", - internal_access_token=dify_config.NOTION_INTERNAL_SECRET, - ) - } - def _build_oauth_server_service( *, @@ -547,7 +522,6 @@ def build_application_services( redis: RedisClientWrapper, ) -> ApplicationServices: installation_state = InstallationStateRepository(session_factory=database_client) - data_source_api_key_auth_bindings = SQLAlchemyDataSourceApiKeyAuthBindingRepository(session_factory=database_client) app_definition_repository = AppDefinitionQueryRepository(session_factory=database_client) app_definitions = AppDefinitionQueryService( definitions=app_definition_repository, @@ -587,6 +561,39 @@ def build_application_services( sources=database_catalog, exporter=RosterAgentPackageExporter() ) workspace_query_repository = WorkspaceQueryRepository(session_factory=database_client) + workspace_member_repository = WorkspaceMemberQueryRepository(session_factory=database_client) + dataset_dependencies = build_dataset_dependencies( + database_client=database_client, + workspace_roles=workspace_member_repository, + ) + datasource_credentials = build_data_source_credentials(database_client=database_client) + data_sources = build_data_source_services( + database_client=database_client, + dataset_access=dataset_dependencies.access, + datasets=dataset_dependencies.datasets, + documents=dataset_dependencies.documents, + actor_credentials=datasource_credentials.actor, + providers=datasource_credentials.providers, + ) + oauth_server = _build_oauth_server_service(database_client=database_client, redis=redis) + apps = build_app_services( + database_client=database_client, + oauth=oauth_server, + recommended_packages=recommended_app_packages, + ) + tags = TagApplicationService(tags=TagRepository(session_factory=database_client)) + knowledge = build_knowledge_services( + database_client=database_client, + dataset_access=dataset_dependencies.access, + datasets=dataset_dependencies.datasets, + documents=dataset_dependencies.documents, + actor_credentials=datasource_credentials.actor, + stored_credentials=datasource_credentials.stored, + providers=datasource_credentials.providers, + redis=redis, + tags=tags, + app_queries=apps.queries, + ) app_scoped_end_user_repository = AppScopedEndUserRepo(session_factory=database_client) file_service = FileService(session_factory=database_client) remote_file_service = RemoteFileService(files=file_service) @@ -598,8 +605,8 @@ def build_application_services( workflow_node_execution_repository = DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository( session_maker=database_client ) - oauth_server = _build_oauth_server_service(database_client=database_client, redis=redis) return ApplicationServices( + credential_queries=CredentialQueryRepository(session_factory=database_client), accounts=AccountServices( access=AccountAccessService( accounts=accounts, @@ -754,11 +761,7 @@ def build_application_services( enabled=dify_config.RBAC_ENABLED, ), ), - apps=build_app_services( - database_client=database_client, - oauth=oauth_server, - recommended_packages=recommended_app_packages, - ), + apps=apps, agent_apps=build_agent_app_services(database_client=database_client), advanced_prompt_templates=AdvancedPromptTemplateService(), app_definitions=app_definitions, @@ -774,7 +777,10 @@ def build_application_services( sites=AppSiteCommandRepository(session_factory=database_client), ), app_api_keys=build_app_api_key_service(database_client=database_client), - dataset_api_keys=build_dataset_api_key_service(database_client=database_client), + dataset_api_keys=build_dataset_api_key_service( + database_client=database_client, + dataset_access=dataset_dependencies.access, + ), app_statistics=AppStatisticQueryRepository(session_factory=database_client), app_tracing_configs=AppTracingConfigService( configs=SQLAlchemyAppTracingConfigRepository(session_factory=database_client), @@ -794,12 +800,8 @@ def build_application_services( redis_client=redis, ), ), - data_source_api_key_auth=DataSourceApiKeyAuthService( - bindings=data_source_api_key_auth_bindings, - validator=ProviderApiKeyAuthCredentialValidator(), - encryptor=TenantApiKeyAuthCredentialEncryptor(), - ), - data_source_oauth=_build_data_source_oauth_services(database_client=database_client), + data_sources=data_sources, + knowledge=knowledge, app_scoped_end_users=AppScopedEndUserServices( commands=AppScopedEndUserService(end_users=app_scoped_end_user_repository), queries=AppScopedEndUserQueryService(end_users=app_scoped_end_user_repository), @@ -905,9 +907,7 @@ def build_application_services( plans=DeploymentWorkspacePlanGateway(), ), workspace_member_queries=WorkspaceMemberQueryService( - members=WorkspaceMemberQueryRepository( - session_factory=database_client, - ), + members=workspace_member_repository, roles=DeploymentWorkspaceMemberRoleResolver(), ), workflow_app_logs=WorkflowAppLogQueryService( @@ -927,9 +927,7 @@ def build_application_services( now=lambda: datetime.now(UTC), access_token_expire_minutes=dify_config.ACCESS_TOKEN_EXPIRE_MINUTES, ), - tags=TagApplicationService( - tags=TagRepository(session_factory=database_client), - ), + tags=tags, workflow_statistics=WorkflowStatisticQueryService( workflow_runs=DifyAPIRepositoryFactory.create_api_workflow_run_repository( session_maker=database_client, diff --git a/api/extensions/ext_celery.py b/api/extensions/ext_celery.py index f4eff8a1012..a91009b4884 100644 --- a/api/extensions/ext_celery.py +++ b/api/extensions/ext_celery.py @@ -176,6 +176,8 @@ def init_app(app: DifyApp) -> Celery: "tasks.trigger_processing_tasks", # async trigger processing "tasks.generate_summary_index_task", # summary index generation "tasks.regenerate_summary_index_task", # summary index regeneration + "tasks.rag_pipeline.rag_pipeline_run_task", # regular RAG pipeline execution + "tasks.rag_pipeline.priority_rag_pipeline_run_task", # priority RAG pipeline execution "tasks.initialize_created_app_rbac_access_task", # app access initialization "tasks.install_default_plugins_task", # tenant default plugin installation "tasks.new_agent_beta_task", # New Agent Beta eligibility checks diff --git a/api/fields/dataset_fields.py b/api/fields/dataset_fields.py index d70cea6d694..e6902f1cab2 100644 --- a/api/fields/dataset_fields.py +++ b/api/fields/dataset_fields.py @@ -1,26 +1,21 @@ -from collections.abc import Sequence -from dataclasses import dataclass, field from datetime import datetime -from typing import Any -import sqlalchemy as sa -from pydantic import Field, field_validator -from sqlalchemy import func, select -from sqlalchemy.orm import Session +from pydantic import BaseModel, Field, field_validator from fields.base import ResponseModel from libs.helper import to_timestamp -from models.account import Account -from models.dataset import ( - AppDatasetJoin, - Dataset, - DatasetMetadata, - Document, - ExternalKnowledgeApis, - ExternalKnowledgeBindings, - Pipeline, -) -from models.model import App, Tag, TagBinding +from services.data_source.entities.notion_import import NotionPageType + + +class NotionEstimatePagePayload(BaseModel): + page_id: str = Field(min_length=1) + page_type: NotionPageType = Field(alias="type") + + +class NotionEstimateWorkspacePayload(BaseModel): + workspace_id: str = Field(min_length=1) + credential_id: str = Field(min_length=1) + pages: list[NotionEstimatePagePayload] = Field(min_length=1) class DatasetMetadataResponse(ResponseModel): @@ -206,270 +201,3 @@ class DatasetDetailResponse(ResponseModel): @classmethod def _expand_null_nested(cls, value: object) -> object: return {} if value is None else value - - -@dataclass(frozen=True) -class DatasetDetailPrefetch: - """Batch-loaded values backing the session-scoped fields of a page of datasets. - - Building a detail response for a single dataset issues about ten queries for - that row alone, so a list endpoint costs one query per dataset per field. - Loading the same values for every dataset on the page up front keeps the - query count bounded regardless of the page size. - """ - - app_counts: dict[str, int] = field(default_factory=dict) - document_counts: dict[str, int] = field(default_factory=dict) - word_counts: dict[str, int] = field(default_factory=dict) - available_document_counts: dict[str, int] = field(default_factory=dict) - author_names: dict[str, str] = field(default_factory=dict) - tags: dict[tuple[str, str], list[Tag]] = field(default_factory=dict) - doc_forms: dict[tuple[str, str], str | None] = field(default_factory=dict) - external_knowledge_infos: dict[tuple[str, str], dict[str, Any]] = field(default_factory=dict) - doc_metadata: dict[str, list[DatasetMetadata]] = field(default_factory=dict) - pipeline_published: dict[str, bool] = field(default_factory=dict) - - -def build_dataset_detail_prefetch(datasets: Sequence[Dataset], *, session: Session) -> DatasetDetailPrefetch: - """Load every session-scoped detail field for the given datasets in a fixed number of queries.""" - dataset_ids = [dataset.id for dataset in datasets] - if not dataset_ids: - return DatasetDetailPrefetch() - - tenant_ids = {dataset.tenant_id for dataset in datasets} - - app_counts = dict( - session.execute( - select(AppDatasetJoin.dataset_id, func.count(AppDatasetJoin.id)) - .where(AppDatasetJoin.dataset_id.in_(dataset_ids), App.id == AppDatasetJoin.app_id) - .group_by(AppDatasetJoin.dataset_id) - ) - .tuples() - .all() - ) - - document_counts: dict[str, int] = {} - word_counts: dict[str, int] = {} - available_document_counts: dict[str, int] = {} - available_document_case = sa.case( - ( - sa.and_( - Document.indexing_status == "completed", - Document.enabled == True, - Document.archived == False, - ), - 1, - ), - else_=0, - ) - document_rows = session.execute( - select( - Document.dataset_id, - func.count(Document.id), - func.coalesce(func.sum(Document.word_count), 0), - func.coalesce(func.sum(available_document_case), 0), - ) - .where(Document.dataset_id.in_(dataset_ids)) - .group_by(Document.dataset_id) - ).all() - for dataset_id, document_count, word_count, available_count in document_rows: - document_counts[dataset_id] = document_count - word_counts[dataset_id] = word_count - available_document_counts[dataset_id] = available_count - - # doc_form only falls back to the documents table for datasets without a - # chunk structure of their own. - doc_forms: dict[tuple[str, str], str | None] = {} - doc_form_dataset_ids = [dataset.id for dataset in datasets if not dataset.chunk_structure] - if doc_form_dataset_ids: - doc_forms = { - (dataset_id, tenant_id): doc_form - for dataset_id, tenant_id, doc_form in session.execute( - select(Document.dataset_id, Document.tenant_id, func.min(Document.doc_form)) - .where(Document.dataset_id.in_(doc_form_dataset_ids), Document.tenant_id.in_(tenant_ids)) - .group_by(Document.dataset_id, Document.tenant_id) - ).all() - } - - author_ids = {dataset.created_by for dataset in datasets if dataset.created_by} - author_names = dict( - session.execute(select(Account.id, Account.name).where(Account.id.in_(author_ids))).tuples().all() - ) - - tags: dict[tuple[str, str], list[Tag]] = {} - tag_rows = session.execute( - select(TagBinding.target_id, TagBinding.tenant_id, Tag) - .join(TagBinding, Tag.id == TagBinding.tag_id) - .where( - TagBinding.target_id.in_(dataset_ids), - TagBinding.tenant_id.in_(tenant_ids), - Tag.tenant_id == TagBinding.tenant_id, - Tag.type == "knowledge", - ) - ).all() - for target_id, tenant_id, tag in tag_rows: - tags.setdefault((target_id, tenant_id), []).append(tag) - - doc_metadata: dict[str, list[DatasetMetadata]] = {} - for dataset_metadata in session.scalars( - select(DatasetMetadata).where(DatasetMetadata.dataset_id.in_(dataset_ids)) - ).all(): - doc_metadata.setdefault(dataset_metadata.dataset_id, []).append(dataset_metadata) - - pipeline_ids = {dataset.pipeline_id for dataset in datasets if dataset.pipeline_id} - pipeline_published: dict[str, bool] = {} - if pipeline_ids: - pipeline_published = dict( - session.execute(select(Pipeline.id, Pipeline.is_published).where(Pipeline.id.in_(pipeline_ids))) - .tuples() - .all() - ) - - return DatasetDetailPrefetch( - app_counts=app_counts, - document_counts=document_counts, - word_counts=word_counts, - available_document_counts=available_document_counts, - author_names=author_names, - tags=tags, - doc_forms=doc_forms, - external_knowledge_infos=_load_external_knowledge_infos(datasets, session=session), - doc_metadata=doc_metadata, - pipeline_published=pipeline_published, - ) - - -def _load_external_knowledge_infos( - datasets: Sequence[Dataset], *, session: Session -) -> dict[tuple[str, str], dict[str, Any]]: - external_datasets = [dataset for dataset in datasets if dataset.provider == "external"] - if not external_datasets: - return {} - - external_dataset_ids = [dataset.id for dataset in external_datasets] - external_tenant_ids = {dataset.tenant_id for dataset in external_datasets} - bindings = session.scalars( - select(ExternalKnowledgeBindings).where( - ExternalKnowledgeBindings.dataset_id.in_(external_dataset_ids), - ExternalKnowledgeBindings.tenant_id.in_(external_tenant_ids), - ) - ).all() - if not bindings: - return {} - - apis = { - (api.id, api.tenant_id): api - for api in session.scalars( - select(ExternalKnowledgeApis).where( - ExternalKnowledgeApis.id.in_({binding.external_knowledge_api_id for binding in bindings}), - ExternalKnowledgeApis.tenant_id.in_(external_tenant_ids), - ) - ).all() - } - - infos: dict[tuple[str, str], dict[str, Any]] = {} - for binding in bindings: - info = Dataset.build_external_knowledge_info( - binding, apis.get((binding.external_knowledge_api_id, binding.tenant_id)) - ) - if info is not None: - infos[(binding.dataset_id, binding.tenant_id)] = info - return infos - - -@dataclass(frozen=True) -class DatasetDetailResponseSource: - """Expose session-backed dataset fields during response validation. - - With a prefetch the values come from the batch load for the whole page; - without one each field is resolved against the session on access. - """ - - dataset: Any - session: Session - prefetch: DatasetDetailPrefetch | None = None - - @property - def _prefetch_key(self) -> tuple[str, str]: - return (self.dataset.id, self.dataset.tenant_id) - - @property - def app_count(self) -> int: - if self.prefetch is not None: - return self.prefetch.app_counts.get(self.dataset.id, 0) - return self.dataset.get_app_count(session=self.session) - - @property - def document_count(self) -> int: - if self.prefetch is not None: - return self.prefetch.document_counts.get(self.dataset.id, 0) - return self.dataset.get_document_count(session=self.session) - - @property - def word_count(self) -> int: - if self.prefetch is not None: - return self.prefetch.word_counts.get(self.dataset.id, 0) - return self.dataset.get_word_count(session=self.session) - - @property - def author_name(self) -> str | None: - if self.prefetch is not None: - return self.prefetch.author_names.get(self.dataset.created_by) - return self.dataset.get_author_name(session=self.session) - - @property - def tags(self) -> Any: - if self.prefetch is not None: - return self.prefetch.tags.get(self._prefetch_key, []) - return self.dataset.get_tags(session=self.session) - - @property - def doc_form(self) -> str | None: - if self.prefetch is not None: - if self.dataset.chunk_structure: - return self.dataset.chunk_structure - return self.prefetch.doc_forms.get(self._prefetch_key) - return self.dataset.get_doc_form(session=self.session) - - @property - def external_knowledge_info(self) -> Any: - if self.prefetch is not None: - if self.dataset.provider != "external": - return None - return self.prefetch.external_knowledge_infos.get(self._prefetch_key) - return self.dataset.get_external_knowledge_info(session=self.session) - - @property - def doc_metadata(self) -> Any: - if self.prefetch is not None: - return self.dataset.build_doc_metadata(self.prefetch.doc_metadata.get(self.dataset.id, [])) - return self.dataset.get_doc_metadata(session=self.session) - - @property - def is_published(self) -> bool: - if self.prefetch is not None: - if not self.dataset.pipeline_id: - return False - return self.prefetch.pipeline_published.get(self.dataset.pipeline_id, False) - return self.dataset.get_is_published(session=self.session) - - @property - def total_documents(self) -> int: - if self.prefetch is not None: - return self.prefetch.document_counts.get(self.dataset.id, 0) - return self.dataset.get_total_documents(session=self.session) - - @property - def total_available_documents(self) -> int: - if self.prefetch is not None: - return self.prefetch.available_document_counts.get(self.dataset.id, 0) - return self.dataset.get_total_available_documents(session=self.session) - - def __getattr__(self, name: str) -> Any: - return getattr(self.dataset, name) # guard-ignore: no-new-getattr -- delegates model fields - - -def dataset_detail_response_source( - dataset: Any, *, session: Session, prefetch: DatasetDetailPrefetch | None = None -) -> DatasetDetailResponseSource: - return DatasetDetailResponseSource(dataset=dataset, session=session, prefetch=prefetch) diff --git a/api/fields/document_fields.py b/api/fields/document_fields.py index 06a69422cf6..d0122c36e08 100644 --- a/api/fields/document_fields.py +++ b/api/fields/document_fields.py @@ -1,16 +1,12 @@ """Response schemas for dataset document endpoints.""" -from collections.abc import Iterable -from dataclasses import dataclass from datetime import datetime from typing import Any from pydantic import Field, field_validator -from sqlalchemy.orm import Session from fields.base import ResponseModel from libs.helper import to_timestamp -from models.dataset import DocMetadataDetailItem, Document def normalize_enum(value: Any) -> Any: @@ -70,39 +66,6 @@ class DocumentResponse(ResponseModel): return to_timestamp(value) -@dataclass(frozen=True) -class DocumentWithSession: - """Expose session-backed document fields during response validation.""" - - document: Document - session: Session - - @property - def data_source_detail_dict(self) -> dict[str, Any]: - return self.document.get_data_source_detail_dict(session=self.session) - - @property - def hit_count(self) -> int: - return self.document.get_hit_count(session=self.session) - - @property - def doc_metadata_details(self) -> list[DocMetadataDetailItem] | None: - return self.document.get_doc_metadata_details(session=self.session) - - def __getattr__(self, name: str) -> Any: - return getattr(self.document, name) # guard-ignore: no-new-getattr -- delegates model fields - - -def document_response(document: Document, *, session: Session) -> DocumentResponse: - return DocumentResponse.model_validate( - DocumentWithSession(document=document, session=session), from_attributes=True - ) - - -def document_responses(documents: Iterable[Document], *, session: Session) -> list[DocumentResponse]: - return [document_response(document, session=session) for document in documents] - - class DocumentListResponse(ResponseModel): data: list[DocumentResponse] has_more: bool diff --git a/api/fields/segment_fields.py b/api/fields/segment_fields.py index ecee5335e53..c4f73c3c08c 100644 --- a/api/fields/segment_fields.py +++ b/api/fields/segment_fields.py @@ -1,10 +1,6 @@ -from collections.abc import Iterable, Mapping -from dataclasses import dataclass from datetime import datetime -from typing import Any from pydantic import field_serializer -from sqlalchemy.orm import Session from fields.base import ResponseModel from libs.helper import to_timestamp @@ -72,40 +68,6 @@ class SegmentResponse(ResponseModel): return to_timestamp(value) -@dataclass(frozen=True) -class SegmentWithSummary: - """Expose session-backed segment relations during response validation.""" - - segment: Any - summary: str | None - session: Session - - @property - def child_chunks(self) -> Any: - return self.segment.get_child_chunks(session=self.session, include_full_doc=False) - - @property - def attachments(self) -> Any: - return self.segment.get_attachments(session=self.session) - - def __getattr__(self, name: str) -> Any: - return getattr(self.segment, name) - - -def segment_response_with_summary(segment: Any, summary: str | None, *, session: Session) -> SegmentResponse: - response_source = SegmentWithSummary(segment=segment, summary=summary, session=session) - return SegmentResponse.model_validate(response_source, from_attributes=True) - - -def segment_responses_with_summaries( - segments: Iterable[Any], - summaries: Mapping[str, str | None], - *, - session: Session, -) -> list[SegmentResponse]: - return [segment_response_with_summary(segment, summaries.get(segment.id), session=session) for segment in segments] - - class SegmentDetailResponse(ResponseModel): data: SegmentResponse doc_form: str diff --git a/api/models/dataset.py b/api/models/dataset.py index 796e3d50ab6..093606796e8 100644 --- a/api/models/dataset.py +++ b/api/models/dataset.py @@ -1,12 +1,5 @@ -import base64 -import hashlib -import hmac import json -import logging -import os import pickle -import re -import time from collections.abc import Sequence from datetime import datetime from json import JSONDecodeError @@ -14,22 +7,17 @@ from typing import Any, ClassVar, TypedDict, cast, override from uuid import uuid4 import sqlalchemy as sa -from sqlalchemy import DateTime, String, func, select -from sqlalchemy.orm import Mapped, Session, mapped_column, scoped_session +from sqlalchemy import DateTime, String, func +from sqlalchemy.orm import Mapped, mapped_column from configs import dify_config -from core.rag.entities import ParentMode, Rule -from core.rag.index_processor.constant.built_in_field import BuiltInField, MetadataDataSource +from core.rag.entities import PreProcessingRuleKey +from core.rag.index_processor.constant.built_in_field import BuiltInField from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType -from core.rag.index_processor.constant.query_type import QueryType from core.rag.retrieval.retrieval_methods import RetrievalMethod -from core.tools.signature import sign_upload_file_preview_url -from extensions.ext_storage import storage from libs.uuid_utils import uuidv7 - -from .account import Account -from .base import Base, TypeBase -from .enums import ( +from models.base import Base, TypeBase +from models.enums import ( CollectionBindingType, CreatorUserRole, DatasetMetadataType, @@ -45,14 +33,11 @@ from .enums import ( SummaryStatus, TidbAuthBindingStatus, ) -from .model import App, Tag, TagBinding, UploadFile -from .types import AdjustedJSON, BinaryData, EnumText, LongText, StringUUID, adjusted_json_index - -logger = logging.getLogger(__name__) +from models.types import AdjustedJSON, BinaryData, EnumText, LongText, StringUUID, adjusted_json_index class PreProcessingRuleItem(TypedDict): - id: str + id: PreProcessingRuleKey enabled: bool @@ -95,6 +80,67 @@ class DatasetBindingItem(TypedDict): name: str +class ExternalKnowledgeApiDict(TypedDict): + id: str + tenant_id: str + name: str + description: str + settings: dict[str, Any] | None + dataset_bindings: list[DatasetBindingItem] + created_by: str + created_at: str + + +class DocumentDict(TypedDict): + id: str + tenant_id: str + dataset_id: str + position: int + data_source_type: str + data_source_info: str | None + dataset_process_rule_id: str | None + batch: str + name: str + created_from: str + created_by: str + created_api_request_id: str | None + created_at: datetime + processing_started_at: datetime | None + file_id: str | None + word_count: int | None + parsing_completed_at: datetime | None + cleaning_completed_at: datetime | None + splitting_completed_at: datetime | None + tokens: int | None + indexing_latency: float | None + completed_at: datetime | None + is_paused: bool | None + paused_by: str | None + paused_at: datetime | None + error: str | None + stopped_at: datetime | None + indexing_status: str + enabled: bool + disabled_at: datetime | None + disabled_by: str | None + archived: bool + archived_reason: str | None + archived_by: str | None + archived_at: datetime | None + updated_at: datetime + doc_type: str | None + doc_metadata: Any + doc_form: IndexStructureType + doc_language: str | None + display_status: str | None + data_source_info_dict: dict[str, Any] + average_segment_length: int + dataset_process_rule: ProcessRuleDict | None + dataset: None + segment_count: int | None + hit_count: int | None + + from models.enums import PermissionEnum # Backward-compatible alias — new code should import PermissionEnum from models.enums @@ -149,25 +195,6 @@ class Dataset(Base): enable_api = mapped_column(sa.Boolean, nullable=False, server_default=sa.true()) is_multimodal = mapped_column(sa.Boolean, default=False, nullable=False) - def get_total_documents(self, *, session: Session) -> int: - return self.get_document_count(session=session) - - def get_total_available_documents(self, *, session: Session) -> int: - return ( - session.scalar( - select(func.count(Document.id)).where( - Document.dataset_id == self.id, - Document.indexing_status == "completed", - Document.enabled == True, - Document.archived == False, - ) - ) - or 0 - ) - - def get_dataset_keyword_table(self, *, session: Session) -> "DatasetKeywordTable | None": - return session.scalar(select(DatasetKeywordTable).where(DatasetKeywordTable.dataset_id == self.id)) - @property def index_struct_dict(self): return json.loads(self.index_struct) if self.index_struct else None @@ -180,78 +207,6 @@ class Dataset(Base): } return self.retrieval_model or default_retrieval_model - def get_created_by_account(self, *, session: Session) -> Account | None: - return session.get(Account, self.created_by) - - def get_author_name(self, *, session: Session) -> str | None: - account = self.get_created_by_account(session=session) - if account: - return account.name - return None - - def get_latest_process_rule(self, *, session: Session) -> "DatasetProcessRule | None": - return session.scalar( - select(DatasetProcessRule) - .where(DatasetProcessRule.dataset_id == self.id) - .order_by(DatasetProcessRule.created_at.desc()) - .limit(1) - ) - - def get_app_count(self, *, session: Session) -> int: - return ( - session.scalar( - select(func.count(AppDatasetJoin.id)).where( - AppDatasetJoin.dataset_id == self.id, App.id == AppDatasetJoin.app_id - ) - ) - or 0 - ) - - def get_document_count(self, *, session: Session) -> int: - return session.scalar(select(func.count(Document.id)).where(Document.dataset_id == self.id)) or 0 - - def get_available_document_count(self, *, session: Session) -> int: - return ( - session.scalar( - select(func.count(Document.id)).where( - Document.dataset_id == self.id, - Document.indexing_status == "completed", - Document.enabled == True, - Document.archived == False, - ) - ) - or 0 - ) - - def get_available_segment_count(self, *, session: Session) -> int: - return ( - session.scalar( - select(func.count(DocumentSegment.id)).where( - DocumentSegment.dataset_id == self.id, - DocumentSegment.status == "completed", - DocumentSegment.enabled == True, - ) - ) - or 0 - ) - - def get_word_count(self, *, session: Session) -> int: - return ( - session.scalar( - select(func.coalesce(func.sum(Document.word_count), 0)).where(Document.dataset_id == self.id) - ) - or 0 - ) - - def get_doc_form(self, *, session: Session) -> str | None: - if self.chunk_structure: - return self.chunk_structure - return session.scalar( - select(Document.doc_form) - .where(Document.dataset_id == self.id, Document.tenant_id == self.tenant_id) - .limit(1) - ) - @property def retrieval_model_dict(self): """Return a normalized retrieval model payload for API responses. @@ -272,39 +227,6 @@ class Dataset(Base): return {**default_retrieval_model, **self.retrieval_model} - def get_tags(self, *, session: Session) -> Sequence[Tag]: - tags = session.scalars( - select(Tag) - .join(TagBinding, Tag.id == TagBinding.tag_id) - .where( - TagBinding.target_id == self.id, - TagBinding.tenant_id == self.tenant_id, - Tag.tenant_id == self.tenant_id, - Tag.type == "knowledge", - ) - ).all() - - return tags or [] - - def get_external_knowledge_info(self, *, session: Session) -> dict[str, Any] | None: - if self.provider != "external": - return None - external_knowledge_binding = session.scalar( - select(ExternalKnowledgeBindings).where( - ExternalKnowledgeBindings.dataset_id == self.id, - ExternalKnowledgeBindings.tenant_id == self.tenant_id, - ) - ) - if not external_knowledge_binding: - return None - external_knowledge_api = session.scalar( - select(ExternalKnowledgeApis).where( - ExternalKnowledgeApis.id == external_knowledge_binding.external_knowledge_api_id, - ExternalKnowledgeApis.tenant_id == self.tenant_id, - ) - ) - return self.build_external_knowledge_info(external_knowledge_binding, external_knowledge_api) - @staticmethod def build_external_knowledge_info( external_knowledge_binding: "ExternalKnowledgeBindings | None", @@ -322,18 +244,6 @@ class Dataset(Base): "external_knowledge_api_endpoint": json.loads(external_knowledge_api.settings).get("endpoint", ""), } - def get_is_published(self, *, session: Session) -> bool: - if self.pipeline_id: - pipeline = session.scalar(select(Pipeline).where(Pipeline.id == self.pipeline_id)) - if pipeline: - return pipeline.is_published - return False - - def get_doc_metadata(self, *, session: Session) -> list[dict[str, str]]: - dataset_metadatas = session.scalars(select(DatasetMetadata).where(DatasetMetadata.dataset_id == self.id)).all() - - return self.build_doc_metadata(dataset_metadatas) - def build_doc_metadata(self, dataset_metadatas: "Sequence[DatasetMetadata]") -> list[dict[str, str]]: """Format metadata rows, appending the built-in fields when they are enabled.""" doc_metadata = [ @@ -407,7 +317,6 @@ class DatasetProcessRule(TypeBase): ) MODES = ["automatic", "custom", "hierarchical"] - PRE_PROCESSING_RULES = ["remove_stopwords", "remove_extra_spaces", "remove_urls_emails"] AUTOMATIC_RULES: ClassVar[AutomaticRulesConfig] = { "pre_processing_rules": [ {"id": "remove_extra_spaces", "enabled": True}, @@ -482,11 +391,11 @@ class Document(Base): paused_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True) # error - error = mapped_column(LongText, nullable=True) + error: Mapped[str | None] = mapped_column(LongText, nullable=True) stopped_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True) # basic fields - indexing_status = mapped_column( + indexing_status: Mapped[str] = mapped_column( EnumText(IndexingStatus, length=255), nullable=False, server_default=sa.text("'waiting'") ) enabled: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=sa.true()) @@ -512,10 +421,10 @@ class Document(Base): @property def display_status(self): status = None - if self.indexing_status == "waiting": - status = "queuing" - elif self.indexing_status not in {"completed", "error", "waiting"} and self.is_paused: + if self.indexing_status not in {"completed", "error"} and self.is_paused: status = "paused" + elif self.indexing_status == "waiting": + status = "queuing" elif self.indexing_status in {"parsing", "cleaning", "splitting", "indexing"}: status = "indexing" elif self.indexing_status == "error": @@ -539,56 +448,6 @@ class Document(Base): return data_source_info_dict return {} - def get_data_source_detail_dict(self, *, session: Session) -> dict[str, Any]: - if self.data_source_info: - if self.data_source_type == "upload_file": - data_source_info_dict: dict[str, Any] = json.loads(self.data_source_info) - file_detail = session.scalar( - select(UploadFile).where(UploadFile.id == data_source_info_dict["upload_file_id"]) - ) - if file_detail: - return { - "upload_file": { - "id": file_detail.id, - "name": file_detail.name, - "size": file_detail.size, - "extension": file_detail.extension, - "mime_type": file_detail.mime_type, - "created_by": file_detail.created_by, - "created_at": file_detail.created_at.timestamp(), - } - } - elif self.data_source_type in {"notion_import", "website_crawl"}: - result: dict[str, Any] = json.loads(self.data_source_info) - return result - return {} - - def get_dataset_process_rule(self, *, session: Session) -> "DatasetProcessRule | None": - if self.dataset_process_rule_id: - return session.get(DatasetProcessRule, self.dataset_process_rule_id) - return None - - def get_dataset(self, *, session: Session) -> Dataset | None: - """Load the owning dataset with the caller-owned database session.""" - return session.get(Dataset, self.dataset_id) - - def get_segment_count(self, *, session: Session) -> int: - return session.scalar(select(func.count(DocumentSegment.id)).where(DocumentSegment.document_id == self.id)) or 0 - - def get_hit_count(self, *, session: Session) -> int: - return ( - session.scalar( - select(func.coalesce(func.sum(DocumentSegment.hit_count), 0)).where( - DocumentSegment.document_id == self.id - ) - ) - or 0 - ) - - def get_uploader(self, *, session: Session) -> str | None: - user = session.scalar(select(Account).where(Account.id == self.created_by)) - return user.name if user else None - @property def upload_date(self): return self.created_at @@ -597,77 +456,57 @@ class Document(Base): def last_update_date(self): return self.updated_at - def get_doc_metadata_details(self, *, session: Session) -> list[DocMetadataDetailItem] | None: - if self.doc_metadata: - document_metadatas = session.scalars( - select(DatasetMetadata) - .join(DatasetMetadataBinding, DatasetMetadataBinding.metadata_id == DatasetMetadata.id) - .where( - DatasetMetadata.tenant_id == self.tenant_id, - DatasetMetadata.dataset_id == self.dataset_id, - DatasetMetadataBinding.tenant_id == self.tenant_id, - DatasetMetadataBinding.dataset_id == self.dataset_id, - DatasetMetadataBinding.document_id == self.id, - ) - ).all() - metadata_list: list[DocMetadataDetailItem] = [] - for metadata in document_metadatas: - metadata_dict: DocMetadataDetailItem = { - "id": metadata.id, - "name": metadata.name, - "type": metadata.type, - "value": self.doc_metadata.get(metadata.name), - } - metadata_list.append(metadata_dict) - # deal built-in fields - metadata_list.extend(self.get_built_in_fields(session=session)) - - return metadata_list - return None - - def get_built_in_fields(self, *, session: Session) -> list[DocMetadataDetailItem]: - built_in_fields: list[DocMetadataDetailItem] = [] - built_in_fields.append( - { - "id": "built-in", - "name": BuiltInField.document_name, - "type": "string", - "value": self.name, - } - ) - built_in_fields.append( - { - "id": "built-in", - "name": BuiltInField.uploader, - "type": "string", - "value": self.get_uploader(session=session), - } - ) - built_in_fields.append( - { - "id": "built-in", - "name": BuiltInField.upload_date, - "type": "time", - "value": str(self.created_at.timestamp()), - } - ) - built_in_fields.append( - { - "id": "built-in", - "name": BuiltInField.last_update_date, - "type": "time", - "value": str(self.updated_at.timestamp()), - } - ) - built_in_fields.append( - { - "id": "built-in", - "name": BuiltInField.source, - "type": "string", - "value": MetadataDataSource[self.data_source_type], - } - ) - return built_in_fields + def to_dict(self, *, segment_count: int, hit_count: int, process_rule: DatasetProcessRule | None) -> DocumentDict: + result: DocumentDict = { + "id": self.id, + "tenant_id": self.tenant_id, + "dataset_id": self.dataset_id, + "position": self.position, + "data_source_type": self.data_source_type, + "data_source_info": self.data_source_info, + "dataset_process_rule_id": self.dataset_process_rule_id, + "batch": self.batch, + "name": self.name, + "created_from": self.created_from, + "created_by": self.created_by, + "created_api_request_id": self.created_api_request_id, + "created_at": self.created_at, + "processing_started_at": self.processing_started_at, + "file_id": self.file_id, + "word_count": self.word_count, + "parsing_completed_at": self.parsing_completed_at, + "cleaning_completed_at": self.cleaning_completed_at, + "splitting_completed_at": self.splitting_completed_at, + "tokens": self.tokens, + "indexing_latency": self.indexing_latency, + "completed_at": self.completed_at, + "is_paused": self.is_paused, + "paused_by": self.paused_by, + "paused_at": self.paused_at, + "error": self.error, + "stopped_at": self.stopped_at, + "indexing_status": self.indexing_status, + "enabled": self.enabled, + "disabled_at": self.disabled_at, + "disabled_by": self.disabled_by, + "archived": self.archived, + "archived_reason": self.archived_reason, + "archived_by": self.archived_by, + "archived_at": self.archived_at, + "updated_at": self.updated_at, + "doc_type": self.doc_type, + "doc_metadata": self.doc_metadata, + "doc_form": self.doc_form, + "doc_language": self.doc_language, + "display_status": self.display_status, + "data_source_info_dict": self.data_source_info_dict, + "average_segment_length": (self.word_count or 0) // segment_count if segment_count else 0, + "dataset_process_rule": process_rule.to_dict() if process_rule else None, + "dataset": None, + "segment_count": segment_count, + "hit_count": hit_count, + } + return result @classmethod def from_dict(cls, data: dict[str, Any]): @@ -761,155 +600,6 @@ class DocumentSegment(TypeBase): stopped_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True, default=None) hit_count: Mapped[int] = mapped_column(sa.Integer, nullable=False, default=0) - def get_dataset(self, *, session: Session) -> Dataset | None: - """Load the owning dataset with the caller-owned database session.""" - return session.get(Dataset, self.dataset_id) - - def get_document(self, *, session: Session) -> Document | None: - """Load the owning document with the caller-owned database session.""" - return session.get(Document, self.document_id) - - def previous_segment(self, session: Session) -> "DocumentSegment | None": - return session.scalar( - select(DocumentSegment).where( - DocumentSegment.document_id == self.document_id, DocumentSegment.position == self.position - 1 - ) - ) - - def next_segment(self, session: Session) -> "DocumentSegment | None": - return session.scalar( - select(DocumentSegment).where( - DocumentSegment.document_id == self.document_id, DocumentSegment.position == self.position + 1 - ) - ) - - def get_child_chunks(self, *, session: Session, include_full_doc: bool = True) -> Sequence["ChildChunk"]: - """Load hierarchical child chunks with the caller-owned database session.""" - document = session.get(Document, self.document_id) - if not document: - return [] - process_rule = document.get_dataset_process_rule(session=session) - if process_rule and process_rule.mode == "hierarchical": - rules_dict = process_rule.rules_dict - if rules_dict: - rules = Rule.model_validate(rules_dict) - if rules.parent_mode and (include_full_doc or rules.parent_mode != ParentMode.FULL_DOC): - child_chunks = session.scalars( - select(ChildChunk).where(ChildChunk.segment_id == self.id).order_by(ChildChunk.position.asc()) - ).all() - return child_chunks or [] - return [] - - @property - def sign_content(self) -> str: - return self.get_sign_content() - - def get_sign_content(self) -> str: - signed_urls: list[tuple[int, int, str]] = [] - text = self.content - - # For data before v0.10.0 - pattern = r"/files/([a-f0-9\-]+)/image-preview(?:\?.*?)?" - matches = re.finditer(pattern, text) - for match in matches: - upload_file_id = match.group(1) - nonce = os.urandom(16).hex() - timestamp = str(int(time.time())) - data_to_sign = f"image-preview|{upload_file_id}|{timestamp}|{nonce}" - secret_key = dify_config.SECRET_KEY.encode() - sign = hmac.new(secret_key, data_to_sign.encode(), hashlib.sha256).digest() - encoded_sign = base64.urlsafe_b64encode(sign).decode() - - params = f"timestamp={timestamp}&nonce={nonce}&sign={encoded_sign}" - base_url = f"/files/{upload_file_id}/image-preview" - signed_url = f"{base_url}?{params}" - signed_urls.append((match.start(), match.end(), signed_url)) - - # For data after v0.10.0 - pattern = r"/files/([a-f0-9\-]+)/file-preview(?:\?.*?)?" - matches = re.finditer(pattern, text) - for match in matches: - upload_file_id = match.group(1) - nonce = os.urandom(16).hex() - timestamp = str(int(time.time())) - data_to_sign = f"file-preview|{upload_file_id}|{timestamp}|{nonce}" - secret_key = dify_config.SECRET_KEY.encode() - sign = hmac.new(secret_key, data_to_sign.encode(), hashlib.sha256).digest() - encoded_sign = base64.urlsafe_b64encode(sign).decode() - - params = f"timestamp={timestamp}&nonce={nonce}&sign={encoded_sign}" - base_url = f"/files/{upload_file_id}/file-preview" - signed_url = f"{base_url}?{params}" - signed_urls.append((match.start(), match.end(), signed_url)) - - # For tools directory - direct file formats (e.g., .png, .jpg, etc.) - # Match URL including any query parameters up to common URL boundaries (space, parenthesis, quotes) - pattern = r"/files/tools/([a-f0-9\-]+)\.([a-zA-Z0-9]+)(?:\?[^\s\)\"\']*)?" - matches = re.finditer(pattern, text) - for match in matches: - upload_file_id = match.group(1) - file_extension = match.group(2) - nonce = os.urandom(16).hex() - timestamp = str(int(time.time())) - data_to_sign = f"file-preview|{upload_file_id}|{timestamp}|{nonce}" - secret_key = dify_config.SECRET_KEY.encode() - sign = hmac.new(secret_key, data_to_sign.encode(), hashlib.sha256).digest() - encoded_sign = base64.urlsafe_b64encode(sign).decode() - - params = f"timestamp={timestamp}&nonce={nonce}&sign={encoded_sign}" - base_url = f"/files/tools/{upload_file_id}.{file_extension}" - signed_url = f"{base_url}?{params}" - signed_urls.append((match.start(), match.end(), signed_url)) - - # Reconstruct the text with signed URLs - offset = 0 - for start, end, signed_url in signed_urls: - text = text[: start + offset] + signed_url + text[end + offset :] - offset += len(signed_url) - (end - start) - - return text - - def get_attachments(self, *, session: Session) -> list[AttachmentItem]: - """Load attachment metadata with the caller-owned database session.""" - # Use JOIN to fetch attachments in a single query instead of two separate queries - attachments_with_bindings = session.execute( - select(SegmentAttachmentBinding, UploadFile) - .join(UploadFile, UploadFile.id == SegmentAttachmentBinding.attachment_id) - .where( - SegmentAttachmentBinding.tenant_id == self.tenant_id, - SegmentAttachmentBinding.dataset_id == self.dataset_id, - SegmentAttachmentBinding.document_id == self.document_id, - SegmentAttachmentBinding.segment_id == self.id, - ) - ).all() - if not attachments_with_bindings: - return [] - attachment_list: list[AttachmentItem] = [] - for _, attachment in attachments_with_bindings: - upload_file_id = attachment.id - nonce = os.urandom(16).hex() - timestamp = str(int(time.time())) - data_to_sign = f"image-preview|{upload_file_id}|{timestamp}|{nonce}" - secret_key = dify_config.SECRET_KEY.encode() - sign = hmac.new(secret_key, data_to_sign.encode(), hashlib.sha256).digest() - encoded_sign = base64.urlsafe_b64encode(sign).decode() - - params = f"timestamp={timestamp}&nonce={nonce}&sign={encoded_sign}" - reference_url = dify_config.FILES_URL or dify_config.CONSOLE_API_URL or "" - base_url = f"{reference_url}/files/{upload_file_id}/image-preview" - source_url = f"{base_url}?{params}" - attachment_list.append( - { - "id": attachment.id, - "name": attachment.name, - "size": attachment.size, - "extension": attachment.extension, - "mime_type": attachment.mime_type, - "source_url": source_url, - } - ) - return attachment_list - class ChildChunk(TypeBase): __tablename__ = "child_chunks" @@ -957,15 +647,6 @@ class ChildChunk(TypeBase): ) error: Mapped[str | None] = mapped_column(LongText, nullable=True, init=False) - def dataset(self, session: Session) -> Dataset | None: - return session.scalar(select(Dataset).where(Dataset.id == self.dataset_id)) - - def document(self, session: Session) -> Document | None: - return session.scalar(select(Document).where(Document.id == self.document_id)) - - def segment(self, session: Session) -> DocumentSegment | None: - return session.scalar(select(DocumentSegment).where(DocumentSegment.id == self.segment_id)) - class AppDatasetJoin(TypeBase): __tablename__ = "app_dataset_joins" @@ -988,9 +669,6 @@ class AppDatasetJoin(TypeBase): DateTime, nullable=False, server_default=sa.func.current_timestamp(), init=False ) - def app(self, session: Session) -> App | None: - return session.get(App, self.app_id) - class DatasetQuery(TypeBase): __tablename__ = "dataset_queries" @@ -1017,37 +695,6 @@ class DatasetQuery(TypeBase): DateTime, nullable=False, server_default=sa.func.current_timestamp(), init=False ) - def get_queries(self, *, session: Session) -> list[dict[str, Any]]: - try: - queries = json.loads(self.content) - if isinstance(queries, list): - for query in queries: - if query["content_type"] == QueryType.IMAGE_QUERY: - file_info = session.scalar(select(UploadFile).where(UploadFile.id == query["content"])) - if file_info: - query["file_info"] = { - "id": file_info.id, - "name": file_info.name, - "size": file_info.size, - "extension": file_info.extension, - "mime_type": file_info.mime_type, - "source_url": sign_upload_file_preview_url(file_info.id, file_info.extension), - } - else: - query["file_info"] = None - - return queries - else: - return [queries] - except JSONDecodeError: - return [ - { - "content_type": QueryType.TEXT_QUERY, - "content": self.content, - "file_info": None, - } - ] - class DatasetKeywordTable(TypeBase): __tablename__ = "dataset_keyword_tables" @@ -1067,40 +714,6 @@ class DatasetKeywordTable(TypeBase): keyword_table: Mapped[str] = mapped_column(LongText, nullable=False) data_source_type: Mapped[str] = mapped_column(String(255), nullable=False, default="database") - def get_keyword_table_dict(self, *, session: Session) -> dict[str, set[Any]] | None: - class SetDecoder(json.JSONDecoder): - def __init__(self, *args: Any, **kwargs: Any) -> None: - def object_hook(dct: Any) -> Any: - if isinstance(dct, dict): - result: dict[str, Any] = {} - items = cast(dict[str, Any], dct).items() - for keyword, node_idxs in items: - if isinstance(node_idxs, list): - result[keyword] = set(node_idxs) - else: - result[keyword] = node_idxs - return result - return dct - - super().__init__(object_hook=object_hook, *args, **kwargs) - - # get dataset - dataset = session.scalar(select(Dataset).where(Dataset.id == self.dataset_id)) - if not dataset: - return None - if self.data_source_type == "database": - return json.loads(self.keyword_table, cls=SetDecoder) if self.keyword_table else None - else: - file_key = "keyword_files/" + dataset.tenant_id + "/" + self.dataset_id + ".txt" - try: - keyword_table_text = storage.load_once(file_key) - if keyword_table_text: - return json.loads(keyword_table_text.decode("utf-8"), cls=SetDecoder) - return None - except Exception: - logger.exception("Failed to load keyword table from file: %s", file_key) - return None - class Embedding(TypeBase): __tablename__ = "embeddings" @@ -1263,6 +876,18 @@ class ExternalKnowledgeApis(TypeBase): DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp(), init=False ) + def to_dict(self, *, dataset_bindings: list[DatasetBindingItem]) -> ExternalKnowledgeApiDict: + return { + "id": self.id, + "tenant_id": self.tenant_id, + "name": self.name, + "description": self.description, + "settings": self.settings_dict, + "dataset_bindings": dataset_bindings, + "created_by": self.created_by, + "created_at": self.created_at.isoformat(), + } + @property def settings_dict(self) -> dict[str, Any] | None: try: @@ -1270,18 +895,6 @@ class ExternalKnowledgeApis(TypeBase): except JSONDecodeError: return None - def get_dataset_bindings(self, *, session: Session) -> list[DatasetBindingItem]: - external_knowledge_bindings = session.scalars( - select(ExternalKnowledgeBindings).where(ExternalKnowledgeBindings.external_knowledge_api_id == self.id) - ).all() - dataset_ids = [binding.dataset_id for binding in external_knowledge_bindings] - datasets = session.scalars(select(Dataset).where(Dataset.id.in_(dataset_ids))).all() - dataset_bindings: list[DatasetBindingItem] = [] - for dataset in datasets: - dataset_bindings.append({"id": dataset.id, "name": dataset.name}) - - return dataset_bindings - class ExternalKnowledgeBindings(TypeBase): __tablename__ = "external_knowledge_bindings" @@ -1467,12 +1080,6 @@ class PipelineCustomizedTemplate(TypeBase): init=False, ) - def created_user_name(self, session: Session) -> str: - account = session.scalar(select(Account).where(Account.id == self.created_by)) - if account: - return account.name - return "" - class Pipeline(TypeBase): __tablename__ = "pipelines" @@ -1500,11 +1107,6 @@ class Pipeline(TypeBase): init=False, ) - def retrieve_dataset(self, session: Session | scoped_session): - return session.scalar( - select(Dataset).where(Dataset.pipeline_id == self.id, Dataset.tenant_id == self.tenant_id) - ) - class DocumentPipelineExecutionLog(TypeBase): __tablename__ = "document_pipeline_execution_logs" diff --git a/api/openapi/markdown/console-openapi.md b/api/openapi/markdown/console-openapi.md index 4f9147765d1..128060a7bd4 100644 --- a/api/openapi/markdown/console-openapi.md +++ b/api/openapi/markdown/console-openapi.md @@ -5083,30 +5083,9 @@ Get compliance document download link ### [GET] /data-source/integrates #### Responses -| Code | Description | -| ---- | ----------- | -| 200 | Success | - -### [PATCH] /data-source/integrates -#### Responses - -| Code | Description | -| ---- | ----------- | -| 200 | Success | - -### [GET] /data-source/integrates/{binding_id}/{action} -#### Parameters - -| Name | Located in | Description | Required | Schema | -| ---- | ---------- | ----------- | -------- | ------ | -| action | path | | Yes | string | -| binding_id | path | | Yes | string (uuid) | - -#### Responses - -| Code | Description | -| ---- | ----------- | -| 200 | Success | +| Code | Description | Schema | +| ---- | ----------- | ------ | +| 200 | Success | **application/json**: [DataSourceIntegrateListResponse](#datasourceintegratelistresponse)
| ### [PATCH] /data-source/integrates/{binding_id}/{action} #### Parameters @@ -5118,9 +5097,9 @@ Get compliance document download link #### Responses -| Code | Description | -| ---- | ----------- | -| 200 | Success | +| Code | Description | Schema | +| ---- | ----------- | ------ | +| 200 | Success | **application/json**: [SimpleResultResponse](#simpleresultresponse)
| ### [GET] /datasets Get list of datasets @@ -5155,7 +5134,7 @@ Create a new dataset | Code | Description | Schema | | ---- | ----------- | ------ | -| 201 | Dataset created successfully | **application/json**: [DatasetDetailResponse](#datasetdetailresponse)
| +| 201 | Dataset created successfully | **application/json**: [DatasetDetailWithPartialMembersResponse](#datasetdetailwithpartialmembersresponse)
| | 400 | Invalid request parameters | | ### [GET] /datasets/api-base-info @@ -5218,25 +5197,6 @@ Delete dataset API key | ---- | ----------- | ------ | | 200 | Batch import status | **application/json**: [SegmentBatchImportStatusResponse](#segmentbatchimportstatusresponse)
| -### [POST] /datasets/batch_import_status/{job_id} -#### Parameters - -| Name | Located in | Description | Required | Schema | -| ---- | ---------- | ----------- | -------- | ------ | -| job_id | path | | Yes | string (uuid) | - -#### Request Body - -| Required | Schema | -| -------- | ------ | -| Yes | **application/json**: [BatchImportPayload](#batchimportpayload)
| - -#### Responses - -| Code | Description | Schema | -| ---- | ----------- | ------ | -| 200 | Batch import started | **application/json**: [SegmentBatchImportStatusResponse](#segmentbatchimportstatusresponse)
| - ### [POST] /datasets/external Create external knowledge dataset @@ -5970,20 +5930,6 @@ Update document processing status (pause/resume) | ---- | ----------- | ------ | | 200 | Segments retrieved successfully | **application/json**: [ConsoleSegmentListResponse](#consolesegmentlistresponse)
| -### [GET] /datasets/{dataset_id}/documents/{document_id}/segments/batch_import -#### Parameters - -| Name | Located in | Description | Required | Schema | -| ---- | ---------- | ----------- | -------- | ------ | -| dataset_id | path | | Yes | string (uuid) | -| document_id | path | | Yes | string (uuid) | - -#### Responses - -| Code | Description | Schema | -| ---- | ----------- | ------ | -| 200 | Batch import status | **application/json**: [SegmentBatchImportStatusResponse](#segmentbatchimportstatusresponse)
| - ### [POST] /datasets/{dataset_id}/documents/{document_id}/segments/batch_import #### Parameters @@ -18102,7 +18048,7 @@ Site token customization strategy | page_id | string | | Yes | | page_name | string | | Yes | | parent_id | string | | Yes | -| type | string | | Yes | +| type | [NotionPageType](#notionpagetype) | | Yes | #### DataSourceIntegrateResponse @@ -21037,15 +20983,30 @@ Coarse node-level status used by Inspector to pick a banner. | notifications | [ [NotificationItemResponse](#notificationitemresponse) ] | | Yes | | should_show | boolean | | Yes | +#### NotionEstimatePagePayload + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| page_id | string | | Yes | +| type | [NotionPageType](#notionpagetype) | | Yes | + #### NotionEstimatePayload | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | | doc_form | string,
**Default:** text_model | | No | | doc_language | string,
**Default:** English | | No | -| notion_info_list | [ object ] | | Yes | +| notion_info_list | [ [NotionEstimateWorkspacePayload](#notionestimateworkspacepayload) ] | | Yes | | process_rule | object | | Yes | +#### NotionEstimateWorkspacePayload + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| credential_id | string | | Yes | +| pages | [ [NotionEstimatePagePayload](#notionestimatepagepayload) ] | | Yes | +| workspace_id | string | | Yes | + #### NotionIcon | Name | Type | Description | Required | @@ -21077,7 +21038,7 @@ Coarse node-level status used by Inspector to pick a banner. | page_id | string | | Yes | | page_name | string | | Yes | | parent_id | string | | Yes | -| type | string | | Yes | +| type | [NotionPageType](#notionpagetype) | | Yes | #### NotionIntegrateWorkspaceResponse @@ -21097,6 +21058,12 @@ Coarse node-level status used by Inspector to pick a banner. | page_name | string | | Yes | | type | string | | Yes | +#### NotionPageType + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| NotionPageType | string | | | + #### OAuthCallbackQuery | Name | Type | Description | Required | @@ -22232,7 +22199,13 @@ Verification of the plugin. | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | | enabled | boolean | Whether this preprocessing rule is enabled. | Yes | -| id | string,
**Available values:** "remove_extra_spaces", "remove_stopwords", "remove_urls_emails" | Rule identifier.
*Enum:* `"remove_extra_spaces"`, `"remove_stopwords"`, `"remove_urls_emails"` | Yes | +| id | [PreProcessingRuleKey](#preprocessingrulekey) | Rule identifier. | Yes | + +#### PreProcessingRuleKey + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| PreProcessingRuleKey | string | | | #### PreviewDetail diff --git a/api/openapi/markdown/service-openapi.md b/api/openapi/markdown/service-openapi.md index 9fd290e93dc..a30ee05de6f 100644 --- a/api/openapi/markdown/service-openapi.md +++ b/api/openapi/markdown/service-openapi.md @@ -3836,7 +3836,13 @@ JSON result for published runs and draft runs using `response_mode: blocking`. | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | | enabled | boolean | Whether this preprocessing rule is enabled. | Yes | -| id | string,
**Available values:** "remove_extra_spaces", "remove_stopwords", "remove_urls_emails" | Rule identifier.
*Enum:* `"remove_extra_spaces"`, `"remove_stopwords"`, `"remove_urls_emails"` | Yes | +| id | [PreProcessingRuleKey](#preprocessingrulekey) | Rule identifier. | Yes | + +#### PreProcessingRuleKey + +| Name | Type | Description | Required | +| ---- | ---- | ----------- | -------- | +| PreProcessingRuleKey | string | | | #### ProcessRule diff --git a/api/repositories/credentials/__init__.py b/api/repositories/credentials/__init__.py new file mode 100644 index 00000000000..944f2df7138 --- /dev/null +++ b/api/repositories/credentials/__init__.py @@ -0,0 +1 @@ +"""Persistence adapters for the credentials domain.""" diff --git a/api/repositories/credentials/query_repository.py b/api/repositories/credentials/query_repository.py new file mode 100644 index 00000000000..8e4d582738a --- /dev/null +++ b/api/repositories/credentials/query_repository.py @@ -0,0 +1,250 @@ +"""Persistence-owned credential queries and visibility filtering shared across domains.""" + +from collections.abc import Sequence +from typing import override + +from sqlalchemy import and_, func, or_, select +from sqlalchemy.orm import InstrumentedAttribute, Session, sessionmaker + +from models.credential_permission import CredentialPermission, CredentialType +from models.enums import PermissionEnum +from models.oauth import DatasourceProvider +from models.provider import ProviderCredential +from models.provider_ids import ModelProviderID +from models.tools import BuiltinToolProvider +from models.trigger import TriggerSubscription, WorkflowPluginTrigger +from services.credentials.query import ( + CredentialQuery, + DatasourceCredentialListItem, + ModelCredentialRecord, + ToolCredentialRecord, + TriggerSubscriptionRecord, +) + + +def apply_credential_visibility_filter_for_actor( + query, + *, + tenant_id: str, + model_id_column: InstrumentedAttribute, + model_user_id_column: InstrumentedAttribute, + model_visibility_column: InstrumentedAttribute, + credential_type: str, + actor_id: str, +): + """Apply the canonical tenant- and actor-scoped credential visibility policy. + + Team credentials and legacy rows without an owner are visible to all tenant + members. Personal credentials are visible to their owner, while partially + shared credentials additionally require an enabled permission row for the + same tenant and credential type. Administrator roles do not bypass this + policy. + """ + partial_member_ids = ( + select(CredentialPermission.credential_id) + .where( + CredentialPermission.tenant_id == tenant_id, + CredentialPermission.credential_type == credential_type, + CredentialPermission.account_id == actor_id, + CredentialPermission.has_permission.is_(True), + ) + .correlate_except(CredentialPermission) + ) + return query.where( + or_( + model_visibility_column == PermissionEnum.ALL_TEAM, + model_user_id_column.is_(None), + model_user_id_column == actor_id, + and_(model_visibility_column == PermissionEnum.PARTIAL_TEAM, model_id_column.in_(partial_member_ids)), + ) + ) + + +class CredentialQueryRepository(CredentialQuery): + """Load complete credential lists without exposing ORM objects or SQL to callers.""" + + def __init__(self, *, session_factory: sessionmaker[Session]) -> None: + self._session_factory = session_factory + + @override + def list_models(self, *, workspace_id: str, provider: str, actor_id: str | None) -> Sequence[ModelCredentialRecord]: + names = [provider] + provider_id = ModelProviderID(provider) + if provider_id.is_langgenius(): + names.append(provider_id.provider_name if "/" in provider else str(provider_id)) + query = ( + select(ProviderCredential) + .where(ProviderCredential.tenant_id == workspace_id, ProviderCredential.provider_name.in_(names)) + .order_by(ProviderCredential.created_at.desc()) + ) + if actor_id is not None: + query = apply_credential_visibility_filter_for_actor( + query, + tenant_id=workspace_id, + model_id_column=ProviderCredential.id, + model_user_id_column=ProviderCredential.user_id, + model_visibility_column=ProviderCredential.visibility, + credential_type=CredentialType.PROVIDER_CREDENTIAL, + actor_id=actor_id, + ) + with self._session_factory() as session: + return [ModelCredentialRecord(row.id, row.credential_name) for row in session.scalars(query)] + + @override + def list_datasources( + self, *, workspace_id: str, provider: str, plugin_id: str, actor_id: str | None + ) -> Sequence[DatasourceCredentialListItem]: + owner = ( + DatasourceProvider.tenant_id == workspace_id, + DatasourceProvider.provider == provider, + DatasourceProvider.plugin_id == plugin_id, + ) + query = select(DatasourceProvider).where(*owner) + if actor_id is not None: + query = apply_credential_visibility_filter_for_actor( + query, + tenant_id=workspace_id, + model_id_column=DatasourceProvider.id, + model_user_id_column=DatasourceProvider.user_id, + model_visibility_column=DatasourceProvider.visibility, + credential_type=CredentialType.DATASOURCE_PROVIDER, + actor_id=actor_id, + ) + with self._session_factory() as session: + rows = session.scalars(query).all() + if not rows: + return [] + # Preserve the configured default even when that credential is hidden from this actor. + default_id = session.scalar( + select(DatasourceProvider.id) + .where(*owner) + .order_by(DatasourceProvider.is_default.desc(), DatasourceProvider.created_at.asc()) + .limit(1) + ) + return [ + DatasourceCredentialListItem( + id=row.id, + name=row.name, + auth_type=row.auth_type, + encrypted_credentials=dict(row.encrypted_credentials), + avatar_url=row.avatar_url, + is_default=row.id == default_id, + ) + for row in rows + ] + + @override + def list_tools( + self, + *, + workspace_id: str, + provider: str, + actor_id: str | None, + include_credential_ids: Sequence[str] = (), + ) -> Sequence[ToolCredentialRecord]: + owner = (BuiltinToolProvider.tenant_id == workspace_id, BuiltinToolProvider.provider == provider) + order = (BuiltinToolProvider.is_default.desc(), BuiltinToolProvider.created_at.asc()) + query = select(BuiltinToolProvider).where(*owner).order_by(*order) + if actor_id is not None: + query = apply_credential_visibility_filter_for_actor( + query, + tenant_id=workspace_id, + model_id_column=BuiltinToolProvider.id, + model_user_id_column=BuiltinToolProvider.user_id, + model_visibility_column=BuiltinToolProvider.visibility, + credential_type=CredentialType.BUILTIN_TOOL_PROVIDER, + actor_id=actor_id, + ) + with self._session_factory() as session: + visible = list(session.scalars(query)) + wanted = set(include_credential_ids) - {row.id for row in visible} - {""} + borrowed: list[BuiltinToolProvider] = ( + list( + session.scalars( + select(BuiltinToolProvider).where(*owner, BuiltinToolProvider.id.in_(wanted)).order_by(*order) + ) + ) + if wanted + else [] + ) + rows = visible + borrowed + partial_ids = [row.id for row in rows if row.visibility == PermissionEnum.PARTIAL_TEAM] + members: dict[str, list[str]] = {} + if partial_ids: + for credential_id, account_id in session.execute( + select(CredentialPermission.credential_id, CredentialPermission.account_id).where( + CredentialPermission.tenant_id == workspace_id, + CredentialPermission.credential_type == CredentialType.BUILTIN_TOOL_PROVIDER, + CredentialPermission.credential_id.in_(partial_ids), + CredentialPermission.has_permission.is_(True), + ) + ): + members.setdefault(credential_id, []).append(account_id) + borrowed_ids = {row.id for row in borrowed} + return [ + ToolCredentialRecord( + id=row.id, + name=row.name, + provider=row.provider, + credential_type=row.credential_type, + credentials=dict(row.credentials), + is_default=(bool(visible) and row.id == visible[0].id) or row.is_default, + visibility=str(row.visibility), + created_by=row.user_id or "", + partial_member_ids=tuple(members.get(row.id, [])), + from_other_member=row.id in borrowed_ids, + ) + for row in rows + ] + + @override + def list_trigger_subscriptions( + self, *, workspace_id: str, provider: str, actor_id: str | None + ) -> Sequence[TriggerSubscriptionRecord]: + query = ( + select(TriggerSubscription) + .where(TriggerSubscription.tenant_id == workspace_id, TriggerSubscription.provider_id == provider) + .order_by(TriggerSubscription.created_at.desc()) + ) + if actor_id is not None: + query = apply_credential_visibility_filter_for_actor( + query, + tenant_id=workspace_id, + model_id_column=TriggerSubscription.id, + model_user_id_column=TriggerSubscription.user_id, + model_visibility_column=TriggerSubscription.visibility, + credential_type=CredentialType.TRIGGER_SUBSCRIPTION, + actor_id=actor_id, + ) + with self._session_factory() as session: + rows = session.scalars(query).all() + if not rows: + return [] + counts = dict( + session.execute( + select( + WorkflowPluginTrigger.subscription_id, func.count(func.distinct(WorkflowPluginTrigger.app_id)) + ) + .where( + WorkflowPluginTrigger.tenant_id == workspace_id, + WorkflowPluginTrigger.subscription_id.in_([row.id for row in rows]), + ) + .group_by(WorkflowPluginTrigger.subscription_id) + ) + .tuples() + .all() + ) + return [ + TriggerSubscriptionRecord( + id=row.id, + name=row.name, + provider=row.provider_id, + credential_type=row.credential_type, + credentials=dict(row.credentials), + endpoint_id=row.endpoint_id, + parameters=dict(row.parameters), + properties=dict(row.properties), + workflows_in_use=counts.get(row.id, 0), + ) + for row in rows + ] diff --git a/api/repositories/data_source/__init__.py b/api/repositories/data_source/__init__.py new file mode 100644 index 00000000000..5b4f34ada4d --- /dev/null +++ b/api/repositories/data_source/__init__.py @@ -0,0 +1 @@ +"""Persistence adapters for the data-source domain.""" diff --git a/api/repositories/data_source_api_key_auth_repository.py b/api/repositories/data_source/api_key_auth_repository.py similarity index 94% rename from api/repositories/data_source_api_key_auth_repository.py rename to api/repositories/data_source/api_key_auth_repository.py index a28ebc16204..9f514d6b3f0 100644 --- a/api/repositories/data_source_api_key_auth_repository.py +++ b/api/repositories/data_source/api_key_auth_repository.py @@ -8,8 +8,8 @@ from sqlalchemy import select from sqlalchemy.orm import Session, sessionmaker from models.source import DataSourceApiKeyAuthBinding -from services.auth.data_source_api_key_auth_service import DataSourceApiKeyAuthBindingRepository -from services.entities.data_source_api_key_auth_entities import ( +from services.data_source.auth.api_key_service import DataSourceApiKeyAuthBindingRepository +from services.data_source.entities.api_key_auth import ( DataSourceApiKeyAuthBindingRecord, DataSourceApiKeyAuthCredentials, ) diff --git a/api/repositories/data_source/credential_repository.py b/api/repositories/data_source/credential_repository.py new file mode 100644 index 00000000000..9862ada014a --- /dev/null +++ b/api/repositories/data_source/credential_repository.py @@ -0,0 +1,252 @@ +"""Actor-aware and trusted-source SQLAlchemy repository for datasource credentials.""" + +import json +from collections.abc import Mapping +from typing import cast + +from pydantic import TypeAdapter, ValidationError +from sqlalchemy import select, update +from sqlalchemy.engine import CursorResult +from sqlalchemy.orm import Session, sessionmaker + +from libs.datetime_utils import naive_utc_now +from models.credential_permission import CredentialType +from models.dataset import Dataset, Document +from models.oauth import DatasourceOauthParamConfig, DatasourceOauthTenantParamConfig, DatasourceProvider +from repositories.credentials.query_repository import apply_credential_visibility_filter_for_actor +from services.data_source.entities.credential import DatasourceCredentialRecord, DatasourceOAuthClientConfigRecord + +_CREDENTIALS_ADAPTER = TypeAdapter(dict[str, object]) + + +def _source_mapping(value: object) -> Mapping[str, object]: + try: + if isinstance(value, str): + return _CREDENTIALS_ADAPTER.validate_python(json.loads(value)) + return _CREDENTIALS_ADAPTER.validate_python(value) + except (TypeError, ValueError, ValidationError): + return {} + + +def _datasource_credential_record(provider: DatasourceProvider) -> DatasourceCredentialRecord: + return DatasourceCredentialRecord( + id=provider.id, + workspace_id=provider.tenant_id, + owner_id=provider.user_id, + name=provider.name, + provider=provider.provider, + plugin_id=provider.plugin_id, + auth_type=provider.auth_type, + encrypted_credentials=_CREDENTIALS_ADAPTER.validate_python(provider.encrypted_credentials), + expires_at=provider.expires_at, + updated_at=provider.updated_at, + ) + + +class SQLAlchemyDatasourceCredentialRepository: + def __init__(self, *, session_factory: sessionmaker[Session]) -> None: + self._session_factory = session_factory + + def get_provider_credential( + self, *, workspace_id: str, provider: str, plugin_id: str, credential_id: str | None + ) -> DatasourceCredentialRecord | None: + with self._session_factory() as session: + rows = _read_provider_credentials( + session, + workspace_id=workspace_id, + provider=provider, + plugin_id=plugin_id, + credential_id=credential_id, + first_only=True, + ) + return _datasource_credential_record(rows[0]) if rows else None + + def list_provider_credentials( + self, *, workspace_id: str, provider: str, plugin_id: str + ) -> list[DatasourceCredentialRecord]: + with self._session_factory() as session: + return [ + _datasource_credential_record(row) + for row in _read_provider_credentials( + session, + workspace_id=workspace_id, + provider=provider, + plugin_id=plugin_id, + ) + ] + + def get_visible( + self, + *, + workspace_id: str, + actor_id: str, + credential_id: str, + provider: str, + plugin_id: str, + ) -> DatasourceCredentialRecord | None: + query = select(DatasourceProvider).where( + DatasourceProvider.tenant_id == workspace_id, + DatasourceProvider.id == credential_id, + DatasourceProvider.provider == provider, + DatasourceProvider.plugin_id == plugin_id, + ) + query = apply_credential_visibility_filter_for_actor( + query, + tenant_id=workspace_id, + model_id_column=DatasourceProvider.id, + model_user_id_column=DatasourceProvider.user_id, + model_visibility_column=DatasourceProvider.visibility, + credential_type=CredentialType.DATASOURCE_PROVIDER, + actor_id=actor_id, + ) + with self._session_factory() as session: + provider_record = session.scalar(query.limit(1)) + return _datasource_credential_record(provider_record) if provider_record is not None else None + + def get_oauth_client_config( + self, + *, + workspace_id: str, + provider: str, + plugin_id: str, + ) -> DatasourceOAuthClientConfigRecord: + """Load OAuth client rows in one bounded read transaction.""" + + with self._session_factory() as session: + return _read_oauth_client_config(session, workspace_id=workspace_id, provider=provider, plugin_id=plugin_id) + + def update_if_unchanged( + self, + *, + record: DatasourceCredentialRecord, + encrypted_credentials: Mapping[str, object], + expires_at: int, + ) -> bool: + with self._session_factory.begin() as session: + return _update_datasource_credentials_if_unchanged( + session, record=record, encrypted_credentials=encrypted_credentials, expires_at=expires_at + ) + + def get_for_stored_document( + self, + *, + workspace_id: str, + dataset_id: str, + document_id: str, + credential_id: str | None, + provider: str, + plugin_id: str, + ) -> DatasourceCredentialRecord | None: + with self._session_factory() as session: + query = ( + select(DatasourceProvider, Document.data_source_info) + .select_from(Document) + .join(Dataset, Dataset.id == Document.dataset_id) + .join(DatasourceProvider, DatasourceProvider.tenant_id == Document.tenant_id) + .where( + Dataset.id == dataset_id, + Dataset.tenant_id == workspace_id, + Document.id == document_id, + Document.dataset_id == dataset_id, + Document.tenant_id == workspace_id, + Document.data_source_type == "notion_import", + DatasourceProvider.tenant_id == workspace_id, + DatasourceProvider.provider == provider, + DatasourceProvider.plugin_id == plugin_id, + ) + ) + if credential_id is not None: + query = query.where(DatasourceProvider.id == credential_id) + else: + query = query.where(DatasourceProvider.is_default.is_(True)).order_by( + DatasourceProvider.created_at.asc() + ) + row = session.execute(query.limit(1)).one_or_none() + if row is None: + return None + datasource_provider, raw_source_info = row + source_info = _source_mapping(raw_source_info) + if credential_id is not None and source_info.get("credential_id") != credential_id: + return None + return _datasource_credential_record(datasource_provider) + + +def _read_oauth_client_config( + session: Session, *, workspace_id: str, provider: str, plugin_id: str +) -> DatasourceOAuthClientConfigRecord: + tenant_config = session.scalar( + select(DatasourceOauthTenantParamConfig) + .where( + DatasourceOauthTenantParamConfig.tenant_id == workspace_id, + DatasourceOauthTenantParamConfig.provider == provider, + DatasourceOauthTenantParamConfig.plugin_id == plugin_id, + DatasourceOauthTenantParamConfig.enabled.is_(True), + ) + .limit(1) + ) + system_config = session.scalar( + select(DatasourceOauthParamConfig) + .where( + DatasourceOauthParamConfig.provider == provider, + DatasourceOauthParamConfig.plugin_id == plugin_id, + ) + .limit(1) + ) + return DatasourceOAuthClientConfigRecord( + encrypted_tenant_params=(dict(tenant_config.client_params) if tenant_config is not None else None), + system_credentials=(dict(system_config.system_credentials) if system_config is not None else None), + ) + + +def _update_datasource_credentials_if_unchanged( + session: Session, + *, + record: DatasourceCredentialRecord, + encrypted_credentials: Mapping[str, object], + expires_at: int, +) -> bool: + credentials = _CREDENTIALS_ADAPTER.validate_python(encrypted_credentials) + result = session.execute( + update(DatasourceProvider) + .where( + DatasourceProvider.id == record.id, + DatasourceProvider.tenant_id == record.workspace_id, + DatasourceProvider.provider == record.provider, + DatasourceProvider.plugin_id == record.plugin_id, + DatasourceProvider.user_id == record.owner_id, + DatasourceProvider.encrypted_credentials + == _CREDENTIALS_ADAPTER.validate_python(record.encrypted_credentials), + DatasourceProvider.expires_at == record.expires_at, + ) + .values( + encrypted_credentials=credentials, + expires_at=expires_at, + updated_at=naive_utc_now(), + ) + ) + return cast(CursorResult[object], result).rowcount == 1 + + +def _read_provider_credentials( + session: Session, + *, + workspace_id: str, + provider: str, + plugin_id: str, + credential_id: str | None = None, + first_only: bool = False, +) -> list[DatasourceProvider]: + statement = ( + select(DatasourceProvider) + .where( + DatasourceProvider.tenant_id == workspace_id, + DatasourceProvider.provider == provider, + DatasourceProvider.plugin_id == plugin_id, + ) + .order_by(DatasourceProvider.is_default.desc(), DatasourceProvider.created_at.asc()) + ) + if credential_id is not None: + statement = statement.where(DatasourceProvider.id == credential_id) + if first_only: + statement = statement.limit(1) + return list(session.scalars(statement)) diff --git a/api/repositories/data_source_oauth_binding_repository.py b/api/repositories/data_source/oauth_binding_repository.py similarity index 56% rename from api/repositories/data_source_oauth_binding_repository.py rename to api/repositories/data_source/oauth_binding_repository.py index 5c52140f13a..98cda3c65d3 100644 --- a/api/repositories/data_source_oauth_binding_repository.py +++ b/api/repositories/data_source/oauth_binding_repository.py @@ -1,24 +1,27 @@ """SQLAlchemy repository for OAuth data-source bindings.""" from collections.abc import Mapping -from typing import override +from typing import cast, override from pydantic import TypeAdapter -from sqlalchemy import select +from sqlalchemy import select, update +from sqlalchemy.engine import CursorResult from sqlalchemy.orm import Session, sessionmaker from libs.datetime_utils import naive_utc_now from models.source import DataSourceOauthBinding -from services.data_source_oauth_service import DataSourceOAuthBindingRepository -from services.entities.data_source_oauth_entities import ( +from services.data_source.binding_application_service import BindingMutationResult, DataSourceBindingStore +from services.data_source.entities.oauth import ( + DataSourceBindingSummary, DataSourceOAuthAuthorization, DataSourceOAuthBindingRecord, ) +from services.data_source.oauth_service import DataSourceOAuthBindingRepository _SOURCE_INFO_ADAPTER = TypeAdapter(dict[str, object]) -class SQLAlchemyDataSourceOAuthBindingRepository(DataSourceOAuthBindingRepository): +class SQLAlchemyDataSourceOAuthBindingRepository(DataSourceOAuthBindingRepository, DataSourceBindingStore): def __init__(self, session_factory: sessionmaker[Session]) -> None: self._session_factory = session_factory @@ -68,7 +71,7 @@ class SQLAlchemyDataSourceOAuthBindingRepository(DataSourceOAuthBindingRepositor DataSourceOauthBinding.tenant_id == workspace_id, DataSourceOauthBinding.provider == provider, DataSourceOauthBinding.id == binding_id, - DataSourceOauthBinding.disabled.is_(False), + DataSourceOauthBinding.disabled.is_not(True), ) ) if binding is None: @@ -96,7 +99,7 @@ class SQLAlchemyDataSourceOAuthBindingRepository(DataSourceOAuthBindingRepositor DataSourceOauthBinding.tenant_id == workspace_id, DataSourceOauthBinding.provider == provider, DataSourceOauthBinding.id == binding_id, - DataSourceOauthBinding.disabled.is_(False), + DataSourceOauthBinding.disabled.is_not(True), ) ) if binding is None: @@ -105,3 +108,59 @@ class SQLAlchemyDataSourceOAuthBindingRepository(DataSourceOAuthBindingRepositor binding.source_info = validated_source_info binding.updated_at = naive_utc_now() return True + + @override + def list_enabled_bindings(self, *, workspace_id: str) -> tuple[DataSourceBindingSummary, ...]: + with self._session_factory() as session: + bindings = session.scalars( + select(DataSourceOauthBinding).where( + DataSourceOauthBinding.tenant_id == workspace_id, + DataSourceOauthBinding.disabled.is_not(True), + ) + ).all() + return tuple( + DataSourceBindingSummary( + id=binding.id, + provider=binding.provider, + created_at=binding.created_at, + disabled=False, + source_info=_SOURCE_INFO_ADAPTER.validate_python(binding.source_info), + ) + for binding in bindings + ) + + @override + def change_disabled_state( + self, + *, + workspace_id: str, + binding_id: str, + disabled: bool, + ) -> BindingMutationResult: + state_predicate = ( + DataSourceOauthBinding.disabled.is_not(True) if disabled else DataSourceOauthBinding.disabled.is_(True) + ) + with self._session_factory.begin() as session: + result = session.execute( + update(DataSourceOauthBinding) + .where( + DataSourceOauthBinding.tenant_id == workspace_id, + DataSourceOauthBinding.id == binding_id, + state_predicate, + ) + .values(disabled=disabled, updated_at=naive_utc_now()) + ) + if cast(CursorResult[object], result).rowcount == 1: + return "updated" + + binding_exists = session.scalar( + select(DataSourceOauthBinding.id) + .where( + DataSourceOauthBinding.tenant_id == workspace_id, + DataSourceOauthBinding.id == binding_id, + ) + .limit(1) + ) + if binding_exists is None: + return "not_found" + return "already_disabled" if disabled else "already_enabled" diff --git a/api/repositories/knowledge/__init__.py b/api/repositories/knowledge/__init__.py index 2ba7ed3f45b..6769eeaf0f9 100644 --- a/api/repositories/knowledge/__init__.py +++ b/api/repositories/knowledge/__init__.py @@ -1 +1 @@ -"""Knowledge base persistence and queries.""" +"""Knowledge-domain persistence adapters and queries.""" diff --git a/api/repositories/knowledge/dataset_api_key_repository.py b/api/repositories/knowledge/dataset_api_key_repository.py index 0dd1264eced..841effc81e5 100644 --- a/api/repositories/knowledge/dataset_api_key_repository.py +++ b/api/repositories/knowledge/dataset_api_key_repository.py @@ -5,8 +5,8 @@ from typing import override from sqlalchemy import Select, func, select from sqlalchemy.orm import Session, sessionmaker -from models.account import Tenant, TenantAccountJoin -from models.dataset import Dataset, DatasetPermission +from models.account import Tenant +from models.dataset import Dataset from models.enums import ApiTokenType from models.model import ApiToken from repositories.knowledge import dataset_api_key_bindings @@ -17,37 +17,12 @@ from services.auth.api_key_contracts import ( ApiKeyResourceNotFoundError, ) from services.knowledge.api_key_service import DatasetApiKeyStore, UnknownDatasetIdsError -from services.knowledge.dataset_access import DatasetAccess class DatasetApiKeyRepository(DatasetApiKeyStore): def __init__(self, *, session_factory: sessionmaker[Session]) -> None: self._session_factory = session_factory - @override - def get_dataset_access(self, workspace_id: str, dataset_id: str, account_id: str) -> DatasetAccess: - with self._session_factory() as session: - dataset = self._get_dataset(session, workspace_id, dataset_id) - role = session.scalar( - select(TenantAccountJoin.role).where( - TenantAccountJoin.tenant_id == workspace_id, TenantAccountJoin.account_id == account_id - ) - ) - has_permission = ( - session.scalar( - select(DatasetPermission.id) - .where( - DatasetPermission.tenant_id == workspace_id, - DatasetPermission.dataset_id == dataset_id, - DatasetPermission.account_id == account_id, - DatasetPermission.has_permission.is_(True), - ) - .limit(1) - ) - is not None - ) - return DatasetAccess(dataset.permission, dataset.maintainer, role, has_permission) - @override def list_keys(self, workspace_id: str, dataset_id: str) -> tuple[ApiKeyRecord, ...]: with self._session_factory() as session: diff --git a/api/repositories/knowledge/dataset_read_repository.py b/api/repositories/knowledge/dataset_read_repository.py new file mode 100644 index 00000000000..a690eee798c --- /dev/null +++ b/api/repositories/knowledge/dataset_read_repository.py @@ -0,0 +1,755 @@ +"""Explicit reads for knowledge models within the caller's transaction. + +These query functions are shared by repositories and existing transaction owners; +they never create, commit, or close a session and perform no external I/O. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass, field +from typing import Any + +import sqlalchemy as sa +from sqlalchemy import func, select, tuple_ +from sqlalchemy.orm import Session, scoped_session + +from core.app.file_access import DatabaseFileAccessController +from core.rag.entities import ParentMode, Rule +from models.account import Account +from models.dataset import ( + AppDatasetJoin, + ChildChunk, + Dataset, + DatasetBindingItem, + DatasetKeywordTable, + DatasetMetadata, + DatasetMetadataBinding, + DatasetProcessRule, + DatasetQuery, + Document, + DocumentSegment, + ExternalKnowledgeApis, + ExternalKnowledgeBindings, + Pipeline, + PipelineCustomizedTemplate, + SegmentAttachmentBinding, +) +from models.model import App, Tag, TagBinding, UploadFile +from models.tools import ToolFile + + +@dataclass(frozen=True) +class DocumentReadBatch: + uploads: dict[tuple[str, str], UploadFile] = field(default_factory=dict) + metadatas: dict[str, list[DatasetMetadata]] = field(default_factory=dict) + built_in_enabled: dict[tuple[str, str], bool] = field(default_factory=dict) + uploader_names: dict[str, str] = field(default_factory=dict) + segment_counts: dict[str, int] = field(default_factory=dict) + hit_counts: dict[str, int] = field(default_factory=dict) + process_rules: dict[str, DatasetProcessRule] = field(default_factory=dict) + + +def load_document_read_batch(documents: Sequence[Document], *, session: Session) -> DocumentReadBatch: + """Load a page's document relations with a bounded number of queries.""" + if not documents: + return DocumentReadBatch() + owners = [(doc.tenant_id, doc.dataset_id, doc.id) for doc in documents] + dataset_owners = list({(doc.tenant_id, doc.dataset_id) for doc in documents}) + upload_owners = list( + { + (doc.tenant_id, file_id) + for doc in documents + if doc.data_source_type == "upload_file" and (file_id := doc.data_source_info_dict.get("upload_file_id")) + } + ) + uploads = ( + { + (file.tenant_id, file.id): file + for file in session.scalars( + select(UploadFile).where(tuple_(UploadFile.tenant_id, UploadFile.id).in_(upload_owners)) + ) + } + if upload_owners + else {} + ) + metadatas: dict[str, list[DatasetMetadata]] = {} + for doc_id, metadata in session.execute( + select(DatasetMetadataBinding.document_id, DatasetMetadata) + .join(DatasetMetadata, DatasetMetadata.id == DatasetMetadataBinding.metadata_id) + .where( + tuple_( + DatasetMetadataBinding.tenant_id, DatasetMetadataBinding.dataset_id, DatasetMetadataBinding.document_id + ).in_(owners), + DatasetMetadata.tenant_id == DatasetMetadataBinding.tenant_id, + DatasetMetadata.dataset_id == DatasetMetadataBinding.dataset_id, + ) + ): + metadatas.setdefault(doc_id, []).append(metadata) + built_in_enabled = { + (tenant_id, dataset_id): enabled + for tenant_id, dataset_id, enabled in session.execute( + select(Dataset.tenant_id, Dataset.id, Dataset.built_in_field_enabled).where( + tuple_(Dataset.tenant_id, Dataset.id).in_(dataset_owners) + ) + ) + } + uploader_names = dict( + session.execute(select(Account.id, Account.name).where(Account.id.in_({doc.created_by for doc in documents}))) + .tuples() + .all() + ) + counts = session.execute( + select( + DocumentSegment.document_id, + func.count(DocumentSegment.id), + func.coalesce(func.sum(DocumentSegment.hit_count), 0), + ) + .where(tuple_(DocumentSegment.tenant_id, DocumentSegment.dataset_id, DocumentSegment.document_id).in_(owners)) + .group_by(DocumentSegment.document_id) + ).all() + rule_ids = {doc.dataset_process_rule_id for doc in documents if doc.dataset_process_rule_id} + process_rules = ( + { + rule.id: rule + for rule in session.scalars( + select(DatasetProcessRule) + .join(Dataset, Dataset.id == DatasetProcessRule.dataset_id) + .where(DatasetProcessRule.id.in_(rule_ids), tuple_(Dataset.tenant_id, Dataset.id).in_(dataset_owners)) + ) + } + if rule_ids + else {} + ) + return DocumentReadBatch( + uploads=uploads, + metadatas=metadatas, + built_in_enabled=built_in_enabled, + uploader_names=uploader_names, + segment_counts={doc_id: count for doc_id, count, _ in counts}, + hit_counts={doc_id: hits for doc_id, _, hits in counts}, + process_rules=process_rules, + ) + + +def get_segment_content_file_ids( + segment: DocumentSegment, *, upload_ids: set[str], tool_ids: set[str], session: Session +) -> tuple[set[str], set[str]]: + access = DatabaseFileAccessController() + uploads = ( + session.scalars( + access.apply_upload_file_filters( + select(UploadFile).where(UploadFile.id.in_(upload_ids), UploadFile.tenant_id == segment.tenant_id) + ) + ).all() + if upload_ids + else [] + ) + tools = ( + session.scalars( + access.apply_tool_file_filters( + select(ToolFile).where(ToolFile.id.in_(tool_ids), ToolFile.tenant_id == segment.tenant_id) + ) + ).all() + if tool_ids + else [] + ) + return {file.id for file in uploads}, {file.id for file in tools} + + +def get_segment_attachment_files(segment: DocumentSegment, *, session: Session) -> Sequence[UploadFile]: + return session.scalars( + select(UploadFile) + .join(SegmentAttachmentBinding, UploadFile.id == SegmentAttachmentBinding.attachment_id) + .where( + SegmentAttachmentBinding.tenant_id == segment.tenant_id, + SegmentAttachmentBinding.dataset_id == segment.dataset_id, + SegmentAttachmentBinding.document_id == segment.document_id, + SegmentAttachmentBinding.segment_id == segment.id, + UploadFile.tenant_id == segment.tenant_id, + ) + ).all() + + +def is_retrieved_segment_owned( + segment: DocumentSegment, *, tenant_id: str, dataset_ids: Sequence[str], session: Session +) -> bool: + """Validate retrieved identifiers against the complete persisted owner chain.""" + if segment.tenant_id != tenant_id or segment.dataset_id not in dataset_ids: + return False + return ( + session.scalar( + select(DocumentSegment.id) + .join(Document, Document.id == DocumentSegment.document_id) + .join(Dataset, Dataset.id == Document.dataset_id) + .where( + Dataset.tenant_id == tenant_id, + Dataset.id == segment.dataset_id, + Document.tenant_id == tenant_id, + Document.dataset_id == segment.dataset_id, + Document.id == segment.document_id, + Document.enabled.is_(True), + Document.archived.is_(False), + DocumentSegment.tenant_id == tenant_id, + DocumentSegment.dataset_id == segment.dataset_id, + DocumentSegment.id == segment.id, + DocumentSegment.enabled.is_(True), + ) + ) + is not None + ) + + +def get_document_upload_file(document: Document, *, session: Session) -> UploadFile | None: + file_id = document.data_source_info_dict.get("upload_file_id") + if not file_id: + return None + return session.scalar( + select(UploadFile).where(UploadFile.id == file_id, UploadFile.tenant_id == document.tenant_id) + ) + + +def get_document_metadata_rows(document: Document, *, session: Session) -> Sequence[DatasetMetadata]: + return session.scalars( + select(DatasetMetadata) + .join(DatasetMetadataBinding, DatasetMetadataBinding.metadata_id == DatasetMetadata.id) + .where( + DatasetMetadata.tenant_id == document.tenant_id, + DatasetMetadata.dataset_id == document.dataset_id, + DatasetMetadataBinding.tenant_id == document.tenant_id, + DatasetMetadataBinding.dataset_id == document.dataset_id, + DatasetMetadataBinding.document_id == document.id, + ) + ).all() + + +def get_query_upload_files(query: DatasetQuery, file_ids: set[str], *, session: Session) -> dict[str, UploadFile]: + if not file_ids: + return {} + rows = session.scalars( + select(UploadFile) + .join(Dataset, Dataset.tenant_id == UploadFile.tenant_id) + .where(Dataset.id == query.dataset_id, UploadFile.id.in_(file_ids)) + ).all() + return {file.id: file for file in rows} + + +def get_external_api_bindings_batch( + apis: Sequence[ExternalKnowledgeApis], *, session: Session +) -> dict[str, list[DatasetBindingItem]]: + if not apis: + return {} + allowed = {(api.id, api.tenant_id) for api in apis} + rows = session.execute( + select( + ExternalKnowledgeBindings.external_knowledge_api_id, + ExternalKnowledgeBindings.tenant_id, + Dataset.id, + Dataset.name, + ) + .join(Dataset, Dataset.id == ExternalKnowledgeBindings.dataset_id) + .where( + tuple_(ExternalKnowledgeBindings.external_knowledge_api_id, ExternalKnowledgeBindings.tenant_id).in_( + allowed + ), + Dataset.tenant_id == ExternalKnowledgeBindings.tenant_id, + ) + .distinct() + ).all() + bindings: dict[str, list[DatasetBindingItem]] = {} + for api_id, tenant_id, dataset_id, name in rows: + if (api_id, tenant_id) in allowed: + bindings.setdefault(api_id, []).append({"id": dataset_id, "name": name}) + return bindings + + +def get_external_api_dataset_bindings(api: ExternalKnowledgeApis, *, session: Session) -> list[DatasetBindingItem]: + return get_external_api_bindings_batch([api], session=session).get(api.id, []) + + +def get_dataset_keyword_table(dataset: Dataset, *, session: Session) -> DatasetKeywordTable | None: + return session.scalar(select(DatasetKeywordTable).where(DatasetKeywordTable.dataset_id == dataset.id)) + + +def get_dataset_creator(dataset: Dataset, *, session: Session) -> Account | None: + return session.get(Account, dataset.created_by) + + +def get_dataset_author_name(dataset: Dataset, *, session: Session) -> str | None: + account = get_dataset_creator(dataset, session=session) + if account: + return account.name + return None + + +def get_latest_dataset_process_rule(dataset: Dataset, *, session: Session) -> DatasetProcessRule | None: + return session.scalar( + select(DatasetProcessRule) + .where(DatasetProcessRule.dataset_id == dataset.id) + .order_by(DatasetProcessRule.created_at.desc()) + .limit(1) + ) + + +def get_dataset_app_count(dataset: Dataset, *, session: Session) -> int: + return ( + session.scalar( + select(func.count(AppDatasetJoin.id)).where( + AppDatasetJoin.dataset_id == dataset.id, + App.id == AppDatasetJoin.app_id, + App.tenant_id == dataset.tenant_id, + ) + ) + or 0 + ) + + +def get_dataset_document_count(dataset: Dataset, *, session: Session) -> int: + return ( + session.scalar( + select(func.count(Document.id)).where( + Document.dataset_id == dataset.id, Document.tenant_id == dataset.tenant_id + ) + ) + or 0 + ) + + +def get_dataset_available_document_count(dataset: Dataset, *, session: Session) -> int: + return ( + session.scalar( + select(func.count(Document.id)).where( + Document.dataset_id == dataset.id, + Document.tenant_id == dataset.tenant_id, + Document.indexing_status == "completed", + Document.enabled == True, + Document.archived == False, + ) + ) + or 0 + ) + + +def get_dataset_available_segment_count(dataset: Dataset, *, session: Session) -> int: + return ( + session.scalar( + select(func.count(DocumentSegment.id)).where( + DocumentSegment.dataset_id == dataset.id, + DocumentSegment.tenant_id == dataset.tenant_id, + DocumentSegment.status == "completed", + DocumentSegment.enabled == True, + ) + ) + or 0 + ) + + +def get_dataset_word_count(dataset: Dataset, *, session: Session) -> int: + return ( + session.scalar( + select(func.coalesce(func.sum(Document.word_count), 0)).where( + Document.dataset_id == dataset.id, Document.tenant_id == dataset.tenant_id + ) + ) + or 0 + ) + + +def get_dataset_doc_form(dataset: Dataset, *, session: Session) -> str | None: + if dataset.chunk_structure: + return dataset.chunk_structure + return session.scalar( + select(Document.doc_form) + .where(Document.dataset_id == dataset.id, Document.tenant_id == dataset.tenant_id) + .limit(1) + ) + + +def get_dataset_tags(dataset: Dataset, *, session: Session) -> Sequence[Tag]: + tags = session.scalars( + select(Tag) + .join(TagBinding, Tag.id == TagBinding.tag_id) + .where( + TagBinding.target_id == dataset.id, + TagBinding.tenant_id == dataset.tenant_id, + Tag.tenant_id == dataset.tenant_id, + Tag.type == "knowledge", + ) + ).all() + return tags or [] + + +def get_external_knowledge_info(dataset: Dataset, *, session: Session) -> dict[str, Any] | None: + if dataset.provider != "external": + return None + external_knowledge_binding = session.scalar( + select(ExternalKnowledgeBindings).where( + ExternalKnowledgeBindings.dataset_id == dataset.id, ExternalKnowledgeBindings.tenant_id == dataset.tenant_id + ) + ) + if not external_knowledge_binding: + return None + external_knowledge_api = session.scalar( + select(ExternalKnowledgeApis).where( + ExternalKnowledgeApis.id == external_knowledge_binding.external_knowledge_api_id, + ExternalKnowledgeApis.tenant_id == dataset.tenant_id, + ) + ) + return dataset.build_external_knowledge_info(external_knowledge_binding, external_knowledge_api) + + +def get_dataset_is_published(dataset: Dataset, *, session: Session) -> bool: + if dataset.pipeline_id: + pipeline = session.scalar( + select(Pipeline).where(Pipeline.id == dataset.pipeline_id, Pipeline.tenant_id == dataset.tenant_id) + ) + if pipeline: + return pipeline.is_published + return False + + +def get_dataset_doc_metadata(dataset: Dataset, *, session: Session) -> list[dict[str, str]]: + dataset_metadatas = session.scalars( + select(DatasetMetadata).where( + DatasetMetadata.dataset_id == dataset.id, DatasetMetadata.tenant_id == dataset.tenant_id + ) + ).all() + return dataset.build_doc_metadata(dataset_metadatas) + + +def get_document_process_rule(document: Document, *, session: Session) -> DatasetProcessRule | None: + if document.dataset_process_rule_id: + return session.scalar( + select(DatasetProcessRule).where( + DatasetProcessRule.id == document.dataset_process_rule_id, + DatasetProcessRule.dataset_id == document.dataset_id, + ) + ) + return None + + +def get_document_dataset(document: Document, *, session: Session) -> Dataset | None: + """Load the owning dataset with the caller-owned database session.""" + return session.scalar( + select(Dataset).where(Dataset.id == document.dataset_id, Dataset.tenant_id == document.tenant_id) + ) + + +def get_document_segment_count(document: Document, *, session: Session) -> int: + return ( + session.scalar( + select(func.count(DocumentSegment.id)).where( + DocumentSegment.document_id == document.id, + DocumentSegment.dataset_id == document.dataset_id, + DocumentSegment.tenant_id == document.tenant_id, + ) + ) + or 0 + ) + + +def get_document_hit_count(document: Document, *, session: Session) -> int: + return ( + session.scalar( + select(func.coalesce(func.sum(DocumentSegment.hit_count), 0)).where( + DocumentSegment.document_id == document.id, + DocumentSegment.dataset_id == document.dataset_id, + DocumentSegment.tenant_id == document.tenant_id, + ) + ) + or 0 + ) + + +def get_document_uploader(document: Document, *, session: Session) -> str | None: + user = session.scalar(select(Account).where(Account.id == document.created_by)) + return user.name if user else None + + +def get_segment_dataset(segment: DocumentSegment, *, session: Session) -> Dataset | None: + """Load the owning dataset with the caller-owned database session.""" + return session.scalar( + select(Dataset).where(Dataset.id == segment.dataset_id, Dataset.tenant_id == segment.tenant_id) + ) + + +def get_segment_document(segment: DocumentSegment, *, session: Session) -> Document | None: + """Load the owning document with the caller-owned database session.""" + return session.scalar( + select(Document).where( + Document.id == segment.document_id, + Document.dataset_id == segment.dataset_id, + Document.tenant_id == segment.tenant_id, + ) + ) + + +def get_previous_segment(segment: DocumentSegment, session: Session) -> DocumentSegment | None: + return session.scalar( + select(DocumentSegment).where( + DocumentSegment.document_id == segment.document_id, + DocumentSegment.dataset_id == segment.dataset_id, + DocumentSegment.tenant_id == segment.tenant_id, + DocumentSegment.position == segment.position - 1, + ) + ) + + +def get_next_segment(segment: DocumentSegment, session: Session) -> DocumentSegment | None: + return session.scalar( + select(DocumentSegment).where( + DocumentSegment.document_id == segment.document_id, + DocumentSegment.dataset_id == segment.dataset_id, + DocumentSegment.tenant_id == segment.tenant_id, + DocumentSegment.position == segment.position + 1, + ) + ) + + +def get_segment_child_chunks( + segment: DocumentSegment, *, session: Session, include_full_doc: bool = True +) -> Sequence[ChildChunk]: + """Load hierarchical child chunks with the caller-owned database session.""" + document = get_segment_document(segment, session=session) + if not document: + return [] + process_rule = get_document_process_rule(document, session=session) + if process_rule and process_rule.mode == "hierarchical": + rules_dict = process_rule.rules_dict + if rules_dict: + rules = Rule.model_validate(rules_dict) + if rules.parent_mode and (include_full_doc or rules.parent_mode != ParentMode.FULL_DOC): + child_chunks = session.scalars( + select(ChildChunk) + .where( + ChildChunk.segment_id == segment.id, + ChildChunk.tenant_id == segment.tenant_id, + ChildChunk.dataset_id == segment.dataset_id, + ChildChunk.document_id == segment.document_id, + ) + .order_by(ChildChunk.position.asc()) + ).all() + return child_chunks or [] + return [] + + +def get_child_dataset(chunk: ChildChunk, session: Session) -> Dataset | None: + return session.scalar(select(Dataset).where(Dataset.id == chunk.dataset_id, Dataset.tenant_id == chunk.tenant_id)) + + +def get_child_document(chunk: ChildChunk, session: Session) -> Document | None: + return session.scalar( + select(Document).where( + Document.id == chunk.document_id, + Document.dataset_id == chunk.dataset_id, + Document.tenant_id == chunk.tenant_id, + ) + ) + + +def get_child_segment(chunk: ChildChunk, session: Session) -> DocumentSegment | None: + return session.scalar( + select(DocumentSegment).where( + DocumentSegment.id == chunk.segment_id, + DocumentSegment.document_id == chunk.document_id, + DocumentSegment.dataset_id == chunk.dataset_id, + DocumentSegment.tenant_id == chunk.tenant_id, + ) + ) + + +def get_joined_app(binding: AppDatasetJoin, session: Session) -> App | None: + return session.get(App, binding.app_id) + + +def get_pipeline_template_creator_name(template: PipelineCustomizedTemplate, session: Session) -> str: + account = session.scalar(select(Account).where(Account.id == template.created_by)) + if account: + return account.name + return "" + + +def get_pipeline_dataset(pipeline: Pipeline, session: Session | scoped_session) -> Dataset | None: + return session.scalar( + select(Dataset).where(Dataset.pipeline_id == pipeline.id, Dataset.tenant_id == pipeline.tenant_id) + ) + + +@dataclass(frozen=True) +class DatasetDetailPrefetch: + """Batch-loaded values backing the session-scoped fields of a page of datasets. + + Building a detail response for a single dataset issues about ten queries for + that row alone, so a list endpoint costs one query per dataset per field. + Loading the same values for every dataset on the page up front keeps the + query count bounded regardless of the page size. + """ + + app_counts: dict[str, int] = field(default_factory=dict) + document_counts: dict[str, int] = field(default_factory=dict) + word_counts: dict[str, int] = field(default_factory=dict) + available_document_counts: dict[str, int] = field(default_factory=dict) + author_names: dict[str, str] = field(default_factory=dict) + tags: dict[tuple[str, str], list[Tag]] = field(default_factory=dict) + doc_forms: dict[tuple[str, str], str | None] = field(default_factory=dict) + external_knowledge_infos: dict[tuple[str, str], dict[str, Any]] = field(default_factory=dict) + doc_metadata: dict[str, list[DatasetMetadata]] = field(default_factory=dict) + pipeline_published: dict[str, bool] = field(default_factory=dict) + + +def build_dataset_detail_prefetch(datasets: Sequence[Dataset], *, session: Session) -> DatasetDetailPrefetch: + """Load every session-scoped detail field for the given datasets in a fixed number of queries.""" + dataset_ids = [dataset.id for dataset in datasets] + if not dataset_ids: + return DatasetDetailPrefetch() + + tenant_ids = {dataset.tenant_id for dataset in datasets} + dataset_owners = [(dataset.tenant_id, dataset.id) for dataset in datasets] + + app_counts = dict( + session.execute( + select(AppDatasetJoin.dataset_id, func.count(AppDatasetJoin.id)) + .where( + AppDatasetJoin.dataset_id.in_(dataset_ids), + App.id == AppDatasetJoin.app_id, + tuple_(App.tenant_id, AppDatasetJoin.dataset_id).in_(dataset_owners), + ) + .group_by(AppDatasetJoin.dataset_id) + ) + .tuples() + .all() + ) + + document_counts: dict[str, int] = {} + word_counts: dict[str, int] = {} + available_document_counts: dict[str, int] = {} + available_document_case = sa.case( + ( + sa.and_( + Document.indexing_status == "completed", + Document.enabled == True, + Document.archived == False, + ), + 1, + ), + else_=0, + ) + document_rows = session.execute( + select( + Document.dataset_id, + func.count(Document.id), + func.coalesce(func.sum(Document.word_count), 0), + func.coalesce(func.sum(available_document_case), 0), + ) + .where(tuple_(Document.tenant_id, Document.dataset_id).in_(dataset_owners)) + .group_by(Document.dataset_id) + ).all() + for dataset_id, document_count, word_count, available_count in document_rows: + document_counts[dataset_id] = document_count + word_counts[dataset_id] = word_count + available_document_counts[dataset_id] = available_count + + # doc_form only falls back to the documents table for datasets without a + # chunk structure of their own. + doc_forms: dict[tuple[str, str], str | None] = {} + doc_form_dataset_ids = [dataset.id for dataset in datasets if not dataset.chunk_structure] + if doc_form_dataset_ids: + doc_forms = { + (dataset_id, tenant_id): doc_form + for dataset_id, tenant_id, doc_form in session.execute( + select(Document.dataset_id, Document.tenant_id, func.min(Document.doc_form)) + .where( + Document.dataset_id.in_(doc_form_dataset_ids), + tuple_(Document.tenant_id, Document.dataset_id).in_(dataset_owners), + ) + .group_by(Document.dataset_id, Document.tenant_id) + ).all() + } + + author_ids = {dataset.created_by for dataset in datasets if dataset.created_by} + author_names = dict( + session.execute(select(Account.id, Account.name).where(Account.id.in_(author_ids))).tuples().all() + ) + + tags: dict[tuple[str, str], list[Tag]] = {} + tag_rows = session.execute( + select(TagBinding.target_id, TagBinding.tenant_id, Tag) + .join(TagBinding, Tag.id == TagBinding.tag_id) + .where( + tuple_(TagBinding.tenant_id, TagBinding.target_id).in_(dataset_owners), + Tag.tenant_id == TagBinding.tenant_id, + Tag.type == "knowledge", + ) + ).all() + for target_id, tenant_id, tag in tag_rows: + tags.setdefault((target_id, tenant_id), []).append(tag) + + doc_metadata: dict[str, list[DatasetMetadata]] = {} + for dataset_metadata in session.scalars( + select(DatasetMetadata).where(tuple_(DatasetMetadata.tenant_id, DatasetMetadata.dataset_id).in_(dataset_owners)) + ).all(): + doc_metadata.setdefault(dataset_metadata.dataset_id, []).append(dataset_metadata) + + pipeline_ids = {dataset.pipeline_id for dataset in datasets if dataset.pipeline_id} + pipeline_published: dict[str, bool] = {} + if pipeline_ids: + pipeline_owners = [(dataset.tenant_id, dataset.pipeline_id) for dataset in datasets if dataset.pipeline_id] + pipeline_published = dict( + session.execute( + select(Pipeline.id, Pipeline.is_published).where( + tuple_(Pipeline.tenant_id, Pipeline.id).in_(pipeline_owners) + ) + ) + .tuples() + .all() + ) + + return DatasetDetailPrefetch( + app_counts=app_counts, + document_counts=document_counts, + word_counts=word_counts, + available_document_counts=available_document_counts, + author_names=author_names, + tags=tags, + doc_forms=doc_forms, + external_knowledge_infos=_load_external_knowledge_infos(datasets, session=session), + doc_metadata=doc_metadata, + pipeline_published=pipeline_published, + ) + + +def _load_external_knowledge_infos( + datasets: Sequence[Dataset], *, session: Session +) -> dict[tuple[str, str], dict[str, Any]]: + external_datasets = [dataset for dataset in datasets if dataset.provider == "external"] + if not external_datasets: + return {} + + external_owners = [(dataset.tenant_id, dataset.id) for dataset in external_datasets] + external_tenant_ids = {dataset.tenant_id for dataset in external_datasets} + bindings = session.scalars( + select(ExternalKnowledgeBindings).where( + tuple_(ExternalKnowledgeBindings.tenant_id, ExternalKnowledgeBindings.dataset_id).in_(external_owners), + ) + ).all() + if not bindings: + return {} + + apis = { + (api.id, api.tenant_id): api + for api in session.scalars( + select(ExternalKnowledgeApis).where( + ExternalKnowledgeApis.id.in_({binding.external_knowledge_api_id for binding in bindings}), + ExternalKnowledgeApis.tenant_id.in_(external_tenant_ids), + ) + ).all() + } + + infos: dict[tuple[str, str], dict[str, Any]] = {} + for binding in bindings: + info = Dataset.build_external_knowledge_info( + binding, apis.get((binding.external_knowledge_api_id, binding.tenant_id)) + ) + if info is not None: + infos[(binding.dataset_id, binding.tenant_id)] = info + return infos diff --git a/api/repositories/knowledge/dataset_repository.py b/api/repositories/knowledge/dataset_repository.py new file mode 100644 index 00000000000..447b4d8a4e8 --- /dev/null +++ b/api/repositories/knowledge/dataset_repository.py @@ -0,0 +1,153 @@ +"""SQLAlchemy repository for tenant-owned dataset state.""" + +from sqlalchemy import and_, exists, literal, select +from sqlalchemy.orm import Session, sessionmaker + +from models.dataset import Dataset, DatasetPermission, Document +from services.knowledge.dataset_access import DatasetAccessRecord, DatasetAccessSnapshot +from services.knowledge.indexing.estimate import DatasetEstimateRecord +from services.knowledge.resource_scope import DatasetRef +from services.knowledge.segments.application import SegmentDatasetRecord, SegmentDocumentRecord, SegmentScope + + +def _dataset_access_record(dataset: Dataset) -> DatasetAccessRecord: + return DatasetAccessRecord( + id=dataset.id, + workspace_id=dataset.tenant_id, + maintainer_id=dataset.maintainer, + permission=str(dataset.permission), + ) + + +def _dataset_estimate_record(dataset: Dataset) -> DatasetEstimateRecord: + return DatasetEstimateRecord( + id=dataset.id, + workspace_id=dataset.tenant_id, + indexing_technique=str(dataset.indexing_technique) if dataset.indexing_technique is not None else None, + ) + + +def _get_dataset(session: Session, dataset_ref: DatasetRef) -> Dataset | None: + """Shared SQL for implementation adapters that already own a transaction.""" + return session.scalar( + select(Dataset) + .where( + Dataset.id == dataset_ref.dataset_id, + Dataset.tenant_id == dataset_ref.tenant_id, + ) + .limit(1) + ) + + +class SQLAlchemyDatasetRepository: + """Own SQL access to tenant-scoped dataset state.""" + + def __init__(self, *, session_factory: sessionmaker[Session]) -> None: + self._session_factory = session_factory + + def get_estimate_record(self, dataset_ref: DatasetRef) -> DatasetEstimateRecord | None: + with self._session_factory() as session: + dataset = _get_dataset(session, dataset_ref) + return _dataset_estimate_record(dataset) if dataset is not None else None + + def is_notion_dataset(self, dataset_ref: DatasetRef) -> bool: + """Return whether the owned dataset is configured for Notion imports.""" + with self._session_factory() as session: + return bool( + session.scalar( + select( + exists().where( + Dataset.id == dataset_ref.dataset_id, + Dataset.tenant_id == dataset_ref.tenant_id, + Dataset.data_source_type == "notion_import", + ) + ) + ) + ) + + def get_access_snapshot( + self, + *, + workspace_id: str, + dataset_id: str, + actor_id: str, + ) -> DatasetAccessSnapshot | None: + with self._session_factory() as session: + dataset = _get_dataset(session, DatasetRef(workspace_id, dataset_id)) + if dataset is None: + return None + + actor_has_partial_access = bool( + session.scalar( + select( + exists().where( + DatasetPermission.tenant_id == workspace_id, + DatasetPermission.dataset_id == dataset_id, + DatasetPermission.account_id == actor_id, + DatasetPermission.has_permission.is_(True), + ) + ) + ) + ) + return DatasetAccessSnapshot( + dataset=_dataset_access_record(dataset), + actor_has_partial_access=actor_has_partial_access, + ) + + def get_segment_scope( + self, + *, + dataset_ref: DatasetRef, + document_id: str, + actor_id: str | None, + ) -> SegmentScope: + """Read ownership, model configuration and the actor grant in one snapshot.""" + grant = ( + exists().where( + DatasetPermission.tenant_id == dataset_ref.tenant_id, + DatasetPermission.dataset_id == dataset_ref.dataset_id, + DatasetPermission.account_id == actor_id, + DatasetPermission.has_permission.is_(True), + ) + if actor_id is not None + else literal(False) + ) + with self._session_factory() as session: + row = session.execute( + select(Dataset, Document, grant) + .outerjoin( + Document, + and_( + Document.id == document_id, + Document.dataset_id == Dataset.id, + Document.tenant_id == Dataset.tenant_id, + ), + ) + .where(Dataset.id == dataset_ref.dataset_id, Dataset.tenant_id == dataset_ref.tenant_id) + ).one_or_none() + if row is None: + return SegmentScope(None, None) + dataset, document, actor_has_access = row + return SegmentScope( + dataset=SegmentDatasetRecord( + id=dataset.id, + workspace_id=dataset.tenant_id, + indexing_technique=str(dataset.indexing_technique), + embedding_model_provider=dataset.embedding_model_provider, + embedding_model=dataset.embedding_model, + ), + document=SegmentDocumentRecord( + id=document.id, + dataset_id=document.dataset_id, + workspace_id=document.tenant_id, + doc_form=str(document.doc_form), + ) + if document is not None + else None, + access_snapshot=DatasetAccessSnapshot( + dataset=_dataset_access_record(dataset), + actor_has_partial_access=bool(actor_has_access), + ) + if actor_id is not None + else None, + ) diff --git a/api/repositories/knowledge/document_repository.py b/api/repositories/knowledge/document_repository.py new file mode 100644 index 00000000000..55ba407c8f3 --- /dev/null +++ b/api/repositories/knowledge/document_repository.py @@ -0,0 +1,241 @@ +"""SQLAlchemy repository for tenant-owned document state.""" + +import json +from collections.abc import Mapping + +from pydantic import TypeAdapter, ValidationError +from sqlalchemy import exists, select, update +from sqlalchemy.orm import Session, sessionmaker + +from libs.datetime_utils import naive_utc_now +from models import Account +from models.dataset import Dataset, DatasetProcessRule, Document, Pipeline +from models.enums import IndexingStatus +from repositories.knowledge.dataset_read_repository import get_pipeline_dataset +from services.knowledge.document_sync import SyncDocumentRecord +from services.knowledge.indexing.errors import DocumentIsDeletedPausedError, DocumentIsPausedError +from services.knowledge.indexing.estimate import EstimateDocumentRecord, StoredSource +from services.knowledge.indexing.execution import IndexingDocument +from services.knowledge.resource_scope import DatasetRef, DocumentRef + +_MAPPING_ADAPTER = TypeAdapter(dict[str, object]) + + +def _mapping(value: object) -> Mapping[str, object] | None: + if value is None: + return None + try: + if isinstance(value, str): + return _MAPPING_ADAPTER.validate_python(json.loads(value)) + return _MAPPING_ADAPTER.validate_python(value) + except (TypeError, ValueError, ValidationError): + return {} + + +def _estimate_document_record(document: Document) -> EstimateDocumentRecord: + return EstimateDocumentRecord( + id=document.id, + workspace_id=document.tenant_id, + dataset_id=document.dataset_id, + data_source_type=str(document.data_source_type), + data_source_info=_mapping(document.data_source_info), + indexing_status=str(document.indexing_status), + doc_form=str(document.doc_form), + doc_language=document.doc_language, + dataset_process_rule_id=document.dataset_process_rule_id, + ) + + +def _get_document(session: Session, document_ref: DocumentRef) -> Document | None: + """Shared SQL for implementation adapters that already own a transaction.""" + return session.scalar( + select(Document) + .where( + Document.id == document_ref.document_id, + Document.dataset_id == document_ref.dataset.dataset_id, + Document.tenant_id == document_ref.dataset.tenant_id, + ) + .limit(1) + ) + + +def require_indexing_document(session: Session, ref: DocumentRef, *, lock: bool = False) -> Document: + """Validate the persisted owner chain before each indexing phase.""" + statement = ( + select(Document) + .join(Dataset, Dataset.id == Document.dataset_id) + .where( + Document.id == ref.document_id, + Document.dataset_id == ref.dataset.dataset_id, + Document.tenant_id == ref.dataset.tenant_id, + Dataset.tenant_id == ref.dataset.tenant_id, + ) + ) + if lock: + statement = statement.with_for_update(of=Document) + document = session.scalar(statement) + if document is None: + raise DocumentIsDeletedPausedError() + if document.is_paused: + raise DocumentIsPausedError() + return document + + +class SQLAlchemyDocumentRepository: + """Own SQL access to documents through their complete owner chain.""" + + def __init__(self, *, session_factory: sessionmaker[Session]) -> None: + self._session_factory = session_factory + + def get_pipeline_dataset(self, pipeline: Pipeline, *, session: Session) -> Dataset | None: + """Read within the pipeline caller's transaction, including uncommitted changes.""" + return get_pipeline_dataset(pipeline, session=session) + + def get_indexing_document(self, ref: DocumentRef) -> IndexingDocument: + with self._session_factory() as session: + document = require_indexing_document(session, ref) + rule = session.scalar( + select(DatasetProcessRule).where( + DatasetProcessRule.id == document.dataset_process_rule_id, + DatasetProcessRule.dataset_id == ref.dataset.dataset_id, + ) + ) + return IndexingDocument( + source=StoredSource.from_document(_estimate_document_record(document)), + processing_rule=rule.to_dict() if rule else {}, + doc_language=document.doc_language or "English", + need_summary=bool(document.need_summary), + ) + + def get_indexing_models(self, ref: DocumentRef) -> tuple[Dataset, Document]: + """Detached legacy processor inputs; never pass them to another thread.""" + with self._session_factory() as session: + document = require_indexing_document(session, ref) + dataset = session.scalar( + select(Dataset).where( + Dataset.id == ref.dataset.dataset_id, + Dataset.tenant_id == ref.dataset.tenant_id, + ) + ) + assert dataset is not None + return dataset, document + + def get_indexing_user(self, ref: DocumentRef) -> Account: + with self._session_factory() as session: + document = require_indexing_document(session, ref) + account = session.get(Account, document.created_by) + if account is None: + raise ValueError("no current user found") + account.set_tenant_id_with_session(ref.dataset.tenant_id, session=session) + return account + + def mark_splitting(self, ref: DocumentRef) -> None: + with self._session_factory.begin() as session: + document = require_indexing_document(session, ref, lock=True) + document.indexing_status = IndexingStatus.SPLITTING + document.parsing_completed_at = naive_utc_now() + + def complete_indexing(self, ref: DocumentRef, *, tokens: int, latency: float) -> None: + with self._session_factory.begin() as session: + document = require_indexing_document(session, ref, lock=True) + document.indexing_status = IndexingStatus.COMPLETED + document.tokens = tokens + document.indexing_latency = latency + document.completed_at = naive_utc_now() + document.error = None + + def fail_indexing(self, ref: DocumentRef, error: str) -> None: + try: + with self._session_factory.begin() as session: + document = require_indexing_document(session, ref, lock=True) + document.indexing_status = IndexingStatus.ERROR + document.error = error + document.stopped_at = naive_utc_now() + except DocumentIsDeletedPausedError: + return + + def exists(self, *, workspace_id: str, dataset_id: str, document_id: str) -> bool: + with self._session_factory() as session: + return bool( + session.scalar( + select( + exists().where( + Document.tenant_id == workspace_id, + Document.dataset_id == dataset_id, + Document.id == document_id, + ) + ) + ) + ) + + def mark_failed(self, *, workspace_id: str, dataset_id: str, document_id: str, error: str) -> None: + with self._session_factory.begin() as session: + session.execute( + update(Document) + .where( + Document.tenant_id == workspace_id, + Document.dataset_id == dataset_id, + Document.id == document_id, + ) + .values(indexing_status=IndexingStatus.ERROR, error=error) + ) + + def get_estimate_document(self, document_ref: DocumentRef) -> EstimateDocumentRecord | None: + with self._session_factory() as session: + document = _get_document(session, document_ref) + return _estimate_document_record(document) if document is not None else None + + def list_estimate_documents_by_batch( + self, dataset_ref: DatasetRef, batch: str + ) -> tuple[EstimateDocumentRecord, ...]: + with self._session_factory() as session: + documents = session.scalars( + select(Document).where( + Document.tenant_id == dataset_ref.tenant_id, + Document.dataset_id == dataset_ref.dataset_id, + Document.batch == batch, + ) + ).all() + return tuple(_estimate_document_record(document) for document in documents) + + def get_sync_document(self, document_ref: DocumentRef) -> SyncDocumentRecord | None: + with self._session_factory() as session: + document = _get_document(session, document_ref) + if document is None: + return None + return SyncDocumentRecord(id=document.id, data_source_type=str(document.data_source_type)) + + def list_active_notion_refs(self, dataset_ref: DatasetRef) -> tuple[DocumentRef, ...]: + with self._session_factory() as session: + return tuple( + dataset_ref.document(document_id) + for document_id in session.scalars( + select(Document.id).where( + Document.tenant_id == dataset_ref.tenant_id, + Document.dataset_id == dataset_ref.dataset_id, + Document.data_source_type == "notion_import", + Document.enabled.is_(True), + Document.archived.is_(False), + ) + ).all() + ) + + def list_bound_notion_page_ids(self, dataset_ref: DatasetRef) -> frozenset[str]: + page_ids: set[str] = set() + with self._session_factory() as session: + documents = session.scalars( + select(Document).where( + Document.tenant_id == dataset_ref.tenant_id, + Document.dataset_id == dataset_ref.dataset_id, + Document.data_source_type == "notion_import", + Document.enabled.is_(True), + ) + ).all() + for document in documents: + source_info = _mapping(document.data_source_info) + if source_info is None: + continue + page_id = source_info.get("notion_page_id") + if isinstance(page_id, str) and page_id: + page_ids.add(page_id) + return frozenset(page_ids) diff --git a/api/repositories/knowledge/keyword_table_repository.py b/api/repositories/knowledge/keyword_table_repository.py new file mode 100644 index 00000000000..e0af43c9be6 --- /dev/null +++ b/api/repositories/knowledge/keyword_table_repository.py @@ -0,0 +1,121 @@ +"""Keyword table serialization, object storage, and tenant-scoped SQL persistence.""" + +import json +from collections.abc import Mapping, Sequence +from typing import Any, Protocol + +import orjson +from sqlalchemy import select +from sqlalchemy.orm import Session + +from models.dataset import Dataset, DatasetKeywordTable, DocumentSegment +from repositories.knowledge.dataset_read_repository import get_dataset_keyword_table + + +class _KeywordStorage(Protocol): + def load_once(self, key: str, /) -> bytes: ... + + def exists(self, key: str, /) -> bool: ... + + def save(self, key: str, data: bytes, /) -> None: ... + + def delete(self, key: str, /) -> None: ... + + +def keyword_file_key(tenant_id: str, dataset_id: str) -> str: + return f"keyword_files/{tenant_id}/{dataset_id}.txt" + + +def _decode_sets(values: dict[str, Any]) -> dict[str, Any]: + return {key: set(value) if isinstance(value, list) else value for key, value in values.items()} + + +def _encode_sets(value: object) -> list[object]: + if isinstance(value, set): + return list(value) + raise TypeError(f"Object of type {type(value).__name__} is not JSON serializable") + + +def load_keyword_table( + *, storage: _KeywordStorage, tenant_id: str, dataset_id: str, storage_type: str, data: str | None +) -> dict[str, Any] | None: + if storage_type == "database": + return json.loads(data, object_hook=_decode_sets) if data else None + file_key = keyword_file_key(tenant_id, dataset_id) + try: + payload = storage.load_once(file_key) + return json.loads(payload.decode("utf-8"), object_hook=_decode_sets) if payload else None + except FileNotFoundError: + return None + + +def load_dataset_keyword_table( + dataset: Dataset, *, storage: _KeywordStorage, session: Session +) -> dict[str, Any] | None: + """Resolve a dataset's stored keyword payload using the caller's session.""" + row = get_dataset_keyword_table(dataset, session=session) + if row is None: + return None + return load_keyword_table( + storage=storage, + tenant_id=dataset.tenant_id, + dataset_id=dataset.id, + storage_type=row.data_source_type, + data=row.keyword_table, + ) + + +def save_keyword_table( + *, storage: _KeywordStorage, tenant_id: str, dataset_id: str, storage_type: str, table: dict[str, set[str]] +) -> str: + data = orjson.dumps( + {"__type__": "keyword_table", "__data__": {"index_id": dataset_id, "summary": None, "table": table}}, + default=_encode_sets, + ) + if storage_type != "database": + file_key = keyword_file_key(tenant_id, dataset_id) + if storage.exists(file_key): + storage.delete(file_key) + storage.save(file_key, data) + return data.decode("utf-8") + + +def persist_keyword_table( + session: Session, + *, + tenant_id: str, + dataset_id: str, + storage_type: str, + data: str, + keywords: Mapping[str, Sequence[str]], +) -> None: + """Write a tenant-owned table and segment keywords in the caller's transaction.""" + result = session.execute( + select(Dataset.id, DatasetKeywordTable) + .outerjoin(DatasetKeywordTable, DatasetKeywordTable.dataset_id == Dataset.id) + .where(Dataset.id == dataset_id, Dataset.tenant_id == tenant_id) + ).one_or_none() + if result is None: + raise LookupError("Dataset no longer exists") + _, row = result + if row is None: + row = DatasetKeywordTable( + dataset_id=dataset_id, + data_source_type=storage_type, + keyword_table=data if storage_type == "database" else "", + ) + session.add(row) + elif storage_type == "database": + row.keyword_table = data + if keywords: + segments = session.scalars( + select(DocumentSegment).where( + DocumentSegment.tenant_id == tenant_id, + DocumentSegment.dataset_id == dataset_id, + DocumentSegment.index_node_id.in_(keywords), + ) + ) + for segment in segments: + assert segment.index_node_id is not None + segment.keywords = list(keywords[segment.index_node_id]) + session.flush() diff --git a/api/repositories/knowledge/metadata_repository.py b/api/repositories/knowledge/metadata_repository.py new file mode 100644 index 00000000000..1610fd3ab0d --- /dev/null +++ b/api/repositories/knowledge/metadata_repository.py @@ -0,0 +1,212 @@ +"""Tenant-scoped metadata operations; each mutation owns one atomic transaction.""" + +from collections.abc import Sequence + +from sqlalchemy import func, select +from sqlalchemy.orm import Session, sessionmaker + +from core.rag.index_processor.constant.built_in_field import BuiltInField, MetadataDataSource +from libs.datetime_utils import naive_utc_now +from models.dataset import Dataset, DatasetMetadata, DatasetMetadataBinding, Document +from repositories.knowledge.dataset_read_repository import get_document_uploader +from services.errors.metadata import MetadataResourceNotFoundError +from services.knowledge.metadata.application import MetadataList, MetadataOperation, MetadataRecord +from services.knowledge.resource_scope import DatasetRef + + +class SQLAlchemyMetadataRepository: + def __init__(self, *, session_factory: sessionmaker[Session]) -> None: + self._sessions = session_factory + + @staticmethod + def _dataset(session: Session, ref: DatasetRef, *, write: bool = False) -> Dataset: + query = select(Dataset).where(Dataset.id == ref.dataset_id, Dataset.tenant_id == ref.tenant_id) + if write: + query = query.with_for_update() + dataset = session.scalar(query) + if dataset is None: + raise MetadataResourceNotFoundError("Dataset not found.") + return dataset + + @staticmethod + def _fields(ref: DatasetRef): + return select(DatasetMetadata).where( + DatasetMetadata.tenant_id == ref.tenant_id, DatasetMetadata.dataset_id == ref.dataset_id + ) + + @staticmethod + def _documents(ref: DatasetRef): + return select(Document).where(Document.tenant_id == ref.tenant_id, Document.dataset_id == ref.dataset_id) + + @staticmethod + def _bindings(ref: DatasetRef): + return select(DatasetMetadataBinding).where( + DatasetMetadataBinding.tenant_id == ref.tenant_id, DatasetMetadataBinding.dataset_id == ref.dataset_id + ) + + @staticmethod + def _record(row: DatasetMetadata) -> MetadataRecord: + return MetadataRecord( + id=row.id, + name=row.name, + type=row.type, + tenant_id=row.tenant_id, + dataset_id=row.dataset_id, + created_by=row.created_by, + created_at=row.created_at, + updated_by=row.updated_by, + updated_at=row.updated_at, + ) + + @classmethod + def _check_name(cls, session: Session, ref: DatasetRef, name: str) -> None: + if session.scalar(cls._fields(ref).where(DatasetMetadata.name == name).limit(1)) is not None: + raise ValueError("Metadata name already exists.") + + def create(self, ref: DatasetRef, *, name: str, field_type: str, actor_id: str) -> MetadataRecord: + with self._sessions.begin() as session: + self._dataset(session, ref, write=True) + self._check_name(session, ref, name) + row = DatasetMetadata( + tenant_id=ref.tenant_id, dataset_id=ref.dataset_id, name=name, type=field_type, created_by=actor_id + ) + session.add(row) + session.flush() + return self._record(row) + + def rename(self, ref: DatasetRef, metadata_id: str, *, name: str, actor_id: str) -> MetadataRecord: + with self._sessions.begin() as session: + self._dataset(session, ref, write=True) + self._check_name(session, ref, name) + row = session.scalar(self._fields(ref).where(DatasetMetadata.id == metadata_id)) + if row is None: + raise MetadataResourceNotFoundError("Metadata not found.") + old_name = row.name + row.name, row.updated_by, row.updated_at = name, actor_id, naive_utc_now() + bindings = session.scalars( + self._bindings(ref).where(DatasetMetadataBinding.metadata_id == metadata_id) + ).all() + for document in session.scalars( + self._documents(ref).where(Document.id.in_([b.document_id for b in bindings])) + ): + values = dict(document.doc_metadata or {}) + values[name] = values.pop(old_name, None) + document.doc_metadata = values + return self._record(row) + + def delete(self, ref: DatasetRef, metadata_id: str) -> MetadataRecord: + with self._sessions.begin() as session: + self._dataset(session, ref, write=True) + row = session.scalar(self._fields(ref).where(DatasetMetadata.id == metadata_id)) + if row is None: + raise MetadataResourceNotFoundError("Metadata not found.") + bindings = session.scalars( + self._bindings(ref).where(DatasetMetadataBinding.metadata_id == metadata_id) + ).all() + for document in session.scalars( + self._documents(ref).where(Document.id.in_([b.document_id for b in bindings])) + ): + values = dict(document.doc_metadata or {}) + values.pop(row.name, None) + document.doc_metadata = values + for binding in bindings: + session.delete(binding) + session.delete(row) + return self._record(row) + + @staticmethod + def _built_in(session: Session, document: Document) -> dict[str, str | int | float | None]: + return { + BuiltInField.document_name: document.name, + BuiltInField.uploader: get_document_uploader(document, session=session), + BuiltInField.upload_date: document.upload_date.timestamp(), + BuiltInField.last_update_date: document.last_update_date.timestamp(), + BuiltInField.source: MetadataDataSource[document.data_source_type], + } + + def set_built_in(self, ref: DatasetRef, *, enabled: bool) -> None: + with self._sessions.begin() as session: + dataset = self._dataset(session, ref, write=True) + if dataset.built_in_field_enabled == enabled: + return + for document in session.scalars( + self._documents(ref).where( + Document.enabled.is_(True), Document.archived.is_(False), Document.indexing_status == "completed" + ) + ): + values = dict(document.doc_metadata or {}) + if enabled: + values.update(self._built_in(session, document)) + else: + for key in BuiltInField: + values.pop(key, None) + document.doc_metadata = values + dataset.built_in_field_enabled = enabled + + def update_documents(self, ref: DatasetRef, operations: Sequence[MetadataOperation], *, actor_id: str) -> None: + with self._sessions.begin() as session: + dataset = self._dataset(session, ref, write=True) + metadata_ids = {value.id for operation in operations for value in operation.metadata_list} + fields = { + row.id: row for row in session.scalars(self._fields(ref).where(DatasetMetadata.id.in_(metadata_ids))) + } + if metadata_ids != set(fields): + raise MetadataResourceNotFoundError("Metadata not found.") + document_ids = {operation.document_id for operation in operations} + documents = { + row.id: row + for row in session.scalars(self._documents(ref).where(Document.id.in_(document_ids)).with_for_update()) + } + if document_ids != set(documents): + raise MetadataResourceNotFoundError("Document not found.") + for operation in operations: + document = documents[operation.document_id] + values = dict(document.doc_metadata or {}) if operation.partial_update else {} + values.update({fields[value.id].name: value.value for value in operation.metadata_list}) + if dataset.built_in_field_enabled: + values.update(self._built_in(session, document)) + document.doc_metadata = values + bindings = session.scalars( + self._bindings(ref).where(DatasetMetadataBinding.document_id == document.id) + ).all() + previous = {binding.metadata_id for binding in bindings} + if not operation.partial_update: + for binding in bindings: + session.delete(binding) + previous.clear() + for value in operation.metadata_list: + if value.id in previous: + continue + session.add( + DatasetMetadataBinding( + tenant_id=ref.tenant_id, + dataset_id=ref.dataset_id, + document_id=document.id, + metadata_id=value.id, + created_by=actor_id, + ) + ) + previous.add(value.id) + + def list_fields(self, ref: DatasetRef) -> MetadataList: + with self._sessions() as session: + dataset = self._dataset(session, ref) + counts = dict( + session.execute( + select(DatasetMetadataBinding.metadata_id, func.count()) + .where( + DatasetMetadataBinding.tenant_id == ref.tenant_id, + DatasetMetadataBinding.dataset_id == ref.dataset_id, + ) + .group_by(DatasetMetadataBinding.metadata_id) + ) + .tuples() + .all() + ) + return { + "doc_metadata": [ + {"id": row.id, "name": row.name, "type": row.type, "count": counts.get(row.id, 0)} + for row in session.scalars(self._fields(ref)) + ], + "built_in_field_enabled": dataset.built_in_field_enabled, + } diff --git a/api/repositories/knowledge/segment_read_adapter.py b/api/repositories/knowledge/segment_read_adapter.py new file mode 100644 index 00000000000..23ba6042b98 --- /dev/null +++ b/api/repositories/knowledge/segment_read_adapter.py @@ -0,0 +1,87 @@ +"""Read adapter for authorized segment content and attachment preview values.""" + +import base64 +import hashlib +import hmac +import os +import re +import time +from collections.abc import Sequence +from typing import Literal + +from sqlalchemy.orm import Session + +from configs import dify_config +from core.tools.signature import bind_file_uri, sign_tool_file_uri +from models.dataset import AttachmentItem, DocumentSegment +from models.model import UploadFile +from repositories.knowledge.dataset_read_repository import get_segment_attachment_files, get_segment_content_file_ids + +# Consume the entire old query string and replace matches in source order. +_FILE_ID = r"[a-fA-F0-9]{8}(?:-[a-fA-F0-9]{4}){3}-[a-fA-F0-9]{12}" +_FILE_LINK = re.compile( + rf"/files/(?:(?P{_FILE_ID})/(?Pimage-preview|file-preview)" + rf"|tools/(?P{_FILE_ID})\.(?P[a-zA-Z0-9]+))" + r"(?:\?[^\s)\"'<>]*)?" +) + + +def _sign_upload_file_uri(upload_file_id: str, *, preview_kind: Literal["image", "file"] = "file") -> str: + """Sign an upload ID whose ownership was checked by this read adapter.""" + timestamp = str(int(time.time())) + nonce = os.urandom(16).hex() + payload = f"{preview_kind}-preview|{upload_file_id}|{timestamp}|{nonce}" + signature = hmac.new(dify_config.SECRET_KEY.encode(), payload.encode(), hashlib.sha256).digest() + encoded_sign = base64.urlsafe_b64encode(signature).decode() + return f"/files/{upload_file_id}/{preview_kind}-preview?timestamp={timestamp}&nonce={nonce}&sign={encoded_sign}" + + +def sign_segment_content(segment: DocumentSegment, *, session: Session) -> str: + """Refresh links only after checking persisted file ownership. + + The segment must come from an authorized knowledge read. File queries also + enforce its tenant and the current execution's file-access scope. Inaccessible + references lose their old signature and never receive a new one. + """ + matches = list(_FILE_LINK.finditer(segment.content)) + if not matches: + return segment.content + upload_ids = {match["upload"].lower() for match in matches if match["upload"]} + tool_ids = {match["tool"].lower() for match in matches if match["tool"]} + allowed_uploads, allowed_tools = get_segment_content_file_ids( + segment, upload_ids=upload_ids, tool_ids=tool_ids, session=session + ) + + def replace(match: re.Match[str]) -> str: + if upload_id := match["upload"]: + upload_id = upload_id.lower() + if upload_id not in allowed_uploads: + return match.group().split("?", 1)[0] + return _sign_upload_file_uri( + upload_id, preview_kind="image" if match["preview"] == "image-preview" else "file" + ) + tool_id = match["tool"].lower() + if tool_id not in allowed_tools: + return match.group().split("?", 1)[0] + return sign_tool_file_uri(tool_id, "." + match["extension"]) + + return _FILE_LINK.sub(replace, segment.content) + + +def format_segment_attachments(files: Sequence[UploadFile]) -> list[AttachmentItem]: + base_url = dify_config.FILES_URL or dify_config.CONSOLE_API_URL or "" + return [ + { + "id": file.id, + "name": file.name, + "size": file.size, + "extension": file.extension, + "mime_type": file.mime_type, + "source_url": bind_file_uri(_sign_upload_file_uri(file.id, preview_kind="image"), base_url), + } + for file in files + ] + + +def get_segment_attachments(segment: DocumentSegment, *, session: Session) -> list[AttachmentItem]: + return format_segment_attachments(get_segment_attachment_files(segment, session=session)) diff --git a/api/repositories/knowledge/segment_repository.py b/api/repositories/knowledge/segment_repository.py new file mode 100644 index 00000000000..47766f46644 --- /dev/null +++ b/api/repositories/knowledge/segment_repository.py @@ -0,0 +1,716 @@ +"""Infrastructure adapters for dataset segment use cases.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import cast + +from sqlalchemy import String, case, delete, func, literal, or_, select, update +from sqlalchemy import cast as sql_cast +from sqlalchemy.dialects.postgresql import JSONB +from sqlalchemy.orm import Session, sessionmaker + +from configs import dify_config +from core.rag.datasource.vdb.vector_factory import Vector +from core.rag.docstore.dataset_docstore import DatasetDocumentStore +from core.rag.index_processor.constant.index_type import IndexStructureType +from core.rag.index_processor.processor.paragraph_index_processor import ParagraphIndexProcessor +from core.rag.models.document import ChildDocument +from core.rag.models.document import Document as IndexDocument +from graphon.file import File +from libs.datetime_utils import naive_utc_now +from libs.helper import escape_like_pattern +from libs.pagination import PaginatedResult, paginate_query +from models.dataset import ( + ChildChunk, + Dataset, + DatasetProcessRule, + Document, + DocumentSegment, + DocumentSegmentSummary, + SegmentAttachmentBinding, +) +from models.enums import IndexingStatus, SegmentStatus, SegmentType, SummaryStatus +from repositories.knowledge.dataset_read_repository import get_dataset_keyword_table, get_segment_child_chunks +from repositories.knowledge.dataset_repository import _get_dataset +from repositories.knowledge.document_repository import _get_document, require_indexing_document +from repositories.knowledge.keyword_table_repository import persist_keyword_table +from repositories.knowledge.segment_read_adapter import get_segment_attachments, sign_segment_content +from services.knowledge.entities.segments import ChildChunkRecord, SegmentRecord +from services.knowledge.indexing.execution import IndexingDocument +from services.knowledge.resource_scope import DatasetRef, DocumentRef, SegmentRef +from services.knowledge.segments.application import ( + ChildChunkListFilter, + ChildChunkPage, + ChildChunkState, + SegmentDetail, + SegmentIndexTarget, + SegmentListFilter, + SegmentPage, + SegmentUpdateState, +) + + +class SQLAlchemySegmentRepository: + """Read and write scoped segment rows in bounded database transactions.""" + + def __init__(self, *, session_factory: sessionmaker[Session]) -> None: + self._session_factory = session_factory + + def clear_for_indexing(self, document: IndexingDocument) -> None: + ref = document.ref + with self._session_factory.begin() as session: + require_indexing_document(session, ref, lock=True) + if document.doc_form == IndexStructureType.PARENT_CHILD_INDEX: + session.execute( + delete(ChildChunk).where( + ChildChunk.tenant_id == ref.dataset.tenant_id, + ChildChunk.dataset_id == ref.dataset.dataset_id, + ChildChunk.document_id == ref.document_id, + ) + ) + session.execute(delete(DocumentSegment).where(*self._indexing_scope(ref))) + + def save_for_indexing( + self, document: IndexingDocument, chunks: list[IndexDocument], token_counts: list[int] + ) -> None: + ref = document.ref + with self._session_factory.begin() as session: + row = require_indexing_document(session, ref, lock=True) + dataset = _get_dataset(session, ref.dataset) + assert dataset is not None + DatasetDocumentStore(dataset=dataset, user_id=row.created_by, document_id=row.id).add_documents( + session=session, + docs=chunks, + token_counts=token_counts, + save_child=document.doc_form == IndexStructureType.PARENT_CHILD_INDEX, + ) + now = naive_utc_now() + row.indexing_status = IndexingStatus.INDEXING + row.cleaning_completed_at = now + row.splitting_completed_at = now + row.word_count = sum(len(chunk.page_content) for chunk in chunks) + session.execute( + update(DocumentSegment) + .where(*self._indexing_scope(ref)) + .values( + status=SegmentStatus.INDEXING, + indexing_at=now, + ) + ) + + def resume_indexing(self, document: IndexingDocument) -> tuple[list[IndexDocument], int]: + ref = document.ref + with self._session_factory() as session: + require_indexing_document(session, ref) + rows = session.scalars( + select(DocumentSegment).where(*self._indexing_scope(ref)).order_by(DocumentSegment.position) + ).all() + chunks = [] + for row in rows: + if row.status == SegmentStatus.COMPLETED: + continue + chunk = IndexDocument( + page_content=row.content, + metadata={ + "doc_id": row.index_node_id, + "doc_hash": row.index_node_hash, + "document_id": ref.document_id, + "dataset_id": ref.dataset.dataset_id, + }, + ) + if document.doc_form == IndexStructureType.PARENT_CHILD_INDEX: + chunk.children = [ + ChildDocument( + page_content=child.content, + metadata={ + "doc_id": child.index_node_id, + "doc_hash": child.index_node_hash, + "document_id": ref.document_id, + "dataset_id": ref.dataset.dataset_id, + }, + ) + for child in get_segment_child_chunks(row, session=session) + ] + chunks.append(chunk) + return chunks, sum(row.tokens for row in rows) + + def complete_indexing_segments(self, ref: DocumentRef, node_ids: Sequence[str]) -> None: + with self._session_factory.begin() as session: + require_indexing_document(session, ref, lock=True) + session.execute( + update(DocumentSegment) + .where( + *self._indexing_scope(ref), + DocumentSegment.index_node_id.in_(node_ids), + DocumentSegment.status == SegmentStatus.INDEXING, + ) + .values(status=SegmentStatus.COMPLETED, enabled=True, completed_at=naive_utc_now()) + ) + + @staticmethod + def _indexing_scope(ref: DocumentRef): + return ( + DocumentSegment.tenant_id == ref.dataset.tenant_id, + DocumentSegment.dataset_id == ref.dataset.dataset_id, + DocumentSegment.document_id == ref.document_id, + ) + + def list_segments(self, document_ref: DocumentRef, query: SegmentListFilter) -> SegmentPage: + with self._session_factory() as session: + statement = ( + select(DocumentSegment) + .where( + DocumentSegment.tenant_id == document_ref.dataset.tenant_id, + DocumentSegment.dataset_id == document_ref.dataset.dataset_id, + DocumentSegment.document_id == document_ref.document_id, + ) + .order_by(DocumentSegment.position.asc()) + ) + if query.statuses: + statement = statement.where(DocumentSegment.status.in_(query.statuses)) + if query.hit_count_gte is not None: + statement = statement.where(DocumentSegment.hit_count >= query.hit_count_gte) + if query.keyword: + escaped_keyword = escape_like_pattern(query.keyword) + if dify_config.SQLALCHEMY_DATABASE_URI_SCHEME == "postgresql": + keywords_jsonb = sql_cast(DocumentSegment.keywords, JSONB) + keywords_array = case( + (func.jsonb_typeof(keywords_jsonb) == "array", keywords_jsonb), + else_=sql_cast(literal("[]"), JSONB), + ) + keywords_condition = func.array_to_string( + func.array( + select(func.jsonb_array_elements_text(keywords_array)) + .correlate(DocumentSegment) + .scalar_subquery() + ), + ",", + ).ilike(f"%{escaped_keyword}%", escape="\\") + else: + keywords_condition = sql_cast(DocumentSegment.keywords, String).ilike( + f"%{escaped_keyword}%", escape="\\" + ) + statement = statement.where( + or_( + DocumentSegment.content.ilike(f"%{escaped_keyword}%", escape="\\"), + keywords_condition, + ) + ) + if query.enabled.lower() == "true": + statement = statement.where(DocumentSegment.enabled.is_(True)) + elif query.enabled.lower() == "false": + statement = statement.where(DocumentSegment.enabled.is_(False)) + + limit = min(query.limit, 100) + segments = paginate_query(statement, session=session, page=query.page, per_page=limit, max_per_page=100) + items = list(segments.items) + summaries = ( + { + summary.chunk_id: summary + for summary in session.scalars( + select(DocumentSegmentSummary).where( + DocumentSegmentSummary.chunk_id.in_([segment.id for segment in items]), + DocumentSegmentSummary.dataset_id == document_ref.dataset.dataset_id, + DocumentSegmentSummary.document_id == document_ref.document_id, + ) + ) + } + if items + else {} + ) + return SegmentPage( + items=tuple( + _segment_data( + segment, + summaries[segment.id].summary_content if segment.id in summaries else None, + session, + ) + for segment in items + ), + total=segments.total, + total_pages=segments.pages, + page=query.page, + limit=limit, + ) + + def list_child_chunks( + self, + segment_ref: SegmentRef, + query: ChildChunkListFilter, + ) -> ChildChunkPage | None: + with self._session_factory() as session: + if _get_segment(session, segment_ref) is None: + return None + limit = min(query.limit, 100) + child_chunks = query_child_chunks( + session, + segment_ref, + page=query.page, + limit=limit, + keyword=query.keyword, + ) + return ChildChunkPage( + items=tuple(_child_chunk_data(chunk) for chunk in child_chunks.items), + total=child_chunks.total, + total_pages=child_chunks.pages, + page=query.page, + limit=limit, + ) + + def get_segment(self, segment_ref: SegmentRef) -> SegmentDetail | None: + with self._session_factory() as session: + segment = _get_segment(session, segment_ref) + if segment is None: + return None + _, document = _require_scope_models(session, segment_ref.document) + return _detail(session, segment, document) + + def get_segment_update_state(self, segment_ref: SegmentRef) -> SegmentUpdateState | None: + with self._session_factory() as session: + segment = _get_segment(session, segment_ref) + if segment is None: + return None + # Legacy rows may contain scalar or malformed JSON keywords. Treat + # those as absent so a write can repair them without building a DTO. + keywords = segment.keywords + return SegmentUpdateState( + content=segment.content, + enabled=segment.enabled, + keywords=tuple(keywords) + if isinstance(keywords, list) and all(isinstance(keyword, str) for keyword in keywords) + else None, + ) + + def get_segments(self, document_ref: DocumentRef, segment_ids: Sequence[str]) -> tuple[SegmentIndexTarget, ...]: + if not segment_ids: + return () + with self._session_factory() as session: + segments = session.scalars( + select(DocumentSegment).where( + DocumentSegment.tenant_id == document_ref.dataset.tenant_id, + DocumentSegment.dataset_id == document_ref.dataset.dataset_id, + DocumentSegment.document_id == document_ref.document_id, + DocumentSegment.id.in_(segment_ids), + ) + ).all() + return tuple(SegmentIndexTarget(segment.id, segment.index_node_id, segment.enabled) for segment in segments) + + def get_children_for_segments( + self, document_ref: DocumentRef, segment_ids: Sequence[str] + ) -> tuple[ChildChunkState, ...]: + if not segment_ids: + return () + with self._session_factory() as session: + return tuple(_child_state(child) for child in session.scalars(_children_query(document_ref, segment_ids))) + + def set_segments(self, document_ref: DocumentRef, segment_ids: Sequence[str], values: Mapping[str, object]) -> None: + if not segment_ids: + return + with self._session_factory.begin() as session: + session.execute( + update(DocumentSegment) + .where( + DocumentSegment.tenant_id == document_ref.dataset.tenant_id, + DocumentSegment.dataset_id == document_ref.dataset.dataset_id, + DocumentSegment.document_id == document_ref.document_id, + DocumentSegment.id.in_(segment_ids), + ) + .values(**values) + ) + + def save_segment(self, segment_ref: SegmentRef, values: Mapping[str, object], *, create: bool = False) -> None: + with self._session_factory.begin() as session: + _, document = _require_scope_models(session, segment_ref.document) + segment = _get_segment(session, segment_ref) + previous_words = segment.word_count if segment is not None else 0 + if create: + if segment is not None: + raise ValueError("Segment already exists") + position = ( + session.scalar( + select(func.max(DocumentSegment.position)).where( + DocumentSegment.tenant_id == segment_ref.document.dataset.tenant_id, + DocumentSegment.dataset_id == segment_ref.document.dataset.dataset_id, + DocumentSegment.document_id == segment_ref.document.document_id, + ) + ) + or 0 + ) + segment = DocumentSegment( + tenant_id=segment_ref.document.dataset.tenant_id, + dataset_id=segment_ref.document.dataset.dataset_id, + document_id=segment_ref.document.document_id, + position=position + 1, + content=cast(str, values["content"]), + word_count=cast(int, values["word_count"]), + tokens=cast(int, values["tokens"]), + created_by=cast(str, values["created_by"]), + ) + segment.id = segment_ref.segment_id + session.add(segment) + if segment is None: + raise LookupError("Segment no longer exists") + for key, value in values.items(): + setattr(segment, key, value) + if "word_count" in values: + document.word_count = max(0, (document.word_count or 0) + segment.word_count - previous_words) + + def delete_segments(self, document_ref: DocumentRef, segment_ids: Sequence[str]) -> None: + if not segment_ids: + return + with self._session_factory.begin() as session: + _, document = _require_scope_models(session, document_ref) + rows = session.scalars( + select(DocumentSegment).where( + DocumentSegment.tenant_id == document_ref.dataset.tenant_id, + DocumentSegment.dataset_id == document_ref.dataset.dataset_id, + DocumentSegment.document_id == document_ref.document_id, + DocumentSegment.id.in_(segment_ids), + ) + ).all() + document.word_count = max(0, (document.word_count or 0) - sum(row.word_count for row in rows)) + for row in rows: + session.delete(row) + + def get_children(self, segment_ref: SegmentRef) -> tuple[ChildChunkState, ...] | None: + with self._session_factory() as session: + if _get_segment(session, segment_ref) is None: + return None + children = session.scalars( + _children_query(segment_ref.document, (segment_ref.segment_id,)).order_by(ChildChunk.position.asc()) + ).all() + return tuple(_child_state(child) for child in children) + + def save_children( + self, + segment_ref: SegmentRef, + *, + added: Sequence[ChildChunkState] = (), + updated: Sequence[ChildChunkState] = (), + deleted: Sequence[ChildChunkState] = (), + ) -> None: + with self._session_factory.begin() as session: + if _get_segment(session, segment_ref) is None: + raise LookupError("Segment no longer exists") + for child in deleted: + row = _get_child_chunk(session, segment_ref, child.data.id) + if row is not None: + session.delete(row) + for child in updated: + row = _get_child_chunk(session, segment_ref, child.data.id) + if row is None: + raise LookupError("Child chunk no longer exists") + row.content = child.data.content + row.word_count = child.data.word_count + row.type = SegmentType(child.data.type) + row.updated_by = child.updated_by + row.updated_at = child.data.updated_at + for child in added: + row = ChildChunk( + **child.data.model_dump(exclude={"id", "created_at", "updated_at"}), + tenant_id=segment_ref.document.dataset.tenant_id, + dataset_id=segment_ref.document.dataset.dataset_id, + document_id=segment_ref.document.document_id, + index_node_id=child.index_node_id, + index_node_hash=child.index_node_hash, + created_by=child.created_by, + ) + row.id = child.data.id + row.created_at = child.data.created_at + row.updated_at = child.data.updated_at + row.updated_by = child.updated_by + session.add(row) + + def get_indexing_snapshot(self, segment_ref: SegmentRef) -> SegmentIndexingSnapshot: + return self.get_indexing_snapshots(segment_ref.document, (segment_ref.segment_id,))[0] + + def get_indexing_snapshots( + self, document_ref: DocumentRef, segment_ids: Sequence[str] + ) -> tuple[SegmentIndexingSnapshot, ...]: + """Load an ordered batch and its owner chain in one bounded read session.""" + if not segment_ids: + return () + with self._session_factory() as session: + dataset, document = _require_scope_models(session, document_ref) + segments = { + segment.id: segment + for segment in session.scalars( + select(DocumentSegment).where( + DocumentSegment.tenant_id == dataset.tenant_id, + DocumentSegment.dataset_id == dataset.id, + DocumentSegment.document_id == document.id, + DocumentSegment.id.in_(segment_ids), + ) + ) + } + if set(segments) != set(segment_ids): + raise LookupError("Segment no longer exists") + rule = ( + session.scalar( + select(DatasetProcessRule).where( + DatasetProcessRule.id == document.dataset_process_rule_id, + DatasetProcessRule.dataset_id == dataset.id, + ) + ) + if document.dataset_process_rule_id + else None + ) + vector_type = ( + Vector.resolve_vector_type(dataset, session=session) + if dataset.indexing_technique == "high_quality" + else None + ) + attachments: dict[str, list[str]] = {segment_id: [] for segment_id in segment_ids} + for binding in session.scalars( + select(SegmentAttachmentBinding).where( + SegmentAttachmentBinding.tenant_id == dataset.tenant_id, + SegmentAttachmentBinding.dataset_id == dataset.id, + SegmentAttachmentBinding.document_id == document.id, + SegmentAttachmentBinding.segment_id.in_(segment_ids), + ) + ): + attachments[binding.segment_id].append(binding.attachment_id) + summaries = { + summary.chunk_id: summary + for summary in session.scalars( + select(DocumentSegmentSummary).where( + DocumentSegmentSummary.dataset_id == dataset.id, + DocumentSegmentSummary.document_id == document.id, + DocumentSegmentSummary.chunk_id.in_(segment_ids), + ) + ) + } + return tuple( + SegmentIndexingSnapshot( + dataset, + document, + segments[segment_id], + rule, + vector_type, + summaries.get(segment_id), + tuple(attachments[segment_id]), + ) + for segment_id in segment_ids + ) + + def replace_attachments(self, segment_ref: SegmentRef, attachment_ids: Sequence[str]) -> None: + with self._session_factory.begin() as session: + if _get_segment(session, segment_ref) is None: + raise LookupError("Segment no longer exists") + session.execute( + delete(SegmentAttachmentBinding).where( + SegmentAttachmentBinding.tenant_id == segment_ref.document.dataset.tenant_id, + SegmentAttachmentBinding.dataset_id == segment_ref.document.dataset.dataset_id, + SegmentAttachmentBinding.document_id == segment_ref.document.document_id, + SegmentAttachmentBinding.segment_id == segment_ref.segment_id, + ) + ) + session.add_all( + [ + SegmentAttachmentBinding( + tenant_id=segment_ref.document.dataset.tenant_id, + dataset_id=segment_ref.document.dataset.dataset_id, + document_id=segment_ref.document.document_id, + segment_id=segment_ref.segment_id, + attachment_id=file_id, + ) + for file_id in attachment_ids + ] + ) + + def save_summary( + self, segment_ref: SegmentRef, content: str | None, *, error: str | None = None + ) -> DocumentSegmentSummary | None: + with self._session_factory.begin() as session: + segment = _get_segment(session, segment_ref) + if segment is None: + raise LookupError("Segment no longer exists") + summary = _summary(session, segment) + if content is None: + if summary is not None: + session.delete(summary) + return None + if summary is None: + summary = DocumentSegmentSummary( + dataset_id=segment.dataset_id, document_id=segment.document_id, chunk_id=segment.id + ) + session.add(summary) + summary.summary_content = content + summary.status = SummaryStatus.ERROR if error else SummaryStatus.GENERATING + summary.error = error + session.flush() + session.expunge(summary) + return summary + + def get_summary_images(self, segment_ref: SegmentRef) -> list[File]: + with self._session_factory() as session: + segment = _get_segment(session, segment_ref) + if segment is None: + raise LookupError("Segment no longer exists") + images = ParagraphIndexProcessor._extract_images_from_segment_attachments( + segment.tenant_id, segment.id, session + ) + return images or ParagraphIndexProcessor._extract_images_from_text( + segment.tenant_id, segment.content, session + ) + + def get_keyword_table(self, dataset_ref: DatasetRef) -> tuple[str, str | None]: + with self._session_factory() as session: + dataset = _get_dataset(session, dataset_ref) + if dataset is None: + raise LookupError("Dataset no longer exists") + row = get_dataset_keyword_table(dataset, session=session) + return (row.data_source_type, row.keyword_table) if row else (dify_config.KEYWORD_DATA_SOURCE_TYPE, None) + + def save_keyword_table( + self, dataset_ref: DatasetRef, *, storage_type: str, data: str, keywords: Mapping[str, Sequence[str]] + ) -> None: + with self._session_factory.begin() as session: + persist_keyword_table( + session, + tenant_id=dataset_ref.tenant_id, + dataset_id=dataset_ref.dataset_id, + storage_type=storage_type, + data=data, + keywords=keywords, + ) + + +@dataclass(frozen=True) +class SegmentIndexingSnapshot: + dataset: Dataset + document: Document + segment: DocumentSegment + process_rule: DatasetProcessRule | None + vector_type: str | None + summary: DocumentSegmentSummary | None + attachment_ids: tuple[str, ...] + + +def _summary(session: Session, segment: DocumentSegment) -> DocumentSegmentSummary | None: + return session.scalar( + select(DocumentSegmentSummary) + .where( + DocumentSegmentSummary.dataset_id == segment.dataset_id, + DocumentSegmentSummary.document_id == segment.document_id, + DocumentSegmentSummary.chunk_id == segment.id, + ) + .limit(1) + ) + + +def _children_query(document_ref: DocumentRef, segment_ids: Sequence[str]): + return select(ChildChunk).where( + ChildChunk.tenant_id == document_ref.dataset.tenant_id, + ChildChunk.dataset_id == document_ref.dataset.dataset_id, + ChildChunk.document_id == document_ref.document_id, + ChildChunk.segment_id.in_(segment_ids), + ) + + +def _child_state(child: ChildChunk) -> ChildChunkState: + return ChildChunkState( + data=_child_chunk_data(child), + index_node_id=child.index_node_id, + index_node_hash=child.index_node_hash, + created_by=child.created_by, + updated_by=child.updated_by, + ) + + +def _detail(session: Session, segment: DocumentSegment, document: Document) -> SegmentDetail: + summary = _summary(session, segment) + return SegmentDetail( + data=_segment_data(segment, summary.summary_content if summary else None, session), + doc_form=str(document.doc_form), + ) + + +def _require_scope_models(session: Session, document_ref: DocumentRef) -> tuple[Dataset, Document]: + dataset = _get_dataset(session, document_ref.dataset) + document = _get_document(session, document_ref) + if dataset is None or document is None: + raise LookupError("Dataset document scope no longer exists") + return dataset, document + + +def _get_segment(session: Session, segment_ref: SegmentRef) -> DocumentSegment | None: + return session.scalar( + select(DocumentSegment) + .where( + DocumentSegment.id == segment_ref.segment_id, + DocumentSegment.document_id == segment_ref.document.document_id, + DocumentSegment.dataset_id == segment_ref.document.dataset.dataset_id, + DocumentSegment.tenant_id == segment_ref.document.dataset.tenant_id, + ) + .limit(1) + ) + + +def _get_child_chunk(session: Session, segment_ref: SegmentRef, child_chunk_id: str) -> ChildChunk | None: + return session.scalar( + select(ChildChunk) + .where( + ChildChunk.id == child_chunk_id, + ChildChunk.segment_id == segment_ref.segment_id, + ChildChunk.document_id == segment_ref.document.document_id, + ChildChunk.dataset_id == segment_ref.document.dataset.dataset_id, + ChildChunk.tenant_id == segment_ref.document.dataset.tenant_id, + ) + .limit(1) + ) + + +def _segment_data(segment: DocumentSegment, summary: str | None, session: Session) -> SegmentRecord: + return SegmentRecord.model_validate( + { + "id": segment.id, + "position": segment.position, + "document_id": segment.document_id, + "content": segment.content, + "answer": segment.answer, + "word_count": segment.word_count, + "tokens": segment.tokens, + "keywords": segment.keywords, + "index_node_id": segment.index_node_id, + "index_node_hash": segment.index_node_hash, + "hit_count": segment.hit_count, + "enabled": segment.enabled, + "disabled_at": segment.disabled_at, + "disabled_by": segment.disabled_by, + "status": segment.status, + "created_by": segment.created_by, + "created_at": segment.created_at, + "updated_at": segment.updated_at, + "updated_by": segment.updated_by, + "indexing_at": segment.indexing_at, + "completed_at": segment.completed_at, + "error": segment.error, + "stopped_at": segment.stopped_at, + "sign_content": sign_segment_content(segment, session=session), + "summary": summary, + "child_chunks": get_segment_child_chunks(segment, session=session, include_full_doc=False), + "attachments": get_segment_attachments(segment, session=session), + } + ) + + +def _child_chunk_data(child_chunk: ChildChunk) -> ChildChunkRecord: + return ChildChunkRecord.model_validate(child_chunk) + + +def query_child_chunks( + session: Session, + segment_ref: SegmentRef, + *, + page: int, + limit: int, + keyword: str | None, +) -> PaginatedResult[ChildChunk]: + """Share the query with callers that already own an explicit session.""" + query = _children_query(segment_ref.document, (segment_ref.segment_id,)).order_by(ChildChunk.position.asc()) + if keyword: + query = query.where(ChildChunk.content.ilike(f"%{escape_like_pattern(keyword)}%", escape="\\")) + return paginate_query(query, session=session, page=page, per_page=limit, max_per_page=100) diff --git a/api/repositories/knowledge/summary_repository.py b/api/repositories/knowledge/summary_repository.py new file mode 100644 index 00000000000..0c2c96b2a3c --- /dev/null +++ b/api/repositories/knowledge/summary_repository.py @@ -0,0 +1,74 @@ +"""Persist summary indexing outcomes using one owned session per attempt.""" + +from collections.abc import Callable + +from sqlalchemy import select +from sqlalchemy.orm import Session + +from libs.datetime_utils import naive_utc_now +from models.dataset import Dataset, Document, DocumentSegment, DocumentSegmentSummary +from models.enums import SummaryStatus +from services.knowledge.summaries.application import SummaryVectorInput, SummaryVectorResult + + +class SQLAlchemySummaryRepository: + def __init__(self, *, new_session: Callable[[], Session]) -> None: + self._new_session = new_session + + @staticmethod + def _record(session: Session, command: SummaryVectorInput) -> DocumentSegmentSummary | None: + ref = command.ref + segment = session.scalar( + select(DocumentSegment) + .join(Document, Document.id == DocumentSegment.document_id) + .join(Dataset, Dataset.id == Document.dataset_id) + .where( + Dataset.id == ref.document.dataset.dataset_id, + Dataset.tenant_id == ref.document.dataset.tenant_id, + Document.id == ref.document.document_id, + Document.tenant_id == ref.document.dataset.tenant_id, + DocumentSegment.id == ref.segment_id, + DocumentSegment.dataset_id == ref.document.dataset.dataset_id, + DocumentSegment.tenant_id == ref.document.dataset.tenant_id, + ) + ) + if segment is None: + raise LookupError("Summary segment no longer exists in this dataset") + return session.scalar( + select(DocumentSegmentSummary) + .where( + DocumentSegmentSummary.dataset_id == ref.document.dataset.dataset_id, + DocumentSegmentSummary.document_id == ref.document.document_id, + DocumentSegmentSummary.chunk_id == ref.segment_id, + ) + .limit(1) + ) + + def completed(self, command: SummaryVectorInput, result: SummaryVectorResult) -> None: + with self._new_session() as session, session.begin(): + row = self._record(session, command) + if row is None: + row = DocumentSegmentSummary( + dataset_id=command.ref.document.dataset.dataset_id, + document_id=command.ref.document.document_id, + chunk_id=command.ref.segment_id, + summary_content=command.content, + enabled=True, + ) + row.id = command.summary_id + session.add(row) + row.summary_content = command.content + row.summary_index_node_id = result.node_id + row.summary_index_node_hash = result.content_hash + row.tokens = result.tokens + row.status = SummaryStatus.COMPLETED + row.error = None + row.updated_at = naive_utc_now() + + def failed(self, command: SummaryVectorInput, error: str) -> None: + with self._new_session() as session, session.begin(): + row = self._record(session, command) + if row is not None: + row.status = SummaryStatus.ERROR + row.error = error + row.updated_at = naive_utc_now() diff --git a/api/repositories/knowledge/upload_file_repository.py b/api/repositories/knowledge/upload_file_repository.py new file mode 100644 index 00000000000..9b9e05fdd2d --- /dev/null +++ b/api/repositories/knowledge/upload_file_repository.py @@ -0,0 +1,55 @@ +"""SQLAlchemy repository for upload files used by knowledge extraction.""" + +from collections.abc import Sequence + +from sqlalchemy import delete, select +from sqlalchemy.orm import Session, sessionmaker + +from core.rag.extractor.entity.extract_setting import UploadFileExtractionInput +from models.model import UploadFile + + +def _query_files(session: Session, *, workspace_id: str, file_ids: Sequence[str]) -> dict[str, UploadFile]: + if not file_ids: + return {} + uploads = session.scalars( + select(UploadFile).where(UploadFile.id.in_(file_ids), UploadFile.tenant_id == workspace_id) + ) + return {upload.id: upload for upload in uploads} + + +def query_upload_extraction_inputs( + session: Session, *, workspace_id: str, file_ids: Sequence[str] +) -> dict[str, UploadFileExtractionInput]: + """Read detached extraction inputs without changing the caller's session lifecycle.""" + return { + upload.id: UploadFileExtractionInput.model_validate(upload) + for upload in _query_files(session, workspace_id=workspace_id, file_ids=file_ids).values() + } + + +class SQLAlchemyKnowledgeUploadRepository: + def __init__(self, *, session_factory: sessionmaker[Session]) -> None: + self._session_factory = session_factory + + def get_by_id(self, *, workspace_id: str, file_id: str) -> UploadFileExtractionInput | None: + return self.get_by_ids(workspace_id=workspace_id, file_ids=(file_id,)).get(file_id) + + def get_by_ids(self, *, workspace_id: str, file_ids: Sequence[str]) -> dict[str, UploadFileExtractionInput]: + with self._session_factory() as session: + return query_upload_extraction_inputs(session, workspace_id=workspace_id, file_ids=file_ids) + + def get_files(self, *, workspace_id: str, file_ids: Sequence[str]) -> dict[str, UploadFile]: + with self._session_factory() as session: + return _query_files(session, workspace_id=workspace_id, file_ids=file_ids) + + def get_file_name(self, *, workspace_id: str, upload_file_id: str) -> str | None: + upload = self.get_files(workspace_id=workspace_id, file_ids=(upload_file_id,)).get(upload_file_id) + return upload.name if upload is not None else None + + def delete_files(self, *, workspace_id: str, file_ids: Sequence[str]) -> None: + """Remove preview records after their storage cleanup has finished.""" + if not file_ids: + return + with self._session_factory.begin() as session: + session.execute(delete(UploadFile).where(UploadFile.tenant_id == workspace_id, UploadFile.id.in_(file_ids))) diff --git a/api/repositories/workspace_member_query_repository.py b/api/repositories/workspace_member_query_repository.py index 1c67ae13e7f..f53e6ad1375 100644 --- a/api/repositories/workspace_member_query_repository.py +++ b/api/repositories/workspace_member_query_repository.py @@ -60,3 +60,17 @@ class WorkspaceMemberQueryRepository(WorkspaceMemberQuery): ) return records + + def get_legacy_role(self, *, workspace_id: str, account_id: str) -> str | None: + """Return one member's role without loading the whole workspace roster.""" + + with self._session_factory() as session: + role = session.scalar( + select(TenantAccountJoin.role) + .where( + TenantAccountJoin.tenant_id == workspace_id, + TenantAccountJoin.account_id == account_id, + ) + .limit(1) + ) + return role.value if role is not None else None diff --git a/api/schedule/clean_unused_datasets_task.py b/api/schedule/clean_unused_datasets_task.py index 02770c83232..ef55110271e 100644 --- a/api/schedule/clean_unused_datasets_task.py +++ b/api/schedule/clean_unused_datasets_task.py @@ -14,6 +14,7 @@ from enums import CloudPlan from extensions.ext_redis import redis_client from libs.pagination import paginate_query from models.dataset import Dataset, DatasetAutoDisableLog, DatasetQuery, Document +from repositories.knowledge.dataset_read_repository import get_dataset_doc_form from services.feature_service import FeatureService @@ -140,7 +141,7 @@ def clean_unused_datasets_task(): # Remove index index_processor = IndexProcessorFactory( - dataset.get_doc_form(session=session) + get_dataset_doc_form(dataset, session=session) ).init_index_processor() index_processor.clean(dataset, None, session=session) diff --git a/api/services/agent/knowledge_datasets.py b/api/services/agent/knowledge_datasets.py index 9db9bbc08c1..a7277e9059c 100644 --- a/api/services/agent/knowledge_datasets.py +++ b/api/services/agent/knowledge_datasets.py @@ -35,7 +35,7 @@ def get_tenant_knowledge_dataset_rows(*, session: Session, tenant_id: str, datas a dataset row, so they are treated as missing instead of breaking the UUID-typed dataset lookup. """ - from services.dataset_service import DatasetService + from services.knowledge.dataset_service import DatasetService valid_ids: list[str] = [] for dataset_id in dataset_ids: diff --git a/api/services/app_service.py b/api/services/app_service.py index 7510d6b3e95..00e926adbfc 100644 --- a/api/services/app_service.py +++ b/api/services/app_service.py @@ -60,7 +60,7 @@ from services.entities.app_entities import ( AppListSortBy, CreateAppParams, ) -from services.model_provider_service import ModelProviderService +from services.model_provider.service import ModelProviderService from services.openapi.visibility import apply_openapi_gate, is_openapi_visible from services.rbac_agent_access_service import initialize_agent_rbac_access from services.system_feature_service import SystemFeatureService diff --git a/api/services/credential_permission_service.py b/api/services/credential_permission_service.py deleted file mode 100644 index d9ce5e7c502..00000000000 --- a/api/services/credential_permission_service.py +++ /dev/null @@ -1,70 +0,0 @@ -from collections.abc import Sequence - -from sqlalchemy import or_, select -from sqlalchemy.orm import InstrumentedAttribute, Session - -from models.account import Account -from models.credential_permission import CredentialPermission -from models.enums import PermissionEnum - - -class CredentialPermissionService: - """ - Shared service for per-credential access control. - Mirrors DatasetPermissionService but supports all credential types - via a credential_type discriminator. - """ - - @classmethod - def get_partial_member_list(cls, credential_id: str, credential_type: str, *, session: Session) -> Sequence[str]: - """Return account_ids that have partial-member access to a credential.""" - return session.scalars( - select(CredentialPermission.account_id).where( - CredentialPermission.credential_id == credential_id, - CredentialPermission.credential_type == credential_type, - ) - ).all() - - @classmethod - def apply_visibility_filter( - cls, - query, - *, - model_id_column: InstrumentedAttribute, - model_user_id_column: InstrumentedAttribute, - model_visibility_column: InstrumentedAttribute, - credential_type: str, - user: Account, - ): - """ - Add WHERE clauses to a SQLAlchemy query so it only returns credentials - visible to the given user. - - - all_team_members: always visible - - only_me: visible only to the creator (user.id matches) - - partial_members: visible to the creator OR users in credential_permissions - - Legacy rows with NULL user_id are treated as all_team_members - - No admin bypass: personal credentials are private regardless of role - """ - # Subquery: credential_ids where user has partial-member permission - partial_subquery = ( - select(CredentialPermission.credential_id) - .where( - CredentialPermission.credential_type == credential_type, - CredentialPermission.account_id == user.id, - ) - .correlate_except(CredentialPermission) - ) - - return query.where( - or_( - # all_team is always visible - model_visibility_column == PermissionEnum.ALL_TEAM, - # legacy rows with NULL user_id treated as all_team - model_user_id_column.is_(None), - # only_me: creator sees their own - (model_user_id_column == user.id), - # partial_members: user is in the permission table - model_id_column.in_(partial_subquery), - ) - ) diff --git a/api/services/auth/firecrawl/__init__.py b/api/services/credentials/__init__.py similarity index 100% rename from api/services/auth/firecrawl/__init__.py rename to api/services/credentials/__init__.py diff --git a/api/services/credentials/query.py b/api/services/credentials/query.py new file mode 100644 index 00000000000..54fb193b107 --- /dev/null +++ b/api/services/credentials/query.py @@ -0,0 +1,79 @@ +"""Consumer-owned read contracts for credential lists and their visibility policy.""" + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Protocol + +from core.plugin.entities.plugin_daemon import CredentialType + + +@dataclass(frozen=True, slots=True) +class ModelCredentialRecord: + id: str + name: str + + +@dataclass(frozen=True, slots=True) +class DatasourceCredentialListItem: + id: str + name: str + auth_type: str + encrypted_credentials: Mapping[str, object] + avatar_url: str | None + is_default: bool + + +@dataclass(frozen=True, slots=True) +class ToolCredentialRecord: + id: str + name: str + provider: str + credential_type: CredentialType + credentials: Mapping[str, object] + is_default: bool + visibility: str + created_by: str + partial_member_ids: tuple[str, ...] + from_other_member: bool + + +@dataclass(frozen=True, slots=True) +class TriggerSubscriptionRecord: + id: str + name: str + provider: str + credential_type: CredentialType + credentials: Mapping[str, object] + endpoint_id: str + parameters: Mapping[str, object] + properties: Mapping[str, object] + workflows_in_use: int + + +class CredentialQuery(Protocol): + """Return loaded values; implementations own SQL and short read sessions. + + An absent actor is reserved for trusted internal callers. Account-scoped + callers must supply their actor ID; administrator roles do not bypass it. + """ + + def list_models( + self, *, workspace_id: str, provider: str, actor_id: str | None + ) -> Sequence[ModelCredentialRecord]: ... + + def list_datasources( + self, *, workspace_id: str, provider: str, plugin_id: str, actor_id: str | None + ) -> Sequence[DatasourceCredentialListItem]: ... + + def list_tools( + self, + *, + workspace_id: str, + provider: str, + actor_id: str | None, + include_credential_ids: Sequence[str] = (), + ) -> Sequence[ToolCredentialRecord]: ... + + def list_trigger_subscriptions( + self, *, workspace_id: str, provider: str, actor_id: str | None + ) -> Sequence[TriggerSubscriptionRecord]: ... diff --git a/api/services/data_source/__init__.py b/api/services/data_source/__init__.py new file mode 100644 index 00000000000..48969ae2282 --- /dev/null +++ b/api/services/data_source/__init__.py @@ -0,0 +1 @@ +"""Application services and adapters for data-source use cases.""" diff --git a/api/services/auth/jina/__init__.py b/api/services/data_source/auth/__init__.py similarity index 100% rename from api/services/auth/jina/__init__.py rename to api/services/data_source/auth/__init__.py diff --git a/api/services/auth/data_source_api_key_auth_gateways.py b/api/services/data_source/auth/api_key_gateways.py similarity index 78% rename from api/services/auth/data_source_api_key_auth_gateways.py rename to api/services/data_source/auth/api_key_gateways.py index d40490c731c..e6db86fddef 100644 --- a/api/services/auth/data_source_api_key_auth_gateways.py +++ b/api/services/data_source/auth/api_key_gateways.py @@ -6,12 +6,12 @@ from typing import Protocol, override import httpx from core.helper import encrypter -from services.auth.data_source_api_key_auth_service import ApiKeyAuthCredentialEncryptor, ApiKeyAuthCredentialValidator from services.auth.errors import ( DataSourceApiKeyAuthProviderUnavailableError, UnsupportedDataSourceApiKeyAuthProviderError, ) -from services.entities.data_source_api_key_auth_entities import DataSourceApiKeyAuthCredentials +from services.data_source.auth.api_key_service import ApiKeyAuthCredentialEncryptor, ApiKeyAuthCredentialValidator +from services.data_source.entities.api_key_auth import DataSourceApiKeyAuthCredentials class _ProviderApiKeyAuthValidator(Protocol): @@ -27,15 +27,15 @@ type _ProviderApiKeyAuthValidatorFactory = Callable[ def _get_provider_validator_factory(provider: str) -> _ProviderApiKeyAuthValidatorFactory: match provider: case "firecrawl": - from services.auth.firecrawl.firecrawl import FirecrawlAuth + from services.data_source.auth.firecrawl.firecrawl import FirecrawlAuth return FirecrawlAuth case "watercrawl": - from services.auth.watercrawl.watercrawl import WatercrawlAuth + from services.data_source.auth.watercrawl.watercrawl import WatercrawlAuth return WatercrawlAuth case "jinareader": - from services.auth.jina.jina import JinaAuth + from services.data_source.auth.jina.jina import JinaAuth return JinaAuth case _: diff --git a/api/services/auth/data_source_api_key_auth_service.py b/api/services/data_source/auth/api_key_service.py similarity index 97% rename from api/services/auth/data_source_api_key_auth_service.py rename to api/services/data_source/auth/api_key_service.py index 585845a4a55..5831177b821 100644 --- a/api/services/auth/data_source_api_key_auth_service.py +++ b/api/services/data_source/auth/api_key_service.py @@ -5,7 +5,7 @@ from dataclasses import replace from typing import Protocol from machinery.context import RequestContext -from services.entities.data_source_api_key_auth_entities import ( +from services.data_source.entities.api_key_auth import ( DataSourceApiKeyAuthBindingCreate, DataSourceApiKeyAuthBindingRecord, DataSourceApiKeyAuthCredentials, diff --git a/api/services/auth/watercrawl/__init__.py b/api/services/data_source/auth/firecrawl/__init__.py similarity index 100% rename from api/services/auth/watercrawl/__init__.py rename to api/services/data_source/auth/firecrawl/__init__.py diff --git a/api/services/auth/firecrawl/firecrawl.py b/api/services/data_source/auth/firecrawl/firecrawl.py similarity index 96% rename from api/services/auth/firecrawl/firecrawl.py rename to api/services/data_source/auth/firecrawl/firecrawl.py index 2261208bf08..fb0b8d8e9aa 100644 --- a/api/services/auth/firecrawl/firecrawl.py +++ b/api/services/data_source/auth/firecrawl/firecrawl.py @@ -9,7 +9,7 @@ from services.auth.errors import ( DataSourceApiKeyAuthProviderUnavailableError, InvalidDataSourceApiKeyAuthCredentialsError, ) -from services.entities.data_source_api_key_auth_entities import DataSourceApiKeyAuthCredentials +from services.data_source.entities.api_key_auth import DataSourceApiKeyAuthCredentials # Explicit bounded timeout for credential-validation requests so a slow or # hanging Firecrawl endpoint cannot block the worker indefinitely. diff --git a/api/services/data_source/auth/jina/__init__.py b/api/services/data_source/auth/jina/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/api/services/auth/jina/jina.py b/api/services/data_source/auth/jina/jina.py similarity index 96% rename from api/services/auth/jina/jina.py rename to api/services/data_source/auth/jina/jina.py index cf201ae0c38..52d1cd58bfd 100644 --- a/api/services/auth/jina/jina.py +++ b/api/services/data_source/auth/jina/jina.py @@ -10,7 +10,7 @@ from services.auth.errors import ( DataSourceApiKeyAuthProviderUnavailableError, InvalidDataSourceApiKeyAuthCredentialsError, ) -from services.entities.data_source_api_key_auth_entities import DataSourceApiKeyAuthCredentials +from services.data_source.entities.api_key_auth import DataSourceApiKeyAuthCredentials _http_client: httpx.Client = get_pooled_http_client( "auth:jina", diff --git a/api/services/data_source/auth/watercrawl/__init__.py b/api/services/data_source/auth/watercrawl/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/api/services/auth/watercrawl/watercrawl.py b/api/services/data_source/auth/watercrawl/watercrawl.py similarity index 96% rename from api/services/auth/watercrawl/watercrawl.py rename to api/services/data_source/auth/watercrawl/watercrawl.py index e4d0951f6df..fe7a9e85f6b 100644 --- a/api/services/auth/watercrawl/watercrawl.py +++ b/api/services/data_source/auth/watercrawl/watercrawl.py @@ -10,7 +10,7 @@ from services.auth.errors import ( DataSourceApiKeyAuthProviderUnavailableError, InvalidDataSourceApiKeyAuthCredentialsError, ) -from services.entities.data_source_api_key_auth_entities import DataSourceApiKeyAuthCredentials +from services.data_source.entities.api_key_auth import DataSourceApiKeyAuthCredentials # Explicit bounded timeout for credential-validation requests so a slow or # hanging WaterCrawl endpoint cannot block the worker indefinitely. diff --git a/api/services/data_source/binding_application_service.py b/api/services/data_source/binding_application_service.py new file mode 100644 index 00000000000..3f11d13738b --- /dev/null +++ b/api/services/data_source/binding_application_service.py @@ -0,0 +1,69 @@ +"""Application service for managing data-source OAuth bindings.""" + +from collections.abc import Sequence +from enum import StrEnum +from typing import Literal, Protocol + +from machinery.context import RequestContext +from services.data_source.entities.oauth import DataSourceBindingSummary + +type BindingMutationResult = Literal["updated", "not_found", "already_enabled", "already_disabled"] + +_SUPPORTED_PROVIDERS = frozenset({"notion"}) + + +class BindingAction(StrEnum): + ENABLE = "enable" + DISABLE = "disable" + + +class DataSourceBindingStore(Protocol): + def list_enabled_bindings(self, *, workspace_id: str) -> Sequence[DataSourceBindingSummary]: ... + + def change_disabled_state( + self, + *, + workspace_id: str, + binding_id: str, + disabled: bool, + ) -> BindingMutationResult: ... + + +class DataSourceBindingError(Exception): + """Base class for framework-neutral binding management failures.""" + + +class DataSourceBindingNotFoundError(DataSourceBindingError): + def __init__(self) -> None: + super().__init__("Data source binding not found") + + +class DataSourceBindingStateError(DataSourceBindingError): + def __init__(self, *, disabled: bool) -> None: + state = "disabled" if disabled else "enabled" + super().__init__(f"Data source is already {state}") + + +class DataSourceBindingApplicationService: + def __init__(self, *, bindings: DataSourceBindingStore) -> None: + self._bindings = bindings + + def list_integrations(self, context: RequestContext) -> tuple[DataSourceBindingSummary, ...]: + return tuple( + binding + for binding in self._bindings.list_enabled_bindings(workspace_id=context.active_workspace_id) + if binding.provider in _SUPPORTED_PROVIDERS + ) + + def change_state(self, context: RequestContext, binding_id: str, action: BindingAction) -> None: + disabled = action == BindingAction.DISABLE + result = self._bindings.change_disabled_state( + workspace_id=context.active_workspace_id, + binding_id=binding_id, + disabled=disabled, + ) + if result == "updated": + return + if result == "not_found": + raise DataSourceBindingNotFoundError() + raise DataSourceBindingStateError(disabled=disabled) diff --git a/api/services/data_source/credential_adapters.py b/api/services/data_source/credential_adapters.py new file mode 100644 index 00000000000..d58c4ea065d --- /dev/null +++ b/api/services/data_source/credential_adapters.py @@ -0,0 +1,197 @@ +"""Plugin and persistence adapters for datasource credential application ports.""" + +from collections.abc import Callable, Mapping, Sequence +from typing import Protocol + +from configs import dify_config +from core.helper import encrypter +from core.helper.provider_cache import NoOpProviderCredentialCache +from core.plugin.entities.plugin_daemon import CredentialType +from core.plugin.impl.datasource import PluginDatasourceManager +from core.plugin.impl.oauth import OAuthHandler +from core.plugin.plugin_service import PluginService +from core.tools.utils.encryption import create_provider_encrypter +from graphon.model_runtime.entities.provider_entities import FormType +from models.provider_ids import DatasourceProviderID +from services.data_source.credential_gateway import ( + DatasourceCredentialError, + DatasourceCredentialRefreshError, + RefreshedDatasourceCredential, +) +from services.data_source.entities.credential import DatasourceCredentialRecord, DatasourceOAuthClientConfigRecord + + +class DatasourceOAuthClientConfigReader(Protocol): + def get_oauth_client_config( + self, + *, + workspace_id: str, + provider: str, + plugin_id: str, + ) -> DatasourceOAuthClientConfigRecord: ... + + +class DatasourceOAuthClientResolver(Protocol): + def resolve(self, *, workspace_id: str, provider_id: DatasourceProviderID) -> dict[str, object]: ... + + +class PluginDatasourceCredentialCodec: + def __init__(self, provider_manager: PluginDatasourceManager) -> None: + self._provider_manager = provider_manager + + def decrypt(self, record: DatasourceCredentialRecord) -> dict[str, object]: + return transform_secret_fields( + record.workspace_id, record.encrypted_credentials, self._secret_variables(record), encrypter.decrypt_token + ) + + def encrypt(self, record: DatasourceCredentialRecord, credentials: Mapping[str, object]) -> dict[str, object]: + return transform_secret_fields( + record.workspace_id, credentials, self._secret_variables(record), encrypter.encrypt_token + ) + + def _secret_variables(self, record: DatasourceCredentialRecord) -> tuple[str, ...]: + return datasource_secret_variables( + self._provider_manager, + record.workspace_id, + f"{record.plugin_id}/{record.provider}", + CredentialType.of(record.auth_type), + ) + + +def transform_secret_fields( + workspace_id: str, + credentials: Mapping[str, object], + secret_variables: Sequence[str], + transform: Callable[[str, str], str], +) -> dict[str, object]: + """Apply the shared token codec only to declared string secrets.""" + result = dict(credentials) + for name in secret_variables: + value = result.get(name) + if isinstance(value, str): + result[name] = transform(workspace_id, value) + return result + + +def datasource_secret_variables( + manager: PluginDatasourceManager, workspace_id: str, provider_id: str, credential_type: CredentialType +) -> tuple[str, ...]: + provider = manager.fetch_datasource_provider(tenant_id=workspace_id, provider_id=provider_id) + if credential_type == CredentialType.API_KEY: + schemas = provider.declaration.credentials_schema + elif credential_type == CredentialType.OAUTH2: + if provider.declaration.oauth_schema is None: + raise DatasourceCredentialError("Datasource provider oauth schema not found") + schemas = provider.declaration.oauth_schema.credentials_schema + else: + raise DatasourceCredentialError(f"Unsupported credential type: {credential_type}") + return tuple(schema.name for schema in schemas if schema.type.value == FormType.SECRET_INPUT) + + +class PluginDatasourceOAuthClientResolver: + """Turn detached OAuth client rows into plugin-ready credentials.""" + + def __init__( + self, + *, + configs: DatasourceOAuthClientConfigReader, + provider_manager: PluginDatasourceManager, + ) -> None: + self._configs = configs + self._provider_manager = provider_manager + + def resolve(self, *, workspace_id: str, provider_id: DatasourceProviderID) -> dict[str, object]: + config = self._configs.get_oauth_client_config( + workspace_id=workspace_id, + provider=provider_id.provider_name, + plugin_id=provider_id.plugin_id, + ) + return resolve_datasource_oauth_client( + config, workspace_id=workspace_id, provider_id=provider_id, provider_manager=self._provider_manager + ) + + +def resolve_datasource_oauth_client( + config: DatasourceOAuthClientConfigRecord, + *, + workspace_id: str, + provider_id: DatasourceProviderID, + provider_manager: PluginDatasourceManager, +) -> dict[str, object]: + provider = provider_manager.fetch_datasource_provider( + tenant_id=workspace_id, + provider_id=str(provider_id), + ) + if config.encrypted_tenant_params is not None: + oauth_schema = provider.declaration.oauth_schema + if oauth_schema is None: + raise DatasourceCredentialError("Datasource provider oauth schema not found") + credential_encrypter, _ = create_provider_encrypter( + tenant_id=workspace_id, + config=[item.to_basic_provider_config() for item in oauth_schema.client_schema], + cache=NoOpProviderCredentialCache(), + ) + return dict(credential_encrypter.decrypt(dict(config.encrypted_tenant_params))) + + if config.system_credentials is not None and PluginService.is_plugin_verified( + workspace_id, provider.plugin_unique_identifier + ): + return dict(config.system_credentials) + raise DatasourceCredentialError(f"OAuth client is not configured for {provider_id}") + + +class OAuthDatasourceCredentialRefresher: + """Refresh datasource credentials through an injected OAuth-client resolver.""" + + def __init__( + self, + *, + oauth_clients: DatasourceOAuthClientResolver, + oauth_handler: OAuthHandler, + ) -> None: + self._oauth_clients = oauth_clients + self._oauth_handler = oauth_handler + + def refresh( + self, + *, + workspace_id: str, + actor_id: str, + record: DatasourceCredentialRecord, + credentials: Mapping[str, object], + ) -> RefreshedDatasourceCredential: + provider_id = DatasourceProviderID(f"{record.plugin_id}/{record.provider}") + try: + return refresh_datasource_credential( + self._oauth_handler, + workspace_id=workspace_id, + actor_id=actor_id, + provider_id=provider_id, + system_credentials=self._oauth_clients.resolve(workspace_id=workspace_id, provider_id=provider_id), + credentials=credentials, + ) + except DatasourceCredentialError: + raise + except Exception as error: + raise DatasourceCredentialRefreshError(record.id) from error + + +def refresh_datasource_credential( + handler: OAuthHandler, + *, + workspace_id: str, + actor_id: str, + provider_id: DatasourceProviderID, + system_credentials: Mapping[str, object], + credentials: Mapping[str, object], +) -> RefreshedDatasourceCredential: + refreshed = handler.refresh_credentials( + tenant_id=workspace_id, + user_id=actor_id, + plugin_id=provider_id.plugin_id, + provider=provider_id.provider_name, + redirect_uri=f"{dify_config.CONSOLE_API_URL}/console/api/oauth/plugin/{provider_id}/datasource/callback", + system_credentials=dict(system_credentials), + credentials=dict(credentials), + ) + return RefreshedDatasourceCredential(credentials=dict(refreshed.credentials), expires_at=refreshed.expires_at) diff --git a/api/services/data_source/credential_gateway.py b/api/services/data_source/credential_gateway.py new file mode 100644 index 00000000000..75dc5cf7cf4 --- /dev/null +++ b/api/services/data_source/credential_gateway.py @@ -0,0 +1,319 @@ +"""Application ports and policies for datasource credential resolution.""" + +import logging +import time +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from typing import Protocol + +from services.data_source.entities.credential import DatasourceCredentialRecord, DatasourceOAuthClientConfigRecord + +logger = logging.getLogger(__name__) + + +class NotionCredentialUnavailableError(ValueError): + pass + + +def resolve_stored_notion_access_token( + *, + credential_id: str | None, + load_token: Callable[[str], str], + integration_token: str | None, +) -> str: + """Use the saved credential, then the configured legacy integration. + + Stored documents never select a tenant default credential: changing that + default must not change the integration used to read an existing document. + Preserve indexing's environment fallback when credential resolution fails, + including refresh failures. Actor-selected imports do not use this policy. + """ + credential_error: Exception | None = None + if credential_id: + try: + token = load_token(credential_id) + if isinstance(token, str) and token: + return token + except Exception as error: + credential_error = error + + if not integration_token: + raise NotionCredentialUnavailableError( + "Must specify `integration_token` or set environment variable `NOTION_INTEGRATION_TOKEN`." + ) from credential_error + + logger.warning("Stored Notion credential %s unavailable; using NOTION_INTEGRATION_TOKEN", credential_id) + return integration_token + + +class DatasourceProviderCredentialStore(Protocol): + """Persisted provider credentials for already-authorized execution paths.""" + + def get_provider_credential( + self, *, workspace_id: str, provider: str, plugin_id: str, credential_id: str | None + ) -> DatasourceCredentialRecord | None: ... + + def list_provider_credentials( + self, *, workspace_id: str, provider: str, plugin_id: str + ) -> Sequence[DatasourceCredentialRecord]: ... + + def get_oauth_client_config( + self, *, workspace_id: str, provider: str, plugin_id: str + ) -> DatasourceOAuthClientConfigRecord: ... + + def update_if_unchanged( + self, *, record: DatasourceCredentialRecord, encrypted_credentials: Mapping[str, object], expires_at: int + ) -> bool: ... + + +class DatasourceCredentialStore(Protocol): + def get_visible( + self, + *, + workspace_id: str, + actor_id: str, + credential_id: str, + provider: str, + plugin_id: str, + ) -> DatasourceCredentialRecord | None: ... + + def update_if_unchanged( + self, + *, + record: DatasourceCredentialRecord, + encrypted_credentials: Mapping[str, object], + expires_at: int, + ) -> bool: ... + + +class TrustedDatasourceCredentialStore(Protocol): + def get_for_stored_document( + self, + *, + workspace_id: str, + dataset_id: str, + document_id: str, + credential_id: str | None, + provider: str, + plugin_id: str, + ) -> DatasourceCredentialRecord | None: ... + + def update_if_unchanged( + self, + *, + record: DatasourceCredentialRecord, + encrypted_credentials: Mapping[str, object], + expires_at: int, + ) -> bool: ... + + +class DatasourceCredentialCodec(Protocol): + def decrypt(self, record: DatasourceCredentialRecord) -> dict[str, object]: ... + + def encrypt( + self, + record: DatasourceCredentialRecord, + credentials: Mapping[str, object], + ) -> dict[str, object]: ... + + +@dataclass(frozen=True, slots=True) +class RefreshedDatasourceCredential: + credentials: Mapping[str, object] + expires_at: int + + +class DatasourceCredentialRefresher(Protocol): + def refresh( + self, + *, + workspace_id: str, + actor_id: str, + record: DatasourceCredentialRecord, + credentials: Mapping[str, object], + ) -> RefreshedDatasourceCredential: ... + + +class ActorDatasourceCredentialResolver(Protocol): + """Resolve a credential selected by an authenticated actor.""" + + def resolve( + self, + *, + workspace_id: str, + actor_id: str, + credential_id: str, + provider: str, + plugin_id: str, + ) -> dict[str, object]: ... + + +class StoredDatasourceCredentialResolver(Protocol): + """Resolve a credential proven by an existing document owner chain.""" + + def resolve_for_document( + self, + *, + workspace_id: str, + dataset_id: str, + document_id: str, + credential_id: str | None, + provider: str, + plugin_id: str, + ) -> dict[str, object]: ... + + +class DatasourceCredentialError(Exception): + """Base class for credential resolution failures.""" + + +class DatasourceCredentialNotFoundError(DatasourceCredentialError): + def __init__(self) -> None: + super().__init__("Credential not found") + + +class DatasourceCredentialConcurrentUpdateError(DatasourceCredentialError): + def __init__(self) -> None: + super().__init__("Credential changed while it was being refreshed") + + +class DatasourceCredentialRefreshError(DatasourceCredentialError): + def __init__(self, credential_id: str) -> None: + super().__init__(f"Failed to refresh datasource credential: {credential_id}") + + +class ActorAwareDatasourceCredentialGateway: + def __init__( + self, + *, + credentials: DatasourceCredentialStore, + codec: DatasourceCredentialCodec, + refresher: DatasourceCredentialRefresher, + now: Callable[[], float] = time.time, + ) -> None: + self._credentials = credentials + self._codec = codec + self._refresher = refresher + self._now = now + + def resolve( + self, + *, + workspace_id: str, + actor_id: str, + credential_id: str, + provider: str, + plugin_id: str, + ) -> dict[str, object]: + def load_record() -> DatasourceCredentialRecord: + record = self._credentials.get_visible( + workspace_id=workspace_id, + actor_id=actor_id, + credential_id=credential_id, + provider=provider, + plugin_id=plugin_id, + ) + if record is None: + raise DatasourceCredentialNotFoundError() + return record + + return resolve_credential_snapshot( + load_record(), + actor_id=actor_id, + codec=self._codec, + refresher=self._refresher, + update=self._credentials.update_if_unchanged, + load_latest=load_record, + now=self._now, + ) + + +class TrustedStoredDatasourceCredentialGateway: + """Resolve a credential only when a stored document proves the full owner chain.""" + + def __init__( + self, + *, + credentials: TrustedDatasourceCredentialStore, + codec: DatasourceCredentialCodec, + refresher: DatasourceCredentialRefresher, + now: Callable[[], float] = time.time, + ) -> None: + self._credentials = credentials + self._codec = codec + self._refresher = refresher + self._now = now + + def resolve_for_document( + self, + *, + workspace_id: str, + dataset_id: str, + document_id: str, + credential_id: str | None, + provider: str, + plugin_id: str, + ) -> dict[str, object]: + def load_record() -> DatasourceCredentialRecord: + record = self._credentials.get_for_stored_document( + workspace_id=workspace_id, + dataset_id=dataset_id, + document_id=document_id, + credential_id=credential_id, + provider=provider, + plugin_id=plugin_id, + ) + if record is None: + raise DatasourceCredentialNotFoundError() + return record + + record = load_record() + return resolve_credential_snapshot( + record, + actor_id=record.owner_id, + codec=self._codec, + refresher=self._refresher, + update=self._credentials.update_if_unchanged, + load_latest=load_record, + now=self._now, + ) + + +def credential_needs_refresh(expires_at: int, *, now: float) -> bool: + return expires_at != -1 and expires_at - 60 < int(now) + + +class CredentialSnapshotUpdate(Protocol): + def __call__( + self, *, record: DatasourceCredentialRecord, encrypted_credentials: Mapping[str, object], expires_at: int + ) -> bool: ... + + +def resolve_credential_snapshot( + record: DatasourceCredentialRecord, + *, + actor_id: str | None, + codec: DatasourceCredentialCodec, + refresher: DatasourceCredentialRefresher, + update: CredentialSnapshotUpdate, + load_latest: Callable[[], DatasourceCredentialRecord], + now: Callable[[], float], +) -> dict[str, object]: + decrypted = codec.decrypt(record) + if not credential_needs_refresh(record.expires_at, now=now()): + return decrypted + if actor_id is None: + raise DatasourceCredentialRefreshError(record.id) + refreshed = refresher.refresh( + workspace_id=record.workspace_id, actor_id=actor_id, record=record, credentials=decrypted + ) + if update( + record=record, + encrypted_credentials=codec.encrypt(record, refreshed.credentials), + expires_at=refreshed.expires_at, + ): + return dict(refreshed.credentials) + latest = load_latest() + if credential_needs_refresh(latest.expires_at, now=now()): + raise DatasourceCredentialConcurrentUpdateError() + return codec.decrypt(latest) diff --git a/api/services/data_source/entities/__init__.py b/api/services/data_source/entities/__init__.py new file mode 100644 index 00000000000..4504a9013aa --- /dev/null +++ b/api/services/data_source/entities/__init__.py @@ -0,0 +1 @@ +"""Shared data contracts for data-source use cases.""" diff --git a/api/services/entities/data_source_api_key_auth_entities.py b/api/services/data_source/entities/api_key_auth.py similarity index 100% rename from api/services/entities/data_source_api_key_auth_entities.py rename to api/services/data_source/entities/api_key_auth.py diff --git a/api/services/data_source/entities/credential.py b/api/services/data_source/entities/credential.py new file mode 100644 index 00000000000..4a06a0d8a66 --- /dev/null +++ b/api/services/data_source/entities/credential.py @@ -0,0 +1,27 @@ +"""Detached data contracts for datasource credentials.""" + +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import datetime + + +@dataclass(frozen=True, slots=True) +class DatasourceCredentialRecord: + id: str + workspace_id: str + owner_id: str | None + name: str + provider: str + plugin_id: str + auth_type: str + encrypted_credentials: Mapping[str, object] + expires_at: int + updated_at: datetime + + +@dataclass(frozen=True, slots=True) +class DatasourceOAuthClientConfigRecord: + """Raw persisted OAuth client configuration detached from its SQLAlchemy session.""" + + encrypted_tenant_params: Mapping[str, object] | None + system_credentials: Mapping[str, object] | None diff --git a/api/services/data_source/entities/notion_import.py b/api/services/data_source/entities/notion_import.py new file mode 100644 index 00000000000..555e1c5a1cb --- /dev/null +++ b/api/services/data_source/entities/notion_import.py @@ -0,0 +1,72 @@ +"""Framework-neutral values used by Notion import use cases.""" + +from collections.abc import Mapping +from dataclasses import dataclass +from enum import StrEnum + + +class NotionPageType(StrEnum): + PAGE = "page" + DATABASE = "database" + + +@dataclass(frozen=True, slots=True) +class NotionPageIcon: + type: str | None = None + url: str | None = None + emoji: str | None = None + + +@dataclass(frozen=True, slots=True) +class AuthorizedNotionPage: + page_id: str + page_name: str + page_icon: NotionPageIcon | None + parent_id: str | None + page_type: NotionPageType + + +@dataclass(frozen=True, slots=True) +class NotionImportPage: + page_id: str + page_name: str + page_icon: NotionPageIcon | None + parent_id: str | None + page_type: NotionPageType + is_bound: bool + + +@dataclass(frozen=True, slots=True) +class NotionWorkspace: + workspace_id: str | None + workspace_name: str | None + workspace_icon: str | None + pages: tuple[AuthorizedNotionPage, ...] + + +@dataclass(frozen=True, slots=True) +class NotionImportWorkspace: + workspace_id: str | None + workspace_name: str | None + workspace_icon: str | None + pages: tuple[NotionImportPage, ...] + + +@dataclass(frozen=True, slots=True) +class NotionImportResult: + workspaces: tuple[NotionImportWorkspace, ...] + + +def notion_page_icon(value: Mapping[str, object] | None) -> NotionPageIcon | None: + """Map a plugin icon payload without leaking the plugin entity type.""" + + if value is None: + return None + icon_type = value.get("type") + url = value.get("url") + emoji = value.get("emoji") + return NotionPageIcon( + type=icon_type if isinstance(icon_type, str) else None, + url=url if isinstance(url, str) else None, + emoji=emoji if isinstance(emoji, str) else None, + ) diff --git a/api/services/entities/data_source_oauth_entities.py b/api/services/data_source/entities/oauth.py similarity index 66% rename from api/services/entities/data_source_oauth_entities.py rename to api/services/data_source/entities/oauth.py index 844cec22847..a65e2126b50 100644 --- a/api/services/entities/data_source_oauth_entities.py +++ b/api/services/data_source/entities/oauth.py @@ -2,6 +2,7 @@ from collections.abc import Mapping from dataclasses import dataclass +from datetime import datetime @dataclass(frozen=True, slots=True) @@ -22,3 +23,14 @@ class DataSourceOAuthBindingRecord: id: str access_token: str source_info: Mapping[str, object] + + +@dataclass(frozen=True, slots=True) +class DataSourceBindingSummary: + """Secret-free management view of a data-source OAuth binding.""" + + id: str + provider: str + created_at: datetime + disabled: bool + source_info: Mapping[str, object] diff --git a/api/services/data_source/notion_import_adapters.py b/api/services/data_source/notion_import_adapters.py new file mode 100644 index 00000000000..6d43d3ccc8d --- /dev/null +++ b/api/services/data_source/notion_import_adapters.py @@ -0,0 +1,146 @@ +"""Notion plugin and extractor adapter for import use cases.""" + +import logging +from collections.abc import Callable, Mapping +from typing import Protocol, cast + +from core.datasource.entities.datasource_entities import DatasourceProviderType +from core.datasource.online_document.online_document_plugin import OnlineDocumentDatasourcePlugin +from core.rag.extractor.notion_extractor import NotionExtractor +from core.rag.models.document import Document as ExtractedDocument +from services.data_source.credential_gateway import ActorDatasourceCredentialResolver, DatasourceCredentialError +from services.data_source.entities.notion_import import ( + AuthorizedNotionPage, + NotionPageType, + NotionWorkspace, + notion_page_icon, +) +from services.data_source.notion_import_application_service import NotionImportCredentialUnavailableError + +_NOTION_PROVIDER = "notion_datasource" +_NOTION_PLUGIN = "langgenius/notion_datasource" +_NOTION_RUNTIME_PROVIDER = f"{_NOTION_PLUGIN}/{_NOTION_PROVIDER}" +logger = logging.getLogger(__name__) + + +class _RuntimeLoader(Protocol): + def __call__( + self, + *, + provider_id: str, + datasource_name: str, + tenant_id: str, + datasource_type: DatasourceProviderType, + ) -> object: ... + + +class _NotionPageExtractor(Protocol): + def extract(self) -> list[ExtractedDocument]: ... + + +class PluginNotionSourceGateway: + def __init__( + self, + *, + credentials: ActorDatasourceCredentialResolver, + runtime_loader: _RuntimeLoader, + extractor_factory: Callable[..., _NotionPageExtractor] = NotionExtractor, + ) -> None: + self._credentials = credentials + self._runtime_loader = runtime_loader + self._extractor_factory = extractor_factory + + def list_authorized_pages( + self, + *, + workspace_id: str, + actor_id: str, + credential_id: str, + ) -> tuple[NotionWorkspace, ...]: + credentials = self._resolve_credentials(workspace_id, actor_id, credential_id) + runtime = cast( + OnlineDocumentDatasourcePlugin, + self._runtime_loader( + provider_id=_NOTION_RUNTIME_PROVIDER, + datasource_name=_NOTION_PROVIDER, + tenant_id=workspace_id, + datasource_type=DatasourceProviderType.ONLINE_DOCUMENT, + ), + ) + runtime.runtime.credentials = dict(credentials) + + workspaces: dict[str | None, tuple[str | None, str | None, list[AuthorizedNotionPage]]] = {} + for message in runtime.get_online_document_pages( + user_id=actor_id, + datasource_parameters={}, + provider_type=runtime.datasource_provider_type(), + ): + for workspace in message.result: + workspace_info = workspaces.get(workspace.workspace_id) + if workspace_info is None: + workspace_info = (workspace.workspace_name, workspace.workspace_icon, []) + workspaces[workspace.workspace_id] = workspace_info + for page in workspace.pages: + try: + page_type = NotionPageType(page.type) + except (TypeError, ValueError): + logger.warning("Skipping unsupported Notion page type %r for page %s", page.type, page.page_id) + continue + workspace_info[2].append( + AuthorizedNotionPage( + page_id=page.page_id, + page_name=page.page_name, + page_icon=notion_page_icon(page.page_icon), + parent_id=page.parent_id, + page_type=page_type, + ) + ) + return tuple( + NotionWorkspace( + workspace_id=workspace_id, + workspace_name=workspace_name, + workspace_icon=workspace_icon, + pages=tuple(pages), + ) + for workspace_id, (workspace_name, workspace_icon, pages) in workspaces.items() + ) + + def preview_page( + self, + *, + workspace_id: str, + actor_id: str, + credential_id: str, + page_id: str, + page_type: NotionPageType, + ) -> str: + credentials = self._resolve_credentials(workspace_id, actor_id, credential_id) + integration_secret = credentials.get("integration_secret") + if not isinstance(integration_secret, str) or not integration_secret: + raise NotionImportCredentialUnavailableError() + extractor = self._extractor_factory( + notion_workspace_id="", + notion_obj_id=page_id, + notion_page_type=page_type.value, + notion_access_token=integration_secret, + notion_token_loader=lambda: integration_secret, + tenant_id=workspace_id, + ) + return "\n".join(document.page_content for document in extractor.extract()) + + def _resolve_credentials( + self, + workspace_id: str, + actor_id: str, + credential_id: str, + ) -> Mapping[str, object]: + try: + return self._credentials.resolve( + workspace_id=workspace_id, + actor_id=actor_id, + credential_id=credential_id, + provider=_NOTION_PROVIDER, + plugin_id=_NOTION_PLUGIN, + ) + except DatasourceCredentialError as error: + raise NotionImportCredentialUnavailableError() from error diff --git a/api/services/data_source/notion_import_application_service.py b/api/services/data_source/notion_import_application_service.py new file mode 100644 index 00000000000..e1454ed71e4 --- /dev/null +++ b/api/services/data_source/notion_import_application_service.py @@ -0,0 +1,129 @@ +"""Application service for listing and previewing importable Notion pages.""" + +from typing import Protocol + +from machinery.context import RequestContext +from services.data_source.entities.notion_import import ( + NotionImportPage, + NotionImportResult, + NotionImportWorkspace, + NotionPageType, + NotionWorkspace, +) +from services.knowledge.dataset_access import DatasetAccess +from services.knowledge.resource_scope import DatasetRef + + +class NotionDocumentBindingReader(Protocol): + def list_bound_notion_page_ids(self, dataset_ref: DatasetRef) -> frozenset[str]: ... + + +class NotionDatasetReader(Protocol): + def is_notion_dataset(self, dataset_ref: DatasetRef) -> bool: ... + + +class NotionSourceGateway(Protocol): + def list_authorized_pages( + self, + *, + workspace_id: str, + actor_id: str, + credential_id: str, + ) -> tuple[NotionWorkspace, ...]: ... + + def preview_page( + self, + *, + workspace_id: str, + actor_id: str, + credential_id: str, + page_id: str, + page_type: NotionPageType, + ) -> str: ... + + +class NotionImportError(Exception): + """Base class for framework-neutral Notion import failures.""" + + +class DatasetIsNotNotionSourceError(NotionImportError): + def __init__(self) -> None: + super().__init__("Dataset is not notion type") + + +class NotionImportCredentialUnavailableError(NotionImportError): + def __init__(self) -> None: + super().__init__("Notion credential is unavailable") + + +class NotionImportApplicationService: + def __init__( + self, + *, + dataset_access: DatasetAccess, + datasets: NotionDatasetReader, + documents: NotionDocumentBindingReader, + source: NotionSourceGateway, + ) -> None: + self._dataset_access = dataset_access + self._datasets = datasets + self._documents = documents + self._source = source + + def list_pages( + self, + context: RequestContext, + *, + credential_id: str, + dataset_id: str | None = None, + ) -> NotionImportResult: + bound_page_ids: frozenset[str] = frozenset() + if dataset_id is not None: + dataset = self._dataset_access.require_accessible(context, dataset_id) + dataset_ref = DatasetRef(dataset.workspace_id, dataset.id) + if not self._datasets.is_notion_dataset(dataset_ref): + raise DatasetIsNotNotionSourceError() + bound_page_ids = self._documents.list_bound_notion_page_ids(dataset_ref) + + workspaces = self._source.list_authorized_pages( + workspace_id=context.active_workspace_id, + actor_id=context.account_id, + credential_id=credential_id, + ) + return NotionImportResult( + workspaces=tuple( + NotionImportWorkspace( + workspace_id=workspace.workspace_id, + workspace_name=workspace.workspace_name, + workspace_icon=workspace.workspace_icon, + pages=tuple( + NotionImportPage( + page_id=page.page_id, + page_name=page.page_name, + page_icon=page.page_icon, + parent_id=page.parent_id, + page_type=page.page_type, + is_bound=page.page_id in bound_page_ids, + ) + for page in workspace.pages + ), + ) + for workspace in workspaces + ) + ) + + def preview_page( + self, + context: RequestContext, + *, + credential_id: str, + page_id: str, + page_type: NotionPageType, + ) -> str: + return self._source.preview_page( + workspace_id=context.active_workspace_id, + actor_id=context.account_id, + credential_id=credential_id, + page_id=page_id, + page_type=page_type, + ) diff --git a/api/services/notion_data_source_gateway.py b/api/services/data_source/notion_oauth_gateway.py similarity index 95% rename from api/services/notion_data_source_gateway.py rename to api/services/data_source/notion_oauth_gateway.py index 5cadbc82fb9..75405f0529a 100644 --- a/api/services/notion_data_source_gateway.py +++ b/api/services/data_source/notion_oauth_gateway.py @@ -2,15 +2,16 @@ import urllib.parse from collections.abc import Mapping -from typing import Any, Literal, TypedDict, override +from typing import Any, TypedDict, override import httpx from pydantic import TypeAdapter, ValidationError from core.helper.ssrf_proxy import MaxRetriesExceededError, SSRFProxy from core.tools.errors import ToolSSRFError -from services.data_source_oauth_service import DataSourceOAuthError, DataSourceProviderGateway -from services.entities.data_source_oauth_entities import DataSourceOAuthAuthorization +from services.data_source.entities.notion_import import NotionPageType +from services.data_source.entities.oauth import DataSourceOAuthAuthorization +from services.data_source.oauth_service import DataSourceOAuthError, DataSourceProviderGateway class _NotionPageSummary(TypedDict): @@ -18,7 +19,7 @@ class _NotionPageSummary(TypedDict): page_name: str page_icon: dict[str, str] | None parent_id: str - type: Literal["page", "database"] + type: NotionPageType class _NotionSourceInfo(TypedDict): @@ -147,7 +148,7 @@ class NotionDataSourceGateway(DataSourceProviderGateway): "page_name": page_name, "page_icon": icon, "parent_id": parent_id, - "type": "page", + "type": NotionPageType.PAGE, } ) ) @@ -180,19 +181,19 @@ class NotionDataSourceGateway(DataSourceProviderGateway): "page_name": page_name, "page_icon": icon, "parent_id": parent_id, - "type": "database", + "type": NotionPageType.DATABASE, } ) ) return pages def notion_page_search(self, access_token: str) -> list[dict[str, Any]]: - return self._search(access_token, object_type="page") + return self._search(access_token, object_type=NotionPageType.PAGE) def notion_database_search(self, access_token: str) -> list[dict[str, Any]]: - return self._search(access_token, object_type="database") + return self._search(access_token, object_type=NotionPageType.DATABASE) - def _search(self, access_token: str, *, object_type: str) -> list[dict[str, Any]]: + def _search(self, access_token: str, *, object_type: NotionPageType) -> list[dict[str, Any]]: results: list[dict[str, Any]] = [] next_cursor: str | None = None has_more = True diff --git a/api/services/data_source_oauth_service.py b/api/services/data_source/oauth_service.py similarity index 98% rename from api/services/data_source_oauth_service.py rename to api/services/data_source/oauth_service.py index 7e91839a9fe..16f5e258957 100644 --- a/api/services/data_source_oauth_service.py +++ b/api/services/data_source/oauth_service.py @@ -4,7 +4,7 @@ from collections.abc import Mapping from typing import Protocol from machinery.context import RequestContext -from services.entities.data_source_oauth_entities import ( +from services.data_source.entities.oauth import ( DataSourceOAuthAuthorization, DataSourceOAuthBindingRecord, DataSourceOAuthCallback, diff --git a/api/services/datasource_provider_service.py b/api/services/data_source/provider_service.py similarity index 79% rename from api/services/datasource_provider_service.py rename to api/services/data_source/provider_service.py index ed394bd0a29..008dc846b31 100644 --- a/api/services/datasource_provider_service.py +++ b/api/services/data_source/provider_service.py @@ -1,7 +1,10 @@ import logging import time from collections.abc import Mapping -from typing import TYPE_CHECKING, Any +from dataclasses import replace +from typing import TYPE_CHECKING, Any, cast + +from services.credentials.query import CredentialQuery if TYPE_CHECKING: from models.account import Account @@ -17,14 +20,25 @@ from core.helper.provider_cache import NoOpProviderCredentialCache from core.plugin.entities.plugin_daemon import CredentialType from core.plugin.impl.datasource import PluginDatasourceManager from core.plugin.impl.oauth import OAuthHandler -from core.plugin.plugin_service import PluginService from core.tools.utils.encryption import ProviderConfigCache, ProviderConfigEncrypter, create_provider_encrypter from extensions.ext_database import db from extensions.ext_redis import redis_client -from graphon.model_runtime.entities.provider_entities import FormType from models.enums import PermissionEnum from models.oauth import DatasourceOauthParamConfig, DatasourceOauthTenantParamConfig, DatasourceProvider from models.provider_ids import DatasourceProviderID +from services.data_source.credential_adapters import ( + datasource_secret_variables, + refresh_datasource_credential, + resolve_datasource_oauth_client, + transform_secret_fields, +) +from services.data_source.credential_gateway import ( + DatasourceCredentialError, + DatasourceProviderCredentialStore, + credential_needs_refresh, + resolve_stored_notion_access_token, +) +from services.data_source.entities.credential import DatasourceCredentialRecord logger = logging.getLogger(__name__) @@ -47,25 +61,47 @@ def get_current_user(): class DatasourceProviderService: """ - Model Provider Service + Datasource provider credentials and authorization management. """ - def __init__(self) -> None: + def __init__(self, *, credentials: DatasourceProviderCredentialStore) -> None: + self._credentials = credentials self.provider_manager = PluginDatasourceManager() + def get_stored_notion_access_token(self, *, tenant_id: str, credential_id: str | None) -> str: + """Resolve a stored Notion credential with the configured integration fallback.""" + + def load_token(saved_id: str) -> str: + credentials = self.get_datasource_credentials( + tenant_id=tenant_id, + credential_id=saved_id, + provider="notion_datasource", + plugin_id="langgenius/notion_datasource", + ) + token = credentials.get("integration_secret") if credentials else None + if not isinstance(token, str) or not token: + raise ValueError(f"No notion credential found for tenant {tenant_id} and credential {saved_id}") + return token + + return resolve_stored_notion_access_token( + credential_id=credential_id, + load_token=load_token, + integration_token=dify_config.NOTION_INTEGRATION_TOKEN, + ) + @staticmethod - def _should_refresh_credentials(datasource_provider: DatasourceProvider, now: int | None = None) -> bool: + def _should_refresh_credentials( + datasource_provider: DatasourceProvider | DatasourceCredentialRecord, now: int | None = None + ) -> bool: current_time = int(time.time()) if now is None else now - if datasource_provider.expires_at == -1: - return False - return (datasource_provider.expires_at - 60) < current_time + return credential_needs_refresh(datasource_provider.expires_at, now=current_time) def _refresh_datasource_credentials( self, tenant_id: str, provider: str, plugin_id: str, - datasource_provider: DatasourceProvider, + datasource_provider: DatasourceProvider | DatasourceCredentialRecord, current_user: Any, ) -> tuple[dict[str, Any], int]: datasource_provider_id = DatasourceProviderID(f"{plugin_id}/{provider}") @@ -90,17 +126,13 @@ class DatasourceProviderService: plugin_id=plugin_id, provider=provider, ) - redirect_uri = ( - f"{dify_config.CONSOLE_API_URL}/console/api/oauth/plugin/{datasource_provider_id}/datasource/callback" - ) system_credentials = self.get_oauth_client(tenant_id, datasource_provider_id) try: - refreshed_credentials = OAuthHandler().refresh_credentials( - tenant_id=tenant_id, - user_id=current_user.id, - plugin_id=datasource_provider_id.plugin_id, - provider=provider_name, - redirect_uri=redirect_uri, + refreshed_credentials = refresh_datasource_credential( + OAuthHandler(), + workspace_id=tenant_id, + actor_id=current_user.id, + provider_id=datasource_provider_id, system_credentials=system_credentials or {}, credentials=decrypted_credentials, ) @@ -157,7 +189,7 @@ class DatasourceProviderService: def decrypt_datasource_provider_credentials( self, tenant_id: str, - datasource_provider: DatasourceProvider, + datasource_provider: DatasourceProvider | DatasourceCredentialRecord, plugin_id: str, provider: str, ) -> dict[str, Any]: @@ -167,11 +199,9 @@ class DatasourceProviderService: provider_id=f"{plugin_id}/{provider}", credential_type=CredentialType.of(datasource_provider.auth_type), ) - decrypted_credentials = encrypted_credentials.copy() - for key, value in decrypted_credentials.items(): - if key in credential_secret_variables: - decrypted_credentials[key] = encrypter.decrypt_token(tenant_id, value) - return decrypted_credentials + return transform_secret_fields( + tenant_id, encrypted_credentials, credential_secret_variables, encrypter.decrypt_token + ) def encrypt_datasource_provider_credentials( self, @@ -179,18 +209,16 @@ class DatasourceProviderService: provider: str, plugin_id: str, raw_credentials: Mapping[str, Any], - datasource_provider: DatasourceProvider, + datasource_provider: DatasourceProvider | DatasourceCredentialRecord, ) -> dict[str, Any]: provider_credential_secret_variables = self.extract_secret_variables( tenant_id=tenant_id, provider_id=f"{plugin_id}/{provider}", credential_type=CredentialType.of(datasource_provider.auth_type), ) - encrypted_credentials = dict(raw_credentials) - for key, value in encrypted_credentials.items(): - if key in provider_credential_secret_variables: - encrypted_credentials[key] = encrypter.encrypt_token(tenant_id, value) - return encrypted_credentials + return transform_secret_fields( + tenant_id, raw_credentials, provider_credential_secret_variables, encrypter.encrypt_token + ) def get_datasource_credentials( self, @@ -205,44 +233,44 @@ class DatasourceProviderService: If the stored credential is expired or about to expire, this method refreshes it through plugin-daemon and persists the refreshed credential before returning. """ - with sessionmaker(bind=db.engine).begin() as session: - if credential_id: - datasource_provider = session.scalar( - select(DatasourceProvider) - .where(DatasourceProvider.tenant_id == tenant_id, DatasourceProvider.id == credential_id) - .limit(1) - ) - else: - datasource_provider = session.scalar( - select(DatasourceProvider) - .where( - DatasourceProvider.tenant_id == tenant_id, - DatasourceProvider.provider == provider, - DatasourceProvider.plugin_id == plugin_id, - ) - .order_by(DatasourceProvider.is_default.desc(), DatasourceProvider.created_at.asc()) - .limit(1) - ) - if not datasource_provider: - return {} - if self._should_refresh_credentials(datasource_provider): - current_user = get_current_user() - encrypted_credentials, expires_at = self._refresh_datasource_credentials( - tenant_id=tenant_id, - provider=provider, - plugin_id=plugin_id, - datasource_provider=datasource_provider, - current_user=current_user, - ) - datasource_provider.encrypted_credentials = encrypted_credentials - datasource_provider.expires_at = expires_at + record = self._credentials.get_provider_credential( + workspace_id=tenant_id, provider=provider, plugin_id=plugin_id, credential_id=credential_id + ) + if record is None: + return {} + return self._resolve_provider_credentials(record) - return self.decrypt_datasource_provider_credentials( - tenant_id=tenant_id, - datasource_provider=datasource_provider, - plugin_id=plugin_id, - provider=provider, + def _resolve_provider_credentials(self, record: DatasourceCredentialRecord) -> dict[str, Any]: + """Refresh a snapshot; the store owns each read and conditional write transaction.""" + if self._should_refresh_credentials(record): + encrypted, expires_at = self._refresh_datasource_credentials( + tenant_id=record.workspace_id, + provider=record.provider, + plugin_id=record.plugin_id, + datasource_provider=record, + current_user=get_current_user(), ) + saved = self._credentials.update_if_unchanged( + record=record, encrypted_credentials=encrypted, expires_at=expires_at + ) + if saved: + record = replace(record, encrypted_credentials=encrypted, expires_at=expires_at) + else: + latest = self._credentials.get_provider_credential( + workspace_id=record.workspace_id, + provider=record.provider, + plugin_id=record.plugin_id, + credential_id=record.id, + ) + if latest is None or self._should_refresh_credentials(latest): + raise ValueError("Credential changed while it was being refreshed") + record = latest + return self.decrypt_datasource_provider_credentials( + tenant_id=record.workspace_id, + datasource_provider=record, + plugin_id=record.plugin_id, + provider=record.provider, + ) def get_all_datasource_credentials_by_provider( self, @@ -256,54 +284,20 @@ class DatasourceProviderService: Expired credentials are refreshed independently. A failed credential refresh is logged and skipped so one broken authorization does not block other credentials. """ - with sessionmaker(bind=db.engine).begin() as session: - datasource_providers = session.scalars( - select(DatasourceProvider) - .where( - DatasourceProvider.tenant_id == tenant_id, - DatasourceProvider.provider == provider, - DatasourceProvider.plugin_id == plugin_id, + datasource_providers = self._credentials.list_provider_credentials( + workspace_id=tenant_id, provider=provider, plugin_id=plugin_id + ) + credentials = [] + for datasource_provider in datasource_providers: + try: + credentials.append(self._resolve_provider_credentials(datasource_provider)) + except Exception: + logger.exception( + "Skipping datasource credentials for provider %s after refresh or decrypt failure", + provider, + extra={"tenant_id": tenant_id, "plugin_id": plugin_id, "credential_id": datasource_provider.id}, ) - .order_by(DatasourceProvider.is_default.desc(), DatasourceProvider.created_at.asc()) - ).all() - if not datasource_providers: - return [] - current_user = get_current_user() - real_credentials_list = [] - for datasource_provider in datasource_providers: - try: - if self._should_refresh_credentials(datasource_provider): - encrypted_credentials, expires_at = self._refresh_datasource_credentials( - tenant_id=tenant_id, - provider=provider, - plugin_id=plugin_id, - datasource_provider=datasource_provider, - current_user=current_user, - ) - datasource_provider.encrypted_credentials = encrypted_credentials - datasource_provider.expires_at = expires_at - real_credentials = self.decrypt_datasource_provider_credentials( - tenant_id=tenant_id, - datasource_provider=datasource_provider, - plugin_id=plugin_id, - provider=provider, - ) - real_credentials_list.append(real_credentials) - except Exception: - logger.exception( - "Skipping datasource credentials for provider %s after refresh or decrypt failure", - provider, - extra={ - "tenant_id": tenant_id, - "plugin_id": plugin_id, - "provider": provider, - "credential_id": getattr(datasource_provider, "id", None), - "credential_name": getattr(datasource_provider, "name", None), - "expires_at": getattr(datasource_provider, "expires_at", None), - }, - ) - - return real_credentials_list + return credentials def update_datasource_provider_name( self, tenant_id: str, datasource_provider_id: DatasourceProviderID, name: str, credential_id: str @@ -516,42 +510,22 @@ class DatasourceProviderService: """ get oauth client """ - provider = datasource_provider_id.provider_name - plugin_id = datasource_provider_id.plugin_id - with Session(db.engine).no_autoflush as session: - # get tenant oauth client params - tenant_oauth_client_params = session.scalar( - select(DatasourceOauthTenantParamConfig) - .where( - DatasourceOauthTenantParamConfig.tenant_id == tenant_id, - DatasourceOauthTenantParamConfig.provider == provider, - DatasourceOauthTenantParamConfig.plugin_id == plugin_id, - DatasourceOauthTenantParamConfig.enabled.is_(True), - ) - .limit(1) + config = self._credentials.get_oauth_client_config( + workspace_id=tenant_id, + provider=datasource_provider_id.provider_name, + plugin_id=datasource_provider_id.plugin_id, + ) + try: + return resolve_datasource_oauth_client( + config, + workspace_id=tenant_id, + provider_id=datasource_provider_id, + provider_manager=self.provider_manager, ) - if tenant_oauth_client_params: - encrypter, _ = self.get_oauth_encrypter(tenant_id, datasource_provider_id) - return dict(encrypter.decrypt(tenant_oauth_client_params.client_params)) - - provider_controller = self.provider_manager.fetch_datasource_provider( - tenant_id=tenant_id, provider_id=str(datasource_provider_id) - ) - is_verified = PluginService.is_plugin_verified(tenant_id, provider_controller.plugin_unique_identifier) - if is_verified: - # fallback to system oauth client params - oauth_client_params = session.scalar( - select(DatasourceOauthParamConfig) - .where( - DatasourceOauthParamConfig.provider == provider, - DatasourceOauthParamConfig.plugin_id == plugin_id, - ) - .limit(1) - ) - if oauth_client_params: - return oauth_client_params.system_credentials - - raise ValueError(f"Please configure oauth client params(system/tenant) for {plugin_id}/{provider}") + except DatasourceCredentialError as error: + raise ValueError( + f"Please configure oauth client params(system/tenant) for {datasource_provider_id}" + ) from error @staticmethod def generate_next_datasource_provider_name( @@ -800,25 +774,12 @@ class DatasourceProviderService: :param credential_form_schemas: :return: """ - datasource_provider = self.provider_manager.fetch_datasource_provider( - tenant_id=tenant_id, provider_id=provider_id - ) - credential_form_schemas = [] - if credential_type == CredentialType.API_KEY: - credential_form_schemas = list(datasource_provider.declaration.credentials_schema) - elif credential_type == CredentialType.OAUTH2: - if not datasource_provider.declaration.oauth_schema: - raise ValueError("Datasource provider oauth schema not found") - credential_form_schemas = list(datasource_provider.declaration.oauth_schema.credentials_schema) - else: + if credential_type not in (CredentialType.API_KEY, CredentialType.OAUTH2): raise ValueError(f"Invalid credential type: {credential_type}") - - secret_input_form_variables = [] - for credential_form_schema in credential_form_schemas: - if credential_form_schema.type.value == FormType.SECRET_INPUT: - secret_input_form_variables.append(credential_form_schema.name) - - return secret_input_form_variables + try: + return list(datasource_secret_variables(self.provider_manager, tenant_id, provider_id, credential_type)) + except DatasourceCredentialError as error: + raise ValueError(str(error)) from error def list_datasource_credentials( self, @@ -827,7 +788,7 @@ class DatasourceProviderService: plugin_id: str, user: "Account | None" = None, *, - session: Session, + credential_query: CredentialQuery, ) -> list[dict]: """ list datasource credentials with obfuscated sensitive fields, @@ -836,41 +797,13 @@ class DatasourceProviderService: :param tenant_id: workspace id :param provider: provider name :param plugin_id: plugin id - :param user: current user (id + admin flag drive the visibility filter) + :param user: current viewer; administrator roles do not bypass credential visibility :return: """ - from models.credential_permission import CredentialType as CredPermType - from services.credential_permission_service import CredentialPermissionService - - # Get all provider configurations of the current workspace - query = select(DatasourceProvider).where( - DatasourceProvider.tenant_id == tenant_id, - DatasourceProvider.provider == provider, - DatasourceProvider.plugin_id == plugin_id, + datasource_providers = credential_query.list_datasources( + workspace_id=tenant_id, provider=provider, plugin_id=plugin_id, actor_id=user.id if user else None ) - if user is not None: - query = CredentialPermissionService.apply_visibility_filter( - query, - model_id_column=DatasourceProvider.id, - model_user_id_column=DatasourceProvider.user_id, - model_visibility_column=DatasourceProvider.visibility, - credential_type=CredPermType.DATASOURCE_PROVIDER, - user=user, - ) - datasource_providers: list[DatasourceProvider] = list(session.scalars(query).all()) - if not datasource_providers: - return [] copy_credentials_list = [] - default_provider = session.execute( - select(DatasourceProvider.id) - .where( - DatasourceProvider.tenant_id == tenant_id, - DatasourceProvider.provider == provider, - DatasourceProvider.plugin_id == plugin_id, - ) - .order_by(DatasourceProvider.is_default.desc(), DatasourceProvider.created_at.asc()) - ).first() - default_provider_id = default_provider.id if default_provider else None for datasource_provider in datasource_providers: encrypted_credentials = datasource_provider.encrypted_credentials # Get provider credential secret variables @@ -881,10 +814,10 @@ class DatasourceProviderService: ) # Obfuscate provider credentials - copy_credentials = encrypted_credentials.copy() + copy_credentials = dict(encrypted_credentials) for key, value in copy_credentials.items(): if key in credential_secret_variables: - copy_credentials[key] = encrypter.obfuscated_token(value) + copy_credentials[key] = encrypter.obfuscated_token(cast(str, value)) copy_credentials_list.append( { "credential": copy_credentials, @@ -892,14 +825,14 @@ class DatasourceProviderService: "name": datasource_provider.name, "avatar_url": datasource_provider.avatar_url, "id": datasource_provider.id, - "is_default": default_provider_id and datasource_provider.id == default_provider_id, + "is_default": datasource_provider.is_default, } ) return copy_credentials_list def get_all_datasource_credentials( - self, tenant_id: str, *, session: Session, user: "Account | None" = None + self, tenant_id: str, *, session: Session, credential_query: CredentialQuery, user: "Account | None" = None ) -> list[dict]: """ get datasource credentials. @@ -921,7 +854,7 @@ class DatasourceProviderService: provider=datasource.provider, plugin_id=datasource.plugin_id, user=user, - session=session, + credential_query=credential_query, ) redirect_uri = ( f"{dify_config.CONSOLE_API_URL}/console/api/oauth/plugin/{datasource_provider_id}/datasource/callback" @@ -965,7 +898,7 @@ class DatasourceProviderService: return datasource_credentials def get_hard_code_datasource_credentials( - self, tenant_id: str, *, session: Session, user: "Account | None" = None + self, tenant_id: str, *, session: Session, credential_query: CredentialQuery, user: "Account | None" = None ) -> list[dict]: """ get hard code datasource credentials. @@ -990,7 +923,7 @@ class DatasourceProviderService: provider=datasource.provider, plugin_id=datasource.plugin_id, user=user, - session=session, + credential_query=credential_query, ) redirect_uri = "{}/console/api/oauth/plugin/{}/datasource/callback".format( dify_config.CONSOLE_API_URL, datasource_provider_id diff --git a/api/services/website_service.py b/api/services/data_source/website_service.py similarity index 88% rename from api/services/website_service.py rename to api/services/data_source/website_service.py index b833c146216..049ae1bb4dc 100644 --- a/api/services/website_service.py +++ b/api/services/data_source/website_service.py @@ -14,7 +14,7 @@ from core.rag.extractor.firecrawl.firecrawl_app import CrawlStatusResponse, Fire from core.rag.extractor.watercrawl.provider import WaterCrawlProvider from extensions.ext_redis import redis_client from extensions.ext_storage import storage -from services.datasource_provider_service import DatasourceProviderService +from services.data_source.provider_service import DatasourceProviderService # Reuse pooled HTTP clients to avoid creating new connections per request and ease testing. # Both clients carry a bounded read/connect timeout so a stalled Jina or @@ -147,8 +147,10 @@ class CrawlStatusDict(TypedDict): class WebsiteService: """Service class for website crawling operations using different providers.""" - @classmethod - def _get_credentials_and_config(cls, tenant_id: str, provider: str) -> tuple[Any, Any]: + def __init__(self, *, providers: DatasourceProviderService) -> None: + self._providers = providers + + def _get_credentials_and_config(self, tenant_id: str, provider: str) -> tuple[Any, Any]: """Get and validate credentials for a provider.""" if provider == "firecrawl": plugin_id = "langgenius/firecrawl_datasource" @@ -158,8 +160,7 @@ class WebsiteService: plugin_id = "langgenius/jina_datasource" else: raise ValueError("Invalid provider") - datasource_provider_service = DatasourceProviderService() - credential = datasource_provider_service.get_datasource_credentials( + credential = self._providers.get_datasource_credentials( tenant_id=tenant_id, provider=provider, plugin_id=plugin_id, @@ -187,19 +188,18 @@ class WebsiteService: except ValueError as e: raise ValueError(f"Invalid arguments: {e}") - @classmethod - def crawl_url(cls, api_request: WebsiteCrawlApiRequest) -> dict[str, Any]: + def crawl_url(self, api_request: WebsiteCrawlApiRequest) -> dict[str, Any]: """Crawl a URL using the specified provider with typed request.""" request = api_request.to_crawl_request() - api_key, config = cls._get_credentials_and_config(current_user.current_tenant_id, request.provider) + api_key, config = self._get_credentials_and_config(current_user.current_tenant_id, request.provider) if request.provider == "firecrawl": - return cls._crawl_with_firecrawl(request=request, api_key=api_key, config=config) + return self._crawl_with_firecrawl(request=request, api_key=api_key, config=config) elif request.provider == "watercrawl": - return cls._crawl_with_watercrawl(request=request, api_key=api_key, config=config) + return self._crawl_with_watercrawl(request=request, api_key=api_key, config=config) elif request.provider == "jinareader": - return cls._crawl_with_jinareader(request=request, api_key=api_key) + return self._crawl_with_jinareader(request=request, api_key=api_key) else: raise ValueError("Invalid provider") @@ -278,23 +278,21 @@ class WebsiteService: raise ValueError("Failed to crawl") return {"status": "active", "job_id": response.json().get("data", {}).get("taskId")} - @classmethod - def get_crawl_status(cls, job_id: str, provider: str) -> CrawlStatusDict: + def get_crawl_status(self, job_id: str, provider: str) -> CrawlStatusDict: """Get crawl status using string parameters.""" api_request = WebsiteCrawlStatusApiRequest(provider=provider, job_id=job_id) - return cls.get_crawl_status_typed(api_request) + return self.get_crawl_status_typed(api_request) - @classmethod - def get_crawl_status_typed(cls, api_request: WebsiteCrawlStatusApiRequest) -> CrawlStatusDict: + def get_crawl_status_typed(self, api_request: WebsiteCrawlStatusApiRequest) -> CrawlStatusDict: """Get crawl status using typed request.""" - api_key, config = cls._get_credentials_and_config(current_user.current_tenant_id, api_request.provider) + api_key, config = self._get_credentials_and_config(current_user.current_tenant_id, api_request.provider) if api_request.provider == "firecrawl": - return cls._get_firecrawl_status(api_request.job_id, api_key, config) + return self._get_firecrawl_status(api_request.job_id, api_key, config) elif api_request.provider == "watercrawl": - return cls._get_watercrawl_status(api_request.job_id, api_key, config) + return self._get_watercrawl_status(api_request.job_id, api_key, config) elif api_request.provider == "jinareader": - return cls._get_jinareader_status(api_request.job_id, api_key) + return self._get_jinareader_status(api_request.job_id, api_key) else: raise ValueError("Invalid provider") @@ -359,16 +357,15 @@ class WebsiteService: crawl_status_data["data"] = formatted_data return crawl_status_data - @classmethod - def get_crawl_url_data(cls, job_id: str, provider: str, url: str, tenant_id: str) -> dict[str, Any] | None: - api_key, config = cls._get_credentials_and_config(tenant_id, provider) + def get_crawl_url_data(self, job_id: str, provider: str, url: str, tenant_id: str) -> dict[str, Any] | None: + api_key, config = self._get_credentials_and_config(tenant_id, provider) if provider == "firecrawl": - return cls._get_firecrawl_url_data(job_id, url, api_key, config) + return self._get_firecrawl_url_data(job_id, url, api_key, config) elif provider == "watercrawl": - return cls._get_watercrawl_url_data(job_id, url, api_key, config) + return self._get_watercrawl_url_data(job_id, url, api_key, config) elif provider == "jinareader": - return cls._get_jinareader_url_data(job_id, url, api_key) + return self._get_jinareader_url_data(job_id, url, api_key) else: raise ValueError("Invalid provider") @@ -435,16 +432,15 @@ class WebsiteService: return dict(item.get("data", {})) return None - @classmethod - def get_scrape_url_data(cls, provider: str, url: str, tenant_id: str, only_main_content: bool) -> dict[str, Any]: + def get_scrape_url_data(self, provider: str, url: str, tenant_id: str, only_main_content: bool) -> dict[str, Any]: request = ScrapeRequest(provider=provider, url=url, tenant_id=tenant_id, only_main_content=only_main_content) - api_key, config = cls._get_credentials_and_config(tenant_id=request.tenant_id, provider=request.provider) + api_key, config = self._get_credentials_and_config(tenant_id=request.tenant_id, provider=request.provider) if request.provider == "firecrawl": - return cls._scrape_with_firecrawl(request=request, api_key=api_key, config=config) + return self._scrape_with_firecrawl(request=request, api_key=api_key, config=config) elif request.provider == "watercrawl": - return cls._scrape_with_watercrawl(request=request, api_key=api_key, config=config) + return self._scrape_with_watercrawl(request=request, api_key=api_key, config=config) else: raise ValueError("Invalid provider") diff --git a/api/services/dataset_ref_service.py b/api/services/dataset_ref_service.py deleted file mode 100644 index 6bd5cf444ff..00000000000 --- a/api/services/dataset_ref_service.py +++ /dev/null @@ -1,65 +0,0 @@ -"""Typed resource references for dataset ownership chains.""" - -from typing import NamedTuple - -from sqlalchemy import select -from sqlalchemy.orm import Session - -from models.dataset import Dataset, Document - - -class DatasetRef(NamedTuple): - """Dataset identifiers used to scope downstream resource lookups.""" - - tenant_id: str - dataset_id: str - - -class DocumentRef(NamedTuple): - """Owner-bound lookup coordinates, not proof that a document is authorized or exists.""" - - dataset: DatasetRef - document_id: str - - -class SegmentRef(NamedTuple): - """Segment identifiers used to scope downstream resource lookups.""" - - document: DocumentRef - segment_id: str - - -class DatasetRefService: - """Build child locators from validated dataset roots and resolve them with owner predicates.""" - - @staticmethod - def create_dataset_ref(dataset: Dataset) -> DatasetRef: - """Create a root ref from a dataset already validated by the caller.""" - return DatasetRef(tenant_id=dataset.tenant_id, dataset_id=dataset.id) - - @staticmethod - def create_document_ref(dataset_ref: DatasetRef, document: Document) -> DocumentRef | None: - if document.tenant_id != dataset_ref.tenant_id or document.dataset_id != dataset_ref.dataset_id: - return None - return DatasetRefService.create_document_ref_from_id(dataset_ref, document.id) - - @staticmethod - def create_document_ref_from_id(dataset_ref: DatasetRef, document_id: str) -> DocumentRef: - """Bind a candidate document ID; ownership is enforced when the ref is consumed.""" - return DocumentRef(dataset=dataset_ref, document_id=document_id) - - @staticmethod - def create_segment_ref(document_ref: DocumentRef, segment_id: str) -> SegmentRef: - """Bind a candidate segment ID; ownership is enforced when the ref is consumed.""" - return SegmentRef(document=document_ref, segment_id=segment_id) - - @staticmethod - def get_document_by_ref(document_ref: DocumentRef, *, session: Session) -> Document | None: - """Resolve a document through its complete tenant and dataset ownership chain.""" - return session.scalar( - select(Document).where( - Document.id == document_ref.document_id, - Document.dataset_id == document_ref.dataset.dataset_id, - Document.tenant_id == document_ref.dataset.tenant_id, - ) - ) diff --git a/api/services/knowledge/__init__.py b/api/services/knowledge/__init__.py index b1ff0725e7a..1ae0c3b5d8b 100644 --- a/api/services/knowledge/__init__.py +++ b/api/services/knowledge/__init__.py @@ -1 +1 @@ -"""Knowledge base application services.""" +"""Knowledge-domain application services and adapters.""" diff --git a/api/services/knowledge/api_key_service.py b/api/services/knowledge/api_key_service.py index f03aad1f1ff..875028fab9c 100644 --- a/api/services/knowledge/api_key_service.py +++ b/api/services/knowledge/api_key_service.py @@ -4,14 +4,16 @@ from collections.abc import Callable from typing import Protocol from machinery.context import RequestContext -from services.auth.api_key_contracts import ApiKeyCache, ApiKeyRecord +from services.auth.api_key_contracts import ApiKeyCache, ApiKeyRecord, ApiKeyResourceNotFoundError from services.errors.account import NoPermissionError -from services.knowledge.dataset_access import DatasetAccess +from services.knowledge.dataset_access import ( + DatasetAccess, + DatasetAccessDeniedError, + DatasetNotFoundError, +) class DatasetApiKeyStore(Protocol): - def get_dataset_access(self, workspace_id: str, dataset_id: str, account_id: str) -> DatasetAccess: ... - def list_keys(self, workspace_id: str, dataset_id: str) -> tuple[ApiKeyRecord, ...]: ... def create_key(self, workspace_id: str, dataset_id: str, *, max_keys: int, prefix: str) -> ApiKeyRecord: ... @@ -35,9 +37,17 @@ class UnknownDatasetIdsError(Exception): class DatasetApiKeyService: MAX_KEYS = 10 - def __init__(self, *, keys: DatasetApiKeyStore, cache: ApiKeyCache, rbac_enabled: Callable[[], bool]) -> None: + def __init__( + self, + *, + keys: DatasetApiKeyStore, + cache: ApiKeyCache, + access: DatasetAccess, + rbac_enabled: Callable[[], bool], + ) -> None: self._keys = keys self._cache = cache + self._access = access self._rbac_enabled = rbac_enabled def list_keys(self, context: RequestContext, dataset_id: str) -> tuple[ApiKeyRecord, ...]: @@ -56,9 +66,12 @@ class DatasetApiKeyService: def _check_access(self, context: RequestContext, dataset_id: str) -> None: # Enterprise RBAC is enforced at admission; community deployments use the dataset ACL. if not self._rbac_enabled(): - access = self._keys.get_dataset_access(context.active_workspace_id, dataset_id, context.account_id) - if not access.allows(context.account_id): - raise NoPermissionError("You do not have permission to access this dataset.") + try: + self._access.require_accessible(context, dataset_id) + except DatasetNotFoundError as error: + raise ApiKeyResourceNotFoundError("Dataset not found.") from error + except DatasetAccessDeniedError as error: + raise NoPermissionError("You do not have permission to access this dataset.") from error def list_workspace_keys(self, context: RequestContext) -> tuple[ApiKeyRecord, ...]: return self._keys.list_workspace_keys(context.active_workspace_id) diff --git a/api/services/knowledge/dataset_access.py b/api/services/knowledge/dataset_access.py index c79f40ac513..4d7fd436293 100644 --- a/api/services/knowledge/dataset_access.py +++ b/api/services/knowledge/dataset_access.py @@ -1,22 +1,123 @@ -"""Dataset ACL policy independent of persistence and request globals.""" +"""Dataset access use case, read model, errors, and consumer-owned ports.""" -from dataclasses import dataclass +from typing import NamedTuple, Protocol + +from machinery.context import RequestContext + +_OWNER_ROLE = "owner" +_ALL_TEAM_MEMBERS = "all_team_members" +_PARTIAL_MEMBERS = "partial_members" -@dataclass(frozen=True, slots=True) -class DatasetAccess: - permission: str +class DatasetAccessRecord(NamedTuple): + """Dataset state consumed only by the access policy.""" + + id: str + workspace_id: str maintainer_id: str | None - role: str | None - has_explicit_permission: bool + permission: str - def allows(self, account_id: str) -> bool: - if self.role is None: - return False - if self.role == "owner" or self.maintainer_id == account_id: - return True - if self.permission == "only_me": - return False - if self.permission == "partial_members": - return self.has_explicit_permission - return self.permission == "all_team_members" + +class AccessibleDataset(NamedTuple): + """Minimal result proving that the request may operate on a dataset.""" + + id: str + workspace_id: str + + +class DatasetAccessSnapshot(NamedTuple): + """Dataset state plus the requesting account's partial-member grant.""" + + dataset: DatasetAccessRecord + actor_has_partial_access: bool + + +class DatasetAccessReader(Protocol): + """Load tenant-scoped dataset state needed by the access policy.""" + + def get_access_snapshot( + self, + *, + workspace_id: str, + dataset_id: str, + actor_id: str, + ) -> DatasetAccessSnapshot | None: ... + + +class WorkspaceRoleReader(Protocol): + """Read the account's legacy role in a workspace.""" + + def get_legacy_role(self, *, workspace_id: str, account_id: str) -> str | None: ... + + +class DatasetAccessError(Exception): + """Base class for framework-neutral dataset access failures.""" + + +class DatasetNotFoundError(DatasetAccessError): + def __init__(self) -> None: + super().__init__("Dataset not found") + + +class DatasetAccessDeniedError(DatasetAccessError): + def __init__(self) -> None: + super().__init__("You do not have permission to access this dataset") + + +def can_access_dataset(snapshot: DatasetAccessSnapshot, *, actor_id: str, workspace_role: str | None) -> bool: + """Evaluate legacy dataset visibility without database or framework access.""" + + if workspace_role is None: + return False + dataset = snapshot.dataset + if workspace_role == _OWNER_ROLE or dataset.maintainer_id == actor_id: + return True + if dataset.permission == _ALL_TEAM_MEMBERS: + return True + if dataset.permission == _PARTIAL_MEMBERS: + return snapshot.actor_has_partial_access + return False + + +class DatasetAccessService: + """Resolve an owned dataset and enforce legacy permissions when configured.""" + + def __init__( + self, + *, + datasets: DatasetAccessReader, + workspace_roles: WorkspaceRoleReader, + legacy_permissions_enabled: bool, + ) -> None: + self._datasets = datasets + self._workspace_roles = workspace_roles + self._legacy_permissions_enabled = legacy_permissions_enabled + + def require_accessible(self, context: RequestContext, dataset_id: str) -> AccessibleDataset: + snapshot = self._datasets.get_access_snapshot( + workspace_id=context.active_workspace_id, + dataset_id=dataset_id, + actor_id=context.account_id, + ) + if snapshot is None: + raise DatasetNotFoundError() + return self.check_access(context, snapshot) + + def check_access(self, context: RequestContext, snapshot: DatasetAccessSnapshot) -> AccessibleDataset: + """Authorize already-loaded state without querying the dataset again.""" + if snapshot.dataset.workspace_id != context.active_workspace_id: + raise DatasetAccessDeniedError() + if self._legacy_permissions_enabled: + workspace_role = self._workspace_roles.get_legacy_role( + workspace_id=context.active_workspace_id, + account_id=context.account_id, + ) + if not can_access_dataset(snapshot, actor_id=context.account_id, workspace_role=workspace_role): + raise DatasetAccessDeniedError() + return AccessibleDataset(id=snapshot.dataset.id, workspace_id=snapshot.dataset.workspace_id) + + +class DatasetAccess(Protocol): + def require_accessible(self, context: RequestContext, dataset_id: str) -> AccessibleDataset: ... + + def check_access(self, context: RequestContext, snapshot: DatasetAccessSnapshot) -> AccessibleDataset: ... diff --git a/api/services/knowledge/dataset_read_service.py b/api/services/knowledge/dataset_read_service.py new file mode 100644 index 00000000000..90b076d4400 --- /dev/null +++ b/api/services/knowledge/dataset_read_service.py @@ -0,0 +1,318 @@ +"""Load knowledge details and materialize data before response serialization.""" + +from __future__ import annotations + +import json +from collections.abc import Iterable, Mapping, Sequence +from typing import Any, cast + +from sqlalchemy.orm import Session + +from core.rag.index_processor.constant.built_in_field import BuiltInField, MetadataDataSource +from core.rag.index_processor.constant.query_type import QueryType +from core.tools.signature import sign_upload_file_preview_url +from models.dataset import Dataset, DatasetMetadata, DatasetQuery, DocMetadataDetailItem, Document, DocumentSegment +from models.model import UploadFile +from repositories.knowledge.dataset_read_repository import ( + DatasetDetailPrefetch, + DocumentReadBatch, + build_dataset_detail_prefetch, + get_document_dataset, + get_document_metadata_rows, + get_document_upload_file, + get_document_uploader, + get_query_upload_files, + get_segment_child_chunks, + load_document_read_batch, +) +from repositories.knowledge.segment_read_adapter import get_segment_attachments, sign_segment_content +from services.knowledge.entities.datasets import DatasetDetailRecord + + +def load_dataset_detail( + dataset: Dataset, *, session: Session, prefetch: DatasetDetailPrefetch | None = None +) -> DatasetDetailRecord: + values = prefetch if prefetch is not None else build_dataset_detail_prefetch([dataset], session=session) + return _dataset_detail_values(dataset, values) + + +def load_dataset_details(datasets: Sequence[Dataset], *, session: Session) -> list[DatasetDetailRecord]: + """Read a page in one batch and return values with no ORM or session references.""" + prefetch = build_dataset_detail_prefetch(datasets, session=session) + return [_dataset_detail_values(dataset, prefetch) for dataset in datasets] + + +def _dataset_detail_values(dataset: Dataset, prefetch: DatasetDetailPrefetch) -> DatasetDetailRecord: + owner = (dataset.id, dataset.tenant_id) + return { + "id": dataset.id, + "name": dataset.name, + "description": dataset.description, + "provider": dataset.provider, + "permission": dataset.permission, + "data_source_type": dataset.data_source_type, + "indexing_technique": dataset.indexing_technique, + "app_count": prefetch.app_counts.get(dataset.id, 0), + "document_count": prefetch.document_counts.get(dataset.id, 0), + "word_count": prefetch.word_counts.get(dataset.id, 0), + "created_by": dataset.created_by, + "author_name": prefetch.author_names.get(dataset.created_by), + "created_at": dataset.created_at, + "updated_by": dataset.updated_by, + "updated_at": dataset.updated_at, + "embedding_model": dataset.embedding_model, + "embedding_model_provider": dataset.embedding_model_provider, + "embedding_available": vars(dataset).get("embedding_available"), + "retrieval_model_dict": dataset.retrieval_model_dict, + "summary_index_setting": dataset.summary_index_setting, + "tags": [{"id": tag.id, "name": tag.name, "type": tag.type} for tag in prefetch.tags.get(owner, [])], + "doc_form": dataset.chunk_structure or prefetch.doc_forms.get(owner), + "external_knowledge_info": prefetch.external_knowledge_infos.get(owner) + if dataset.provider == "external" + else None, + "external_retrieval_model": dataset.external_retrieval_model, + "doc_metadata": dataset.build_doc_metadata(prefetch.doc_metadata.get(dataset.id, [])), + "built_in_field_enabled": dataset.built_in_field_enabled, + "pipeline_id": dataset.pipeline_id, + "runtime_mode": dataset.runtime_mode, + "chunk_structure": dataset.chunk_structure, + "icon_info": dataset.icon_info, + "is_published": prefetch.pipeline_published.get(dataset.pipeline_id, False) if dataset.pipeline_id else False, + "total_documents": prefetch.document_counts.get(dataset.id, 0), + "total_available_documents": prefetch.available_document_counts.get(dataset.id, 0), + "enable_api": dataset.enable_api, + "is_multimodal": dataset.is_multimodal, + "permission_keys": vars(dataset).get("permission_keys", []), + "maintainer": dataset.maintainer, + } + + +def load_document_detail(document: Document, *, session: Session) -> dict[str, Any]: + return load_document_details([document], session=session)[0] + + +def load_document_details(documents: Iterable[Document], *, session: Session) -> list[dict[str, Any]]: + """Materialize document attributes and relations while the caller's session is open.""" + items = list(documents) + batch = load_document_read_batch(items, session=session) + return [ + { + "id": document.id, + "position": document.position, + "data_source_type": document.data_source_type, + "data_source_info_dict": document.data_source_info_dict, + "dataset_process_rule_id": document.dataset_process_rule_id, + "name": document.name, + "created_from": document.created_from, + "created_by": document.created_by, + "created_at": document.created_at, + "tokens": document.tokens, + "indexing_status": document.indexing_status, + "error": document.error, + "enabled": document.enabled, + "disabled_at": document.disabled_at, + "disabled_by": document.disabled_by, + "archived": document.archived, + "display_status": document.display_status, + "word_count": document.word_count, + "doc_form": document.doc_form, + "summary_index_status": vars(document).get("summary_index_status"), + "need_summary": document.need_summary, + "completed_segments": vars(document).get("completed_segments"), + "total_segments": vars(document).get("total_segments"), + **document_response_values(document, batch), + } + for document in items + ] + + +def load_segment_detail(segment: DocumentSegment, summary: str | None, *, session: Session) -> dict[str, Any]: + """Load relations and authorize file links before handing plain values to the response model.""" + child_chunks = get_segment_child_chunks(segment, session=session, include_full_doc=False) + return { + "id": segment.id, + "position": segment.position, + "document_id": segment.document_id, + "content": segment.content, + "sign_content": sign_segment_content(segment, session=session), + "answer": segment.answer, + "word_count": segment.word_count, + "tokens": segment.tokens, + "keywords": segment.keywords, + "index_node_id": segment.index_node_id, + "index_node_hash": segment.index_node_hash, + "hit_count": segment.hit_count, + "enabled": segment.enabled, + "disabled_at": segment.disabled_at, + "disabled_by": segment.disabled_by, + "status": segment.status, + "created_by": segment.created_by, + "created_at": segment.created_at, + "updated_at": segment.updated_at, + "updated_by": segment.updated_by, + "indexing_at": segment.indexing_at, + "completed_at": segment.completed_at, + "error": segment.error, + "stopped_at": segment.stopped_at, + "child_chunks": [ + { + "id": chunk.id, + "segment_id": chunk.segment_id, + "content": chunk.content, + "position": chunk.position, + "word_count": chunk.word_count, + "type": chunk.type, + "created_at": chunk.created_at, + "updated_at": chunk.updated_at, + } + for chunk in child_chunks + ], + "attachments": get_segment_attachments(segment, session=session), + "summary": summary, + } + + +def load_segment_details( + segments: Iterable[DocumentSegment], summaries: Mapping[str, str | None], *, session: Session +) -> list[dict[str, Any]]: + return [load_segment_detail(segment, summaries.get(segment.id), session=session) for segment in segments] + + +def format_document_source_detail(document: Document, file: UploadFile | None) -> dict[str, Any]: + if document.data_source_type == "upload_file" and file is not None: + return { + "upload_file": { + "id": file.id, + "name": file.name, + "size": file.size, + "extension": file.extension, + "mime_type": file.mime_type, + "created_by": file.created_by, + "created_at": file.created_at.timestamp(), + } + } + if document.data_source_type in {"notion_import", "website_crawl"}: + return document.data_source_info_dict + return {} + + +def format_document_metadata( + document: Document, metadatas: Sequence[DatasetMetadata], *, built_in_enabled: bool, uploader_name: str | None +) -> list[DocMetadataDetailItem] | None: + values = document.doc_metadata or {} + details: list[DocMetadataDetailItem] = [ + {"id": metadata.id, "name": metadata.name, "type": metadata.type, "value": values.get(metadata.name)} + for metadata in metadatas + ] + if built_in_enabled: + details.extend(format_document_built_in_fields(document, uploader_name=uploader_name)) + return details or None + + +def document_response_values(document: Document, batch: DocumentReadBatch) -> dict[str, Any]: + """Build response-only values from a loaded snapshot without database access.""" + file_id = document.data_source_info_dict.get("upload_file_id") + process_rule = batch.process_rules.get(document.dataset_process_rule_id) + if process_rule is not None and process_rule.dataset_id != document.dataset_id: + process_rule = None + return { + "data_source_detail_dict": format_document_source_detail( + document, batch.uploads.get((document.tenant_id, file_id)) if isinstance(file_id, str) else None + ), + "hit_count": batch.hit_counts.get(document.id, 0), + "doc_metadata_details": format_document_metadata( + document, + batch.metadatas.get(document.id, []), + built_in_enabled=batch.built_in_enabled.get((document.tenant_id, document.dataset_id), False), + uploader_name=batch.uploader_names.get(document.created_by), + ), + "process_rule_dict": process_rule.to_dict() if process_rule else None, + } + + +def get_document_source_detail(document: Document, *, session: Session) -> dict[str, Any]: + file = get_document_upload_file(document, session=session) if document.data_source_type == "upload_file" else None + return format_document_source_detail(document, file) + + +def get_document_metadata_details(document: Document, *, session: Session) -> list[DocMetadataDetailItem] | None: + dataset = get_document_dataset(document, session=session) + built_in_enabled = bool(dataset is not None and dataset.built_in_field_enabled) + return format_document_metadata( + document, + get_document_metadata_rows(document, session=session) if document.doc_metadata else [], + built_in_enabled=built_in_enabled, + uploader_name=get_document_uploader(document, session=session) if built_in_enabled else None, + ) + + +def get_document_built_in_fields(document: Document, *, session: Session) -> list[DocMetadataDetailItem]: + return format_document_built_in_fields(document, uploader_name=get_document_uploader(document, session=session)) + + +def format_document_built_in_fields(document: Document, *, uploader_name: str | None) -> list[DocMetadataDetailItem]: + built_in_fields: list[DocMetadataDetailItem] = [] + built_in_fields.append( + {"id": "built-in", "name": BuiltInField.document_name, "type": "string", "value": document.name} + ) + built_in_fields.append({"id": "built-in", "name": BuiltInField.uploader, "type": "string", "value": uploader_name}) + built_in_fields.append( + { + "id": "built-in", + "name": BuiltInField.upload_date, + "type": "time", + "value": str(document.created_at.timestamp()), + } + ) + built_in_fields.append( + { + "id": "built-in", + "name": BuiltInField.last_update_date, + "type": "time", + "value": str(document.updated_at.timestamp()), + } + ) + built_in_fields.append( + { + "id": "built-in", + "name": BuiltInField.source, + "type": "string", + "value": MetadataDataSource[document.data_source_type], + } + ) + return built_in_fields + + +def get_dataset_queries(query: DatasetQuery, *, session: Session) -> list[dict[str, Any]]: + """Read structured audit content, preserving legacy plain-text queries.""" + try: + decoded = json.loads(query.content) + except (ValueError, TypeError): + decoded = None + raw_records = decoded if isinstance(decoded, list) else [decoded] + if not raw_records or not all( + isinstance(item, dict) + and item.get("content_type") in {QueryType.TEXT_QUERY, QueryType.IMAGE_QUERY} + and isinstance(item.get("content"), str) + for item in raw_records + ): + return [{"content_type": QueryType.TEXT_QUERY, "content": query.content, "file_info": None}] + records = cast(list[dict[str, Any]], raw_records) + files = get_query_upload_files( + query, {item["content"] for item in records if item["content_type"] == QueryType.IMAGE_QUERY}, session=session + ) + result: list[dict[str, Any]] = [] + for item in records: + record: dict[str, Any] = {"content_type": item["content_type"], "content": item["content"], "file_info": None} + file = files.get(item["content"]) if item["content_type"] == QueryType.IMAGE_QUERY else None + if file is not None: + record["file_info"] = { + "id": file.id, + "name": file.name, + "size": file.size, + "extension": file.extension, + "mime_type": file.mime_type, + "source_url": sign_upload_file_preview_url(file.id, file.extension), + } + result.append(record) + return result diff --git a/api/services/dataset_service.py b/api/services/knowledge/dataset_service.py similarity index 74% rename from api/services/dataset_service.py rename to api/services/knowledge/dataset_service.py index 2c6628f6c43..65e5152fd64 100644 --- a/api/services/dataset_service.py +++ b/api/services/knowledge/dataset_service.py @@ -5,12 +5,12 @@ import logging import secrets import time import uuid -from collections import Counter from collections.abc import Sequence -from typing import Annotated, Any, Literal, TypedDict, cast +from contextlib import nullcontext +from dataclasses import replace +from typing import Any, Literal, TypedDict, cast import sqlalchemy as sa -from pydantic import BaseModel, ConfigDict, Field, ValidationError, field_validator from redis.exceptions import LockNotOwnedError from sqlalchemy import ColumnElement, case, delete, exists, func, select, tuple_, update from sqlalchemy.orm import Session @@ -49,7 +49,6 @@ from models.dataset import ( DocumentSegment, ExternalKnowledgeBindings, Pipeline, - SegmentAttachmentBinding, ) from models.enums import ( DatasetRuntimeMode, @@ -58,25 +57,18 @@ from models.enums import ( IndexingStatus, ProcessRuleMode, SegmentStatus, - SegmentType, ) from models.model import UploadFile from models.provider_ids import ModelProviderID from models.source import DataSourceOauthBinding from models.workflow import Workflow from repositories.knowledge import dataset_api_key_bindings -from services.dataset_ref_service import DatasetRef, DatasetRefService, SegmentRef +from repositories.knowledge.dataset_read_repository import get_dataset_doc_form, get_latest_dataset_process_rule +from repositories.knowledge.segment_repository import query_child_chunks from services.document_indexing_proxy.document_indexing_task_proxy import DocumentIndexingTaskProxy from services.document_indexing_proxy.duplicate_document_indexing_task_proxy import DuplicateDocumentIndexingTaskProxy from services.enterprise import rbac_service as enterprise_rbac_service from services.entities.feature_entities import FeatureModel -from services.entities.knowledge_entities.knowledge_entities import ( - ChildChunkUpdateArgs, - KnowledgeConfig, - RerankingModel, - RetrievalModel, - SegmentUpdateArgs, -) from services.entities.knowledge_entities.rag_pipeline_entities import ( KnowledgeConfiguration, RagPipelineDatasetCreateEntity, @@ -86,23 +78,35 @@ from services.errors.chunk import ChildChunkDeleteIndexError, ChildChunkIndexing from services.errors.dataset import DatasetNameDuplicateError from services.errors.document import DocumentIndexingError from services.errors.file import FileNotExistsError -from services.external_knowledge_service import ExternalDatasetService from services.feature_service import FeatureService from services.file_service import FileService -from services.knowledge.dataset_access import DatasetAccess +from services.knowledge.dataset_access import DatasetAccessRecord, DatasetAccessSnapshot, can_access_dataset +from services.knowledge.entities.document_creation import DocumentCreationDataset, DocumentIndexingJobs +from services.knowledge.entities.knowledge_entities import ( + KnowledgeConfig, + RerankingModel, + RetrievalModel, +) +from services.knowledge.entities.segments import SegmentUpdateArgs +from services.knowledge.external.service import ExternalDatasetService +from services.knowledge.resource_scope import DatasetRef, SegmentRef +from services.knowledge.segments.application import ( + ChildChunkDeleteIndexApplicationError, + ChildChunkIndexingApplicationError, + SegmentDatasetRecord, + SegmentDocumentRecord, + SegmentMutationScope, + SegmentMutationService, + validate_segment_values, +) from services.rag_pipeline.rag_pipeline import RagPipelineService from services.tag_application_service import TagTargetQuery -from services.vector_service import VectorService from tasks.add_document_to_index_task import add_document_to_index_task from tasks.batch_clean_document_task import batch_clean_document_task from tasks.clean_notion_document_task import clean_notion_document_task from tasks.deal_dataset_index_update_task import deal_dataset_index_update_task from tasks.deal_dataset_vector_index_task import deal_dataset_vector_index_task -from tasks.delete_segment_from_index_task import delete_segment_from_index_task -from tasks.disable_segment_from_index_task import disable_segment_from_index_task -from tasks.disable_segments_from_index_task import disable_segments_from_index_task from tasks.document_indexing_update_task import document_indexing_update_task -from tasks.enable_segments_to_index_task import enable_segments_to_index_task from tasks.recover_document_indexing_task import recover_document_indexing_task from tasks.regenerate_summary_index_task import regenerate_summary_index_task from tasks.remove_document_from_index_task import remove_document_from_index_task @@ -122,121 +126,6 @@ class AutoDisableLogsDict(TypedDict): count: int -class _EstimatePreProcessingRule(BaseModel): - id: str = Field(min_length=1) - enabled: bool - - @field_validator("id") - @classmethod - def _validate_id(cls, v: str) -> str: - if v not in DatasetProcessRule.PRE_PROCESSING_RULES: - raise ValueError("Process rule pre_processing_rules id is invalid") - return v - - -class _EstimateSegmentation(BaseModel): - separator: str = Field(min_length=1) - max_tokens: int = Field(gt=0) - - -class _EstimateRules(BaseModel): - pre_processing_rules: list[_EstimatePreProcessingRule] - segmentation: _EstimateSegmentation - - @field_validator("pre_processing_rules") - @classmethod - def _deduplicate(cls, v: list[_EstimatePreProcessingRule]) -> list[_EstimatePreProcessingRule]: - seen: dict[str, _EstimatePreProcessingRule] = {} - for rule in v: - seen[rule.id] = rule - return list(seen.values()) - - -class _EstimateHierarchicalRules(_EstimateRules): - parent_mode: Literal["full-doc", "paragraph"] | None = None - subchunk_segmentation: _EstimateSegmentation | None = None - - -class _SummaryIndexSettingDisabled(BaseModel): - enable: Literal[False] = False - - -class _SummaryIndexSettingEnabled(BaseModel): - enable: Literal[True] - model_name: str = Field(min_length=1) - model_provider_name: str = Field(min_length=1) - - -_SummaryIndexSetting = Annotated[ - _SummaryIndexSettingDisabled | _SummaryIndexSettingEnabled, - Field(discriminator="enable"), -] - - -class _AutomaticProcessRule(BaseModel): - model_config = ConfigDict(extra="allow") - - mode: Literal[ProcessRuleMode.AUTOMATIC] - summary_index_setting: _SummaryIndexSetting | None = None - - @field_validator("summary_index_setting", mode="before") - @classmethod - def _normalize_summary_index_setting(cls, v: Any) -> Any: - """Treat dicts with enable=None (or missing enable) as None (#36602).""" - if v is None: - return None - if isinstance(v, dict) and v.get("enable") is None: - return None - return v - - -class _CustomProcessRule(BaseModel): - model_config = ConfigDict(extra="allow") - - mode: Literal[ProcessRuleMode.CUSTOM] - rules: _EstimateRules - summary_index_setting: _SummaryIndexSetting | None = None - - @field_validator("summary_index_setting", mode="before") - @classmethod - def _normalize_summary_index_setting(cls, v: Any) -> Any: - """Treat dicts with enable=None (or missing enable) as None (#36602).""" - if v is None: - return None - if isinstance(v, dict) and v.get("enable") is None: - return None - return v - - -class _HierarchicalProcessRule(BaseModel): - model_config = ConfigDict(extra="allow") - - mode: Literal[ProcessRuleMode.HIERARCHICAL] - rules: _EstimateHierarchicalRules - summary_index_setting: _SummaryIndexSetting | None = None - - @field_validator("summary_index_setting", mode="before") - @classmethod - def _normalize_summary_index_setting(cls, v: Any) -> Any: - """Treat dicts with enable=None (or missing enable) as None (#36602).""" - if v is None: - return None - if isinstance(v, dict) and v.get("enable") is None: - return None - return v - - -_EstimateProcessRule = Annotated[ - _AutomaticProcessRule | _CustomProcessRule | _HierarchicalProcessRule, - Field(discriminator="mode"), -] - - -class _EstimateArgs(BaseModel): - info_list: dict[str, Any] - process_rule: _EstimateProcessRule - - class DatasetService: @staticmethod def _can_manage_all_datasets(tenant_id: str, account_id: str, *, session: Session) -> bool: @@ -561,7 +450,7 @@ class DatasetService: @staticmethod def check_doc_form(dataset: Dataset, doc_form: str, *, session: Session): - dataset_doc_form = dataset.get_doc_form(session=session) + dataset_doc_form = get_dataset_doc_form(dataset, session=session) if dataset_doc_form and doc_form != dataset_doc_form: raise ValueError("doc_form is different from the dataset doc_form.") @@ -1403,8 +1292,16 @@ class DatasetService: ) is not None ) - access = DatasetAccess(dataset.permission, dataset.maintainer, user.current_role, has_permission) - if not access.allows(user.id): + snapshot = DatasetAccessSnapshot( + dataset=DatasetAccessRecord( + id=dataset.id, + workspace_id=dataset.tenant_id, + maintainer_id=dataset.maintainer, + permission=dataset.permission, + ), + actor_has_partial_access=has_permission, + ) + if not can_access_dataset(snapshot, actor_id=user.id, workspace_role=user.current_role): logger.debug("User %s does not have permission to access dataset %s", user.id, dataset.id) raise NoPermissionError("You do not have permission to access this dataset.") @@ -1775,9 +1672,9 @@ class DocumentService: # Calculate summary_index_status for documents that need summary (only if dataset summary index is enabled) summary_status_map: dict[str, str | None] = {} if has_summary_index and document_ids_need_summary: - from services.summary_index_service import SummaryIndexService + from services.knowledge.summaries.adapters import SummaryIndexAdapter - summary_status_map = SummaryIndexService.get_documents_summary_index_status( + summary_status_map = SummaryIndexAdapter.get_documents_summary_index_status( document_ids=document_ids_need_summary, dataset_id=dataset.id, tenant_id=tenant_id, @@ -2122,7 +2019,7 @@ class DocumentService: return document @staticmethod - def pause_document(document, session: Session): + def pause_document(document, session: Session, *, actor_id: str): if document.indexing_status not in { IndexingStatus.WAITING, IndexingStatus.PARSING, @@ -2132,9 +2029,8 @@ class DocumentService: }: raise DocumentIndexingError() # update document to be paused - assert current_user is not None document.is_paused = True - document.paused_by = current_user.id + document.paused_by = actor_id document.paused_at = naive_utc_now() session.add(document) @@ -2161,13 +2057,13 @@ class DocumentService: recover_document_indexing_task.delay(document.dataset_id, document.id) @staticmethod - def retry_document(dataset_id: str, documents: list[Document], session: Session): + def retry_document(dataset_id: str, documents: list[Document], session: Session, *, actor_id: str): """Reserve the whole retry batch before changing any document state. Redis lock acquisition is intentionally coupled to this bounded status transaction so a concurrent request cannot partially admit the batch. """ - if not current_user or not current_user.id: + if not actor_id: raise ValueError("Current user or current user id not found") unique_documents = list({document.id: document for document in documents}.values()) @@ -2202,12 +2098,11 @@ class DocumentService: raise document_ids = [document.id for document in unique_documents] - retry_document_indexing_task.delay(dataset_id, document_ids, current_user.id) + retry_document_indexing_task.delay(dataset_id, document_ids, actor_id) @staticmethod def sync_website_document(dataset: Dataset, document: Document, session: Session): - dataset_ref = DatasetRefService.create_dataset_ref(dataset) - if DatasetRefService.create_document_ref(dataset_ref, document) is None: + if document.tenant_id != dataset.tenant_id or document.dataset_id != dataset.id: raise ValueError("Document not found.") # add sync flag @@ -2248,14 +2143,47 @@ class DocumentService: *, session: Session, ) -> tuple[list[Document], str]: - # check doc_form + # Compatibility entry point for existing callers that own their session. DatasetService.check_doc_form(dataset, knowledge_config.doc_form, session=session) - # check document limit - assert isinstance(current_user, Account) - assert current_user.current_tenant_id is not None + assert isinstance(account, Account) + assert account.current_tenant_id is not None + config = DocumentService.prepare_document_creation( + DocumentCreationDataset( + dataset.tenant_id, dataset.indexing_technique, dataset.embedding_model, dataset.embedding_model_provider + ), + knowledge_config, + ) + documents: list[Document] = [] + batch = "" + lock = ( + nullcontext() + if config.original_document_id + else redis_client.lock(f"add_document_lock_dataset_id_{dataset.id}", timeout=600) + ) + try: + with lock: + documents, batch, jobs = DocumentService.save_prepared_documents( + dataset, config, account, dataset_process_rule, created_from, session=session + ) + DocumentService.dispatch_document_indexing(jobs) + except LockNotOwnedError: + if config.original_document_id: + raise + if not batch: + batch = DocumentService.generate_document_batch() + return documents, batch + @staticmethod + def generate_document_batch() -> str: + return time.strftime("%Y%m%d%H%M%S") + str(100000 + secrets.randbelow(exclusive_upper_bound=900000)) + + @staticmethod + def prepare_document_creation( + dataset: DocumentCreationDataset, knowledge_config: KnowledgeConfig + ) -> KnowledgeConfig: + """Resolve remote quotas and models from values; new callers close their read transaction first.""" if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD: - features = FeatureService.get_features(current_user.current_tenant_id, exclude_vector_space=True) + features = FeatureService.get_features(dataset.tenant_id, exclude_vector_space=True) if not knowledge_config.original_document_id: count = 0 if knowledge_config.data_source: @@ -2274,6 +2202,59 @@ class DocumentService: count = len(website_info.urls) DocumentService.check_document_creation_limits(count, features) + config = knowledge_config + effective_dataset = dataset + if not dataset.indexing_technique: + if config.indexing_technique not in Dataset.INDEXING_TECHNIQUE_LIST: + raise ValueError("Indexing technique is invalid") + if config.indexing_technique == IndexTechniqueType.HIGH_QUALITY: + if not config.embedding_model or not config.embedding_model_provider: + model = ModelManager.for_tenant(tenant_id=dataset.tenant_id).get_default_model_instance( + tenant_id=dataset.tenant_id, model_type=ModelType.TEXT_EMBEDDING + ) + config = config.model_copy( + update={"embedding_model": model.model_name, "embedding_model_provider": model.provider} + ) + effective_dataset = replace( + dataset, + indexing_technique=config.indexing_technique, + embedding_model=config.embedding_model, + embedding_model_provider=config.embedding_model_provider, + ) + if config.original_document_id: + DatasetService.check_dataset_model_setting(effective_dataset) + return config + + @staticmethod + def dispatch_document_indexing(jobs: DocumentIndexingJobs) -> None: + """Publish only identifiers after the corresponding database writes commit.""" + ref = jobs.dataset + if jobs.removed_notion: + clean_notion_document_task.delay(list(jobs.removed_notion), ref.dataset_id) + if jobs.created: + DocumentIndexingTaskProxy(ref.tenant_id, ref.dataset_id, list(jobs.created)).delay() + if jobs.duplicated: + DuplicateDocumentIndexingTaskProxy(ref.tenant_id, ref.dataset_id, list(jobs.duplicated)).delay() + if jobs.updated: + document_indexing_update_task.delay(ref.dataset_id, jobs.updated) + + @staticmethod + def save_prepared_documents( + dataset: Dataset, + knowledge_config: KnowledgeConfig, + account: Account | Any, + dataset_process_rule: DatasetProcessRule | None = None, + created_from: str = DocumentCreatedFrom.WEB, + *, + session: Session, + ) -> tuple[list[Document], str, DocumentIndexingJobs]: + """Persist prepared input without remote calls, Redis locks, or task dispatch. + + The caller must resolve quotas/models with prepare_document_creation and + hold the dataset's creation lock for new documents. ORM values stay within + this persistence adapter; publish the returned jobs after closing its session. + """ + DatasetService.check_doc_form(dataset, knowledge_config.doc_form, session=session) # if dataset is empty, update dataset data_source_type if not dataset.data_source_type and knowledge_config.data_source: dataset.data_source_type = knowledge_config.data_source.info_list.data_source_type @@ -2284,16 +2265,10 @@ class DocumentService: dataset.indexing_technique = IndexTechniqueType(knowledge_config.indexing_technique) if knowledge_config.indexing_technique == IndexTechniqueType.HIGH_QUALITY: - model_manager = ModelManager.for_tenant(tenant_id=current_user.current_tenant_id) - if knowledge_config.embedding_model and knowledge_config.embedding_model_provider: - dataset_embedding_model = knowledge_config.embedding_model - dataset_embedding_model_provider = knowledge_config.embedding_model_provider - else: - embedding_model = model_manager.get_default_model_instance( - tenant_id=current_user.current_tenant_id, model_type=ModelType.TEXT_EMBEDDING - ) - dataset_embedding_model = embedding_model.model_name - dataset_embedding_model_provider = embedding_model.provider + assert knowledge_config.embedding_model + assert knowledge_config.embedding_model_provider + dataset_embedding_model = knowledge_config.embedding_model + dataset_embedding_model_provider = knowledge_config.embedding_model_provider dataset.embedding_model = dataset_embedding_model dataset.embedding_model_provider = dataset_embedding_model_provider dataset_collection_binding = DatasetCollectionBindingService.get_dataset_collection_binding( @@ -2315,19 +2290,19 @@ class DocumentService: else default_retrieval_model ) - documents = [] + documents: list[Document] = [] + jobs = DocumentIndexingJobs(DatasetRef(dataset.tenant_id, dataset.id)) if knowledge_config.original_document_id: - document = DocumentService.update_document_with_dataset_id( - dataset, knowledge_config, account, session=session - ) + document = DocumentService.update_prepared_document(dataset, knowledge_config, account, session=session) documents.append(document) batch = document.batch + jobs = replace(jobs, updated=document.id) else: # When creating new documents, data_source must be provided if not knowledge_config.data_source: raise ValueError("Data source is required when creating new documents") - batch = time.strftime("%Y%m%d%H%M%S") + str(100000 + secrets.randbelow(exclusive_upper_bound=900000)) + batch = DocumentService.generate_document_batch() # save process rule if not dataset_process_rule: process_rule = knowledge_config.process_rule @@ -2341,7 +2316,7 @@ class DocumentService: created_by=account.id, ) else: - dataset_process_rule = dataset.get_latest_process_rule(session=session) + dataset_process_rule = get_latest_dataset_process_rule(dataset, session=session) if not dataset_process_rule: raise ValueError("No process rule found.") elif process_rule.mode == ProcessRuleMode.AUTOMATIC: @@ -2356,14 +2331,14 @@ class DocumentService: "Invalid process rule mode: %s, can not find dataset process rule", process_rule.mode, ) - return [], "" + return [], "", jobs session.add(dataset_process_rule) session.flush() else: # Fallback when no process_rule provided in knowledge_config: # 1) reuse the dataset's latest process rule if present # 2) otherwise create an automatic rule - dataset_process_rule = dataset.get_latest_process_rule(session=session) + dataset_process_rule = get_latest_dataset_process_rule(dataset, session=session) if not dataset_process_rule: dataset_process_rule = DatasetProcessRule( dataset_id=dataset.id, @@ -2373,170 +2348,127 @@ class DocumentService: ) session.add(dataset_process_rule) session.flush() - lock_name = f"add_document_lock_dataset_id_{dataset.id}" - try: - with redis_client.lock(lock_name, timeout=600): - assert dataset_process_rule - position = DocumentService.get_documents_position(dataset.id, session) - document_ids = [] - duplicate_document_ids = [] - if knowledge_config.data_source.info_list.data_source_type == "upload_file": - if not knowledge_config.data_source.info_list.file_info_list: - raise ValueError("File source info is required") - upload_file_list = knowledge_config.data_source.info_list.file_info_list.file_ids - # Issue #41735: under MySQL's default - # ``REPEATABLE READ`` isolation, the SELECT snapshot - # for this transaction is pinned by the first - # read done earlier in ``save_document_with_dataset_id`` - # (e.g. ``check_doc_form``). The caller has just - # committed the upload in a different session - # (``FileService.upload_text`` / - # ``FileService.upload_file``), so the rows exist - # but are invisible to *this* session until its - # snapshot is refreshed. End the current transaction - # so the next SELECT starts with a fresh snapshot. - # The pending changes on ``dataset`` are also - # flushed, which is what we want — they must be - # persisted by this point anyway. PostgreSQL's - # default ``READ COMMITTED`` rebuilds the read view - # per statement so it isn't affected. - session.commit() - files = list( - session.scalars( - select(UploadFile).where( - UploadFile.tenant_id == dataset.tenant_id, - UploadFile.id.in_(upload_file_list), - ) - ).all() + assert dataset_process_rule + position = DocumentService.get_documents_position(dataset.id, session) + document_ids = [] + duplicate_document_ids = [] + removed_notion_ids: list[str] = [] + if knowledge_config.data_source.info_list.data_source_type == "upload_file": + if not knowledge_config.data_source.info_list.file_info_list: + raise ValueError("File source info is required") + upload_file_list = knowledge_config.data_source.info_list.file_info_list.file_ids + # Issue #41735: under MySQL's default + # ``REPEATABLE READ`` isolation, the SELECT snapshot + # for this transaction is pinned by the first + # read done earlier in ``save_document_with_dataset_id`` + # (e.g. ``check_doc_form``). The caller has just + # committed the upload in a different session + # (``FileService.upload_text`` / + # ``FileService.upload_file``), so the rows exist + # but are invisible to *this* session until its + # snapshot is refreshed. End the current transaction + # so the next SELECT starts with a fresh snapshot. + # The pending changes on ``dataset`` are also + # flushed, which is what we want — they must be + # persisted by this point anyway. PostgreSQL's + # default ``READ COMMITTED`` rebuilds the read view + # per statement so it isn't affected. + session.commit() + files = list( + session.scalars( + select(UploadFile).where( + UploadFile.tenant_id == dataset.tenant_id, + UploadFile.id.in_(upload_file_list), ) - if len(files) != len(set(upload_file_list)): - raise FileNotExistsError("One or more files not found.") + ).all() + ) + if len(files) != len(set(upload_file_list)): + raise FileNotExistsError("One or more files not found.") - file_names = [file.name for file in files] - db_documents = list( - session.scalars( - select(Document).where( - Document.dataset_id == dataset.id, - Document.tenant_id == current_user.current_tenant_id, - Document.data_source_type == DataSourceType.UPLOAD_FILE, - Document.enabled == True, - Document.name.in_(file_names), - ) - ).all() + file_names = [file.name for file in files] + db_documents = list( + session.scalars( + select(Document).where( + Document.dataset_id == dataset.id, + Document.tenant_id == account.current_tenant_id, + Document.data_source_type == DataSourceType.UPLOAD_FILE, + Document.enabled == True, + Document.name.in_(file_names), ) - documents_map = {document.name: document for document in db_documents} - for file in files: - data_source_info: dict[str, object] = { - "upload_file_id": file.id, - } - document = documents_map.get(file.name) - if knowledge_config.duplicate and document: - document.dataset_process_rule_id = dataset_process_rule.id - document.updated_at = naive_utc_now() - document.created_from = created_from - document.doc_form = IndexStructureType(knowledge_config.doc_form) - document.doc_language = knowledge_config.doc_language - document.data_source_info = json.dumps(data_source_info) - document.batch = batch - document.indexing_status = IndexingStatus.WAITING - session.add(document) - documents.append(document) - duplicate_document_ids.append(document.id) - continue - else: - document = DocumentService.build_document( - dataset, - dataset_process_rule.id, - knowledge_config.data_source.info_list.data_source_type, - knowledge_config.doc_form, - knowledge_config.doc_language, - data_source_info, - created_from, - position, - account, - file.name, - batch, - ) - session.add(document) - session.flush() - document_ids.append(document.id) - documents.append(document) - position += 1 - elif knowledge_config.data_source.info_list.data_source_type == "notion_import": - notion_info_list = knowledge_config.data_source.info_list.notion_info_list # type: ignore - if not notion_info_list: - raise ValueError("No notion info list found.") - exist_page_ids = [] - exist_document = {} - documents = list( - session.scalars( - select(Document).where( - Document.dataset_id == dataset.id, - Document.tenant_id == current_user.current_tenant_id, - Document.data_source_type == DataSourceType.NOTION_IMPORT, - Document.enabled == True, - ) - ).all() + ).all() + ) + documents_map = {document.name: document for document in db_documents} + for file in files: + data_source_info: dict[str, object] = { + "upload_file_id": file.id, + } + document = documents_map.get(file.name) + if knowledge_config.duplicate and document: + document.dataset_process_rule_id = dataset_process_rule.id + document.updated_at = naive_utc_now() + document.created_from = created_from + document.doc_form = IndexStructureType(knowledge_config.doc_form) + document.doc_language = knowledge_config.doc_language + document.data_source_info = json.dumps(data_source_info) + document.batch = batch + document.indexing_status = IndexingStatus.WAITING + session.add(document) + documents.append(document) + duplicate_document_ids.append(document.id) + continue + else: + document = DocumentService.build_document( + dataset, + dataset_process_rule.id, + knowledge_config.data_source.info_list.data_source_type, + knowledge_config.doc_form, + knowledge_config.doc_language, + data_source_info, + created_from, + position, + account, + file.name, + batch, ) - if documents: - for document in documents: - data_source_info = json.loads(document.data_source_info) - exist_page_ids.append(data_source_info["notion_page_id"]) - exist_document[data_source_info["notion_page_id"]] = document.id - for notion_info in notion_info_list: - workspace_id = notion_info.workspace_id - for page in notion_info.pages: - if page.page_id not in exist_page_ids: - data_source_info = { - "credential_id": notion_info.credential_id, - "notion_workspace_id": workspace_id, - "notion_page_id": page.page_id, - "notion_page_icon": page.page_icon.model_dump() if page.page_icon else None, # type: ignore - "type": page.type, - } - # Truncate page name to 255 characters to prevent DB field length errors - truncated_page_name = page.page_name[:255] if page.page_name else "nopagename" - document = DocumentService.build_document( - dataset, - dataset_process_rule.id, - knowledge_config.data_source.info_list.data_source_type, - knowledge_config.doc_form, - knowledge_config.doc_language, - data_source_info, - created_from, - position, - account, - truncated_page_name, - batch, - ) - document.id = str(uuid.uuid4()) - session.add(document) - document_ids.append(document.id) - documents.append(document) - position += 1 - else: - exist_document.pop(page.page_id) + session.add(document) session.flush() - # delete not selected documents - if len(exist_document) > 0: - clean_notion_document_task.delay(list(exist_document.values()), dataset.id) - elif knowledge_config.data_source.info_list.data_source_type == "website_crawl": - website_info = knowledge_config.data_source.info_list.website_info_list - if not website_info: - raise ValueError("No website info list found.") - urls = website_info.urls - for url in urls: + document_ids.append(document.id) + documents.append(document) + position += 1 + elif knowledge_config.data_source.info_list.data_source_type == "notion_import": + notion_info_list = knowledge_config.data_source.info_list.notion_info_list # type: ignore + if not notion_info_list: + raise ValueError("No notion info list found.") + exist_page_ids = [] + exist_document = {} + documents = list( + session.scalars( + select(Document).where( + Document.dataset_id == dataset.id, + Document.tenant_id == account.current_tenant_id, + Document.data_source_type == DataSourceType.NOTION_IMPORT, + Document.enabled == True, + ) + ).all() + ) + if documents: + for document in documents: + data_source_info = json.loads(document.data_source_info) + exist_page_ids.append(data_source_info["notion_page_id"]) + exist_document[data_source_info["notion_page_id"]] = document.id + for notion_info in notion_info_list: + workspace_id = notion_info.workspace_id + for page in notion_info.pages: + if page.page_id not in exist_page_ids: data_source_info = { - "url": url, - "provider": website_info.provider, - "job_id": website_info.job_id, - "only_main_content": website_info.only_main_content, - "mode": "crawl", + "credential_id": notion_info.credential_id, + "notion_workspace_id": workspace_id, + "notion_page_id": page.page_id, + "notion_page_icon": page.page_icon.model_dump() if page.page_icon else None, # type: ignore + "type": page.type, } - if len(url) > 255: - document_name = url[:200] + "..." - else: - document_name = url + # Truncate page name to 255 characters to prevent DB field length errors + truncated_page_name = page.page_name[:255] if page.page_name else "nopagename" document = DocumentService.build_document( dataset, dataset_process_rule.id, @@ -2547,29 +2479,64 @@ class DocumentService: created_from, position, account, - document_name, + truncated_page_name, batch, ) + document.id = str(uuid.uuid4()) session.add(document) - session.flush() document_ids.append(document.id) documents.append(document) position += 1 - session.commit() + else: + exist_document.pop(page.page_id) + session.flush() + # delete not selected documents + if len(exist_document) > 0: + removed_notion_ids = list(exist_document.values()) + elif knowledge_config.data_source.info_list.data_source_type == "website_crawl": + website_info = knowledge_config.data_source.info_list.website_info_list + if not website_info: + raise ValueError("No website info list found.") + urls = website_info.urls + for url in urls: + data_source_info = { + "url": url, + "provider": website_info.provider, + "job_id": website_info.job_id, + "only_main_content": website_info.only_main_content, + "mode": "crawl", + } + if len(url) > 255: + document_name = url[:200] + "..." + else: + document_name = url + document = DocumentService.build_document( + dataset, + dataset_process_rule.id, + knowledge_config.data_source.info_list.data_source_type, + knowledge_config.doc_form, + knowledge_config.doc_language, + data_source_info, + created_from, + position, + account, + document_name, + batch, + ) + session.add(document) + session.flush() + document_ids.append(document.id) + documents.append(document) + position += 1 + session.commit() - # trigger async task - if document_ids: - DocumentIndexingTaskProxy(dataset.tenant_id, dataset.id, document_ids).delay() - if duplicate_document_ids: - DuplicateDocumentIndexingTaskProxy( - dataset.tenant_id, dataset.id, duplicate_document_ids - ).delay() - # Note: Summary index generation is triggered in document_indexing_task after indexing completes - # to ensure segments are available. See tasks/document_indexing_task.py - except LockNotOwnedError: - pass - - return documents, batch + jobs = replace( + jobs, + created=tuple(document_ids), + duplicated=tuple(duplicate_document_ids), + removed_notion=tuple(removed_notion_ids), + ) + return documents, batch, jobs # @staticmethod # def save_document_with_dataset_id( @@ -2945,9 +2912,27 @@ class DocumentService: *, session: Session, ): - assert isinstance(current_user, Account) - DatasetService.check_dataset_model_setting(dataset) + document = DocumentService.update_prepared_document( + dataset, document_data, account, dataset_process_rule, created_from, session=session + ) + DocumentService.dispatch_document_indexing( + DocumentIndexingJobs(DatasetRef(dataset.tenant_id, dataset.id), updated=document.id) + ) + return document + + @staticmethod + def update_prepared_document( + dataset: Dataset, + document_data: KnowledgeConfig, + account: Account, + dataset_process_rule: DatasetProcessRule | None = None, + created_from: str = DocumentCreatedFrom.WEB, + *, + session: Session, + ): + assert isinstance(account, Account) + document = DocumentService.get_document(dataset.id, document_data.original_document_id, session=session) if document is None: raise NotFound("Document not found") @@ -3007,7 +2992,7 @@ class DocumentService: select(DataSourceOauthBinding) .where( sa.and_( - DataSourceOauthBinding.tenant_id == current_user.current_tenant_id, + DataSourceOauthBinding.tenant_id == account.current_tenant_id, DataSourceOauthBinding.provider == "notion", DataSourceOauthBinding.disabled == False, DataSourceOauthBinding.source_info["workspace_id"] == f'"{workspace_id}"', @@ -3064,20 +3049,18 @@ class DocumentService: .values(status=SegmentStatus.RE_SEGMENT) ) session.commit() - # trigger async task - document_indexing_update_task.delay(document.dataset_id, document.id) return document @staticmethod def save_document_without_dataset_id( tenant_id: str, knowledge_config: KnowledgeConfig, account: Account, session: Session ): - assert isinstance(current_user, Account) - assert current_user.current_tenant_id is not None + assert isinstance(account, Account) + assert account.current_tenant_id is not None assert knowledge_config.data_source if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD: - features = FeatureService.get_features(current_user.current_tenant_id, exclude_vector_space=True) + features = FeatureService.get_features(account.current_tenant_id, exclude_vector_space=True) count = 0 if knowledge_config.data_source.info_list.data_source_type == "upload_file": upload_file_list = ( @@ -3097,6 +3080,28 @@ class DocumentService: count = len(website_info.urls) DocumentService.check_document_creation_limits(count, features) + dataset = DocumentService.build_dataset_for_documents( + DatasetRef(tenant_id, str(uuid.uuid4())), knowledge_config, account, session + ) + + documents, batch = DocumentService.save_document_with_dataset_id( + dataset, knowledge_config, account, session=session + ) + + cut_length = 18 + cut_name = documents[0].name[:cut_length] + dataset.name = cut_name + "..." + dataset.description = "useful for when you want to answer queries about the " + documents[0].name + session.flush() + + return dataset, documents, batch + + @staticmethod + def build_dataset_for_documents( + ref: DatasetRef, knowledge_config: KnowledgeConfig, account: Account, session: Session + ) -> Dataset: + """Create a dataset inside the caller's SQL-only persistence phase.""" + assert knowledge_config.data_source dataset_collection_binding_id = None retrieval_model = None if knowledge_config.indexing_technique == IndexTechniqueType.HIGH_QUALITY: @@ -3120,7 +3125,8 @@ class DocumentService: ) # save dataset dataset = Dataset( - tenant_id=tenant_id, + id=ref.dataset_id, + tenant_id=ref.tenant_id, name="", data_source_type=knowledge_config.data_source.info_list.data_source_type, indexing_technique=IndexTechniqueType(knowledge_config.indexing_technique), @@ -3137,17 +3143,7 @@ class DocumentService: session.add(dataset) session.flush() - documents, batch = DocumentService.save_document_with_dataset_id( - dataset, knowledge_config, account, session=session - ) - - cut_length = 18 - cut_name = documents[0].name[:cut_length] - dataset.name = cut_name + "..." - dataset.description = "useful for when you want to answer queries about the " + documents[0].name - session.flush() - - return dataset, documents, batch + return dataset @classmethod def document_create_args_validate(cls, knowledge_config: KnowledgeConfig): @@ -3235,23 +3231,6 @@ class DocumentService: if not isinstance(knowledge_config.process_rule.rules.segmentation.max_tokens, int): raise ValueError("Process rule segmentation max_tokens is invalid") - @classmethod - def estimate_args_validate(cls, args: dict[str, Any]): - try: - validated = _EstimateArgs.model_validate(args) - except ValidationError as e: - first = e.errors()[0] - original = first.get("ctx", {}).get("error") - raise ValueError(str(original) if isinstance(original, ValueError) else first["msg"]) from e - process_rule_dict = validated.process_rule.model_dump(exclude_none=True) - if validated.process_rule.mode == ProcessRuleMode.AUTOMATIC: - process_rule_dict["rules"] = {} - elif validated.process_rule.mode == ProcessRuleMode.HIERARCHICAL: - rules = process_rule_dict.get("rules") - if isinstance(rules, dict) and not rules.get("parent_mode"): - rules["parent_mode"] = "paragraph" - args["process_rule"] = process_rule_dict - @staticmethod def batch_update_document_status( dataset: Dataset, @@ -3448,147 +3427,125 @@ class DocumentService: return update_info +def _segment_mutation_scope(dataset: Dataset, document: Document) -> SegmentMutationScope: + """Materialize already-authorized rows before ending their read transaction.""" + if document.tenant_id != dataset.tenant_id or document.dataset_id != dataset.id: + raise ValueError("Document does not belong to the dataset") + return SegmentMutationScope( + dataset=SegmentDatasetRecord( + id=dataset.id, + workspace_id=dataset.tenant_id, + indexing_technique=dataset.indexing_technique or "economy", + embedding_model_provider=dataset.embedding_model_provider, + embedding_model=dataset.embedding_model, + ), + document=SegmentDocumentRecord( + id=document.id, + dataset_id=document.dataset_id, + workspace_id=document.tenant_id, + doc_form=document.doc_form, + ), + ) + + +def _segment_ref_for_mutation(scope: SegmentMutationScope, segment: DocumentSegment) -> SegmentRef: + document_ref = scope.document.ref + if ( + segment.tenant_id != document_ref.dataset.tenant_id + or segment.dataset_id != document_ref.dataset.dataset_id + or segment.document_id != document_ref.document_id + ): + raise ValueError("Segment does not belong to the document") + return document_ref.segment(segment.id) + + +def _check_child_chunk_scope(child: ChildChunk, segment_ref: SegmentRef) -> None: + if ( + child.tenant_id != segment_ref.document.dataset.tenant_id + or child.dataset_id != segment_ref.document.dataset.dataset_id + or child.document_id != segment_ref.document.document_id + or child.segment_id != segment_ref.segment_id + ): + raise ValueError("Child chunk does not belong to the segment") + + class SegmentService: @classmethod def segment_create_args_validate(cls, args: dict[str, Any], document: Document): - if document.doc_form == IndexStructureType.QA_INDEX: - if "answer" not in args or not args["answer"]: - raise ValueError("Answer is required") - if not args["answer"].strip(): - raise ValueError("Answer is empty") - if "content" not in args or not args["content"] or not args["content"].strip(): - raise ValueError("Content is empty") - - if args.get("attachment_ids"): - if not isinstance(args["attachment_ids"], list): - raise ValueError("Attachment IDs is invalid") - single_chunk_attachment_limit = dify_config.SINGLE_CHUNK_ATTACHMENT_LIMIT - if len(args["attachment_ids"]) > single_chunk_attachment_limit: - raise ValueError(f"Exceeded maximum attachment limit of {single_chunk_attachment_limit}") + validate_segment_values( + args, doc_form=document.doc_form, attachment_limit=dify_config.SINGLE_CHUNK_ATTACHMENT_LIMIT + ) @classmethod - def create_segment(cls, args: dict[str, Any], document: Document, dataset: Dataset, session: Session): + def multi_create_segment( + cls, + segments: list, + document: Document, + dataset: Dataset, + session: Session, + *, + mutations: SegmentMutationService, + ): + """Commit a Service API batch before dispatching its scoped index update.""" assert isinstance(current_user, Account) assert current_user.current_tenant_id is not None - - content = args["content"] - doc_id = str(uuid.uuid4()) - segment_hash = helper.generate_text_hash(content) - tokens = 0 - if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY: - model_manager = ModelManager.for_tenant(tenant_id=current_user.current_tenant_id) - embedding_model = model_manager.get_model_instance( - tenant_id=current_user.current_tenant_id, - provider=dataset.embedding_model_provider, - model_type=ModelType.TEXT_EMBEDDING, - model=dataset.embedding_model, - ) - # calc embedding use tokens - tokens = embedding_model.get_text_embedding_num_tokens(texts=[content])[0] - lock_name = f"add_segment_lock_document_id_{document.id}" - try: - with redis_client.lock(lock_name, timeout=600): - max_position = session.scalar( - select(func.max(DocumentSegment.position)).where(DocumentSegment.document_id == document.id) - ) - segment_document = DocumentSegment( - tenant_id=current_user.current_tenant_id, - dataset_id=document.dataset_id, - document_id=document.id, - index_node_id=doc_id, - index_node_hash=segment_hash, - position=max_position + 1 if max_position else 1, - content=content, - word_count=len(content), - tokens=tokens, - status=SegmentStatus.COMPLETED, - indexing_at=naive_utc_now(), - completed_at=naive_utc_now(), - created_by=current_user.id, - ) - if document.doc_form == IndexStructureType.QA_INDEX: - segment_document.word_count += len(args["answer"]) - segment_document.answer = args["answer"] - - session.add(segment_document) - # update document word count - assert document.word_count is not None - document.word_count += segment_document.word_count - session.add(document) - session.commit() - - if args["attachment_ids"]: - for attachment_id in args["attachment_ids"]: - binding = SegmentAttachmentBinding( - tenant_id=current_user.current_tenant_id, - dataset_id=document.dataset_id, - document_id=document.id, - segment_id=segment_document.id, - attachment_id=attachment_id, - ) - session.add(binding) - session.commit() - - # save vector index - try: - keywords = args.get("keywords") - keywords_list = [keywords] if keywords is not None else None - VectorService.create_segments_vector( - keywords_list, [segment_document], dataset, document.doc_form, session=session - ) - except Exception as e: - logger.exception("create segment index failed") - segment_document.enabled = False - segment_document.disabled_at = naive_utc_now() - segment_document.status = SegmentStatus.ERROR - segment_document.error = str(e) - session.commit() - segment = session.get(DocumentSegment, segment_document.id) - return segment - except LockNotOwnedError: - pass - - @classmethod - def multi_create_segment(cls, segments: list, document: Document, dataset: Dataset, session: Session): - assert isinstance(current_user, Account) - assert current_user.current_tenant_id is not None - - lock_name = f"multi_add_segment_lock_document_id_{document.id}" + scope = _segment_mutation_scope(dataset, document) + if current_user.current_tenant_id != scope.dataset.workspace_id: + raise ValueError("Dataset does not belong to the current workspace") + actor_id = current_user.id + document_ref = scope.document.ref + session.commit() + lock_name = f"multi_add_segment_lock_document_id_{document_ref.document_id}" increment_word_count = 0 try: with redis_client.lock(lock_name, timeout=600): embedding_model = None - if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY: - model_manager = ModelManager.for_tenant(tenant_id=current_user.current_tenant_id) + if scope.dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY: + model_manager = ModelManager.for_tenant(tenant_id=scope.dataset.workspace_id) embedding_model = model_manager.get_model_instance( - tenant_id=current_user.current_tenant_id, - provider=dataset.embedding_model_provider, + tenant_id=scope.dataset.workspace_id, + provider=scope.dataset.embedding_model_provider or "", model_type=ModelType.TEXT_EMBEDDING, - model=dataset.embedding_model, + model=scope.dataset.embedding_model or "", + ) + token_counts = [ + embedding_model.get_text_embedding_num_tokens( + texts=[ + item["content"] + item["answer"] + if scope.document.doc_form == IndexStructureType.QA_INDEX + else item["content"] + ] + )[0] + if embedding_model + else 0 + for item in segments + ] + persisted_document = session.scalar( + select(Document).where( + Document.id == document_ref.document_id, + Document.dataset_id == document_ref.dataset.dataset_id, + Document.tenant_id == document_ref.dataset.tenant_id, + ) + ) + if persisted_document is None: + raise ValueError("Document no longer exists") + document = persisted_document + max_position = session.scalar( + select(func.max(DocumentSegment.position)).where( + DocumentSegment.tenant_id == document_ref.dataset.tenant_id, + DocumentSegment.dataset_id == document_ref.dataset.dataset_id, + DocumentSegment.document_id == document_ref.document_id, ) - max_position = session.scalar( - select(func.max(DocumentSegment.position)).where(DocumentSegment.document_id == document.id) ) - pre_segment_data_list = [] segment_data_list = [] keywords_list = [] position = max_position + 1 if max_position else 1 - for segment_item in segments: + for segment_item, tokens in zip(segments, token_counts, strict=True): content = segment_item["content"] doc_id = str(uuid.uuid4()) segment_hash = helper.generate_text_hash(content) - tokens = 0 - if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY and embedding_model: - # calc embedding use tokens - if document.doc_form == IndexStructureType.QA_INDEX: - tokens = embedding_model.get_text_embedding_num_tokens( - texts=[content + segment_item["answer"]] - )[0] - else: - tokens = embedding_model.get_text_embedding_num_tokens(texts=[content])[0] - segment_document = DocumentSegment( - tenant_id=current_user.current_tenant_id, + tenant_id=document_ref.dataset.tenant_id, dataset_id=document.dataset_id, document_id=document.id, index_node_id=doc_id, @@ -3601,7 +3558,7 @@ class SegmentService: status=SegmentStatus.COMPLETED, indexing_at=naive_utc_now(), completed_at=naive_utc_now(), - created_by=current_user.id, + created_by=actor_id, ) if document.doc_form == IndexStructureType.QA_INDEX: segment_document.answer = segment_item["answer"] @@ -3611,7 +3568,6 @@ class SegmentService: segment_data_list.append(segment_document) position += 1 - pre_segment_data_list.append(segment_document) if "keywords" in segment_item: keywords_list.append(segment_item["keywords"]) else: @@ -3620,19 +3576,16 @@ class SegmentService: assert document.word_count is not None document.word_count += increment_word_count session.add(document) - try: - # save vector index - VectorService.create_segments_vector( - keywords_list, pre_segment_data_list, dataset, document.doc_form, session=session - ) - except Exception as e: - logger.exception("create segment index failed") - for segment_document in segment_data_list: - segment_document.enabled = False - segment_document.disabled_at = naive_utc_now() - segment_document.status = SegmentStatus.ERROR - segment_document.error = str(e) + session.flush() + segment_ids = [segment.id for segment in segment_data_list] session.commit() + try: + mutations.index_segments(document_ref, segment_ids=segment_ids, keywords_list=keywords_list) + except Exception: + logger.exception("create segment index failed") + session.expire_all() + for segment in segment_data_list: + session.refresh(segment) return segment_data_list except LockNotOwnedError: pass @@ -3645,446 +3598,37 @@ class SegmentService: document: Document, dataset: Dataset, session: Session, - ): - assert isinstance(current_user, Account) - assert current_user.current_tenant_id is not None - - indexing_cache_key = f"segment_{segment.id}_indexing" - cache_result = redis_client.get(indexing_cache_key) - if cache_result is not None: - raise ValueError("Segment is indexing, please try again later") - if args.enabled is not None: - action = args.enabled - if segment.enabled != action: - if not action: - segment.enabled = action - segment.disabled_at = naive_utc_now() - segment.disabled_by = current_user.id - session.add(segment) - session.commit() - # Set cache to prevent indexing the same segment multiple times - redis_client.setex(indexing_cache_key, 600, 1) - disable_segment_from_index_task.delay(segment.id) - return segment - if not segment.enabled: - if args.enabled is not None: - if not args.enabled: - raise ValueError("Can't update disabled segment") - else: - raise ValueError("Can't update disabled segment") - try: - word_count_change = segment.word_count - content = args.content or segment.content - if segment.content == content: - segment.word_count = len(content) - if document.doc_form == IndexStructureType.QA_INDEX: - segment.answer = args.answer - segment.word_count += len(args.answer) if args.answer else 0 - word_count_change = segment.word_count - word_count_change - keyword_changed = False - if args.keywords: - if Counter(segment.keywords) != Counter(args.keywords): - segment.keywords = args.keywords - keyword_changed = True - segment.enabled = True - segment.disabled_at = None - segment.disabled_by = None - session.add(segment) - session.commit() - # update document word count - if word_count_change != 0: - assert document.word_count is not None - document.word_count = max(0, document.word_count + word_count_change) - session.add(document) - # update segment index task - if document.doc_form == IndexStructureType.PARENT_CHILD_INDEX and args.regenerate_child_chunks: - # regenerate child chunks - # get embedding model instance - if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY: - # check embedding model setting - model_manager = ModelManager.for_tenant(tenant_id=dataset.tenant_id) - - if dataset.embedding_model_provider: - embedding_model_instance = model_manager.get_model_instance( - tenant_id=dataset.tenant_id, - provider=dataset.embedding_model_provider, - model_type=ModelType.TEXT_EMBEDDING, - model=dataset.embedding_model, - ) - else: - embedding_model_instance = model_manager.get_default_model_instance( - tenant_id=dataset.tenant_id, - model_type=ModelType.TEXT_EMBEDDING, - ) - else: - raise ValueError("The knowledge base index technique is not high quality!") - # get the process rule - processing_rule = session.get(DatasetProcessRule, document.dataset_process_rule_id) - if processing_rule: - VectorService.generate_child_chunks( - segment, document, dataset, embedding_model_instance, processing_rule, True, session=session - ) - elif document.doc_form in (IndexStructureType.PARAGRAPH_INDEX, IndexStructureType.QA_INDEX): - if args.enabled or keyword_changed: - # update segment vector index - VectorService.update_segment_vector(args.keywords, segment, dataset, session=session) - # update summary index if summary is provided and has changed - if args.summary is not None: - # When user manually provides summary, allow saving even if summary_index_setting doesn't exist - # summary_index_setting is only needed for LLM generation, not for manual summary vectorization - # Vectorization uses dataset.embedding_model, which doesn't require summary_index_setting - if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY: - # Query existing summary from database - from models.dataset import DocumentSegmentSummary - - existing_summary = session.scalar( - select(DocumentSegmentSummary) - .where( - DocumentSegmentSummary.chunk_id == segment.id, - DocumentSegmentSummary.dataset_id == dataset.id, - ) - .limit(1) - ) - - # Check if summary has changed - existing_summary_content = existing_summary.summary_content if existing_summary else None - if existing_summary_content != args.summary: - # Summary has changed, update it - from services.summary_index_service import SummaryIndexService - - try: - SummaryIndexService.update_summary_for_segment( - segment, - dataset, - args.summary, - session=session, - ) - except Exception: - logger.exception("Failed to update summary for segment %s", segment.id) - # Don't fail the entire update if summary update fails - else: - segment_hash = helper.generate_text_hash(content) - tokens = 0 - if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY: - model_manager = ModelManager.for_tenant(tenant_id=current_user.current_tenant_id) - embedding_model = model_manager.get_model_instance( - tenant_id=current_user.current_tenant_id, - provider=dataset.embedding_model_provider, - model_type=ModelType.TEXT_EMBEDDING, - model=dataset.embedding_model, - ) - - # calc embedding use tokens - if document.doc_form == IndexStructureType.QA_INDEX: - segment.answer = args.answer - tokens = embedding_model.get_text_embedding_num_tokens(texts=[content + segment.answer])[0] # type: ignore - else: - tokens = embedding_model.get_text_embedding_num_tokens(texts=[content])[0] - segment.content = content - segment.index_node_hash = segment_hash - segment.word_count = len(content) - segment.tokens = tokens - segment.status = SegmentStatus.COMPLETED - segment.indexing_at = naive_utc_now() - segment.completed_at = naive_utc_now() - segment.updated_by = current_user.id - segment.updated_at = naive_utc_now() - segment.enabled = True - segment.disabled_at = None - segment.disabled_by = None - if document.doc_form == IndexStructureType.QA_INDEX: - segment.answer = args.answer - segment.word_count += len(args.answer) if args.answer else 0 - word_count_change = segment.word_count - word_count_change - # update document word count - if word_count_change != 0: - assert document.word_count is not None - document.word_count = max(0, document.word_count + word_count_change) - session.add(document) - session.add(segment) - session.commit() - if document.doc_form == IndexStructureType.PARENT_CHILD_INDEX and args.regenerate_child_chunks: - # get embedding model instance - if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY: - # check embedding model setting - model_manager = ModelManager.for_tenant(tenant_id=dataset.tenant_id) - - if dataset.embedding_model_provider: - embedding_model_instance = model_manager.get_model_instance( - tenant_id=dataset.tenant_id, - provider=dataset.embedding_model_provider, - model_type=ModelType.TEXT_EMBEDDING, - model=dataset.embedding_model, - ) - else: - embedding_model_instance = model_manager.get_default_model_instance( - tenant_id=dataset.tenant_id, - model_type=ModelType.TEXT_EMBEDDING, - ) - else: - raise ValueError("The knowledge base index technique is not high quality!") - # get the process rule - processing_rule = session.get(DatasetProcessRule, document.dataset_process_rule_id) - if processing_rule: - VectorService.generate_child_chunks( - segment, document, dataset, embedding_model_instance, processing_rule, True, session=session - ) - elif document.doc_form in (IndexStructureType.PARAGRAPH_INDEX, IndexStructureType.QA_INDEX): - # update segment vector index - VectorService.update_segment_vector(args.keywords, segment, dataset, session=session) - # Handle summary index when content changed - if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY: - from models.dataset import DocumentSegmentSummary - - existing_summary = session.scalar( - select(DocumentSegmentSummary) - .where( - DocumentSegmentSummary.chunk_id == segment.id, - DocumentSegmentSummary.dataset_id == dataset.id, - ) - .limit(1) - ) - - if args.summary is None: - # User didn't provide summary, auto-regenerate if segment previously had summary - # Auto-regeneration only happens if summary_index_setting exists and enable is True - if ( - existing_summary - and dataset.summary_index_setting - and dataset.summary_index_setting.get("enable") is True - ): - # Segment previously had summary, regenerate it with new content - from services.summary_index_service import SummaryIndexService - - try: - SummaryIndexService.generate_and_vectorize_summary( - segment, - dataset, - dataset.summary_index_setting, - session=session, - ) - logger.info("Auto-regenerated summary for segment %s after content change", segment.id) - except Exception: - logger.exception("Failed to auto-regenerate summary for segment %s", segment.id) - # Don't fail the entire update if summary regeneration fails - else: - # User provided summary, check if it has changed - # Manual summary updates are allowed even if summary_index_setting doesn't exist - existing_summary_content = existing_summary.summary_content if existing_summary else None - if existing_summary_content != args.summary: - # Summary has changed, use user-provided summary - from services.summary_index_service import SummaryIndexService - - try: - SummaryIndexService.update_summary_for_segment( - segment, - dataset, - args.summary, - session=session, - ) - logger.info("Updated summary for segment %s with user-provided content", segment.id) - except Exception: - logger.exception("Failed to update summary for segment %s", segment.id) - # Don't fail the entire update if summary update fails - else: - # Summary hasn't changed, regenerate based on new content - # Auto-regeneration only happens if summary_index_setting exists and enable is True - if ( - existing_summary - and dataset.summary_index_setting - and dataset.summary_index_setting.get("enable") is True - ): - from services.summary_index_service import SummaryIndexService - - try: - SummaryIndexService.generate_and_vectorize_summary( - segment, - dataset, - dataset.summary_index_setting, - session=session, - ) - logger.info( - "Regenerated summary for segment %s after content change (summary unchanged)", - segment.id, - ) - except Exception: - logger.exception("Failed to regenerate summary for segment %s", segment.id) - # Don't fail the entire update if summary regeneration fails - # update multimodel vector index - if args.attachment_ids is not None: - VectorService.update_multimodel_vector(segment, args.attachment_ids, dataset, session=session) - except Exception as e: - logger.exception("update segment index failed") - segment.enabled = False - segment.disabled_at = naive_utc_now() - segment.status = SegmentStatus.ERROR - segment.error = str(e) - session.commit() - new_segment = session.get(DocumentSegment, segment.id) - if not new_segment: + *, + actor_id: str, + mutations: SegmentMutationService, + ) -> DocumentSegment: + """Delegate an authorized Service API patch to the shared mutation owner.""" + scope = _segment_mutation_scope(dataset, document) + segment_ref = _segment_ref_for_mutation(scope, segment) + # Finish the caller's read transaction before the mutation owner performs + # external indexing I/O in between its own bounded write transactions. + session.commit() + mutations.update_segment(scope, actor_id=actor_id, segment_id=segment_ref.segment_id, args=args) + session.expire_all() + updated = cls.get_segment_by_ref(segment_ref, session=session) + if updated is None: raise ValueError("new_segment is not found") - return new_segment + return updated @classmethod - def delete_segment(cls, segment: DocumentSegment, document: Document, dataset: Dataset, session: Session): - indexing_cache_key = f"segment_{segment.id}_delete_indexing" - cache_result = redis_client.get(indexing_cache_key) - if cache_result is not None: - raise ValueError("Segment is deleting.") - - # enabled segment need to delete index - if segment.enabled: - # send delete segment index task - redis_client.setex(indexing_cache_key, 600, 1) - - # Get child chunk IDs before parent segment is deleted - child_node_ids = [] - if segment.index_node_id: - child_node_ids = list( - session.scalars( - select(ChildChunk.index_node_id).where( - ChildChunk.segment_id == segment.id, - ChildChunk.dataset_id == dataset.id, - ) - ).all() - ) - - delete_segment_from_index_task.delay( - [segment.index_node_id], dataset.id, document.id, [segment.id], child_node_ids - ) - - session.delete(segment) - # update document word count - assert document.word_count is not None - document.word_count -= segment.word_count - session.add(document) - session.commit() - - @classmethod - def delete_segments(cls, segment_ids: list, document: Document, dataset: Dataset, session: Session): - assert current_user is not None - # Check if segment_ids is not empty to avoid WHERE false condition - if not segment_ids or len(segment_ids) == 0: - return - segments_info = session.execute( - select(DocumentSegment.index_node_id, DocumentSegment.id, DocumentSegment.word_count).where( - DocumentSegment.id.in_(segment_ids), - DocumentSegment.dataset_id == dataset.id, - DocumentSegment.document_id == document.id, - DocumentSegment.tenant_id == current_user.current_tenant_id, - ) - ).all() - - if not segments_info: - return - - index_node_ids = [info[0] for info in segments_info] - segment_db_ids = [info[1] for info in segments_info] - total_words = sum(info[2] for info in segments_info if info[2] is not None) - - # Get child chunk IDs before parent segments are deleted - child_node_ids = [] - if index_node_ids: - child_node_ids = [ - nid - for nid in session.scalars( - select(ChildChunk.index_node_id).where( - ChildChunk.segment_id.in_(segment_db_ids), - ChildChunk.dataset_id == dataset.id, - ) - ).all() - if nid - ] - - # Start async cleanup with both parent and child node IDs - if index_node_ids or child_node_ids: - delete_segment_from_index_task.delay( - index_node_ids, dataset.id, document.id, segment_db_ids, child_node_ids - ) - - if document.word_count is None: - document.word_count = 0 - else: - document.word_count = max(0, document.word_count - total_words) - - session.add(document) - - # Delete database records - session.execute( - delete(DocumentSegment).where( - DocumentSegment.id.in_(segment_db_ids), - DocumentSegment.dataset_id == dataset.id, - DocumentSegment.document_id == document.id, - DocumentSegment.tenant_id == current_user.current_tenant_id, - ) - ) - session.commit() - - @classmethod - def update_segments_status( + def delete_segment( cls, - segment_ids: list, - action: Literal["enable", "disable"], - dataset: Dataset, + segment: DocumentSegment, document: Document, + dataset: Dataset, session: Session, - ): - assert current_user is not None - - # Check if segment_ids is not empty to avoid WHERE false condition - if not segment_ids or len(segment_ids) == 0: - return - match action: - case "enable": - segments = session.scalars( - select(DocumentSegment).where( - DocumentSegment.id.in_(segment_ids), - DocumentSegment.dataset_id == dataset.id, - DocumentSegment.document_id == document.id, - DocumentSegment.enabled == False, - ) - ).all() - if not segments: - return - real_deal_segment_ids = [] - for segment in segments: - indexing_cache_key = f"segment_{segment.id}_indexing" - cache_result = redis_client.get(indexing_cache_key) - if cache_result is not None: - continue - segment.enabled = True - segment.disabled_at = None - segment.disabled_by = None - session.add(segment) - real_deal_segment_ids.append(segment.id) - session.commit() - - enable_segments_to_index_task.delay(real_deal_segment_ids, dataset.id, document.id) - case "disable": - segments = session.scalars( - select(DocumentSegment).where( - DocumentSegment.id.in_(segment_ids), - DocumentSegment.dataset_id == dataset.id, - DocumentSegment.document_id == document.id, - DocumentSegment.enabled == True, - ) - ).all() - if not segments: - return - real_deal_segment_ids = [] - for segment in segments: - indexing_cache_key = f"segment_{segment.id}_indexing" - cache_result = redis_client.get(indexing_cache_key) - if cache_result is not None: - continue - segment.enabled = False - segment.disabled_at = naive_utc_now() - segment.disabled_by = current_user.id - session.add(segment) - real_deal_segment_ids.append(segment.id) - session.commit() - - disable_segments_from_index_task.delay(real_deal_segment_ids, dataset.id, document.id) + *, + mutations: SegmentMutationService, + ) -> None: + segment_ref = _segment_ref_for_mutation(_segment_mutation_scope(dataset, document), segment) + session.commit() + mutations.delete_segment(segment_ref) + session.expire_all() @classmethod def create_child_chunk( @@ -4094,126 +3638,21 @@ class SegmentService: document: Document, dataset: Dataset, session: Session, + *, + actor_id: str, + mutations: SegmentMutationService, ) -> ChildChunk: - assert isinstance(current_user, Account) - - lock_name = f"add_child_lock_{segment.id}" - with redis_client.lock(lock_name, timeout=20): - index_node_id = str(uuid.uuid4()) - index_node_hash = helper.generate_text_hash(content) - max_position = session.scalar( - select(func.max(ChildChunk.position)).where( - ChildChunk.tenant_id == current_user.current_tenant_id, - ChildChunk.dataset_id == dataset.id, - ChildChunk.document_id == document.id, - ChildChunk.segment_id == segment.id, - ) - ) - assert current_user.current_tenant_id - child_chunk = ChildChunk( - tenant_id=current_user.current_tenant_id, - dataset_id=dataset.id, - document_id=document.id, - segment_id=segment.id, - position=max_position + 1 if max_position else 1, - index_node_id=index_node_id, - index_node_hash=index_node_hash, - content=content, - word_count=len(content), - type=SegmentType.CUSTOMIZED, - created_by=current_user.id, - ) - session.add(child_chunk) - # save vector index - try: - VectorService.create_child_chunk_vector(child_chunk, dataset, session=session) - except Exception as e: - logger.exception("create child chunk index failed") - session.rollback() - raise ChildChunkIndexingError(str(e)) - session.commit() - - return child_chunk - - @classmethod - def update_child_chunks( - cls, - child_chunks_update_args: list[ChildChunkUpdateArgs], - segment: DocumentSegment, - document: Document, - dataset: Dataset, - session: Session, - ) -> list[ChildChunk]: - assert isinstance(current_user, Account) - child_chunks = session.scalars( - select(ChildChunk).where( - ChildChunk.dataset_id == dataset.id, - ChildChunk.document_id == document.id, - ChildChunk.segment_id == segment.id, - ) - ).all() - child_chunks_map = {chunk.id: chunk for chunk in child_chunks} - - new_child_chunks, update_child_chunks, delete_child_chunks, new_child_chunks_args = [], [], [], [] - - for child_chunk_update_args in child_chunks_update_args: - if child_chunk_update_args.id: - child_chunk = child_chunks_map.pop(child_chunk_update_args.id, None) - if child_chunk: - if child_chunk.content != child_chunk_update_args.content: - child_chunk.content = child_chunk_update_args.content - child_chunk.word_count = len(child_chunk.content) - child_chunk.updated_by = current_user.id - child_chunk.updated_at = naive_utc_now() - child_chunk.type = SegmentType.CUSTOMIZED - update_child_chunks.append(child_chunk) - else: - new_child_chunks_args.append(child_chunk_update_args) - if child_chunks_map: - delete_child_chunks = list(child_chunks_map.values()) + segment_ref = _segment_ref_for_mutation(_segment_mutation_scope(dataset, document), segment) + session.commit() try: - if update_child_chunks: - session.bulk_save_objects(update_child_chunks) - - if delete_child_chunks: - for child_chunk in delete_child_chunks: - session.delete(child_chunk) - if new_child_chunks_args: - child_chunk_count = len(child_chunks) - for position, args in enumerate(new_child_chunks_args, start=child_chunk_count + 1): - assert current_user.current_tenant_id - index_node_id = str(uuid.uuid4()) - index_node_hash = helper.generate_text_hash(args.content) - child_chunk = ChildChunk( - tenant_id=current_user.current_tenant_id, - dataset_id=dataset.id, - document_id=document.id, - segment_id=segment.id, - position=position, - index_node_id=index_node_id, - index_node_hash=index_node_hash, - content=args.content, - word_count=len(args.content), - type=SegmentType.CUSTOMIZED, - created_by=current_user.id, - ) - - session.add(child_chunk) - session.flush() - new_child_chunks.append(child_chunk) - VectorService.update_child_chunk_vector( - new_child_chunks, - update_child_chunks, - delete_child_chunks, - dataset, - session=session, - ) - session.commit() - except Exception as e: - logger.exception("update child chunk index failed") - session.rollback() - raise ChildChunkIndexingError(str(e)) - return sorted(new_child_chunks + update_child_chunks, key=lambda x: x.position) + created = mutations.create_child_chunk(segment_ref, content=content, actor_id=actor_id) + except ChildChunkIndexingApplicationError as error: + raise ChildChunkIndexingError(str(error)) from error + session.expire_all() + child = cls.get_child_chunk_by_segment_ref(created.id, segment_ref, session=session) + if child is None: + raise ValueError("Child chunk no longer exists") + return child @classmethod def update_child_chunk( @@ -4224,34 +3663,46 @@ class SegmentService: document: Document, dataset: Dataset, session: Session, + *, + actor_id: str, + mutations: SegmentMutationService, ) -> ChildChunk: - assert current_user is not None - + segment_ref = _segment_ref_for_mutation(_segment_mutation_scope(dataset, document), segment) + _check_child_chunk_scope(child_chunk, segment_ref) + child_id = child_chunk.id + session.commit() try: - child_chunk.content = content - child_chunk.word_count = len(content) - child_chunk.updated_by = current_user.id - child_chunk.updated_at = naive_utc_now() - child_chunk.type = SegmentType.CUSTOMIZED - session.add(child_chunk) - VectorService.update_child_chunk_vector([], [child_chunk], [], dataset, session=session) - session.commit() - except Exception as e: - logger.exception("update child chunk index failed") - session.rollback() - raise ChildChunkIndexingError(str(e)) - return child_chunk + mutations.update_child_chunk(segment_ref, child_chunk_id=child_id, content=content, actor_id=actor_id) + except ChildChunkIndexingApplicationError as error: + raise ChildChunkIndexingError(str(error)) from error + session.expire_all() + updated = cls.get_child_chunk_by_segment_ref(child_id, segment_ref, session=session) + if updated is None: + raise ValueError("Child chunk no longer exists") + return updated @classmethod - def delete_child_chunk(cls, child_chunk: ChildChunk, dataset: Dataset, session: Session): - session.delete(child_chunk) - try: - VectorService.delete_child_chunk_vector(child_chunk, dataset, session=session) - except Exception as e: - logger.exception("delete child chunk index failed") - session.rollback() - raise ChildChunkDeleteIndexError(str(e)) + def delete_child_chunk( + cls, + child_chunk: ChildChunk, + dataset: Dataset, + session: Session, + *, + mutations: SegmentMutationService, + ) -> None: + segment_ref = ( + DatasetRef(tenant_id=dataset.tenant_id, dataset_id=dataset.id) + .document(child_chunk.document_id) + .segment(child_chunk.segment_id) + ) + _check_child_chunk_scope(child_chunk, segment_ref) + child_id = child_chunk.id session.commit() + try: + mutations.delete_child_chunk(segment_ref, child_chunk_id=child_id) + except ChildChunkDeleteIndexApplicationError as error: + raise ChildChunkDeleteIndexError(str(error)) from error + session.expire_all() @classmethod def get_child_chunks( @@ -4267,20 +3718,14 @@ class SegmentService: ): assert isinstance(current_user, Account) - query = ( - select(ChildChunk) - .filter_by( - tenant_id=current_user.current_tenant_id, - dataset_id=dataset_id, - document_id=document_id, - segment_id=segment_id, - ) - .order_by(ChildChunk.position.asc()) + assert current_user.current_tenant_id is not None + return query_child_chunks( + session, + DatasetRef(current_user.current_tenant_id, dataset_id).document(document_id).segment(segment_id), + page=page, + limit=limit, + keyword=keyword, ) - if keyword: - escaped_keyword = helper.escape_like_pattern(keyword) - query = query.where(ChildChunk.content.ilike(f"%{escaped_keyword}%", escape="\\")) - return paginate_query(query, session=session, page=page, per_page=limit, max_per_page=100) @classmethod def get_child_chunk_by_id(cls, child_chunk_id: str, tenant_id: str, session: Session) -> ChildChunk | None: @@ -4371,13 +3816,16 @@ class SegmentService: session: Session, status: str | None = None, enabled: bool | None = None, + *, + tenant_id: str, ) -> Sequence[DocumentSegment]: """ - Get segments for a document in a dataset with optional filtering. + Get segments through their tenant, dataset, and document owner chain. Args: document_id: Document ID dataset_id: Dataset ID + tenant_id: Owning tenant ID status: Optional status filter (e.g., "completed") enabled: Optional enabled filter (True/False) @@ -4385,6 +3833,7 @@ class SegmentService: Sequence of DocumentSegment instances """ query = select(DocumentSegment).where( + DocumentSegment.tenant_id == tenant_id, DocumentSegment.document_id == document_id, DocumentSegment.dataset_id == dataset_id, ) diff --git a/api/services/knowledge/datasets/__init__.py b/api/services/knowledge/datasets/__init__.py new file mode 100644 index 00000000000..697c750e8a8 --- /dev/null +++ b/api/services/knowledge/datasets/__init__.py @@ -0,0 +1 @@ +"""Dataset administration use cases and infrastructure adapters.""" diff --git a/api/services/knowledge/datasets/adapters.py b/api/services/knowledge/datasets/adapters.py new file mode 100644 index 00000000000..592b43ed1c4 --- /dev/null +++ b/api/services/knowledge/datasets/adapters.py @@ -0,0 +1,328 @@ +"""Session-owning dataset administration adapters.""" + +from collections.abc import Generator, Mapping +from contextlib import contextmanager +from typing import Any + +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker + +from configs import dify_config +from core.plugin.impl.model_runtime_factory import create_plugin_provider_manager +from graphon.model_runtime.entities.model_entities import ModelType +from libs.pagination import clamp_pagination +from machinery.context import RequestContext +from models import Account, Dataset, Document +from models.dataset import DatasetPermission, DatasetPermissionEnum +from models.provider_ids import ModelProviderID +from repositories.knowledge.dataset_repository import _get_dataset +from services.app.query_service import AppQueryService +from services.enterprise import rbac_service +from services.errors.account import NoPermissionError +from services.knowledge.dataset_access import DatasetAccessDeniedError, DatasetNotFoundError +from services.knowledge.dataset_read_service import get_dataset_queries, load_dataset_detail, load_dataset_details +from services.knowledge.dataset_service import DatasetPermissionService, DatasetService, DocumentService +from services.knowledge.datasets.application import ( + DatasetListFilter, + DatasetVisibility, +) +from services.knowledge.entities.datasets import DatasetDetailRecord, DatasetPage +from services.knowledge.resource_scope import DatasetRef +from services.tag_application_service import TagTargetQuery + + +@contextmanager +def _translate_permissions() -> Generator[None, None, None]: + try: + yield + except NoPermissionError as error: + raise DatasetAccessDeniedError() from error + + +def require_dataset(session: Session, ref: DatasetRef) -> Dataset: + dataset = _get_dataset(session, ref) + if dataset is None: + raise DatasetNotFoundError() + return dataset + + +def load_actor(session: Session, context: RequestContext) -> Account: + account = session.get(Account, context.account_id) + if account is None: + raise DatasetAccessDeniedError() + account.set_tenant_id_with_session(context.active_workspace_id, session=session) + return account + + +def _normalize_provider(item: DatasetDetailRecord) -> DatasetDetailRecord: + if item["indexing_technique"] == "high_quality" and item["embedding_model_provider"]: + item["embedding_model_provider"] = str(ModelProviderID(item["embedding_model_provider"])) + return item + + +def _status(document: Document, counts: tuple[int, int] | None = None) -> dict[str, Any]: + result = { + "id": document.id, + "indexing_status": document.indexing_status, + "processing_started_at": document.processing_started_at, + "parsing_completed_at": document.parsing_completed_at, + "cleaning_completed_at": document.cleaning_completed_at, + "splitting_completed_at": document.splitting_completed_at, + "completed_at": document.completed_at, + "paused_at": document.paused_at, + "error": document.error, + "stopped_at": document.stopped_at, + } + if counts is not None: + result.update(completed_segments=counts[0], total_segments=counts[1]) + return result + + +class SQLAlchemyDatasetOperations: + def __init__( + self, *, session_factory: sessionmaker[Session], tags: TagTargetQuery, app_queries: AppQueryService + ) -> None: + self._sessions = session_factory + self._tags = tags + self._app_queries = app_queries + + def visibility(self, context: RequestContext) -> DatasetVisibility: + with self._sessions() as session: + permissions = rbac_service.RBACService.MyPermissions.get( + context.active_workspace_id, context.account_id, session=session + ) + scope = ( + rbac_service.RBACService.DatasetAccess.whitelist_resources(context.active_workspace_id, context.account_id) + if dify_config.RBAC_ENABLED + else None + ) + return DatasetVisibility( + default_permissions=permissions.dataset.default_permission_keys, + workspace_permissions=permissions.workspace.permission_keys, + overrides={item.resource_id: item.permission_keys for item in permissions.dataset.overrides}, + unrestricted=scope is None or scope.unrestricted, + whitelist_ids=scope.resource_ids if scope is not None else (), + ) + + def embedding_models(self, workspace_id: str) -> set[str]: + manager = create_plugin_provider_manager(tenant_id=workspace_id) + models = manager.get_configurations(tenant_id=workspace_id).get_models( + model_type=ModelType.TEXT_EMBEDDING, only_active=True + ) + return {f"{model.model}:{model.provider.provider}" for model in models} + + def list_datasets( + self, context: RequestContext, query: DatasetListFilter, accessible_ids: list[str] | None, include_own: bool + ) -> DatasetPage: + page, limit = clamp_pagination(query.page, query.limit, 100) + with self._sessions() as session: + account = load_actor(session, context) + if query.ids: + datasets, total = DatasetService.get_datasets_by_ids( + list(query.ids), + context.active_workspace_id, + user=account, + accessible_dataset_ids=accessible_ids, + include_own_datasets=include_own, + session=session, + ) + else: + datasets, total = DatasetService.get_datasets( + page, + limit, + session, + context.active_workspace_id, + account, + query.keyword, + list(query.tag_ids), + query.include_all, + tags=self._tags, + accessible_dataset_ids=accessible_ids, + include_own_datasets=include_own, + ) + data = load_dataset_details(datasets, session=session) + partial_ids = [item["id"] for item in data if item["permission"] == "partial_members"] + members: dict[str, list[str]] = {} + if partial_ids: + for dataset_id, account_id in session.execute( + select(DatasetPermission.dataset_id, DatasetPermission.account_id).where( + DatasetPermission.dataset_id.in_(partial_ids), + DatasetPermission.tenant_id == context.active_workspace_id, + ) + ): + members.setdefault(dataset_id, []).append(account_id) + for item in data: + _normalize_provider(item) + item["partial_member_list"] = members.get(item["id"], []) + return { + "data": data, + "has_more": not query.ids and page * limit < total, + "total": total, + "page": page, + "limit": limit, + } + + @staticmethod + def _detail(session: Session, context: RequestContext, dataset: Dataset) -> DatasetDetailRecord: + result = load_dataset_detail(dataset, session=session) + permissions = rbac_service.RBACService.DatasetPermissions.batch_get( + context.active_workspace_id, context.account_id, [dataset.id], session=session + ) + result["permission_keys"] = permissions.get(dataset.id, []) + return result + + def get_dataset(self, context: RequestContext, ref: DatasetRef) -> DatasetDetailRecord: + with self._sessions() as session: + dataset = require_dataset(session, ref) + result = load_dataset_detail(dataset, session=session) + permissions = rbac_service.RBACService.MyPermissions.get( + context.active_workspace_id, context.account_id, dataset_id=ref.dataset_id, session=session + ) + result["permission_keys"] = permissions.dataset.permission_keys_by_resource_ids([dataset.id]).get( + dataset.id, [] + ) + if dataset.permission == DatasetPermissionEnum.PARTIAL_TEAM: + result["partial_member_list"] = self._members(session, ref) + return _normalize_provider(result) + + def create_dataset(self, context: RequestContext, values: Mapping[str, Any]) -> DatasetDetailRecord: + with _translate_permissions(), self._sessions() as session: + dataset = DatasetService.create_empty_dataset( + session=session, + tenant_id=context.active_workspace_id, + account=load_actor(session, context), + **dict(values), + ) + result = self._detail(session, context, dataset) + result["partial_member_list"] = None + dataset_id = dataset.id + session.commit() + if dify_config.RBAC_ENABLED: + rbac_service.RBACService.DatasetAccess.replace_whitelist( + context.active_workspace_id, + context.account_id, + dataset_id, + rbac_service.ReplaceMemberBindings(automatic_include_workspace_members=False), + ) + return result + + def update_dataset( + self, context: RequestContext, ref: DatasetRef, values: Mapping[str, Any] + ) -> DatasetDetailRecord: + data = dict(values) + if ( + data.get("indexing_technique") == "high_quality" + and data.get("embedding_model_provider") is not None + and data.get("embedding_model") is not None + ): + data["is_multimodal"] = DatasetService.check_is_multimodal_model( + ref.tenant_id, data["embedding_model_provider"], data["embedding_model"] + ) + with _translate_permissions(), self._sessions() as session: + dataset = require_dataset(session, ref) + account = load_actor(session, context) + if not dify_config.RBAC_ENABLED: + DatasetPermissionService.check_permission( + account, dataset, data.get("permission"), data.get("partial_member_list"), session=session + ) + # The legacy updater consumes fields from its input; retain the + # member selection for the permission update in this transaction. + dataset = DatasetService.update_dataset(ref.dataset_id, dict(data), account, session=session) + if dataset is None: + raise DatasetNotFoundError() + result = self._detail(session, context, dataset) + if ( + data.get("partial_member_list") is not None + and data.get("permission") == DatasetPermissionEnum.PARTIAL_TEAM + ): + DatasetPermissionService.update_partial_member_list( + ref.tenant_id, ref.dataset_id, data["partial_member_list"], session + ) + elif data.get("permission") in {DatasetPermissionEnum.ONLY_ME, DatasetPermissionEnum.ALL_TEAM}: + DatasetPermissionService.clear_partial_member_list(ref.dataset_id, session) + result["partial_member_list"] = self._members(session, ref) + session.commit() + return result + + def delete_dataset(self, context: RequestContext, ref: DatasetRef) -> None: + with _translate_permissions(), self._sessions() as session: + require_dataset(session, ref) + if not DatasetService.delete_dataset(ref.dataset_id, load_actor(session, context), session): + raise DatasetNotFoundError() + DatasetPermissionService.clear_partial_member_list(ref.dataset_id, session) + session.commit() + + def is_in_use(self, ref: DatasetRef) -> bool: + with self._sessions() as session: + require_dataset(session, ref) + return DatasetService.dataset_use_check(ref, session) + + def queries(self, ref: DatasetRef, *, page: int, limit: int) -> dict[str, Any]: + page, limit = clamp_pagination(page, limit, 100) + with self._sessions() as session: + require_dataset(session, ref) + queries, total = DatasetService.get_dataset_queries( + dataset_id=ref.dataset_id, page=page, per_page=limit, session=session + ) + data = [ + { + "id": query.id, + "queries": get_dataset_queries(query, session=session), + "source": query.source, + "source_app_id": query.source_app_id, + "created_by_role": query.created_by_role, + "created_by": query.created_by, + "created_at": query.created_at, + } + for query in queries + ] + return {"data": data, "has_more": page * limit < total, "total": total, "page": page, "limit": limit} + + def related_apps(self, ref: DatasetRef) -> dict[str, Any]: + with self._sessions() as session: + require_dataset(session, ref) + joins = DatasetService.get_related_apps(ref.dataset_id, session) + app_ids = [join.app_id for join in joins] + apps = self._app_queries.related_apps(ref.tenant_id, app_ids) + return {"data": apps, "total": len(apps)} + + def indexing_status(self, ref: DatasetRef) -> dict[str, Any]: + with self._sessions() as session: + require_dataset(session, ref) + documents = session.scalars( + select(Document).where(Document.dataset_id == ref.dataset_id, Document.tenant_id == ref.tenant_id) + ).all() + counts = DocumentService.get_document_segment_counts(documents, session=session) + return {"data": [_status(document, counts.get(document.id, (0, 0))) for document in documents]} + + def error_documents(self, ref: DatasetRef) -> dict[str, Any]: + with self._sessions() as session: + require_dataset(session, ref) + documents = DocumentService.get_error_documents_by_dataset_ref(ref, session) + return {"data": [_status(document) for document in documents], "total": len(documents)} + + @staticmethod + def _members(session: Session, ref: DatasetRef) -> list[str]: + return list( + session.scalars( + select(DatasetPermission.account_id).where( + DatasetPermission.dataset_id == ref.dataset_id, DatasetPermission.tenant_id == ref.tenant_id + ) + ).all() + ) + + def partial_members(self, ref: DatasetRef) -> list[str]: + with self._sessions() as session: + require_dataset(session, ref) + return self._members(session, ref) + + def auto_disable_logs(self, ref: DatasetRef) -> dict[str, Any]: + with self._sessions() as session: + return dict(DatasetService.get_dataset_auto_disable_logs(ref, session)) + + def set_api_enabled(self, context: RequestContext, ref: DatasetRef, enabled: bool) -> None: + with self._sessions() as session: + DatasetService.update_dataset_api_status( + require_dataset(session, ref), enabled, load_actor(session, context), session + ) + session.commit() diff --git a/api/services/knowledge/datasets/application.py b/api/services/knowledge/datasets/application.py new file mode 100644 index 00000000000..279bac2438d --- /dev/null +++ b/api/services/knowledge/datasets/application.py @@ -0,0 +1,167 @@ +"""Dataset administration policies independent of HTTP and persistence.""" + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from typing import Any, Protocol + +from libs.url_utils import normalize_api_base_url +from machinery.context import RequestContext +from services.knowledge.dataset_access import DatasetAccess +from services.knowledge.datasets.retrieval import retrieval_methods +from services.knowledge.entities.datasets import DatasetDetailRecord, DatasetPage +from services.knowledge.resource_scope import DatasetRef + + +@dataclass(frozen=True) +class DatasetListFilter: + page: int = 1 + limit: int = 20 + keyword: str | None = None + include_all: bool = False + ids: Sequence[str] = () + tag_ids: Sequence[str] = () + + +@dataclass(frozen=True) +class DatasetVisibility: + default_permissions: Sequence[str] = () + workspace_permissions: Sequence[str] = () + overrides: Mapping[str, Sequence[str]] = field(default_factory=dict) + unrestricted: bool = True + whitelist_ids: Sequence[str] = () + + def list_scope(self, *, rbac_enabled: bool) -> tuple[list[str] | None, bool]: + if not rbac_enabled: + return None, False + # A restricted whitelist always takes precedence over default grants, + # resource overrides and the creator's own-dataset fallback. + if not self.unrestricted: + return sorted(set(self.whitelist_ids)), False + read_keys = {"dataset.preview", "dataset.acl.preview", "dataset.full_access"} + default_read = bool( + read_keys.intersection(self.default_permissions) or read_keys.intersection(self.workspace_permissions) + ) + ids = ( + None + if default_read + else sorted(key for key, grants in self.overrides.items() if read_keys.intersection(grants)) + ) + return ids, "dataset.create_and_management" in self.workspace_permissions + + +class DatasetOperations(Protocol): + """Materialize owned values and complete writes before returning.""" + + def visibility(self, context: RequestContext) -> DatasetVisibility: ... + def list_datasets( + self, context: RequestContext, query: DatasetListFilter, accessible_ids: list[str] | None, include_own: bool + ) -> DatasetPage: ... + def embedding_models(self, workspace_id: str) -> set[str]: ... + def get_dataset(self, context: RequestContext, ref: DatasetRef) -> DatasetDetailRecord: ... + def create_dataset(self, context: RequestContext, values: Mapping[str, Any]) -> DatasetDetailRecord: ... + def update_dataset( + self, context: RequestContext, ref: DatasetRef, values: Mapping[str, Any] + ) -> DatasetDetailRecord: ... + def delete_dataset(self, context: RequestContext, ref: DatasetRef) -> None: ... + def is_in_use(self, ref: DatasetRef) -> bool: ... + def queries(self, ref: DatasetRef, *, page: int, limit: int) -> dict[str, Any]: ... + def related_apps(self, ref: DatasetRef) -> dict[str, Any]: ... + def indexing_status(self, ref: DatasetRef) -> dict[str, Any]: ... + def error_documents(self, ref: DatasetRef) -> dict[str, Any]: ... + def partial_members(self, ref: DatasetRef) -> list[str]: ... + def auto_disable_logs(self, ref: DatasetRef) -> dict[str, Any]: ... + def set_api_enabled(self, context: RequestContext, ref: DatasetRef, enabled: bool) -> None: ... + + +class DatasetApplicationService: + def __init__( + self, + *, + dataset_access: DatasetAccess, + operations: DatasetOperations, + rbac_enabled: bool, + service_api_url: str, + vector_store: str | None, + tidb_fulltext: bool, + ) -> None: + self._dataset_access = dataset_access + self._operations = operations + self._rbac_enabled = rbac_enabled + self._service_api_url = service_api_url + self._vector_store = vector_store + self._tidb_fulltext = tidb_fulltext + + def _dataset(self, context: RequestContext, dataset_id: str) -> DatasetRef: + dataset = self._dataset_access.require_accessible(context, dataset_id) + return DatasetRef(dataset.workspace_id, dataset.id) + + @staticmethod + def _embedding_available(item: DatasetDetailRecord, models: set[str], *, listing: bool) -> None: + high_quality = item["indexing_technique"] == "high_quality" + item["embedding_available"] = ( + not high_quality + or (listing and not item["embedding_model_provider"]) + or f"{item['embedding_model']}:{item['embedding_model_provider']}" in models + ) + + def list_datasets(self, context: RequestContext, query: DatasetListFilter) -> DatasetPage: + visibility = self._operations.visibility(context) + ids, include_own = visibility.list_scope(rbac_enabled=self._rbac_enabled) + result = self._operations.list_datasets(context, query, ids, include_own) + models = self._operations.embedding_models(context.active_workspace_id) + for item in result["data"]: + item["permission_keys"] = list(visibility.overrides.get(item["id"], visibility.default_permissions)) + self._embedding_available(item, models, listing=True) + return result + + def get_dataset(self, context: RequestContext, *, dataset_id: str) -> DatasetDetailRecord: + result = self._operations.get_dataset(context, self._dataset(context, dataset_id)) + self._embedding_available(result, self._operations.embedding_models(context.active_workspace_id), listing=False) + return result + + def create_dataset(self, context: RequestContext, *, values: Mapping[str, Any]) -> DatasetDetailRecord: + settings = dict(values) + settings["permission"] = "all_team_members" if self._rbac_enabled else settings.get("permission") or "only_me" + return self._operations.create_dataset(context, settings) + + def update_dataset( + self, context: RequestContext, *, dataset_id: str, values: Mapping[str, Any] + ) -> DatasetDetailRecord: + return self._operations.update_dataset(context, self._dataset(context, dataset_id), values) + + def delete_dataset(self, context: RequestContext, *, dataset_id: str) -> None: + self._operations.delete_dataset(context, self._dataset(context, dataset_id)) + + def is_in_use(self, context: RequestContext, *, dataset_id: str) -> bool: + return self._operations.is_in_use(self._dataset(context, dataset_id)) + + def queries(self, context: RequestContext, *, dataset_id: str, page: int, limit: int) -> dict[str, Any]: + return self._operations.queries(self._dataset(context, dataset_id), page=page, limit=limit) + + def related_apps(self, context: RequestContext, *, dataset_id: str) -> dict[str, Any]: + return self._operations.related_apps(self._dataset(context, dataset_id)) + + def indexing_status(self, context: RequestContext, *, dataset_id: str) -> dict[str, Any]: + return self._operations.indexing_status(self._dataset(context, dataset_id)) + + def error_documents(self, context: RequestContext, *, dataset_id: str) -> dict[str, Any]: + return self._operations.error_documents(self._dataset(context, dataset_id)) + + def partial_members(self, context: RequestContext, *, dataset_id: str) -> list[str]: + return self._operations.partial_members(self._dataset(context, dataset_id)) + + def auto_disable_logs(self, context: RequestContext, *, dataset_id: str) -> dict[str, Any]: + return self._operations.auto_disable_logs(self._dataset(context, dataset_id)) + + def set_api_enabled(self, context: RequestContext, *, dataset_id: str, status: str) -> None: + self._operations.set_api_enabled(context, self._dataset(context, dataset_id), status == "enable") + + def api_base_url(self, context: RequestContext, *, request_base_url: str) -> str: + return normalize_api_base_url(self._service_api_url or request_base_url) + + def retrieval_settings( + self, context: RequestContext, *, vector_type: str | None = None, is_mock: bool = False + ) -> dict[str, list[str]]: + return retrieval_methods( + vector_type if is_mock else self._vector_store, is_mock=is_mock, tidb_fulltext=self._tidb_fulltext + ) diff --git a/api/services/knowledge/datasets/retrieval.py b/api/services/knowledge/datasets/retrieval.py new file mode 100644 index 00000000000..f07fc477a5d --- /dev/null +++ b/api/services/knowledge/datasets/retrieval.py @@ -0,0 +1,85 @@ +"""Supported dataset retrieval methods for each vector store.""" + +from core.rag.datasource.vdb.vector_type import VectorType +from core.rag.retrieval.retrieval_methods import RetrievalMethod + + +def retrieval_methods( + vector_type: str | None, *, is_mock: bool = False, tidb_fulltext: bool = False +) -> dict[str, list[str]]: + """ + Get supported retrieval methods based on vector database type. + + Args: + vector_type: Vector database type, can be None + is_mock: Whether this is a Mock API, affects MILVUS handling + tidb_fulltext: Whether TiDB full-text retrieval is enabled + + Returns: + Dictionary containing supported retrieval methods + + Raises: + ValueError: If vector_type is None or unsupported + """ + if vector_type is None: + raise ValueError("Vector store type is not configured.") + + # Define vector database types that only support semantic search + semantic_only_types = { + VectorType.RELYT, + VectorType.CHROMA, + VectorType.PGVECTO_RS, + VectorType.VIKINGDB, + VectorType.UPSTASH, + } + + # Define vector database types that support all retrieval methods + full_search_types = { + VectorType.QDRANT, + VectorType.WEAVIATE, + VectorType.OPENSEARCH, + VectorType.ANALYTICDB, + VectorType.MYSCALE, + VectorType.ORACLE, + VectorType.ELASTICSEARCH, + VectorType.ELASTICSEARCH_JA, + VectorType.PGVECTOR, + VectorType.VASTBASE, + VectorType.TIDB_ON_QDRANT, + VectorType.LINDORM, + VectorType.COUCHBASE, + VectorType.OPENGAUSS, + VectorType.OCEANBASE, + VectorType.SEEKDB, + VectorType.TABLESTORE, + VectorType.HUAWEI_CLOUD, + VectorType.TENCENT, + VectorType.MATRIXONE, + VectorType.CLICKZETTA, + VectorType.BAIDU, + VectorType.ALIBABACLOUD_MYSQL, + VectorType.IRIS, + VectorType.HOLOGRES, + } + + semantic_methods = {"retrieval_method": [RetrievalMethod.SEMANTIC_SEARCH.value]} + full_methods = { + "retrieval_method": [ + RetrievalMethod.SEMANTIC_SEARCH.value, + RetrievalMethod.FULL_TEXT_SEARCH.value, + RetrievalMethod.HYBRID_SEARCH.value, + ] + } + + if vector_type == VectorType.MILVUS: + return semantic_methods if is_mock else full_methods + + if vector_type == VectorType.TIDB_VECTOR: + return full_methods if tidb_fulltext else semantic_methods + + if vector_type in semantic_only_types: + return semantic_methods + elif vector_type in full_search_types: + return full_methods + else: + raise ValueError(f"Unsupported vector db type {vector_type}.") diff --git a/api/services/knowledge/document_sync.py b/api/services/knowledge/document_sync.py new file mode 100644 index 00000000000..a22bf695f08 --- /dev/null +++ b/api/services/knowledge/document_sync.py @@ -0,0 +1,71 @@ +"""Request-side document synchronization use cases.""" + +from typing import NamedTuple, Protocol + +from machinery.context import RequestContext +from services.knowledge.dataset_access import DatasetAccess +from services.knowledge.resource_scope import DatasetRef, DocumentRef + + +class SyncDocumentRecord(NamedTuple): + """Document state required to authorize and dispatch a sync request.""" + + id: str + data_source_type: str + + +class DocumentSyncReader(Protocol): + def get_sync_document(self, document_ref: DocumentRef) -> SyncDocumentRecord | None: ... + + def list_active_notion_refs(self, dataset_ref: DatasetRef) -> tuple[DocumentRef, ...]: ... + + +class DocumentSyncDispatcher(Protocol): + def dispatch(self, *, dataset_id: str, document_id: str) -> None: ... + + +class DocumentSyncError(Exception): + """Base class for document synchronization failures.""" + + +class SyncDocumentNotFoundError(DocumentSyncError): + def __init__(self) -> None: + super().__init__("Document not found") + + +class SyncDocumentSourceError(DocumentSyncError): + def __init__(self) -> None: + super().__init__("Document is not notion type") + + +class DocumentSyncApplicationService: + """Authorize synchronization requests and dispatch the worker use case.""" + + def __init__( + self, + *, + dataset_access: DatasetAccess, + documents: DocumentSyncReader, + dispatcher: DocumentSyncDispatcher, + ) -> None: + self._dataset_access = dataset_access + self._documents = documents + self._dispatcher = dispatcher + + def sync_dataset(self, context: RequestContext, dataset_id: str) -> int: + dataset = self._dataset_access.require_accessible(context, dataset_id) + dataset_ref = DatasetRef(dataset.workspace_id, dataset.id) + document_refs = self._documents.list_active_notion_refs(dataset_ref) + for document_ref in document_refs: + self._dispatcher.dispatch(dataset_id=dataset_ref.dataset_id, document_id=document_ref.document_id) + return len(document_refs) + + def sync_document(self, context: RequestContext, dataset_id: str, document_id: str) -> None: + dataset = self._dataset_access.require_accessible(context, dataset_id) + document_ref = DatasetRef(dataset.workspace_id, dataset.id).document(document_id) + document = self._documents.get_sync_document(document_ref) + if document is None: + raise SyncDocumentNotFoundError() + if document.data_source_type != "notion_import": + raise SyncDocumentSourceError() + self._dispatcher.dispatch(dataset_id=dataset.id, document_id=document.id) diff --git a/api/services/knowledge/document_sync_adapters.py b/api/services/knowledge/document_sync_adapters.py new file mode 100644 index 00000000000..c110fe22faf --- /dev/null +++ b/api/services/knowledge/document_sync_adapters.py @@ -0,0 +1,11 @@ +"""Infrastructure adapters for document-sync requests.""" + +from collections.abc import Callable + + +class CeleryDocumentSyncDispatcher: + def __init__(self, *, delay: Callable[[str, str], object]) -> None: + self._delay = delay + + def dispatch(self, *, dataset_id: str, document_id: str) -> None: + self._delay(dataset_id, document_id) diff --git a/api/services/knowledge/documents/__init__.py b/api/services/knowledge/documents/__init__.py new file mode 100644 index 00000000000..5c5a2feed80 --- /dev/null +++ b/api/services/knowledge/documents/__init__.py @@ -0,0 +1 @@ +"""Document application use cases and their infrastructure adapters.""" diff --git a/api/services/knowledge/documents/adapters.py b/api/services/knowledge/documents/adapters.py new file mode 100644 index 00000000000..1c2bc398dd5 --- /dev/null +++ b/api/services/knowledge/documents/adapters.py @@ -0,0 +1,580 @@ +"""Session-owning adapters for document queries and existing document operations.""" + +import copy +import json +from collections.abc import Generator, Mapping, Sequence +from contextlib import contextmanager, nullcontext +from typing import Any, Literal +from uuid import uuid4 + +from redis.exceptions import LockNotOwnedError +from sqlalchemy import asc, desc, func, select, update +from sqlalchemy.orm import Session, sessionmaker +from werkzeug.exceptions import NotFound + +from configs import dify_config +from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError +from core.model_manager import ModelManager +from core.rag.index_processor.constant.built_in_field import BuiltInField +from core.rag.index_processor.constant.index_type import IndexTechniqueType +from extensions.ext_redis import redis_client +from graphon.model_runtime.entities.model_entities import ModelType +from graphon.model_runtime.errors.invoke import InvokeAuthorizationError +from libs.pagination import paginate_query +from machinery.context import RequestContext +from models import Account, Dataset, Document, DocumentSegment, UploadFile +from models.dataset import DatasetPermissionEnum, DocumentPipelineExecutionLog +from models.enums import IndexingStatus +from repositories.knowledge.dataset_read_repository import ( + get_dataset_doc_form, + get_document_hit_count, + get_document_process_rule, + get_document_segment_count, + get_latest_dataset_process_rule, +) +from repositories.knowledge.dataset_repository import _get_dataset +from repositories.knowledge.document_repository import _get_document +from services.enterprise import rbac_service +from services.errors.account import NoPermissionError +from services.errors.document import DocumentIndexingError +from services.file_service import FileService +from services.knowledge.dataset_access import DatasetAccessDeniedError, DatasetNotFoundError +from services.knowledge.dataset_read_service import ( + get_document_metadata_details, + get_document_source_detail, + load_document_detail, + load_document_details, +) +from services.knowledge.dataset_service import DatasetService, DocumentService +from services.knowledge.documents.application import ( + DocumentIndexingStateError, + DocumentInvalidActionError, + DocumentListFilter, + DocumentNotFoundError, + DocumentProviderError, + DocumentState, + DocumentZip, +) +from services.knowledge.entities.document_creation import DocumentCreationDataset, DocumentIndexingJobs +from services.knowledge.entities.knowledge_entities import KnowledgeConfig +from services.knowledge.resource_scope import DatasetRef, DocumentRef +from services.knowledge.summaries.adapters import SummaryIndexAdapter +from services.vector_space_admission_service import get_vector_space_admission_error_fields +from tasks.generate_summary_index_task import generate_summary_index_task + + +@contextmanager +def _translate_errors() -> Generator[None, None, None]: + try: + yield + except NotFound as error: + raise DocumentNotFoundError(error.description) from error + except NoPermissionError as error: + raise DatasetAccessDeniedError() from error + except DocumentIndexingError as error: + raise DocumentIndexingStateError(str(error)) from error + except (ProviderTokenNotInitError, InvokeAuthorizationError) as error: + raise DocumentProviderError(str(error), kind="uninitialized") from error + except QuotaExceededError as error: + raise DocumentProviderError(str(error), kind="quota") from error + except ModelCurrentlyNotSupportError as error: + raise DocumentProviderError(str(error), kind="unsupported") from error + + +def _require_dataset(session: Session, ref: DatasetRef) -> Dataset: + dataset = _get_dataset(session, ref) + if dataset is None: + raise DatasetNotFoundError() + return dataset + + +def _require_document(session: Session, ref: DocumentRef) -> Document: + document = _get_document(session, ref) + if document is None: + raise DocumentNotFoundError("Document not found.") + return document + + +def _state(document: Document) -> DocumentState: + return DocumentState( + id=document.id, + archived=bool(document.archived), + indexing_status=str(document.indexing_status), + is_paused=bool(document.is_paused), + data_source_type=str(document.data_source_type), + doc_form=str(document.doc_form), + need_summary=bool(document.need_summary), + ) + + +def _created_response(dataset: Dataset, documents: Sequence[Document], batch: str, session: Session) -> dict[str, Any]: + return { + "dataset": { + "id": dataset.id, + "name": dataset.name, + "description": dataset.description, + "permission": dataset.permission, + "data_source_type": dataset.data_source_type, + "indexing_technique": dataset.indexing_technique, + "created_by": dataset.created_by, + "created_at": dataset.created_at, + }, + "documents": load_document_details(documents, session=session), + "batch": batch, + } + + +def _indexing_status(document: Document, counts: tuple[int, int]) -> dict[str, Any]: + return { + "id": document.id, + "processing_started_at": document.processing_started_at, + "parsing_completed_at": document.parsing_completed_at, + "cleaning_completed_at": document.cleaning_completed_at, + "splitting_completed_at": document.splitting_completed_at, + "completed_at": document.completed_at, + "paused_at": document.paused_at, + "error": document.error, + "stopped_at": document.stopped_at, + "indexing_status": IndexingStatus.PAUSED if document.is_paused else document.indexing_status, + **get_vector_space_admission_error_fields(document.error), + "completed_segments": counts[0], + "total_segments": counts[1], + } + + +class SQLAlchemyDocumentOperations: + def __init__(self, *, session_factory: sessionmaker[Session]) -> None: + self._sessions = session_factory + + def find_document(self, *, workspace_id: str, document_id: str) -> DocumentRef | None: + with self._sessions() as session: + document = session.scalar( + select(Document).where(Document.id == document_id, Document.tenant_id == workspace_id) + ) + return DatasetRef(workspace_id, document.dataset_id).document(document.id) if document else None + + def get_state(self, ref: DocumentRef) -> DocumentState | None: + with self._sessions() as session: + document = _get_document(session, ref) + return _state(document) if document else None + + def get_states(self, ref: DatasetRef, document_ids: Sequence[str]) -> Sequence[DocumentState]: + with self._sessions() as session: + return tuple( + _state(document) for document in DocumentService.get_documents_by_ids(ref, document_ids, session) + ) + + def get_process_rule(self, ref: DatasetRef | None) -> dict[str, Any]: + result = copy.deepcopy(DocumentService.DEFAULT_RULES) + if ref is not None: + with self._sessions() as session: + rule = get_latest_dataset_process_rule(_require_dataset(session, ref), session=session) + if rule is not None: + result.update(mode=rule.mode, rules=rule.rules_dict) + return result + + def list_documents(self, ref: DatasetRef, query: DocumentListFilter) -> dict[str, Any]: + with self._sessions() as session: + dataset = _require_dataset(session, ref) + statement = select(Document).where( + Document.dataset_id == ref.dataset_id, Document.tenant_id == ref.tenant_id + ) + statement = DocumentService.apply_display_status_filter(statement, query.status) + if query.search: + statement = statement.where(Document.name.like(f"%{query.search}%")) + direction = desc if query.sort.startswith("-") else asc + sort = query.sort.removeprefix("-") + if sort == "hit_count": + hits = ( + select(DocumentSegment.document_id, func.sum(DocumentSegment.hit_count).label("total_hit_count")) + .where(DocumentSegment.dataset_id == ref.dataset_id, DocumentSegment.tenant_id == ref.tenant_id) + .group_by(DocumentSegment.document_id) + .subquery() + ) + statement = statement.outerjoin(hits, hits.c.document_id == Document.id).order_by( + direction(func.coalesce(hits.c.total_hit_count, 0)), direction(Document.position) + ) + elif sort == "created_at": + statement = statement.order_by(direction(Document.created_at), direction(Document.position)) + else: + statement = statement.order_by(desc(Document.created_at), desc(Document.position)) + page = paginate_query(statement, session=session, page=query.page, per_page=query.limit, max_per_page=100) + DocumentService.enrich_documents_with_summary_index_status( + documents=page.items, dataset=dataset, tenant_id=ref.tenant_id, session=session + ) + if query.fetch: + counts = DocumentService.get_document_segment_counts(page.items, session=session) + for document in page.items: + document.completed_segments, document.total_segments = counts.get(document.id, (0, 0)) + return { + "data": load_document_details(page.items, session=session), + "has_more": page.has_next, + "limit": page.per_page, + "total": page.total, + "page": page.page, + } + + def get_detail(self, ref: DocumentRef, *, metadata_only: bool) -> dict[str, Any]: + with self._sessions() as session: + document = _require_document(session, ref) + result = { + "id": document.id, + "doc_type": document.doc_type, + "doc_metadata": get_document_metadata_details(document, session=session), + } + if metadata_only: + return result + rule = get_document_process_rule(document, session=session) + segment_count = get_document_segment_count(document, session=session) + result.update( + { + "position": document.position, + "data_source_type": document.data_source_type, + "dataset_process_rule_id": document.dataset_process_rule_id, + "name": document.name, + "created_from": document.created_from, + "created_by": document.created_by, + "tokens": document.tokens, + "indexing_status": document.indexing_status, + "indexing_latency": document.indexing_latency, + "error": document.error, + "enabled": document.enabled, + "disabled_by": document.disabled_by, + "archived": document.archived, + "display_status": document.display_status, + "doc_form": document.doc_form, + "doc_language": document.doc_language, + } + ) + result.update( + { + "created_at": int(document.created_at.timestamp()) if document.created_at else None, + "completed_at": int(document.completed_at.timestamp()) if document.completed_at else None, + "updated_at": int(document.updated_at.timestamp()) if document.updated_at else None, + "disabled_at": int(document.disabled_at.timestamp()) if document.disabled_at else None, + } + ) + result.update( + data_source_info=document.data_source_info_dict, + data_source_detail_dict=get_document_source_detail(document, session=session), + dataset_process_rule=DatasetService.get_process_rules(ref.dataset.dataset_id, session), + document_process_rule=dict(rule.to_dict()) if rule else {}, + segment_count=segment_count, + average_segment_length=(document.word_count or 0) // segment_count if segment_count else 0, + hit_count=get_document_hit_count(document, session=session), + need_summary=bool(document.need_summary), + ) + return result + + def get_indexing_status(self, ref: DocumentRef) -> dict[str, Any]: + with self._sessions() as session: + document = _require_document(session, ref) + counts = DocumentService.get_document_segment_counts([document], session=session) + return _indexing_status(document, counts.get(document.id, (0, 0))) + + def get_batch_indexing_status(self, ref: DatasetRef, batch: str) -> dict[str, Any]: + with self._sessions() as session: + documents = session.scalars( + select(Document).where( + Document.tenant_id == ref.tenant_id, Document.dataset_id == ref.dataset_id, Document.batch == batch + ) + ).all() + if not documents: + raise DocumentNotFoundError("Documents not found.") + counts = DocumentService.get_document_segment_counts(documents, session=session) + return {"data": [_indexing_status(document, counts.get(document.id, (0, 0))) for document in documents]} + + def get_execution_log(self, ref: DocumentRef) -> dict[str, Any]: + with self._sessions() as session: + _require_document(session, ref) + log = session.scalar( + select(DocumentPipelineExecutionLog) + .where(DocumentPipelineExecutionLog.document_id == ref.document_id) + .order_by(DocumentPipelineExecutionLog.created_at.desc()) + .limit(1) + ) + if log is None: + return {} + return { + "datasource_info": json.loads(log.datasource_info), + "datasource_type": log.datasource_type, + "input_data": log.input_data, + "datasource_node_id": log.datasource_node_id, + } + + def get_summary_status(self, ref: DocumentRef) -> dict[str, Any]: + with self._sessions() as session: + _require_document(session, ref) + return dict( + SummaryIndexAdapter.get_document_summary_status_detail( + tenant_id=ref.dataset.tenant_id, + document_id=ref.document_id, + dataset_id=ref.dataset.dataset_id, + session=session, + ) + ) + + def get_download_url(self, ref: DocumentRef) -> str: + with _translate_errors(), self._sessions() as session: + return DocumentService.get_document_download_url(_require_document(session, ref), session) + + @contextmanager + def build_download_zip(self, ref: DatasetRef, document_ids: Sequence[str]) -> Generator[DocumentZip, None, None]: + with _translate_errors(), self._sessions() as session: + files_by_document = DocumentService._get_upload_files_by_document_id_for_zip_download( + dataset_id=ref.dataset_id, document_ids=document_ids, tenant_id=ref.tenant_id, session=session + ) + files = [files_by_document[document_id] for document_id in document_ids] + filename = DocumentService._generate_document_batch_download_zip_filename() + # The read transaction is closed before reading storage or streaming the response. + with FileService.build_upload_files_zip_tempfile(upload_files=files) as path: + yield DocumentZip(path=path, filename=filename) + + @staticmethod + def _actor(session: Session, context: RequestContext) -> Account: + account = session.get(Account, context.account_id) + if account is None: + raise DatasetAccessDeniedError() + account.set_tenant_id_with_session(context.active_workspace_id, session=session) + return account + + def _creation_dataset( + self, context: RequestContext, ref: DatasetRef, config: KnowledgeConfig + ) -> DocumentCreationDataset: + with self._sessions() as session: + dataset = _require_dataset(session, ref) + if not dataset.indexing_technique and not config.indexing_technique: + raise ValueError("indexing_technique is required.") + self._actor(session, context) + DatasetService.check_doc_form(dataset, config.doc_form, session=session) + return DocumentCreationDataset( + dataset.tenant_id, dataset.indexing_technique, dataset.embedding_model, dataset.embedding_model_provider + ) + + def _persist_creation( + self, context: RequestContext, ref: DatasetRef, config: KnowledgeConfig, *, initialize: bool + ) -> tuple[dict[str, Any], DocumentIndexingJobs]: + """Keep ORM values local and materialize the response before closing the session.""" + jobs: DocumentIndexingJobs | None = None + try: + with self._sessions() as session: + account = self._actor(session, context) + dataset = ( + DocumentService.build_dataset_for_documents(ref, config, account, session) + if initialize + else _require_dataset(session, ref) + ) + documents, batch, jobs = DocumentService.save_prepared_documents( + dataset, config, account, session=session + ) + if initialize: + dataset.name = documents[0].name[:18] + "..." + dataset.description = "useful for when you want to answer queries about the " + documents[0].name + dataset.permission = ( + DatasetPermissionEnum.ALL_TEAM if dify_config.RBAC_ENABLED else DatasetPermissionEnum.ONLY_ME + ) + session.flush() + result = _created_response(dataset, documents, batch, session) + session.commit() + return result, jobs + except Exception: + # The persistence helper commits documents before loading response + # details. Preserve their indexing even if that later read fails. + if jobs is not None: + DocumentService.dispatch_document_indexing(jobs) + raise + + def _run_creation( + self, context: RequestContext, ref: DatasetRef, config: KnowledgeConfig, *, initialize: bool + ) -> dict[str, Any]: + lock = ( + nullcontext() + if config.original_document_id + else redis_client.lock(f"add_document_lock_dataset_id_{ref.dataset_id}", timeout=600) + ) + result = None + try: + with lock: + result, jobs = self._persist_creation(context, ref, config, initialize=initialize) + # No session or ORM entity crosses the task boundary. Keep the + # creation lock until committed documents have been dispatched. + DocumentService.dispatch_document_indexing(jobs) + except LockNotOwnedError: + if config.original_document_id: + raise + # Preserve successful creation when the Redis lease expires on release. + if result is None: + if initialize: + raise + with self._sessions() as session: + result = _created_response( + _require_dataset(session, ref), [], DocumentService.generate_document_batch(), session + ) + assert result is not None + return result + + def create_documents(self, context: RequestContext, ref: DatasetRef, settings: Mapping[str, Any]) -> dict[str, Any]: + config = KnowledgeConfig.model_validate(settings) + DocumentService.document_create_args_validate(config) + with _translate_errors(): + snapshot = self._creation_dataset(context, ref, config) + config = DocumentService.prepare_document_creation(snapshot, config) + return self._run_creation(context, ref, config, initialize=False) + + def initialize_dataset(self, context: RequestContext, settings: Mapping[str, Any]) -> dict[str, Any]: + config = KnowledgeConfig.model_validate(settings) + with _translate_errors(): + if config.indexing_technique == IndexTechniqueType.HIGH_QUALITY: + if config.embedding_model is None or config.embedding_model_provider is None: + raise ValueError( + "embedding model and embedding model provider are required for high quality indexing." + ) + ModelManager.for_tenant(tenant_id=context.active_workspace_id).get_model_instance( + tenant_id=context.active_workspace_id, + provider=config.embedding_model_provider, + model_type=ModelType.TEXT_EMBEDDING, + model=config.embedding_model, + ) + config = config.model_copy( + update={ + "is_multimodal": DatasetService.check_is_multimodal_model( + context.active_workspace_id, config.embedding_model_provider, config.embedding_model + ) + } + ) + DocumentService.document_create_args_validate(config) + with self._sessions() as session: + self._actor(session, context) + config = DocumentService.prepare_document_creation( + DocumentCreationDataset( + context.active_workspace_id, + config.indexing_technique, + config.embedding_model, + config.embedding_model_provider, + ), + config, + ) + dataset_id = str(uuid4()) + result = self._run_creation( + context, DatasetRef(context.active_workspace_id, dataset_id), config, initialize=True + ) + if dify_config.RBAC_ENABLED: + rbac_service.RBACService.DatasetAccess.replace_whitelist( + context.active_workspace_id, + context.account_id, + dataset_id, + rbac_service.ReplaceMemberBindings(automatic_include_workspace_members=False), + ) + rbac_service.try_sync_creator_access_policy_member_bindings( + context.active_workspace_id, + context.account_id, + rbac_service.RBACResourceType.DATASET, + dataset_id, + ) + return result + + def _check_model(self, ref: DatasetRef) -> None: + with self._sessions() as session: + dataset = _require_dataset(session, ref) + DatasetService.check_dataset_model_setting(dataset) + + def delete_document(self, ref: DocumentRef) -> None: + self._check_model(ref.dataset) + with _translate_errors(), self._sessions() as session: + DocumentService.delete_document(_require_document(session, ref), session) + + def delete_documents(self, ref: DatasetRef, document_ids: Sequence[str]) -> None: + with _translate_errors(), self._sessions() as session: + dataset = _require_dataset(session, ref) + DocumentService.delete_documents( + ref, list(document_ids), get_dataset_doc_form(dataset, session=session), session + ) + + def update_document(self, ref: DocumentRef, values: Mapping[str, object]) -> None: + with self._sessions.begin() as session: + document = _require_document(session, ref) + for name, value in values.items(): + setattr(document, name, value) + + def rename_document(self, ref: DocumentRef, name: str) -> dict[str, Any]: + with self._sessions.begin() as session: + dataset = _require_dataset(session, ref.dataset) + document = _require_document(session, ref) + if dataset.built_in_field_enabled and document.doc_metadata: + document.doc_metadata = {**document.doc_metadata, BuiltInField.document_name: name} + document.name = name + if document.data_source_info_dict and (file_id := document.data_source_info_dict.get("upload_file_id")): + session.execute( + update(UploadFile) + .where(UploadFile.id == file_id, UploadFile.tenant_id == ref.dataset.tenant_id) + .values(name=name) + ) + session.flush() + return load_document_detail(document, session=session) + + def pause_document(self, ref: DocumentRef, *, actor_id: str) -> None: + with _translate_errors(), self._sessions() as session: + DocumentService.pause_document(_require_document(session, ref), session, actor_id=actor_id) + + def recover_document(self, ref: DocumentRef) -> None: + with _translate_errors(), self._sessions() as session: + DocumentService.recover_document(_require_document(session, ref), session) + + def retry_documents(self, ref: DatasetRef, document_ids: Sequence[str], *, actor_id: str) -> None: + with _translate_errors(), self._sessions() as session: + documents = DocumentService.get_documents_by_ids(ref, document_ids, session) + DocumentService.retry_document(ref.dataset_id, list(documents), session, actor_id=actor_id) + + def change_status( + self, + context: RequestContext, + ref: DatasetRef, + document_ids: Sequence[str], + action: Literal["enable", "disable", "archive", "un_archive"], + ) -> None: + self._check_model(ref) + with _translate_errors(), self._sessions() as session: + dataset = _require_dataset(session, ref) + account = self._actor(session, context) + documents = DocumentService.get_documents_by_ids(ref, document_ids, session) + try: + DocumentService.batch_update_document_status( + dataset, [doc.id for doc in documents], action, account, session + ) + except (ValueError, DocumentIndexingError) as error: + raise DocumentInvalidActionError(str(error)) from error + + def sync_website(self, ref: DocumentRef) -> None: + with _translate_errors(), self._sessions() as session: + DocumentService.sync_website_document( + _require_dataset(session, ref.dataset), _require_document(session, ref), session + ) + + def require_summary_enabled(self, ref: DatasetRef) -> None: + with self._sessions() as session: + dataset = _require_dataset(session, ref) + if dataset.indexing_technique != IndexTechniqueType.HIGH_QUALITY: + raise ValueError( + "Summary generation is only available for 'high_quality' indexing technique. " + f"Current indexing technique: {dataset.indexing_technique}" + ) + if not dataset.summary_index_setting or not dataset.summary_index_setting.get("enable"): + raise ValueError( + "Summary index is not enabled for this dataset. Please enable it in the dataset settings." + ) + + def enable_summary(self, ref: DatasetRef, document_ids: Sequence[str]) -> None: + if document_ids: + with self._sessions.begin() as session: + session.execute( + update(Document) + .where( + Document.tenant_id == ref.tenant_id, + Document.dataset_id == ref.dataset_id, + Document.id.in_(document_ids), + ) + .values(need_summary=True) + ) + + def dispatch_summary(self, ref: DocumentRef) -> None: + generate_summary_index_task.delay(ref.dataset.dataset_id, ref.document_id) diff --git a/api/services/knowledge/documents/application.py b/api/services/knowledge/documents/application.py new file mode 100644 index 00000000000..604193de7bc --- /dev/null +++ b/api/services/knowledge/documents/application.py @@ -0,0 +1,310 @@ +"""Console document use cases, independent of HTTP and persistence frameworks.""" + +from collections.abc import Mapping, Sequence +from contextlib import AbstractContextManager +from dataclasses import dataclass +from typing import Any, Literal, Protocol + +from libs.datetime_utils import naive_utc_now +from machinery.context import RequestContext +from services.knowledge.dataset_access import DatasetAccess +from services.knowledge.resource_scope import DatasetRef, DocumentRef + + +@dataclass(frozen=True) +class DocumentListFilter: + page: int = 1 + limit: int = 20 + search: str = "" + sort: str = "-created_at" + status: str | None = None + fetch: bool = False + + +@dataclass(frozen=True) +class DocumentState: + id: str + archived: bool + indexing_status: str + is_paused: bool + data_source_type: str + doc_form: str + need_summary: bool + + +@dataclass(frozen=True) +class DocumentZip: + path: str + filename: str + + +class DocumentApplicationError(Exception): + """Base class for document use-case failures mapped by the transport.""" + + +class DocumentNotFoundError(DocumentApplicationError): + pass + + +class DocumentArchivedError(DocumentApplicationError): + pass + + +class DocumentIndexingStateError(DocumentApplicationError): + pass + + +class DocumentInvalidActionError(DocumentApplicationError): + pass + + +class DocumentProviderError(DocumentApplicationError): + def __init__(self, description: str, *, kind: Literal["uninitialized", "quota", "unsupported"]) -> None: + super().__init__(description) + self.kind = kind + + +class DocumentOperations(Protocol): + """Owned document operations; adapters return materialized values and own their sessions.""" + + def find_document(self, *, workspace_id: str, document_id: str) -> DocumentRef | None: ... + def get_state(self, ref: DocumentRef) -> DocumentState | None: ... + def get_states(self, ref: DatasetRef, document_ids: Sequence[str]) -> Sequence[DocumentState]: ... + def list_documents(self, ref: DatasetRef, query: DocumentListFilter) -> dict[str, Any]: ... + def get_detail(self, ref: DocumentRef, *, metadata_only: bool) -> dict[str, Any]: ... + def get_process_rule(self, ref: DatasetRef | None) -> dict[str, Any]: ... + def get_indexing_status(self, ref: DocumentRef) -> dict[str, Any]: ... + def get_batch_indexing_status(self, ref: DatasetRef, batch: str) -> dict[str, Any]: ... + def get_execution_log(self, ref: DocumentRef) -> dict[str, Any]: ... + def get_summary_status(self, ref: DocumentRef) -> dict[str, Any]: ... + def get_download_url(self, ref: DocumentRef) -> str: ... + def build_download_zip( + self, ref: DatasetRef, document_ids: Sequence[str] + ) -> AbstractContextManager[DocumentZip]: ... + def create_documents( + self, context: RequestContext, ref: DatasetRef, settings: Mapping[str, Any] + ) -> dict[str, Any]: ... + def initialize_dataset(self, context: RequestContext, settings: Mapping[str, Any]) -> dict[str, Any]: ... + def delete_document(self, ref: DocumentRef) -> None: ... + def delete_documents(self, ref: DatasetRef, document_ids: Sequence[str]) -> None: ... + def update_document(self, ref: DocumentRef, values: Mapping[str, object]) -> None: ... + def rename_document(self, ref: DocumentRef, name: str) -> dict[str, Any]: ... + def pause_document(self, ref: DocumentRef, *, actor_id: str) -> None: ... + def recover_document(self, ref: DocumentRef) -> None: ... + def retry_documents(self, ref: DatasetRef, document_ids: Sequence[str], *, actor_id: str) -> None: ... + def change_status( + self, + context: RequestContext, + ref: DatasetRef, + document_ids: Sequence[str], + action: Literal["enable", "disable", "archive", "un_archive"], + ) -> None: ... + def sync_website(self, ref: DocumentRef) -> None: ... + def require_summary_enabled(self, ref: DatasetRef) -> None: ... + def enable_summary(self, ref: DatasetRef, document_ids: Sequence[str]) -> None: ... + def dispatch_summary(self, ref: DocumentRef) -> None: ... + + +class DatasetDocumentApplicationService: + """Authorize document use cases and apply document state transitions.""" + + def __init__( + self, + *, + dataset_access: DatasetAccess, + operations: DocumentOperations, + metadata_schema: Mapping[str, Mapping[str, type]], + ) -> None: + self._dataset_access = dataset_access + self._operations = operations + self._metadata_schema = metadata_schema + + def _dataset(self, context: RequestContext, dataset_id: str) -> DatasetRef: + dataset = self._dataset_access.require_accessible(context, dataset_id) + return DatasetRef(tenant_id=dataset.workspace_id, dataset_id=dataset.id) + + def _document( + self, context: RequestContext, dataset_id: str, document_id: str + ) -> tuple[DocumentRef, DocumentState]: + ref = self._dataset(context, dataset_id).document(document_id) + state = self._operations.get_state(ref) + if state is None: + raise DocumentNotFoundError("Document not found.") + return ref, state + + @staticmethod + def _require_mutable(state: DocumentState) -> None: + if state.archived: + raise DocumentArchivedError() + + def get_process_rule(self, context: RequestContext, *, document_id: str | None) -> dict[str, Any]: + ref = None + if document_id: + document_ref = self._operations.find_document( + workspace_id=context.active_workspace_id, document_id=document_id + ) + if document_ref is None: + raise DocumentNotFoundError("Document not found.") + ref = self._dataset(context, document_ref.dataset.dataset_id) + return self._operations.get_process_rule(ref) + + def list_documents(self, context: RequestContext, *, dataset_id: str, query: DocumentListFilter) -> dict[str, Any]: + return self._operations.list_documents(self._dataset(context, dataset_id), query) + + def get_document( + self, context: RequestContext, *, dataset_id: str, document_id: str, metadata_only: bool = False + ) -> dict[str, Any]: + ref, _ = self._document(context, dataset_id, document_id) + return self._operations.get_detail(ref, metadata_only=metadata_only) + + def get_indexing_status(self, context: RequestContext, *, dataset_id: str, document_id: str) -> dict[str, Any]: + ref, _ = self._document(context, dataset_id, document_id) + return self._operations.get_indexing_status(ref) + + def get_batch_indexing_status(self, context: RequestContext, *, dataset_id: str, batch: str) -> dict[str, Any]: + return self._operations.get_batch_indexing_status(self._dataset(context, dataset_id), batch) + + def get_execution_log(self, context: RequestContext, *, dataset_id: str, document_id: str) -> dict[str, Any]: + ref, _ = self._document(context, dataset_id, document_id) + return self._operations.get_execution_log(ref) + + def get_summary_status(self, context: RequestContext, *, dataset_id: str, document_id: str) -> dict[str, Any]: + ref, _ = self._document(context, dataset_id, document_id) + return self._operations.get_summary_status(ref) + + def get_download_url(self, context: RequestContext, *, dataset_id: str, document_id: str) -> str: + ref, _ = self._document(context, dataset_id, document_id) + return self._operations.get_download_url(ref) + + def build_download_zip( + self, context: RequestContext, *, dataset_id: str, document_ids: Sequence[str] + ) -> AbstractContextManager[DocumentZip]: + return self._operations.build_download_zip(self._dataset(context, dataset_id), document_ids) + + def create_documents( + self, context: RequestContext, *, dataset_id: str, settings: Mapping[str, Any] + ) -> dict[str, Any]: + return self._operations.create_documents(context, self._dataset(context, dataset_id), settings) + + def initialize_dataset(self, context: RequestContext, *, settings: Mapping[str, Any]) -> dict[str, Any]: + return self._operations.initialize_dataset(context, settings) + + def delete_document(self, context: RequestContext, *, dataset_id: str, document_id: str) -> None: + ref, _ = self._document(context, dataset_id, document_id) + self._operations.delete_document(ref) + + def delete_documents(self, context: RequestContext, *, dataset_id: str, document_ids: Sequence[str]) -> None: + self._operations.delete_documents(self._dataset(context, dataset_id), document_ids) + + def update_processing( + self, + context: RequestContext, + *, + dataset_id: str, + document_id: str, + action: Literal["pause", "resume"], + ) -> None: + ref, state = self._document(context, dataset_id, document_id) + if action == "pause": + if state.indexing_status != "indexing": + raise DocumentInvalidActionError("Document not in indexing state.") + values: dict[str, object] = { + "paused_by": context.account_id, + "paused_at": naive_utc_now(), + "is_paused": True, + } + elif action == "resume": + if state.indexing_status not in {"paused", "error"}: + raise DocumentInvalidActionError("Document not in paused or error state.") + values = {"paused_by": None, "paused_at": None, "is_paused": False} + else: + raise DocumentInvalidActionError("Invalid action.") + self._operations.update_document(ref, values) + + def update_metadata( + self, + context: RequestContext, + *, + dataset_id: str, + document_id: str, + doc_type: str | None, + doc_metadata: dict[str, Any] | None, + ) -> None: + ref, _ = self._document(context, dataset_id, document_id) + if doc_type is None or doc_metadata is None: + raise ValueError("Both doc_type and doc_metadata must be provided.") + if doc_type not in self._metadata_schema: + raise ValueError("Invalid doc_type.") + metadata = ( + doc_metadata + if doc_type == "others" + else { + key: doc_metadata[key] + for key, value_type in self._metadata_schema[doc_type].items() + if doc_metadata.get(key) is not None and isinstance(doc_metadata[key], value_type) + } + ) + self._operations.update_document( + ref, {"doc_type": doc_type, "doc_metadata": metadata, "updated_at": naive_utc_now()} + ) + + def change_status( + self, + context: RequestContext, + *, + dataset_id: str, + document_ids: Sequence[str], + action: Literal["enable", "disable", "archive", "un_archive"], + ) -> None: + self._operations.change_status(context, self._dataset(context, dataset_id), document_ids, action) + + def pause_document(self, context: RequestContext, *, dataset_id: str, document_id: str) -> None: + ref, state = self._document(context, dataset_id, document_id) + self._require_mutable(state) + if state.indexing_status not in {"waiting", "parsing", "cleaning", "splitting", "indexing"}: + raise DocumentIndexingStateError("Cannot pause completed document.") + self._operations.pause_document(ref, actor_id=context.account_id) + + def recover_document(self, context: RequestContext, *, dataset_id: str, document_id: str) -> None: + ref, state = self._document(context, dataset_id, document_id) + self._require_mutable(state) + if not state.is_paused: + raise DocumentIndexingStateError("Document is not in paused status.") + self._operations.recover_document(ref) + + def retry_documents(self, context: RequestContext, *, dataset_id: str, document_ids: Sequence[str]) -> None: + ref = self._dataset(context, dataset_id) + states = {state.id: state for state in self._operations.get_states(ref, document_ids)} + retry_ids = [ + document_id + for document_id in dict.fromkeys(document_ids) + if (state := states.get(document_id)) is not None + and not state.archived + and state.indexing_status != "completed" + ] + self._operations.retry_documents(ref, retry_ids, actor_id=context.account_id) + + def rename_document( + self, context: RequestContext, *, dataset_id: str, document_id: str, name: str + ) -> dict[str, Any]: + ref, _ = self._document(context, dataset_id, document_id) + return self._operations.rename_document(ref, name) + + def sync_website(self, context: RequestContext, *, dataset_id: str, document_id: str) -> None: + ref, state = self._document(context, dataset_id, document_id) + if state.data_source_type != "website_crawl": + raise ValueError("Document is not a website document.") + self._require_mutable(state) + self._operations.sync_website(ref) + + def generate_summary(self, context: RequestContext, *, dataset_id: str, document_ids: Sequence[str]) -> None: + ref = self._dataset(context, dataset_id) + self._operations.require_summary_enabled(ref) + states = self._operations.get_states(ref, document_ids) + if len(states) != len(document_ids): + missing_ids = set(document_ids) - {state.id for state in states} + raise DocumentNotFoundError(f"Some documents not found: {sorted(missing_ids)}") + targets = [state for state in states if state.doc_form != "qa_model"] + self._operations.enable_summary(ref, [state.id for state in targets if not state.need_summary]) + for state in targets: + self._operations.dispatch_summary(ref.document(state.id)) diff --git a/api/services/knowledge/entities/__init__.py b/api/services/knowledge/entities/__init__.py new file mode 100644 index 00000000000..bb315666d2c --- /dev/null +++ b/api/services/knowledge/entities/__init__.py @@ -0,0 +1 @@ +"""Shared data contracts for knowledge use cases.""" diff --git a/api/services/knowledge/entities/datasets.py b/api/services/knowledge/entities/datasets.py new file mode 100644 index 00000000000..21d5d659b99 --- /dev/null +++ b/api/services/knowledge/entities/datasets.py @@ -0,0 +1,69 @@ +"""Materialized dataset read contracts, independent of persistence and HTTP. + +Timestamps remain datetimes until transport serialization. JSON configuration +objects retain their stored keys; this boundary does not apply response defaults. +""" + +from collections.abc import Mapping, Sequence +from datetime import datetime +from typing import NotRequired, TypedDict + +from pydantic import JsonValue + + +class DatasetTagRecord(TypedDict): + id: str + name: str + type: str + + +class DatasetDetailRecord(TypedDict): + id: str + name: str + description: str | None + provider: str + permission: str + data_source_type: str | None + indexing_technique: str | None + app_count: int + document_count: int + word_count: int + created_by: str + author_name: str | None + created_at: datetime + updated_by: str | None + updated_at: datetime + embedding_model: str | None + embedding_model_provider: str | None + # None means availability has not been evaluated by the application service. + embedding_available: bool | None + retrieval_model_dict: Mapping[str, JsonValue] + summary_index_setting: Mapping[str, JsonValue] | None + tags: list[DatasetTagRecord] + doc_form: str | None + external_knowledge_info: Mapping[str, JsonValue] | None + external_retrieval_model: Mapping[str, JsonValue] | None + doc_metadata: Sequence[Mapping[str, str]] + built_in_field_enabled: bool + pipeline_id: str | None + runtime_mode: str | None + chunk_structure: str | None + icon_info: Mapping[str, JsonValue] | None + is_published: bool + total_documents: int + total_available_documents: int + enable_api: bool + is_multimodal: bool + permission_keys: list[str] + maintainer: str | None + # Detail/create responses may omit membership or explicitly return null; + # listing and updates materialize a list, including an empty one. + partial_member_list: NotRequired[list[str] | None] + + +class DatasetPage(TypedDict): + data: list[DatasetDetailRecord] + has_more: bool + total: int + page: int + limit: int diff --git a/api/services/knowledge/entities/document_creation.py b/api/services/knowledge/entities/document_creation.py new file mode 100644 index 00000000000..fb7bb474a61 --- /dev/null +++ b/api/services/knowledge/entities/document_creation.py @@ -0,0 +1,22 @@ +"""Owned values passed between document creation's database and I/O phases.""" + +from dataclasses import dataclass + +from services.knowledge.resource_scope import DatasetRef + + +@dataclass(frozen=True) +class DocumentCreationDataset: + tenant_id: str + indexing_technique: str | None + embedding_model: str | None + embedding_model_provider: str | None + + +@dataclass(frozen=True) +class DocumentIndexingJobs: + dataset: DatasetRef + created: tuple[str, ...] = () + duplicated: tuple[str, ...] = () + removed_notion: tuple[str, ...] = () + updated: str | None = None diff --git a/api/services/knowledge/entities/indexing_estimate.py b/api/services/knowledge/entities/indexing_estimate.py new file mode 100644 index 00000000000..437a16b7671 --- /dev/null +++ b/api/services/knowledge/entities/indexing_estimate.py @@ -0,0 +1,159 @@ +"""Commands and validation shared by knowledge indexing-estimate use cases.""" + +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import Annotated, Literal + +from pydantic import BaseModel, ConfigDict, Field, ValidationError, field_validator + +from core.rag.entities import PreProcessingRuleKey +from services.data_source.entities.notion_import import NotionPageType + + +class EstimateValidationError(ValueError): + """Raised when an estimate process rule is invalid.""" + + +class _PreProcessingRule(BaseModel): + id: PreProcessingRuleKey + enabled: bool + + +class _Segmentation(BaseModel): + separator: str = Field(min_length=1) + max_tokens: int = Field(gt=0) + + +class _Rules(BaseModel): + pre_processing_rules: list[_PreProcessingRule] + segmentation: _Segmentation + + @field_validator("pre_processing_rules") + @classmethod + def _deduplicate(cls, value: list[_PreProcessingRule]) -> list[_PreProcessingRule]: + by_id = {rule.id: rule for rule in value} + return list(by_id.values()) + + +class _HierarchicalRules(_Rules): + parent_mode: Literal["full-doc", "paragraph"] | None = None + subchunk_segmentation: _Segmentation | None = None + + +class _SummaryIndexSettingDisabled(BaseModel): + enable: Literal[False] = False + + +class _SummaryIndexSettingEnabled(BaseModel): + enable: Literal[True] + model_name: str = Field(min_length=1) + model_provider_name: str = Field(min_length=1) + + +_SummaryIndexSetting = Annotated[ + _SummaryIndexSettingDisabled | _SummaryIndexSettingEnabled, + Field(discriminator="enable"), +] + + +class _ProcessRuleBase(BaseModel): + model_config = ConfigDict(extra="allow") + + summary_index_setting: _SummaryIndexSetting | None = None + + @field_validator("summary_index_setting", mode="before") + @classmethod + def _normalize_summary_index_setting(cls, value: object) -> object: + """Treat mappings with a missing/None enable flag as absent (#36602).""" + + if value is None: + return None + if isinstance(value, Mapping) and value.get("enable") is None: + return None + return value + + +class _AutomaticProcessRule(_ProcessRuleBase): + mode: Literal["automatic"] + + +class _CustomProcessRule(_ProcessRuleBase): + mode: Literal["custom"] + rules: _Rules + + +class _HierarchicalProcessRule(_ProcessRuleBase): + mode: Literal["hierarchical"] + rules: _HierarchicalRules + + +_ProcessRule = Annotated[ + _AutomaticProcessRule | _CustomProcessRule | _HierarchicalProcessRule, + Field(discriminator="mode"), +] + + +class _ProcessRuleEnvelope(BaseModel): + process_rule: _ProcessRule + + +def _validation_message(error: ValidationError) -> str: + first = error.errors()[0] + original = first.get("ctx", {}).get("error") + return str(original) if isinstance(original, ValueError) else str(first["msg"]) + + +def normalize_process_rule(process_rule: Mapping[str, object]) -> dict[str, object]: + """Validate and normalize a process rule without mutating caller-owned input.""" + + try: + validated = _ProcessRuleEnvelope.model_validate({"process_rule": process_rule}) + except ValidationError as error: + raise EstimateValidationError(_validation_message(error)) from error + + result = validated.process_rule.model_dump(exclude_none=True) + if validated.process_rule.mode == "automatic": + result["rules"] = {} + elif validated.process_rule.mode == "hierarchical": + rules = result.get("rules") + if isinstance(rules, dict) and not rules.get("parent_mode"): + rules["parent_mode"] = "paragraph" + return result + + +@dataclass(frozen=True, slots=True) +class UploadFileEstimateSource: + file_id: str + source_type: Literal["upload_file"] = field(default="upload_file", init=False) + + +@dataclass(frozen=True, slots=True) +class NotionEstimateSource: + workspace_id: str + page_id: str + page_type: NotionPageType + credential_id: str + source_type: Literal["notion_import"] = field(default="notion_import", init=False) + + +@dataclass(frozen=True, slots=True) +class WebsiteEstimateSource: + provider: str + job_id: str + url: str + mode: str = "crawl" + only_main_content: bool = False + source_type: Literal["website_crawl"] = field(default="website_crawl", init=False) + + +type NewEstimateSource = UploadFileEstimateSource | NotionEstimateSource | WebsiteEstimateSource + + +@dataclass(frozen=True, slots=True) +class NewSourcesEstimateCommand: + sources: tuple[NewEstimateSource, ...] + process_rule: Mapping[str, object] + doc_form: str = "text_model" + doc_language: str = "English" + dataset_id: str | None = None + indexing_technique: str = "economy" diff --git a/api/services/entities/knowledge_entities/knowledge_entities.py b/api/services/knowledge/entities/knowledge_entities.py similarity index 91% rename from api/services/entities/knowledge_entities/knowledge_entities.py rename to api/services/knowledge/entities/knowledge_entities.py index f24283073b9..759fa7328be 100644 --- a/api/services/entities/knowledge_entities/knowledge_entities.py +++ b/api/services/knowledge/entities/knowledge_entities.py @@ -2,8 +2,7 @@ from typing import Annotated, Any, Literal from pydantic import BaseModel, Field, WithJsonSchema, field_validator -from core.rag.entities import Rule -from core.rag.entities.metadata_entities import MetadataFilteringCondition +from core.rag.entities import MetadataFilteringCondition, Rule from core.rag.index_processor.constant.index_type import IndexStructureType from core.rag.retrieval.retrieval_methods import RetrievalMethod from libs.helper import UUIDStr @@ -250,24 +249,6 @@ class SegmentCreateArgs(BaseModel): attachment_ids: list[str] | None = Field(default=None, description="Attachment file IDs.") -class SegmentUpdateArgs(BaseModel): - content: str | None = Field(default=None, description="Updated chunk text content.") - answer: str | None = Field(default=None, description="Updated answer content for QA mode.") - keywords: list[str] | None = Field(default=None, description="Updated keywords for the chunk.") - regenerate_child_chunks: bool = Field( - default=False, - description="Whether to regenerate child chunks after updating a parent chunk.", - ) - enabled: bool | None = Field(default=None, description="Whether the chunk is enabled.") - attachment_ids: list[str] | None = Field(default=None, description="Attachment file IDs.") - summary: str | None = Field(default=None, description="Summary content for summary index.") - - -class ChildChunkUpdateArgs(BaseModel): - id: str | None = Field(default=None, description="Existing child chunk ID. Omit to create a new child chunk.") - content: str = Field(description="Child chunk text content.") - - class MetadataArgs(BaseModel): type: Literal["string", "number", "time"] = Field( description="`string` for text values, `number` for numeric values, `time` for date/time values." diff --git a/api/services/knowledge/entities/segments.py b/api/services/knowledge/entities/segments.py new file mode 100644 index 00000000000..8431265a4b7 --- /dev/null +++ b/api/services/knowledge/entities/segments.py @@ -0,0 +1,79 @@ +"""Detached knowledge segment values shared by application ports.""" + +from datetime import datetime + +from pydantic import BaseModel, ConfigDict, Field + + +class ChildChunkUpdateArgs(BaseModel): + id: str | None = Field(default=None, description="Existing child chunk ID. Omit to create a new child chunk.") + content: str = Field(description="Child chunk text content.") + + +class SegmentAttachmentRecord(BaseModel): + model_config = ConfigDict(from_attributes=True, frozen=True) + + id: str + name: str + size: int + extension: str + mime_type: str | None + source_url: str + + +class ChildChunkRecord(BaseModel): + model_config = ConfigDict(from_attributes=True, frozen=True) + + id: str + segment_id: str + content: str + position: int + word_count: int + type: str + created_at: datetime + updated_at: datetime + + +class SegmentRecord(BaseModel): + model_config = ConfigDict(from_attributes=True, frozen=True) + + id: str + position: int + document_id: str + content: str + sign_content: str + answer: str | None + word_count: int + tokens: int + keywords: tuple[str, ...] | None + index_node_id: str | None + index_node_hash: str | None + hit_count: int + enabled: bool + disabled_at: datetime | None + disabled_by: str | None + status: str + created_by: str + created_at: datetime + updated_at: datetime + updated_by: str | None + indexing_at: datetime | None + completed_at: datetime | None + error: str | None + stopped_at: datetime | None + child_chunks: tuple[ChildChunkRecord, ...] + attachments: tuple[SegmentAttachmentRecord, ...] + summary: str | None + + +class SegmentUpdateArgs(BaseModel): + content: str | None = Field(default=None, description="Updated chunk text content.") + answer: str | None = Field(default=None, description="Updated answer content for QA mode.") + keywords: list[str] | None = Field(default=None, description="Updated keywords for the chunk.") + regenerate_child_chunks: bool = Field( + default=False, + description="Whether to regenerate child chunks after updating a parent chunk.", + ) + enabled: bool | None = Field(default=None, description="Whether the chunk is enabled.") + attachment_ids: list[str] | None = Field(default=None, description="Attachment file IDs.") + summary: str | None = Field(default=None, description="Summary content for summary index.") diff --git a/api/services/knowledge/external/__init__.py b/api/services/knowledge/external/__init__.py new file mode 100644 index 00000000000..ed656469c7a --- /dev/null +++ b/api/services/knowledge/external/__init__.py @@ -0,0 +1 @@ +"""External knowledge connection and retrieval use cases.""" diff --git a/api/services/knowledge/external/adapters.py b/api/services/knowledge/external/adapters.py new file mode 100644 index 00000000000..3add04e9425 --- /dev/null +++ b/api/services/knowledge/external/adapters.py @@ -0,0 +1,153 @@ +"""External knowledge persistence and network adapters with bounded sessions.""" + +from collections.abc import Mapping +from typing import Any + +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker + +from core.rag.entities import MetadataFilteringCondition +from libs.pagination import clamp_pagination +from machinery.context import RequestContext +from models.dataset import DatasetQuery, ExternalKnowledgeApis +from models.enums import CreatorUserRole, DatasetQuerySource +from repositories.knowledge.dataset_read_repository import ( + get_external_api_bindings_batch, + get_external_api_dataset_bindings, +) +from services.enterprise import rbac_service +from services.entities.external_knowledge_entities.external_knowledge_entities import ExternalDatasetCreatePayload +from services.hit_testing_service import HitTestingService +from services.knowledge.dataset_read_service import load_dataset_detail +from services.knowledge.datasets.adapters import require_dataset +from services.knowledge.entities.datasets import DatasetDetailRecord +from services.knowledge.external.application import ExternalTemplateNotFoundError +from services.knowledge.external.service import ExternalDatasetService +from services.knowledge.resource_scope import DatasetRef + + +class SQLAlchemyExternalKnowledgeOperations: + def __init__(self, *, session_factory: sessionmaker[Session]) -> None: + self._sessions = session_factory + + @staticmethod + def _template(session: Session, workspace_id: str, template_id: str) -> ExternalKnowledgeApis: + template = session.scalar( + select(ExternalKnowledgeApis).where( + ExternalKnowledgeApis.id == template_id, ExternalKnowledgeApis.tenant_id == workspace_id + ) + ) + if template is None: + raise ExternalTemplateNotFoundError("API template not found.") + return template + + @staticmethod + def _values(session: Session, template: ExternalKnowledgeApis) -> dict[str, Any]: + return dict(template.to_dict(dataset_bindings=get_external_api_dataset_bindings(template, session=session))) + + def list_templates(self, workspace_id: str, *, page: int, limit: int, keyword: str | None) -> dict[str, Any]: + page, limit = clamp_pagination(page, limit, 100) + with self._sessions() as session: + templates, total = ExternalDatasetService.get_external_knowledge_apis( + page, limit, workspace_id, keyword, session=session + ) + bindings = get_external_api_bindings_batch(templates, session=session) + return { + "data": [item.to_dict(dataset_bindings=bindings.get(item.id, [])) for item in templates], + "has_more": page * limit < total, + "total": total, + "page": page, + "limit": limit, + } + + def get_template(self, workspace_id: str, template_id: str) -> dict[str, Any]: + with self._sessions() as session: + return self._values(session, self._template(session, workspace_id, template_id)) + + def create_template(self, context: RequestContext, *, name: str, settings: Mapping[str, Any]) -> dict[str, Any]: + # The legacy creator probes the endpoint before its first SQL statement. + # An unopened session keeps that external request outside a transaction. + with self._sessions() as session: + template = ExternalDatasetService.create_external_knowledge_api( + context.active_workspace_id, + context.account_id, + {"name": name, "settings": dict(settings)}, + session=session, + ) + result = self._values(session, template) + session.commit() + return result + + def update_template( + self, context: RequestContext, template_id: str, *, name: str, settings: Mapping[str, Any] + ) -> dict[str, Any]: + with self._sessions.begin() as session: + self._template(session, context.active_workspace_id, template_id) + template = ExternalDatasetService.update_external_knowledge_api( + context.active_workspace_id, + context.account_id, + template_id, + {"name": name, "settings": dict(settings)}, + session=session, + ) + return self._values(session, template) + + def delete_template(self, workspace_id: str, template_id: str) -> None: + with self._sessions.begin() as session: + session.delete(self._template(session, workspace_id, template_id)) + + def template_usage(self, workspace_id: str, template_id: str) -> tuple[bool, int]: + with self._sessions() as session: + return ExternalDatasetService.external_knowledge_api_use_check(template_id, workspace_id, session=session) + + def create_dataset(self, context: RequestContext, payload: ExternalDatasetCreatePayload) -> DatasetDetailRecord: + with self._sessions() as session: + self._template(session, context.active_workspace_id, payload.external_knowledge_api_id) + dataset = ExternalDatasetService.create_external_dataset( + context.active_workspace_id, context.account_id, payload, session=session + ) + permissions = rbac_service.RBACService.DatasetPermissions.batch_get( + context.active_workspace_id, context.account_id, [dataset.id], session=session + ) + result = load_dataset_detail(dataset, session=session) + result["permission_keys"] = permissions.get(dataset.id, []) + session.commit() + return result + + def retrieve( + self, + context: RequestContext, + ref: DatasetRef, + *, + query: str, + retrieval_model: dict[str, Any] | None, + metadata_filters: dict[str, Any] | None, + ) -> dict[str, Any]: + with self._sessions() as session: + dataset = require_dataset(session, ref) + if dataset.provider != "external": + return {"query": {"content": query}, "records": []} + request = ExternalDatasetService.prepare_external_knowledge_retrieval( + tenant_id=ref.tenant_id, + dataset_id=ref.dataset_id, + query=HitTestingService.escape_query_for_search(query), + external_retrieval_parameters=retrieval_model or {}, + metadata_condition=MetadataFilteringCondition.model_validate(metadata_filters) + if metadata_filters + else None, + session=session, + ) + documents = ExternalDatasetService.execute_external_knowledge_retrieval(request) + with self._sessions.begin() as session: + require_dataset(session, ref) + session.add( + DatasetQuery( + dataset_id=ref.dataset_id, + content=query, + source=DatasetQuerySource.HIT_TESTING, + source_app_id=None, + created_by_role=CreatorUserRole.ACCOUNT, + created_by=context.account_id, + ) + ) + return dict(HitTestingService.compact_external_retrieve_response(dataset, query, documents)) diff --git a/api/services/knowledge/external/application.py b/api/services/knowledge/external/application.py new file mode 100644 index 00000000000..8fa758e1a54 --- /dev/null +++ b/api/services/knowledge/external/application.py @@ -0,0 +1,104 @@ +"""External knowledge use cases without transport or persistence dependencies.""" + +from collections.abc import Mapping +from typing import Any, Protocol + +from machinery.context import RequestContext +from services.entities.external_knowledge_entities.external_knowledge_entities import ExternalDatasetCreatePayload +from services.knowledge.dataset_access import DatasetAccess +from services.knowledge.entities.datasets import DatasetDetailRecord +from services.knowledge.resource_scope import DatasetRef + + +class ExternalTemplateNotFoundError(Exception): + pass + + +class ExternalHitTestingError(Exception): + pass + + +class ExternalKnowledgeOperations(Protocol): + def list_templates(self, workspace_id: str, *, page: int, limit: int, keyword: str | None) -> dict[str, Any]: ... + def get_template(self, workspace_id: str, template_id: str) -> dict[str, Any]: ... + def create_template(self, context: RequestContext, *, name: str, settings: Mapping[str, Any]) -> dict[str, Any]: ... + def update_template( + self, context: RequestContext, template_id: str, *, name: str, settings: Mapping[str, Any] + ) -> dict[str, Any]: ... + def delete_template(self, workspace_id: str, template_id: str) -> None: ... + def template_usage(self, workspace_id: str, template_id: str) -> tuple[bool, int]: ... + def create_dataset(self, context: RequestContext, payload: ExternalDatasetCreatePayload) -> DatasetDetailRecord: ... + def retrieve( + self, + context: RequestContext, + ref: DatasetRef, + *, + query: str, + retrieval_model: dict[str, Any] | None, + metadata_filters: dict[str, Any] | None, + ) -> dict[str, Any]: ... + + +class ExternalKnowledgeApplicationService: + def __init__(self, *, dataset_access: DatasetAccess, operations: ExternalKnowledgeOperations) -> None: + self._dataset_access = dataset_access + self._operations = operations + + @staticmethod + def _validate_settings(settings: Mapping[str, Any]) -> None: + if not settings: + raise ValueError("api list is empty") + if not settings.get("endpoint"): + raise ValueError("endpoint is required") + if not settings.get("api_key"): + raise ValueError("api_key is required") + + def list_templates(self, context: RequestContext, *, page: int, limit: int, keyword: str | None) -> dict[str, Any]: + return self._operations.list_templates(context.active_workspace_id, page=page, limit=limit, keyword=keyword) + + def get_template(self, context: RequestContext, *, template_id: str) -> dict[str, Any]: + return self._operations.get_template(context.active_workspace_id, template_id) + + def create_template(self, context: RequestContext, *, name: str, settings: Mapping[str, Any]) -> dict[str, Any]: + self._validate_settings(settings) + return self._operations.create_template(context, name=name, settings=settings) + + def update_template( + self, context: RequestContext, *, template_id: str, name: str, settings: Mapping[str, Any] + ) -> dict[str, Any]: + self._validate_settings(settings) + return self._operations.update_template(context, template_id, name=name, settings=settings) + + def delete_template(self, context: RequestContext, *, template_id: str) -> None: + self._operations.delete_template(context.active_workspace_id, template_id) + + def template_usage(self, context: RequestContext, *, template_id: str) -> tuple[bool, int]: + return self._operations.template_usage(context.active_workspace_id, template_id) + + def create_dataset(self, context: RequestContext, *, payload: ExternalDatasetCreatePayload) -> DatasetDetailRecord: + return self._operations.create_dataset(context, payload) + + def hit_testing( + self, + context: RequestContext, + *, + dataset_id: str, + query: str, + retrieval_model: dict[str, Any] | None, + metadata_filters: dict[str, Any] | None, + ) -> dict[str, Any]: + dataset = self._dataset_access.require_accessible(context, dataset_id) + if not query: + raise ValueError("Query or attachment_ids is required") + if len(query) > 250: + raise ValueError("Query cannot exceed 250 characters") + try: + return self._operations.retrieve( + context, + DatasetRef(dataset.workspace_id, dataset.id), + query=query, + retrieval_model=retrieval_model, + metadata_filters=metadata_filters, + ) + except Exception as error: + raise ExternalHitTestingError(str(error)) from error diff --git a/api/services/external_knowledge_service.py b/api/services/knowledge/external/service.py similarity index 93% rename from api/services/external_knowledge_service.py rename to api/services/knowledge/external/service.py index 76620c48f75..8d7d5855800 100644 --- a/api/services/external_knowledge_service.py +++ b/api/services/knowledge/external/service.py @@ -358,6 +358,22 @@ class ExternalDatasetService: the inner knowledge retrieval API—can consistently expose ``502 external_knowledge_failed``. """ + request = ExternalDatasetService.prepare_external_knowledge_retrieval( + tenant_id, dataset_id, query, external_retrieval_parameters, metadata_condition, session=session + ) + return ExternalDatasetService.execute_external_knowledge_retrieval(request) + + @staticmethod + def prepare_external_knowledge_retrieval( + tenant_id: str, + dataset_id: str, + query: str, + external_retrieval_parameters: dict[str, Any], + metadata_condition: MetadataFilteringCondition | None = None, + *, + session: Session, + ) -> ExternalKnowledgeApiSetting: + """Resolve owned configuration into a request that outlives the read session.""" external_knowledge_binding = session.scalar( select(ExternalKnowledgeBindings) .where(ExternalKnowledgeBindings.dataset_id == dataset_id, ExternalKnowledgeBindings.tenant_id == tenant_id) @@ -393,14 +409,19 @@ class ExternalDatasetService: "metadata_condition": metadata_condition.model_dump() if metadata_condition else None, } + return ExternalKnowledgeApiSetting( + url=f"{settings.get('endpoint')}/retrieval", + request_method="post", + headers=headers, + params=request_params, + ) + + @staticmethod + def execute_external_knowledge_retrieval(request: ExternalKnowledgeApiSetting) -> list[Any]: + """Execute and validate a prepared request without holding a database session.""" try: response = ExternalDatasetService.process_external_api( - ExternalKnowledgeApiSetting( - url=f"{settings.get('endpoint')}/retrieval", - request_method="post", - headers=headers, - params=request_params, - ), + request, None, ) except ExternalKnowledgeRetrievalError: diff --git a/api/services/knowledge/indexing/__init__.py b/api/services/knowledge/indexing/__init__.py new file mode 100644 index 00000000000..0629046e587 --- /dev/null +++ b/api/services/knowledge/indexing/__init__.py @@ -0,0 +1 @@ +"""Use-case-centered stored-document indexing slice.""" diff --git a/api/services/knowledge/indexing/adapters/__init__.py b/api/services/knowledge/indexing/adapters/__init__.py new file mode 100644 index 00000000000..f8596026c1d --- /dev/null +++ b/api/services/knowledge/indexing/adapters/__init__.py @@ -0,0 +1 @@ +"""Infrastructure adapters for the knowledge indexing slice.""" diff --git a/api/services/knowledge/indexing/adapters/cleanup.py b/api/services/knowledge/indexing/adapters/cleanup.py new file mode 100644 index 00000000000..e6ef60a5cb4 --- /dev/null +++ b/api/services/knowledge/indexing/adapters/cleanup.py @@ -0,0 +1,120 @@ +"""Delete document indexes between bounded database transactions.""" + +from collections.abc import Callable, Mapping, Sequence + +from sqlalchemy import delete, select +from sqlalchemy.orm import Session + +from configs import dify_config +from core.rag.datasource.keyword.jieba.jieba import Jieba +from core.rag.datasource.vdb.vector_factory import Vector +from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType +from extensions.ext_redis import redis_client +from models.dataset import ChildChunk, Dataset, DocumentSegment, DocumentSegmentSummary +from repositories.knowledge.dataset_read_repository import get_dataset_keyword_table +from repositories.knowledge.keyword_table_repository import persist_keyword_table + + +def clean_document_indexes( + *, dataset_id: str, document_ids: Sequence[str], doc_form: str | None, new_session: Callable[[], Session] +) -> str | None: + """Clean indexes even when the owning document rows have already been deleted. + + Snapshot the dataset-owned rows, close the read session, remove external + indexes, then delete the selected summary and child rows in one transaction. + A failed vector deletion leaves those rows available for a retry. + Empty document selections never mean deleting the entire dataset index. + """ + if not document_ids: + return None + with new_session() as session: + dataset = session.scalar(select(Dataset).where(Dataset.id == dataset_id)) + if dataset is None: + return None + tenant_id = dataset.tenant_id + segments = session.scalars( + select(DocumentSegment).where( + DocumentSegment.tenant_id == tenant_id, + DocumentSegment.dataset_id == dataset_id, + DocumentSegment.document_id.in_(document_ids), + ) + ).all() + summaries = session.scalars( + select(DocumentSegmentSummary).where( + DocumentSegmentSummary.dataset_id == dataset_id, + DocumentSegmentSummary.document_id.in_(document_ids), + ) + ).all() + children = session.scalars( + select(ChildChunk).where( + ChildChunk.tenant_id == tenant_id, + ChildChunk.dataset_id == dataset_id, + ChildChunk.document_id.in_(document_ids), + ) + ).all() + summary_ids = [summary.id for summary in summaries] + child_ids = [child.id for child in children] + summary_node_ids = [summary.summary_index_node_id for summary in summaries if summary.summary_index_node_id] + body_node_ids = ( + [child.index_node_id for child in children if child.index_node_id] + if doc_form == IndexStructureType.PARENT_CHILD_INDEX + else [segment.index_node_id for segment in segments if segment.index_node_id] + ) + high_quality = dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY + node_ids = list(dict.fromkeys([*summary_node_ids, *body_node_ids])) + vector_type = Vector.resolve_vector_type(dataset, session=session) if high_quality and node_ids else None + session.expunge(dataset) + + if not node_ids and not summary_ids and not child_ids: + return None + + if high_quality: + if node_ids: + Vector(dataset, session=None, vector_type=vector_type).delete_by_ids(node_ids) + elif body_node_ids: + # The keyword adapter owns its lock and storage I/O; callbacks only do SQL. + def read() -> tuple[str, str | None]: + with new_session() as session: + row = get_dataset_keyword_table(dataset, session=session) + return ( + (row.data_source_type, row.keyword_table) if row else (dify_config.KEYWORD_DATA_SOURCE_TYPE, None) + ) + + def write(storage_type: str, data: str, keywords: Mapping[str, Sequence[str]]) -> None: + with new_session() as session, session.begin(): + persist_keyword_table( + session, + tenant_id=tenant_id, + dataset_id=dataset_id, + storage_type=storage_type, + data=data, + keywords=keywords, + ) + + Jieba(dataset).update_texts( + [], + read=read, + write=write, + lock=redis_client.lock(f"keyword_indexing_lock_{dataset_id}", timeout=600), + deleted_ids=body_node_ids, + ) + + with new_session() as session, session.begin(): + if summary_ids: + session.execute( + delete(DocumentSegmentSummary).where( + DocumentSegmentSummary.dataset_id == dataset_id, + DocumentSegmentSummary.document_id.in_(document_ids), + DocumentSegmentSummary.id.in_(summary_ids), + ) + ) + if child_ids: + session.execute( + delete(ChildChunk).where( + ChildChunk.tenant_id == tenant_id, + ChildChunk.dataset_id == dataset_id, + ChildChunk.document_id.in_(document_ids), + ChildChunk.id.in_(child_ids), + ) + ) + return tenant_id diff --git a/api/services/knowledge/indexing/adapters/estimate.py b/api/services/knowledge/indexing/adapters/estimate.py new file mode 100644 index 00000000000..c09b8c42438 --- /dev/null +++ b/api/services/knowledge/indexing/adapters/estimate.py @@ -0,0 +1,187 @@ +"""Infrastructure adapters for indexing estimates.""" + +import json +import logging +from collections.abc import Mapping +from typing import cast + +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker + +from configs import dify_config +from core.credit_usage import CreditUsageCreatedBy +from core.entities.knowledge_entities import IndexingEstimate, PreviewDetail, QAPreviewDetail +from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError +from core.model_context import with_credit_usage_created_by +from core.model_manager import ModelInstance, ModelManager +from core.plugin.impl.exc import PluginDaemonClientSideError +from core.rag.extractor.entity.extract_setting import ExtractSetting +from core.rag.index_processor.constant.index_type import IndexTechniqueType +from core.rag.index_processor.index_processor_base import SummaryIndexSettingDict +from core.rag.index_processor.index_processor_factory import IndexProcessorFactory +from core.tools.utils.web_reader_tool import get_image_upload_file_ids +from enums import DeploymentEdition +from extensions.ext_storage import storage +from graphon.model_runtime.entities.model_entities import ModelType +from models.dataset import Dataset, DatasetProcessRule +from repositories.knowledge.dataset_repository import _get_dataset +from repositories.knowledge.upload_file_repository import SQLAlchemyKnowledgeUploadRepository +from services.knowledge.indexing.estimate import ( + IndexingEstimateExecutionError, + IndexingEstimateProviderUnavailableError, +) +from services.knowledge.resource_scope import DatasetRef + +logger = logging.getLogger(__name__) + + +class SQLAlchemyProcessRuleReader: + """Load one tenant-owned process rule in a bounded read session.""" + + def __init__(self, *, session_factory: sessionmaker[Session]) -> None: + self._session_factory = session_factory + + def get_by_id(self, *, dataset_ref: DatasetRef, process_rule_id: str) -> Mapping[str, object] | None: + with self._session_factory() as session: + process_rule = session.scalar( + select(DatasetProcessRule) + .join(Dataset, Dataset.id == DatasetProcessRule.dataset_id) + .where( + DatasetProcessRule.id == process_rule_id, + DatasetProcessRule.dataset_id == dataset_ref.dataset_id, + Dataset.tenant_id == dataset_ref.tenant_id, + ) + .limit(1) + ) + if process_rule is None: + return None + rules = json.loads(process_rule.rules) if process_rule.rules else {} + return {"mode": str(process_rule.mode), "rules": rules} + + +class IndexingEstimateAdapter: + """Extract previews without depending on the document indexing execution flow.""" + + def __init__( + self, + *, + session_factory: sessionmaker[Session], + ) -> None: + self._session_factory = session_factory + self._uploads = SQLAlchemyKnowledgeUploadRepository(session_factory=session_factory) + + @with_credit_usage_created_by(CreditUsageCreatedBy.KNOWLEDGE_INDEXING) + def run( + self, + tenant_id: str, + extract_settings: list[ExtractSetting], + tmp_processing_rule: Mapping[str, object], + doc_form: str | None = None, + doc_language: str = "English", + dataset_id: str | None = None, + indexing_technique: str = "economy", + ) -> IndexingEstimate: + try: + return self._estimate( + tenant_id, extract_settings, tmp_processing_rule, doc_form, doc_language, dataset_id, indexing_technique + ) + except (LLMBadRequestError, ProviderTokenNotInitError, PluginDaemonClientSideError) as error: + raise IndexingEstimateProviderUnavailableError(error.description) from error + except Exception as error: + raise IndexingEstimateExecutionError(str(error)) from error + + def _estimate( + self, + tenant_id: str, + extract_settings: list[ExtractSetting], + processing_rule: Mapping[str, object], + doc_form: str | None, + doc_language: str, + dataset_id: str | None, + indexing_technique: str, + ) -> IndexingEstimate: + if ( + dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD + and len(extract_settings) > dify_config.BATCH_UPLOAD_LIMIT + ): + raise ValueError(f"You have reached the batch upload limit of {dify_config.BATCH_UPLOAD_LIMIT}.") + + model = self._embedding_model(tenant_id, dataset_id, indexing_technique) + processor = IndexProcessorFactory(doc_form).init_index_processor() + previews: list[PreviewDetail] = [] + qa_previews: list[QAPreviewDetail] = [] + image_ids: set[str] = set() + total_segments = 0 + # Legacy processors need a local session for extraction artifacts. No + # document-state transaction is shared with this preview operation. + with self._session_factory() as session: + for setting in extract_settings: + texts = processor.extract(setting, process_rule_mode=processing_rule["mode"], session=session) + documents = processor.transform( + texts, + current_user=None, + embedding_model_instance=model, + process_rule={"mode": processing_rule["mode"], "rules": processing_rule.get("rules")}, + tenant_id=tenant_id, + doc_language=doc_language, + preview=True, + session=session, + ) + total_segments += len(documents) + for document in documents: + if doc_form == "qa_model": + if len(qa_previews) < 10: + qa_previews.append( + QAPreviewDetail( + question=document.page_content, answer=document.metadata.get("answer") or "" + ) + ) + elif len(previews) < 10: + preview = PreviewDetail(content=document.page_content) + if document.children: + preview.child_chunks = [child.page_content for child in document.children] + previews.append(preview) + image_ids.update(get_image_upload_file_ids(document.page_content)) + session.commit() + + # Storage I/O takes place after extraction commits and outside the + # upload-record deletion transaction. Foreign tenant files are excluded. + uploads = self._uploads.get_files(workspace_id=tenant_id, file_ids=tuple(image_ids)) + for upload in uploads.values(): + try: + storage.delete(upload.key) + except Exception: + logger.exception("Delete image file failed while estimating indexing, upload_file_id: %s", upload.id) + self._uploads.delete_files(workspace_id=tenant_id, file_ids=tuple(uploads)) + + if doc_form == "qa_model": + return IndexingEstimate(total_segments=total_segments * 20, qa_preview=qa_previews, preview=[]) + summary = processing_rule.get("summary_index_setting") + if isinstance(summary, Mapping) and summary.get("enable") and previews: + with self._session_factory() as session: + previews = processor.generate_summary_preview( + tenant_id, previews, cast(SummaryIndexSettingDict, summary), doc_language, session=session + ) + return IndexingEstimate(total_segments=total_segments, preview=previews) + + def _embedding_model(self, tenant_id: str, dataset_id: str | None, technique: str) -> ModelInstance | None: + provider = None + model_name = None + if dataset_id: + with self._session_factory() as session: + dataset = _get_dataset(session, DatasetRef(tenant_id, dataset_id)) + if dataset is None: + raise ValueError("Dataset not found.") + if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY: + technique = IndexTechniqueType.HIGH_QUALITY + provider, model_name = dataset.embedding_model_provider, dataset.embedding_model + if technique != IndexTechniqueType.HIGH_QUALITY: + return None + manager = ModelManager.for_tenant(tenant_id=tenant_id) + if provider: + if not model_name: + raise ValueError("Embedding model is not configured.") + return manager.get_model_instance( + tenant_id=tenant_id, provider=provider, model_type=ModelType.TEXT_EMBEDDING, model=model_name + ) + return manager.get_default_model_instance(tenant_id=tenant_id, model_type=ModelType.TEXT_EMBEDDING) diff --git a/api/services/knowledge/indexing/adapters/execution.py b/api/services/knowledge/indexing/adapters/execution.py new file mode 100644 index 00000000000..c5578cb97ef --- /dev/null +++ b/api/services/knowledge/indexing/adapters/execution.py @@ -0,0 +1,247 @@ +"""Adapt legacy processors to committed document-indexing phases.""" + +from concurrent.futures import ThreadPoolExecutor +from typing import cast + +from flask import Flask, current_app +from sqlalchemy.orm import Session, sessionmaker +from werkzeug.exceptions import HTTPException +from werkzeug.local import LocalProxy + +from core.errors.error import ( + AppInvokeQuotaExceededError, + InvokeRateLimitError, + LLMError, + ModelCurrentlyNotSupportError, + ProviderTokenNotInitError, + QuotaExceededError, +) +from core.model_manager import ModelManager +from core.plugin.impl.exc import PluginDaemonError +from core.rag.datasource.keyword.jieba.jieba import Jieba +from core.rag.datasource.vdb.vector_factory import Vector +from core.rag.embedding.token_counter import calculate_segment_token_counts +from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType +from core.rag.index_processor.index_processor_factory import IndexProcessorFactory +from core.rag.models.document import Document +from extensions.application_services.data_sources import build_data_source_credentials +from extensions.ext_redis import redis_client +from extensions.otel import propagate_context +from graphon.model_runtime.entities.model_entities import ModelType +from libs import helper +from models.enums import DataSourceType +from repositories.knowledge.document_repository import SQLAlchemyDocumentRepository +from repositories.knowledge.segment_repository import SQLAlchemySegmentRepository +from repositories.knowledge.upload_file_repository import SQLAlchemyKnowledgeUploadRepository +from services.knowledge.indexing.adapters.sources import ( + CompositeStoredSourceResolver, + FileSourceAdapter, + NotionSourceResolver, + WebsiteSourceAdapter, +) +from services.knowledge.indexing.errors import DocumentIsPausedError +from services.knowledge.indexing.execution import DocumentIndexingService, IndexingDocument +from services.knowledge.resource_scope import DocumentRef +from services.vector_space_admission_service import VectorSpaceAdmissionService + + +class IndexingExecutionAdapter: + def __init__( + self, + *, + session_factory: sessionmaker[Session], + documents: SQLAlchemyDocumentRepository, + segments: SQLAlchemySegmentRepository, + uploads: SQLAlchemyKnowledgeUploadRepository, + sources: CompositeStoredSourceResolver, + ) -> None: + self._session_factory = session_factory + self._documents = documents + self._segments = segments + self._uploads = uploads + self._sources = sources + + def describe_error(self, error: Exception) -> str: + """Preserve provider and extractor descriptions without their transport prefixes.""" + if isinstance( + error, + ( + LLMError, + ProviderTokenNotInitError, + QuotaExceededError, + AppInvokeQuotaExceededError, + ModelCurrentlyNotSupportError, + InvokeRateLimitError, + PluginDaemonError, + HTTPException, + ), + ): + return str(error.description) + return str(error) + + def check_paused(self, ref: DocumentRef) -> None: + if redis_client.get(f"document_{ref.document_id}_is_paused"): + raise DocumentIsPausedError(f"Document paused, document id: {ref.document_id}") + + def extract(self, document: IndexingDocument) -> list[Document]: + if not document.processing_rule: + raise ValueError("no process rule found") + setting = self._sources.resolve_for_indexing(document.source) + processor = IndexProcessorFactory(document.doc_form).init_index_processor() + # Legacy extractors persist downloaded images and Notion source metadata. + # Their session is independent of document-state write transactions. + with self._session_factory() as session: + texts = processor.extract(setting, process_rule_mode=document.processing_rule["mode"], session=session) + session.commit() + for text in texts: + text.metadata.update(document_id=document.ref.document_id, dataset_id=document.ref.dataset.dataset_id) + return texts + + def transform(self, document: IndexingDocument, texts: list[Document]) -> list[Document]: + dataset, _ = self._documents.get_indexing_models(document.ref) + user = self._documents.get_indexing_user(document.ref) + model = None + if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY: + manager = ModelManager.for_tenant(tenant_id=dataset.tenant_id) + if dataset.embedding_model_provider: + model = manager.get_model_instance( + tenant_id=dataset.tenant_id, + provider=dataset.embedding_model_provider, + model_type=ModelType.TEXT_EMBEDDING, + model=dataset.embedding_model, + ) + else: + model = manager.get_default_model_instance( + tenant_id=dataset.tenant_id, model_type=ModelType.TEXT_EMBEDDING + ) + processor = IndexProcessorFactory(document.doc_form).init_index_processor() + with self._session_factory() as session: + chunks = processor.transform( + texts, + user, + embedding_model_instance=model, + process_rule=document.processing_rule, + tenant_id=dataset.tenant_id, + doc_language=document.doc_language, + session=session, + ) + session.commit() + return chunks + + def count_tokens(self, document: IndexingDocument, chunks: list[Document]) -> list[int]: + dataset, _ = self._documents.get_indexing_models(document.ref) + return calculate_segment_token_counts(dataset=dataset, documents=chunks) + + def ensure_admission(self, document: IndexingDocument, chunks: list[Document]) -> None: + dataset, _ = self._documents.get_indexing_models(document.ref) + with self._session_factory() as session: + VectorSpaceAdmissionService().ensure_document_can_be_indexed( + dataset=dataset, + document_id=document.ref.document_id, + doc_form=document.doc_form, + documents=chunks, + include_summaries=document.need_summary, + session=session, + ) + + def load(self, document: IndexingDocument, chunks: list[Document]) -> None: + dataset, _ = self._documents.get_indexing_models(document.ref) + app = cast(LocalProxy[Flask], current_app)._get_current_object() + if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY: + groups: list[list[Document]] = [[] for _ in range(10)] + for chunk in chunks: + groups[int(helper.generate_text_hash(chunk.page_content), 16) % 10].append(chunk) + with ThreadPoolExecutor(max_workers=10) as executor: + futures = [ + executor.submit( + propagate_context(self._process_chunk), + app, + document.ref, + group, + False, + ) + for group in groups + if group + ] + for future in futures: + future.result() + elif ( + dataset.indexing_technique == IndexTechniqueType.ECONOMY + and document.doc_form != IndexStructureType.PARENT_CHILD_INDEX + ): + # A future propagates keyword failures; Thread.join() used to hide + # them and incorrectly let the document become completed. + with ThreadPoolExecutor(max_workers=1) as executor: + executor.submit(propagate_context(self._process_chunk), app, document.ref, chunks, True).result() + + def _process_chunk(self, app: Flask, ref: DocumentRef, chunks: list[Document], keywords: bool) -> None: + with app.app_context(): + self.check_paused(ref) + # Workers load their own inputs and own their sessions. No ORM + # instance or session crosses the executor boundary. + dataset, document = self._documents.get_indexing_models(ref) + if keywords: + # Keep the read/modify/write lock while each repository call + # closes its transaction before keyword extraction or storage I/O. + Jieba(dataset).update_texts( + chunks, + read=lambda: self._segments.get_keyword_table(ref.dataset), + write=lambda storage_type, data, selected: self._segments.save_keyword_table( + ref.dataset, storage_type=storage_type, data=data, keywords=selected + ), + lock=redis_client.lock(f"keyword_indexing_lock_{dataset.id}", timeout=600), + ) + elif dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY: + with self._session_factory() as session: + vector_type = Vector.resolve_vector_type(dataset, session=session) + attachments = ( + [attachment for chunk in chunks for attachment in chunk.attachments or []] + if dataset.is_multimodal + else [] + ) + files = self._uploads.get_files( + workspace_id=ref.dataset.tenant_id, + file_ids=[attachment.metadata["doc_id"] for attachment in attachments], + ) + # All database inputs are materialized before initializing the + # vector backend, embedding text/images, or reading object storage. + vector = Vector(dataset, session=None, vector_type=vector_type) + if document.doc_form == IndexStructureType.PARENT_CHILD_INDEX: + for chunk in chunks: + if chunk.children: + vector.create([Document.model_validate(child.model_dump()) for child in chunk.children]) + else: + vector.create(chunks) + if attachments: + vector.create_multimodal(attachments, upload_files=files) + self.check_paused(ref) + self._segments.complete_indexing_segments(ref, [chunk.metadata["doc_id"] for chunk in chunks]) + + +def build_document_indexing_service( + *, + session_factory: sessionmaker[Session], + enforce_vector_space_admission: bool = False, +) -> DocumentIndexingService: + """Assemble a fresh indexing use case from explicit database dependencies.""" + documents = SQLAlchemyDocumentRepository(session_factory=session_factory) + segments = SQLAlchemySegmentRepository(session_factory=session_factory) + uploads = SQLAlchemyKnowledgeUploadRepository(session_factory=session_factory) + credentials = build_data_source_credentials(database_client=session_factory) + sources = CompositeStoredSourceResolver( + adapters={ + DataSourceType.UPLOAD_FILE: FileSourceAdapter(uploads=uploads), + DataSourceType.NOTION_IMPORT: NotionSourceResolver( + actor_credentials=credentials.actor, stored_credentials=credentials.stored + ), + DataSourceType.WEBSITE_CRAWL: WebsiteSourceAdapter(), + } + ) + return DocumentIndexingService( + documents=documents, + segments=segments, + backend=IndexingExecutionAdapter( + session_factory=session_factory, documents=documents, segments=segments, uploads=uploads, sources=sources + ), + enforce_vector_space_admission=enforce_vector_space_admission, + ) diff --git a/api/services/knowledge/indexing/adapters/sources.py b/api/services/knowledge/indexing/adapters/sources.py new file mode 100644 index 00000000000..08a2966829a --- /dev/null +++ b/api/services/knowledge/indexing/adapters/sources.py @@ -0,0 +1,282 @@ +"""Datasource-specific extraction input resolvers.""" + +from collections.abc import Mapping, Sequence +from typing import Protocol + +from configs import dify_config +from core.rag.extractor.entity.datasource_type import DatasourceType +from core.rag.extractor.entity.extract_setting import ( + ExtractSetting, + NotionInfo, + StoredDocumentExtractionInput, + UploadFileExtractionInput, + WebsiteInfo, +) +from services.data_source.credential_gateway import ( + ActorDatasourceCredentialResolver, + NotionCredentialUnavailableError, + StoredDatasourceCredentialResolver, + resolve_stored_notion_access_token, +) +from services.data_source.entities.notion_import import NotionPageType +from services.knowledge.indexing.errors import ( + IndexingInputSourceError, + SourceCredentialUnavailableError, + UnsupportedStoredSourceError, +) +from services.knowledge.indexing.estimate import StoredSource + +_NOTION_PROVIDER = "notion_datasource" +_NOTION_PLUGIN = "langgenius/notion_datasource" + + +class StoredSourceAdapter(Protocol): + """Resolve one persisted datasource kind.""" + + def resolve(self, source: StoredSource) -> ExtractSetting: ... + + +class KnowledgeUploadCatalog(Protocol): + def get_by_ids(self, *, workspace_id: str, file_ids: Sequence[str]) -> Mapping[str, UploadFileExtractionInput]: ... + + +class FileSourceAdapter: + def __init__(self, *, uploads: KnowledgeUploadCatalog) -> None: + self._uploads = uploads + + def resolve(self, source: StoredSource) -> ExtractSetting: + info = _required_info(source, "no upload file found") + file_id = _required_string(info, "upload_file_id", "no upload file found") + return self.resolve_selection( + workspace_id=source.document_ref.dataset.tenant_id, + file_id=file_id, + document_model=source.document_model, + ) + + def resolve_selection( + self, + *, + workspace_id: str, + file_id: str, + document_model: str, + ) -> ExtractSetting: + return self.resolve_selections(workspace_id=workspace_id, file_ids=(file_id,), document_model=document_model)[ + file_id + ] + + def resolve_selections( + self, *, workspace_id: str, file_ids: Sequence[str], document_model: str + ) -> dict[str, ExtractSetting]: + if not file_ids: + return {} + uploads = self._uploads.get_by_ids(workspace_id=workspace_id, file_ids=file_ids) + settings: dict[str, ExtractSetting] = {} + for file_id in file_ids: + upload = uploads.get(file_id) + if upload is None: + raise IndexingInputSourceError(f"no upload file found: {file_id}") + settings[file_id] = ExtractSetting( + datasource_type=DatasourceType.FILE, upload_file=upload, document_model=document_model + ) + return settings + + +class WebsiteSourceAdapter: + def resolve(self, source: StoredSource) -> ExtractSetting: + info = _required_info(source, "no website crawl info found") + return self.resolve_selection( + workspace_id=source.document_ref.dataset.tenant_id, + provider=_required_string(info, "provider"), + job_id=_required_string(info, "job_id"), + url=_required_string(info, "url"), + mode=_required_string(info, "mode"), + only_main_content=bool(info.get("only_main_content", False)), + document_model=source.document_model, + ) + + def resolve_selection( + self, + *, + workspace_id: str, + provider: str, + job_id: str, + url: str, + mode: str, + only_main_content: bool, + document_model: str, + ) -> ExtractSetting: + return ExtractSetting( + datasource_type=DatasourceType.WEBSITE, + website_info=WebsiteInfo( + provider=provider, + job_id=job_id, + url=url, + mode=mode, + tenant_id=workspace_id, + only_main_content=only_main_content, + ), + document_model=document_model, + ) + + +class NotionSourceResolver: + """Resolve actor-selected and trusted-stored Notion extraction inputs.""" + + def __init__( + self, + *, + actor_credentials: ActorDatasourceCredentialResolver, + stored_credentials: StoredDatasourceCredentialResolver, + integration_token: str | None = None, + ) -> None: + self._actor_credentials = actor_credentials + self._stored_credentials = stored_credentials + self._integration_token = integration_token + + def resolve(self, source: StoredSource) -> ExtractSetting: + info = _required_info(source, "no notion import info found") + document_ref = source.document_ref + credential_id = _string_value(info, "credential_id") + token = self._stored_token(source, credential_id) + return self._prepared( + tenant_id=document_ref.dataset.tenant_id, + notion_workspace_id=_required_string(info, "notion_workspace_id"), + page_id=_required_string(info, "notion_page_id"), + page_type=_notion_page_type(info), + document_model=source.document_model, + credential_id=credential_id, + token=token, + ) + + def _stored_token(self, source: StoredSource, credential_id: str | None) -> str: + document_ref = source.document_ref + + def load_token(saved_id: str) -> str: + credentials = self._stored_credentials.resolve_for_document( + workspace_id=document_ref.dataset.tenant_id, + dataset_id=document_ref.dataset.dataset_id, + document_id=document_ref.document_id, + credential_id=saved_id, + provider=_NOTION_PROVIDER, + plugin_id=_NOTION_PLUGIN, + ) + return self._required_token(credentials) + + try: + return resolve_stored_notion_access_token( + credential_id=credential_id, + load_token=load_token, + integration_token=self._integration_token or dify_config.NOTION_INTEGRATION_TOKEN, + ) + except NotionCredentialUnavailableError as error: + raise SourceCredentialUnavailableError("Notion credential is unavailable") from error + + def resolve_selection( + self, + *, + workspace_id: str, + actor_id: str, + credential_id: str, + notion_workspace_id: str, + page_id: str, + page_type: NotionPageType, + document_model: str, + ) -> ExtractSetting: + credentials = self._actor_credentials.resolve( + workspace_id=workspace_id, + actor_id=actor_id, + credential_id=credential_id, + provider=_NOTION_PROVIDER, + plugin_id=_NOTION_PLUGIN, + ) + return self._prepared( + tenant_id=workspace_id, + notion_workspace_id=notion_workspace_id, + page_id=page_id, + page_type=page_type, + document_model=document_model, + credential_id=credential_id, + token=self._required_token(credentials), + ) + + def _prepared( + self, + *, + tenant_id: str, + notion_workspace_id: str, + page_id: str, + page_type: NotionPageType, + document_model: str, + credential_id: str | None, + token: str, + ) -> ExtractSetting: + return ExtractSetting( + datasource_type=DatasourceType.NOTION, + notion_info=NotionInfo( + credential_id=credential_id, + notion_workspace_id=notion_workspace_id, + notion_obj_id=page_id, + notion_page_type=page_type.value, + tenant_id=tenant_id, + notion_access_token=token, + ), + document_model=document_model, + ) + + @staticmethod + def _required_token(credentials: Mapping[str, object]) -> str: + token = credentials.get("integration_secret") + if isinstance(token, str) and token: + return token + raise SourceCredentialUnavailableError("Notion credential is unavailable") + + +class CompositeStoredSourceResolver: + """Route persisted sources to datasource-specific adapters.""" + + def __init__(self, *, adapters: Mapping[str, StoredSourceAdapter]) -> None: + self._adapters = dict(adapters) + + def resolve(self, source: StoredSource) -> ExtractSetting: + adapter = self._adapters.get(source.source_type) + if adapter is None: + raise UnsupportedStoredSourceError(source.source_type) + return adapter.resolve(source) + + def resolve_for_indexing(self, source: StoredSource) -> ExtractSetting: + setting = self.resolve(source) + if setting.notion_info is not None: + ref = source.document_ref + setting.notion_info.document = StoredDocumentExtractionInput( + id=ref.document_id, + tenant_id=ref.dataset.tenant_id, + dataset_id=ref.dataset.dataset_id, + data_source_info_dict=dict(source.source_info) if source.source_info is not None else None, + ) + return setting + + +def _required_info(source: StoredSource, message: str) -> Mapping[str, object]: + if source.source_info is None: + raise IndexingInputSourceError(message) + return source.source_info + + +def _string_value(info: Mapping[str, object], key: str) -> str | None: + value = info.get(key) + return value if isinstance(value, str) and value else None + + +def _required_string(info: Mapping[str, object], key: str, message: str | None = None) -> str: + value = _string_value(info, key) + if value is None: + raise IndexingInputSourceError(message or f"invalid data source field: {key}") + return value + + +def _notion_page_type(info: Mapping[str, object]) -> NotionPageType: + value = _required_string(info, "type") + try: + return NotionPageType(value) + except ValueError as error: + raise IndexingInputSourceError("invalid data source field: type") from error diff --git a/api/services/knowledge/indexing/errors.py b/api/services/knowledge/indexing/errors.py new file mode 100644 index 00000000000..ab16e980ae0 --- /dev/null +++ b/api/services/knowledge/indexing/errors.py @@ -0,0 +1,27 @@ +"""Application failures raised while preparing or executing indexing.""" + + +class IndexingInputError(Exception): + """Base class for failures while resolving stored extraction inputs.""" + + +class IndexingInputSourceError(IndexingInputError): + """Raised when a persisted source cannot form a valid extraction input.""" + + +class UnsupportedStoredSourceError(IndexingInputSourceError): + def __init__(self, source_type: str) -> None: + self.source_type = source_type + super().__init__(f"unsupported data source type: {source_type}") + + +class SourceCredentialUnavailableError(IndexingInputSourceError): + """Raised when a datasource cannot resolve its required credential policy.""" + + +class DocumentIsPausedError(Exception): + """Indexing stopped because the owned document was paused.""" + + +class DocumentIsDeletedPausedError(Exception): + """Indexing stopped because its document owner chain no longer exists.""" diff --git a/api/services/knowledge/indexing/estimate.py b/api/services/knowledge/indexing/estimate.py new file mode 100644 index 00000000000..6e75705e27c --- /dev/null +++ b/api/services/knowledge/indexing/estimate.py @@ -0,0 +1,345 @@ +"""Indexing-estimate application service and its consumer-owned ports.""" + +from collections.abc import Mapping, Sequence +from typing import NamedTuple, Never, Protocol + +from core.entities.knowledge_entities import IndexingEstimate +from core.rag.extractor.entity.extract_setting import ExtractSetting +from machinery.context import RequestContext +from services.data_source.credential_gateway import DatasourceCredentialError +from services.data_source.entities.notion_import import NotionPageType +from services.knowledge.dataset_access import DatasetAccess, DatasetNotFoundError +from services.knowledge.entities.indexing_estimate import ( + NewEstimateSource, + NewSourcesEstimateCommand, + NotionEstimateSource, + UploadFileEstimateSource, + WebsiteEstimateSource, + normalize_process_rule, +) +from services.knowledge.indexing.errors import ( + IndexingInputSourceError, + SourceCredentialUnavailableError, + UnsupportedStoredSourceError, +) +from services.knowledge.resource_scope import DatasetRef, DocumentRef + + +class DatasetEstimateRecord(NamedTuple): + """Dataset state required to estimate stored documents.""" + + id: str + workspace_id: str + indexing_technique: str | None + + +class EstimateDocumentRecord(NamedTuple): + """Document state required by the stored-document estimate use case.""" + + id: str + workspace_id: str + dataset_id: str + data_source_type: str + data_source_info: Mapping[str, object] | None + indexing_status: str + doc_form: str + doc_language: str | None + dataset_process_rule_id: str | None + + @property + def ref(self) -> DocumentRef: + return DatasetRef(self.workspace_id, self.dataset_id).document(self.id) + + +class StoredSource(NamedTuple): + """A persisted extraction source bound to its owning document.""" + + document_ref: DocumentRef + source_type: str + source_info: Mapping[str, object] | None + document_model: str + + @classmethod + def from_document(cls, document: EstimateDocumentRecord) -> "StoredSource": + return cls( + document_ref=document.ref, + source_type=document.data_source_type, + source_info=document.data_source_info, + document_model=document.doc_form, + ) + + +class EstimateDocumentReader(Protocol): + def get_estimate_document(self, document_ref: DocumentRef) -> EstimateDocumentRecord | None: ... + + def list_estimate_documents_by_batch( + self, dataset_ref: DatasetRef, batch: str + ) -> tuple[EstimateDocumentRecord, ...]: ... + + +class EstimateDatasetReader(Protocol): + def get_estimate_record(self, dataset_ref: DatasetRef) -> DatasetEstimateRecord | None: ... + + +class UploadEstimateSourceResolver(Protocol): + def resolve_selections( + self, *, workspace_id: str, file_ids: Sequence[str], document_model: str + ) -> Mapping[str, ExtractSetting]: ... + + +class NotionEstimateSourceResolver(Protocol): + def resolve_selection( + self, + *, + workspace_id: str, + actor_id: str, + credential_id: str, + notion_workspace_id: str, + page_id: str, + page_type: NotionPageType, + document_model: str, + ) -> ExtractSetting: ... + + +class WebsiteEstimateSourceResolver(Protocol): + def resolve_selection( + self, + *, + workspace_id: str, + provider: str, + job_id: str, + url: str, + mode: str, + only_main_content: bool, + document_model: str, + ) -> ExtractSetting: ... + + +class StoredEstimateSourceResolver(Protocol): + def resolve(self, source: StoredSource) -> ExtractSetting: ... + + +class ProcessRuleReader(Protocol): + def get_by_id(self, *, dataset_ref: DatasetRef, process_rule_id: str) -> Mapping[str, object] | None: ... + + +class IndexingEstimateRunner(Protocol): + def run( + self, + tenant_id: str, + extract_settings: list[ExtractSetting], + tmp_processing_rule: Mapping[str, object], + doc_form: str | None = None, + doc_language: str = "English", + dataset_id: str | None = None, + indexing_technique: str = "economy", + ) -> IndexingEstimate: ... + + +class IndexingEstimateApplicationError(Exception): + """Base class for framework-neutral estimate failures.""" + + +class EstimateSourceNotFoundError(IndexingEstimateApplicationError): + def __init__(self, source_id: str) -> None: + super().__init__(f"Estimate source not found: {source_id}") + + +class UnsupportedEstimateSourceError(IndexingEstimateApplicationError): + def __init__(self, source_type: str) -> None: + super().__init__(f"Data source type not supported: {source_type}") + + +class IndexingEstimateCredentialUnavailableError(IndexingEstimateApplicationError): + """Raised when a datasource credential required by an estimate is unavailable.""" + + +class IndexingEstimateProviderUnavailableError(IndexingEstimateApplicationError): + """Raised when a configured model or plugin provider is unavailable.""" + + +class IndexingEstimateExecutionError(IndexingEstimateApplicationError): + """Raised when the estimate processor cannot complete an estimate.""" + + +class EstimateDocumentNotFoundError(IndexingEstimateApplicationError): + def __init__(self) -> None: + super().__init__("Document not found") + + +class EstimateDocumentAlreadyFinishedError(IndexingEstimateApplicationError): + def __init__(self) -> None: + super().__init__("Document already finished") + + +class IndexingEstimateApplicationService: + def __init__( + self, + *, + dataset_access: DatasetAccess, + datasets: EstimateDatasetReader, + documents: EstimateDocumentReader, + files: UploadEstimateSourceResolver, + notion: NotionEstimateSourceResolver, + websites: WebsiteEstimateSourceResolver, + stored_sources: StoredEstimateSourceResolver, + process_rules: ProcessRuleReader, + runner: IndexingEstimateRunner, + ) -> None: + self._dataset_access = dataset_access + self._datasets = datasets + self._documents = documents + self._files = files + self._notion = notion + self._websites = websites + self._stored_sources = stored_sources + self._process_rules = process_rules + self._runner = runner + + def estimate_new_sources( + self, + context: RequestContext, + command: NewSourcesEstimateCommand, + ) -> IndexingEstimate: + dataset_id = command.dataset_id or None + if dataset_id is not None: + dataset = self._dataset_access.require_accessible(context, dataset_id) + dataset_id = dataset.id + process_rule = normalize_process_rule(command.process_rule) + + try: + file_settings = self._files.resolve_selections( + workspace_id=context.active_workspace_id, + file_ids=tuple( + dict.fromkeys( + source.file_id for source in command.sources if isinstance(source, UploadFileEstimateSource) + ) + ), + document_model=command.doc_form, + ) + extract_settings = [ + file_settings[source.file_id] + if isinstance(source, UploadFileEstimateSource) + else self._resolve_new_source(context, source, document_model=command.doc_form) + for source in command.sources + ] + if not extract_settings: + return IndexingEstimate(total_segments=0, preview=[]) + return self._runner.run( + tenant_id=context.active_workspace_id, + extract_settings=extract_settings, + tmp_processing_rule=process_rule, + doc_form=command.doc_form, + doc_language=command.doc_language, + dataset_id=dataset_id, + indexing_technique=command.indexing_technique, + ) + except Exception as error: + self._translate_error(error) + + def estimate_document( + self, + context: RequestContext, + *, + dataset_id: str, + document_id: str, + ) -> IndexingEstimate: + accessible = self._dataset_access.require_accessible(context, dataset_id) + dataset_ref = DatasetRef(accessible.workspace_id, accessible.id) + dataset = self._require_estimate_dataset(dataset_ref) + document = self._documents.get_estimate_document(dataset_ref.document(document_id)) + if document is None: + raise EstimateDocumentNotFoundError() + return self._estimate_document(context, dataset, document) + + def estimate_batch(self, context: RequestContext, *, dataset_id: str, batch: str) -> IndexingEstimate: + accessible = self._dataset_access.require_accessible(context, dataset_id) + dataset_ref = DatasetRef(accessible.workspace_id, accessible.id) + dataset = self._require_estimate_dataset(dataset_ref) + documents = self._documents.list_estimate_documents_by_batch(dataset_ref, batch) + if not documents: + raise EstimateDocumentNotFoundError() + # A batch preview uses its first document, as the existing API does. + return self._estimate_document(context, dataset, documents[0]) + + def _estimate_document( + self, + context: RequestContext, + dataset: DatasetEstimateRecord, + document: EstimateDocumentRecord, + ) -> IndexingEstimate: + if document.indexing_status in {"completed", "error"}: + raise EstimateDocumentAlreadyFinishedError() + try: + extract_setting = self._stored_sources.resolve(StoredSource.from_document(document)) + process_rule = self._load_process_rule(dataset, document) + return self._runner.run( + tenant_id=context.active_workspace_id, + extract_settings=[extract_setting], + tmp_processing_rule=process_rule, + doc_form=document.doc_form, + doc_language=document.doc_language or "English", + dataset_id=dataset.id, + indexing_technique=dataset.indexing_technique or "economy", + ) + except Exception as error: + self._translate_error(error) + + def _resolve_new_source( + self, context: RequestContext, source: NewEstimateSource, *, document_model: str + ) -> ExtractSetting: + if isinstance(source, NotionEstimateSource): + return self._notion.resolve_selection( + workspace_id=context.active_workspace_id, + actor_id=context.account_id, + credential_id=source.credential_id, + notion_workspace_id=source.workspace_id, + page_id=source.page_id, + page_type=source.page_type, + document_model=document_model, + ) + if isinstance(source, WebsiteEstimateSource): + return self._websites.resolve_selection( + workspace_id=context.active_workspace_id, + provider=source.provider, + job_id=source.job_id, + url=source.url, + mode=source.mode, + only_main_content=source.only_main_content, + document_model=document_model, + ) + raise UnsupportedEstimateSourceError(source.source_type) + + def _load_process_rule( + self, + dataset: DatasetEstimateRecord, + document: EstimateDocumentRecord, + ) -> Mapping[str, object]: + process_rule_id = document.dataset_process_rule_id + if not process_rule_id: + return {"mode": "automatic", "rules": {}} + process_rule = self._process_rules.get_by_id( + dataset_ref=DatasetRef(dataset.workspace_id, dataset.id), + process_rule_id=process_rule_id, + ) + if process_rule is None: + raise EstimateSourceNotFoundError(process_rule_id) + return process_rule + + @staticmethod + def _translate_error(error: Exception) -> Never: + if isinstance(error, IndexingEstimateApplicationError): + raise error + if isinstance(error, UnsupportedStoredSourceError): + raise UnsupportedEstimateSourceError(error.source_type) from error + if isinstance(error, (DatasourceCredentialError, SourceCredentialUnavailableError)): + raise IndexingEstimateCredentialUnavailableError(str(error)) from error + if isinstance(error, IndexingInputSourceError): + raise EstimateSourceNotFoundError(str(error)) from error + raise IndexingEstimateExecutionError(str(error)) from error + + def _require_estimate_dataset(self, dataset_ref: DatasetRef) -> DatasetEstimateRecord: + dataset = self._datasets.get_estimate_record(dataset_ref) + if dataset is None: + raise DatasetNotFoundError() + return dataset diff --git a/api/services/knowledge/indexing/execution.py b/api/services/knowledge/indexing/execution.py new file mode 100644 index 00000000000..dba53de6dbb --- /dev/null +++ b/api/services/knowledge/indexing/execution.py @@ -0,0 +1,138 @@ +"""Document indexing phases and their application-owned persistence and execution ports.""" + +import logging +import time +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from typing import Literal, Protocol + +from core.credit_usage import CreditUsageCreatedBy +from core.model_context import with_credit_usage_created_by +from core.rag.models.document import Document +from services.knowledge.indexing.errors import DocumentIsDeletedPausedError, DocumentIsPausedError +from services.knowledge.indexing.estimate import StoredSource +from services.knowledge.resource_scope import DocumentRef + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class IndexingDocument: + source: StoredSource + processing_rule: Mapping[str, object] + doc_language: str + need_summary: bool + + @property + def ref(self) -> DocumentRef: + return self.source.document_ref + + @property + def doc_form(self) -> str: + return self.source.document_model + + +class IndexingDocumentStore(Protocol): + def get_indexing_document(self, ref: DocumentRef) -> IndexingDocument | None: ... + + def mark_splitting(self, ref: DocumentRef) -> None: ... + + def complete_indexing(self, ref: DocumentRef, *, tokens: int, latency: float) -> None: ... + + def fail_indexing(self, ref: DocumentRef, error: str) -> None: ... + + +class IndexingSegmentStore(Protocol): + def clear_for_indexing(self, document: IndexingDocument) -> None: ... + + def save_for_indexing( + self, document: IndexingDocument, chunks: list[Document], token_counts: list[int] + ) -> None: ... + + def resume_indexing(self, document: IndexingDocument) -> tuple[list[Document], int]: ... + + +class IndexingBackend(Protocol): + def describe_error(self, error: Exception) -> str: + """Translate backend-specific exceptions into the persisted error message.""" + ... + + def check_paused(self, ref: DocumentRef) -> None: ... + + def extract(self, document: IndexingDocument) -> list[Document]: ... + + def transform(self, document: IndexingDocument, texts: list[Document]) -> list[Document]: ... + + def ensure_admission(self, document: IndexingDocument, chunks: list[Document]) -> None: ... + + def count_tokens(self, document: IndexingDocument, chunks: list[Document]) -> list[int]: ... + + def load(self, document: IndexingDocument, chunks: list[Document]) -> None: ... + + +class DocumentIndexingService: + """Coordinate committed phases; no caller session is held across execution.""" + + def __init__( + self, + *, + documents: IndexingDocumentStore, + segments: IndexingSegmentStore, + backend: IndexingBackend, + enforce_vector_space_admission: bool = False, + clock: Callable[[], float] = time.perf_counter, + ) -> None: + self._documents = documents + self._segments = segments + self._backend = backend + self._enforce_vector_space_admission = enforce_vector_space_admission + self._clock = clock + + @with_credit_usage_created_by(CreditUsageCreatedBy.KNOWLEDGE_INDEXING) + def run(self, document_refs: Sequence[DocumentRef]) -> None: + for ref in document_refs: + self._run(ref, start="parsing") + + @with_credit_usage_created_by(CreditUsageCreatedBy.KNOWLEDGE_INDEXING) + def run_in_splitting_status(self, ref: DocumentRef) -> None: + self._run(ref, start="splitting") + + @with_credit_usage_created_by(CreditUsageCreatedBy.KNOWLEDGE_INDEXING) + def run_in_indexing_status(self, ref: DocumentRef) -> None: + self._run(ref, start="indexing") + + def _run(self, ref: DocumentRef, *, start: Literal["parsing", "splitting", "indexing"]) -> None: + try: + document = self._documents.get_indexing_document(ref) + if document is None: + logger.info("Document not found, skipping indexing: %s", ref) + return + self._backend.check_paused(ref) + if start == "indexing": + chunks, total_tokens = self._segments.resume_indexing(document) + else: + if start == "splitting": + self._segments.clear_for_indexing(document) + texts = self._backend.extract(document) + self._documents.mark_splitting(ref) + chunks = self._backend.transform(document, texts) + # Recovery retains its existing admission policy; only normal + # indexing callers that explicitly opt in reserve vector space. + if start == "parsing" and self._enforce_vector_space_admission: + self._backend.ensure_admission(document, chunks) + token_counts = self._backend.count_tokens(document, chunks) + total_tokens = sum(token_counts) + self._segments.save_for_indexing(document, chunks, token_counts) + + started = self._clock() + self._backend.check_paused(ref) + self._backend.load(document, chunks) + self._backend.check_paused(ref) + self._documents.complete_indexing(ref, tokens=total_tokens, latency=self._clock() - started) + except DocumentIsPausedError: + raise + except DocumentIsDeletedPausedError: + logger.info("Document deleted during indexing: %s", ref) + except Exception as error: + logger.exception("Document indexing failed: %s", ref) + self._documents.fail_indexing(ref, self._backend.describe_error(error)) diff --git a/api/services/knowledge/metadata/__init__.py b/api/services/knowledge/metadata/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/api/services/knowledge/metadata/application.py b/api/services/knowledge/metadata/application.py new file mode 100644 index 00000000000..6a39918cea6 --- /dev/null +++ b/api/services/knowledge/metadata/application.py @@ -0,0 +1,125 @@ +"""Metadata use cases with explicit ownership and persistence contracts.""" + +from collections.abc import Sequence +from dataclasses import dataclass +from datetime import datetime +from typing import Protocol, TypedDict + +from core.rag.index_processor.constant.built_in_field import BuiltInField +from machinery.context import RequestContext +from services.knowledge.dataset_access import DatasetAccess +from services.knowledge.resource_scope import DatasetRef + + +class MetadataInput(Protocol): + @property + def name(self) -> str: ... + @property + def type(self) -> str: ... + + +class MetadataValue(Protocol): + @property + def id(self) -> str: ... + @property + def value(self) -> str | int | float | None: ... + + +class MetadataOperation(Protocol): + @property + def document_id(self) -> str: ... + @property + def metadata_list(self) -> Sequence[MetadataValue]: ... + @property + def partial_update(self) -> bool: ... + + +class MetadataOperationsInput(Protocol): + @property + def operation_data(self) -> Sequence[MetadataOperation]: ... + + +@dataclass(frozen=True) +class MetadataRecord: + id: str + name: str + type: str + tenant_id: str + dataset_id: str + created_by: str + created_at: datetime + updated_by: str | None + updated_at: datetime | None + + +class MetadataFieldSummary(TypedDict): + id: str + name: str + type: str + count: int + + +class MetadataList(TypedDict): + doc_metadata: list[MetadataFieldSummary] + built_in_field_enabled: bool + + +class MetadataStore(Protocol): + """Own bounded transactions and recheck the complete dataset owner chain.""" + + def create(self, ref: DatasetRef, *, name: str, field_type: str, actor_id: str) -> MetadataRecord: ... + def rename(self, ref: DatasetRef, metadata_id: str, *, name: str, actor_id: str) -> MetadataRecord: ... + def delete(self, ref: DatasetRef, metadata_id: str) -> MetadataRecord: ... + def set_built_in(self, ref: DatasetRef, *, enabled: bool) -> None: ... + def update_documents(self, ref: DatasetRef, operations: Sequence[MetadataOperation], *, actor_id: str) -> None: ... + def list_fields(self, ref: DatasetRef) -> MetadataList: ... + + +class MetadataService: + def __init__(self, *, store: MetadataStore, dataset_access: DatasetAccess) -> None: + self._store = store + self._access = dataset_access + + def require_dataset(self, context: RequestContext, dataset_id: str) -> DatasetRef: + dataset = self._access.require_accessible(context, dataset_id) + return DatasetRef(dataset.workspace_id, dataset.id) + + @staticmethod + def _validate_name(name: str) -> None: + if len(name) > 255: + raise ValueError("Metadata name cannot exceed 255 characters.") + if name in BuiltInField: + raise ValueError("Metadata name already exists in Built-in fields.") + + def create_metadata(self, ref: DatasetRef, args: MetadataInput, *, actor_id: str) -> MetadataRecord: + self._validate_name(args.name) + return self._store.create(ref, name=args.name, field_type=args.type, actor_id=actor_id) + + def update_metadata_name(self, ref: DatasetRef, metadata_id: str, name: str, *, actor_id: str) -> MetadataRecord: + self._validate_name(name) + return self._store.rename(ref, metadata_id, name=name, actor_id=actor_id) + + def delete_metadata(self, ref: DatasetRef, metadata_id: str) -> MetadataRecord: + return self._store.delete(ref, metadata_id) + + @staticmethod + def get_built_in_fields() -> list[dict[str, str]]: + return [ + { + "name": field.value, + "type": "time" if field in (BuiltInField.upload_date, BuiltInField.last_update_date) else "string", + } + for field in BuiltInField + ] + + def enable_built_in_field(self, ref: DatasetRef) -> None: + self._store.set_built_in(ref, enabled=True) + + def disable_built_in_field(self, ref: DatasetRef) -> None: + self._store.set_built_in(ref, enabled=False) + + def update_documents_metadata(self, ref: DatasetRef, args: MetadataOperationsInput, *, actor_id: str) -> None: + self._store.update_documents(ref, args.operation_data, actor_id=actor_id) + + def get_dataset_metadatas(self, ref: DatasetRef) -> MetadataList: + return self._store.list_fields(ref) diff --git a/api/services/knowledge/resource_scope.py b/api/services/knowledge/resource_scope.py new file mode 100644 index 00000000000..2a29f9dedba --- /dev/null +++ b/api/services/knowledge/resource_scope.py @@ -0,0 +1,32 @@ +"""Typed owner chains used by knowledge application and persistence ports.""" + +from __future__ import annotations + +from typing import NamedTuple + + +class DatasetRef(NamedTuple): + """Identify a dataset inside its owning workspace.""" + + tenant_id: str + dataset_id: str + + def document(self, document_id: str) -> DocumentRef: + return DocumentRef(dataset=self, document_id=document_id) + + +class DocumentRef(NamedTuple): + """Identify a document through its complete dataset owner chain.""" + + dataset: DatasetRef + document_id: str + + def segment(self, segment_id: str) -> SegmentRef: + return SegmentRef(document=self, segment_id=segment_id) + + +class SegmentRef(NamedTuple): + """Identify a segment through its complete document owner chain.""" + + document: DocumentRef + segment_id: str diff --git a/api/services/knowledge/retrieval/__init__.py b/api/services/knowledge/retrieval/__init__.py new file mode 100644 index 00000000000..675f7c2640b --- /dev/null +++ b/api/services/knowledge/retrieval/__init__.py @@ -0,0 +1 @@ +"""Knowledge retrieval authorization and application operations.""" diff --git a/api/services/knowledge/retrieval/attachments.py b/api/services/knowledge/retrieval/attachments.py new file mode 100644 index 00000000000..9c5ae792bec --- /dev/null +++ b/api/services/knowledge/retrieval/attachments.py @@ -0,0 +1,33 @@ +"""Grant execution-local file access only for validated knowledge retrieval results.""" + +from collections.abc import Sequence + +from sqlalchemy.orm import Session + +from core.app.file_access import ( + get_current_file_access_scope, + grant_retriever_segment_access, + grant_upload_file_access, +) +from models.dataset import DocumentSegment +from models.model import UploadFile +from repositories.knowledge.dataset_read_repository import get_segment_attachment_files, is_retrieved_segment_owned + + +def authorize_retrieved_segment( + segment: DocumentSegment, *, tenant_id: str, dataset_ids: Sequence[str], session: Session +) -> Sequence[UploadFile] | None: + """Authorize the owner chain and attachments before content links are signed. + + A valid segment with no attachments returns an empty sequence; an invalid + retrieval result returns None and must not be included in the response. + """ + scope = get_current_file_access_scope() + if scope is not None and scope.tenant_id != tenant_id: + return None + if not is_retrieved_segment_owned(segment, tenant_id=tenant_id, dataset_ids=dataset_ids, session=session): + return None + attachments = get_segment_attachment_files(segment, session=session) + grant_retriever_segment_access([segment.id]) + grant_upload_file_access(file.id for file in attachments) + return attachments diff --git a/api/services/knowledge/segments/__init__.py b/api/services/knowledge/segments/__init__.py new file mode 100644 index 00000000000..45c7098779b --- /dev/null +++ b/api/services/knowledge/segments/__init__.py @@ -0,0 +1 @@ +"""Dataset segment use cases.""" diff --git a/api/services/knowledge/segments/adapters.py b/api/services/knowledge/segments/adapters.py new file mode 100644 index 00000000000..ae501e3b8d3 --- /dev/null +++ b/api/services/knowledge/segments/adapters.py @@ -0,0 +1,82 @@ +"""Infrastructure adapters for dataset segment use cases.""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Any, Protocol + +from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError +from core.model_manager import ModelManager +from graphon.model_runtime.entities.model_entities import ModelType +from services.knowledge.segments.application import ( + SegmentDatasetRecord, + SegmentModelProviderError, +) + + +class RedisSegmentClient(Protocol): + def get(self, name: str | bytes) -> Any: ... + + def setnx(self, name: str | bytes, value: Any) -> Any: ... + def setex(self, name: str, time: int, value: Any) -> Any: ... + def lock(self, name: str, timeout: int) -> Any: ... + + +class ModelManagerSegmentGuard: + def check(self, dataset: SegmentDatasetRecord) -> None: + try: + ModelManager.for_tenant(tenant_id=dataset.workspace_id).get_model_instance( + tenant_id=dataset.workspace_id, + provider=dataset.embedding_model_provider or "", + model_type=ModelType.TEXT_EMBEDDING, + model=dataset.embedding_model or "", + ) + except LLMBadRequestError as error: + raise SegmentModelProviderError(kind="bad_request", description=str(error)) from error + except ProviderTokenNotInitError as error: + raise SegmentModelProviderError(kind="token", description=error.description) from error + + +class RedisSegmentIndexingState: + def __init__(self, redis: RedisSegmentClient) -> None: + self._redis = redis + + def lock(self, name: str, *, timeout: int): + return self._redis.lock(name, timeout=timeout) + + def is_segment_indexing(self, segment_id: str, *, deleting: bool = False) -> bool: + suffix = "delete_indexing" if deleting else "indexing" + return self._redis.get(f"segment_{segment_id}_{suffix}") is not None + + def mark_segment_indexing(self, segment_id: str, *, deleting: bool = False) -> None: + suffix = "delete_indexing" if deleting else "indexing" + self._redis.setex(f"segment_{segment_id}_{suffix}", 600, 1) + + def is_document_indexing(self, document_id: str) -> bool: + return self._redis.get(f"document_{document_id}_indexing") is not None + + def set_batch_waiting(self, job_id: str) -> None: + self._redis.setnx(f"segment_batch_import_{job_id}", "waiting") + + def get_batch_status(self, job_id: str) -> str | None: + value = self._redis.get(f"segment_batch_import_{job_id}") + if value is None: + return None + return value.decode() if isinstance(value, bytes) else value + + +class CelerySegmentBatchImportDispatcher: + def __init__(self, *, delay: Callable[..., object]) -> None: + self._delay = delay + + def dispatch( + self, + *, + job_id: str, + upload_file_id: str, + dataset_id: str, + document_id: str, + workspace_id: str, + actor_id: str, + ) -> None: + self._delay(job_id, upload_file_id, dataset_id, document_id, workspace_id, actor_id) diff --git a/api/services/knowledge/segments/application.py b/api/services/knowledge/segments/application.py new file mode 100644 index 00000000000..3600468c16c --- /dev/null +++ b/api/services/knowledge/segments/application.py @@ -0,0 +1,958 @@ +"""Application service for dataset segment and child-chunk HTTP use cases.""" + +from __future__ import annotations + +import logging +from collections import Counter +from collections.abc import Callable, Mapping, Sequence +from contextlib import AbstractContextManager +from dataclasses import dataclass, replace +from typing import Literal, Protocol +from uuid import uuid4 + +from libs.datetime_utils import naive_utc_now +from machinery.context import RequestContext +from services.knowledge.dataset_access import DatasetAccess, DatasetAccessDeniedError, DatasetAccessSnapshot +from services.knowledge.entities.segments import ( + ChildChunkRecord, + ChildChunkUpdateArgs, + SegmentRecord, + SegmentUpdateArgs, +) +from services.knowledge.resource_scope import DatasetRef, DocumentRef, SegmentRef + +_HIGH_QUALITY = "high_quality" +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True, slots=True) +class SegmentDatasetRecord: + id: str + workspace_id: str + indexing_technique: str + embedding_model_provider: str | None + embedding_model: str | None + + +@dataclass(frozen=True, slots=True) +class SegmentDocumentRecord: + id: str + dataset_id: str + workspace_id: str + doc_form: str + + @property + def ref(self) -> DocumentRef: + return DatasetRef(self.workspace_id, self.dataset_id).document(self.id) + + +@dataclass(frozen=True, slots=True) +class SegmentScope: + dataset: SegmentDatasetRecord | None + document: SegmentDocumentRecord | None + access_snapshot: DatasetAccessSnapshot | None = None + + +@dataclass(frozen=True, slots=True) +class SegmentListFilter: + page: int = 1 + limit: int = 20 + statuses: tuple[str, ...] = () + hit_count_gte: int | None = None + enabled: str = "all" + keyword: str | None = None + + +@dataclass(frozen=True, slots=True) +class ChildChunkListFilter: + page: int = 1 + limit: int = 20 + keyword: str | None = None + + +@dataclass(frozen=True, slots=True) +class SegmentPage: + items: tuple[SegmentRecord, ...] + total: int + total_pages: int + page: int + limit: int + + +@dataclass(frozen=True, slots=True) +class SegmentDetail: + data: SegmentRecord + doc_form: str + + +@dataclass(frozen=True, slots=True) +class ChildChunkPage: + items: tuple[ChildChunkRecord, ...] + total: int + total_pages: int + page: int + limit: int + + +@dataclass(frozen=True, slots=True) +class SegmentBatchImport: + job_id: str + job_status: str + + +class SegmentScopeReader(Protocol): + def get_segment_scope( + self, + *, + dataset_ref: DatasetRef, + document_id: str, + actor_id: str | None, + ) -> SegmentScope: ... + + +@dataclass(frozen=True, slots=True) +class ChildChunkState: + data: ChildChunkRecord + index_node_id: str | None + index_node_hash: str | None + created_by: str + updated_by: str | None = None + + +@dataclass(frozen=True, slots=True) +class SegmentIndexTarget: + id: str + index_node_id: str | None + enabled: bool + + +@dataclass(frozen=True, slots=True) +class SegmentUpdateState: + """Persisted fields needed to decide a segment update, without presentation data.""" + + content: str + enabled: bool + keywords: tuple[str, ...] | None + + +class SegmentStore(Protocol): + """Owner-scoped persistence; each operation owns a short session.""" + + def list_segments(self, document_ref: DocumentRef, query: SegmentListFilter) -> SegmentPage: ... + def get_segments(self, document_ref: DocumentRef, segment_ids: Sequence[str]) -> tuple[SegmentIndexTarget, ...]: ... + def get_segment(self, segment_ref: SegmentRef) -> SegmentDetail | None: ... + def get_segment_update_state(self, segment_ref: SegmentRef) -> SegmentUpdateState | None: ... + def save_segment(self, segment_ref: SegmentRef, values: Mapping[str, object], *, create: bool = False) -> None: ... + def delete_segments(self, document_ref: DocumentRef, segment_ids: Sequence[str]) -> None: ... + def list_child_chunks(self, segment_ref: SegmentRef, query: ChildChunkListFilter) -> ChildChunkPage | None: ... + def get_children_for_segments( + self, document_ref: DocumentRef, segment_ids: Sequence[str] + ) -> tuple[ChildChunkState, ...]: ... + def set_segments( + self, document_ref: DocumentRef, segment_ids: Sequence[str], values: Mapping[str, object] + ) -> None: ... + def get_children(self, segment_ref: SegmentRef) -> tuple[ChildChunkState, ...] | None: ... + def save_children( + self, + segment_ref: SegmentRef, + *, + added: Sequence[ChildChunkState] = (), + updated: Sequence[ChildChunkState] = (), + deleted: Sequence[ChildChunkState] = (), + ) -> None: ... + + +class SegmentIndex(Protocol): + def count_tokens(self, dataset: SegmentDatasetRecord, text: str) -> int: ... + def create_many( + self, + document_ref: DocumentRef, + *, + segment_ids: Sequence[str], + keywords_list: Sequence[Sequence[str] | None] | None, + ) -> None: ... + def create( + self, segment_ref: SegmentRef, *, keywords: Sequence[str] | None, attachment_ids: Sequence[str] + ) -> None: ... + def update(self, segment_ref: SegmentRef, *, keywords: Sequence[str] | None, regenerate_children: bool) -> None: ... + def update_attachments(self, segment_ref: SegmentRef, attachment_ids: Sequence[str]) -> None: ... + def update_summary(self, segment_ref: SegmentRef, *, summary: str | None, content_changed: bool) -> None: ... + def update_children( + self, + segment_ref: SegmentRef, + *, + added: Sequence[ChildChunkState] = (), + updated: Sequence[ChildChunkState] = (), + deleted: Sequence[ChildChunkState] = (), + ) -> None: ... + def delete( + self, document_ref: DocumentRef, segments: Sequence[SegmentIndexTarget], children: Sequence[ChildChunkState] + ) -> None: ... + def change_status( + self, document_ref: DocumentRef, segment_ids: Sequence[str], action: Literal["enable", "disable"] + ) -> None: ... + + +class SegmentLimits(Protocol): + SINGLE_CHUNK_ATTACHMENT_LIMIT: int + + +class SegmentModelGuard(Protocol): + def check(self, dataset: SegmentDatasetRecord) -> None: ... + + +class SegmentUploadCatalog(Protocol): + def get_file_name(self, *, workspace_id: str, upload_file_id: str) -> str | None: ... + + +class SegmentIndexingState(Protocol): + def lock(self, name: str, *, timeout: int) -> AbstractContextManager[object]: ... + def is_segment_indexing(self, segment_id: str, *, deleting: bool = False) -> bool: ... + def mark_segment_indexing(self, segment_id: str, *, deleting: bool = False) -> None: ... + + def is_document_indexing(self, document_id: str) -> bool: ... + + def set_batch_waiting(self, job_id: str) -> None: ... + + def get_batch_status(self, job_id: str) -> str | None: ... + + +class SegmentBatchImportDispatcher(Protocol): + def dispatch( + self, + *, + job_id: str, + upload_file_id: str, + dataset_id: str, + document_id: str, + workspace_id: str, + actor_id: str, + ) -> None: ... + + +class SegmentApplicationError(Exception): + """Base class for framework-neutral segment use-case errors.""" + + +class SegmentDatasetNotFoundError(SegmentApplicationError): + pass + + +class SegmentDocumentNotFoundError(SegmentApplicationError): + pass + + +class SegmentNotFoundError(SegmentApplicationError): + pass + + +class ChildChunkNotFoundError(SegmentApplicationError): + pass + + +class SegmentUploadFileNotFoundError(SegmentApplicationError): + pass + + +class SegmentPermissionDeniedError(SegmentApplicationError): + pass + + +class SegmentInvalidFileTypeError(SegmentApplicationError): + pass + + +class SegmentDocumentIndexingError(SegmentApplicationError): + pass + + +class SegmentStatusUpdateError(SegmentApplicationError): + pass + + +class SegmentBatchImportDispatchError(SegmentApplicationError): + pass + + +class SegmentBatchImportNotFoundError(SegmentApplicationError): + pass + + +class SegmentModelProviderError(Exception): + """Adapter error raised when the configured embedding model cannot be loaded.""" + + def __init__(self, *, kind: Literal["bad_request", "token"], description: str) -> None: + super().__init__(description) + self.kind = kind + self.description = description + + +class SegmentDatasetModelUnavailableError(SegmentApplicationError): + pass + + +class SegmentEmbeddingModelUnavailableError(SegmentApplicationError): + pass + + +class ChildChunkIndexingApplicationError(SegmentApplicationError): + pass + + +class ChildChunkDeleteIndexApplicationError(SegmentApplicationError): + pass + + +class DatasetSegmentApplicationService: + def __init__( + self, + *, + dataset_access: DatasetAccess, + scopes: SegmentScopeReader, + store: SegmentStore, + index: SegmentIndex, + limits: SegmentLimits, + text_hash: Callable[[str], str], + uploads: SegmentUploadCatalog, + model_guard: SegmentModelGuard, + indexing_state: SegmentIndexingState, + batch_dispatcher: SegmentBatchImportDispatcher, + job_id_factory: Callable[[], str], + ) -> None: + self._dataset_access = dataset_access + self._scopes = scopes + self._store = store + self._index = index + self._limits = limits + self._uploads = uploads + self._model_guard = model_guard + self._indexing_state = indexing_state + self._batch_dispatcher = batch_dispatcher + self._job_id_factory = job_id_factory + self.mutations = SegmentMutationService( + store=store, index=index, text_hash=text_hash, indexing_state=indexing_state + ) + + def list_segments( + self, + context: RequestContext, + *, + dataset_id: str, + document_id: str, + query: SegmentListFilter, + ) -> SegmentPage: + scope = self._require_scope(context, dataset_id=dataset_id, document_id=document_id) + return self._store.list_segments(scope.document.ref, query) + + def delete_segments( + self, + context: RequestContext, + *, + dataset_id: str, + document_id: str, + segment_ids: Sequence[str], + ) -> None: + scope = self._require_scope(context, dataset_id=dataset_id, document_id=document_id) + self._check_dataset_model(scope.dataset) + segments = self._store.get_segments(scope.document.ref, segment_ids) + children = self._store.get_children_for_segments(scope.document.ref, [segment.id for segment in segments]) + if segments: + self._index.delete(scope.document.ref, segments, children) + self._store.delete_segments(scope.document.ref, [segment.id for segment in segments]) + + def change_segment_status( + self, + context: RequestContext, + *, + dataset_id: str, + document_id: str, + segment_ids: Sequence[str], + action: Literal["enable", "disable"], + ) -> None: + scope = self._require_scope(context, dataset_id=dataset_id, document_id=document_id) + self._check_dataset_model(scope.dataset) + self._check_embedding_model(scope.dataset) + if self._indexing_state.is_document_indexing(scope.document.id): + raise SegmentDocumentIndexingError("Document is being indexed, please try again later") + try: + segments = self._store.get_segments(scope.document.ref, segment_ids) + ids = [] + for segment in segments: + if segment.enabled == (action == "enable") or self._indexing_state.is_segment_indexing(segment.id): + continue + ids.append(segment.id) + if ids: + self._store.set_segments( + scope.document.ref, + ids, + { + "enabled": action == "enable", + "disabled_at": naive_utc_now() if action == "disable" else None, + "disabled_by": context.account_id if action == "disable" else None, + }, + ) + self._index.change_status(scope.document.ref, ids, action) + except Exception as error: + raise SegmentStatusUpdateError(str(error)) from error + + def create_segment( + self, + context: RequestContext, + *, + dataset_id: str, + document_id: str, + values: Mapping[str, object], + ) -> SegmentDetail: + scope = self._require_scope(context, dataset_id=dataset_id, document_id=document_id) + self._check_embedding_model(scope.dataset) + args = self._validate_values(values, scope.document.doc_form) + segment_ref = self.mutations.create_segment(scope, actor_id=context.account_id, args=args) + return self._detail(segment_ref) + + def update_segment( + self, + context: RequestContext, + *, + dataset_id: str, + document_id: str, + segment_id: str, + values: Mapping[str, object], + ) -> SegmentDetail: + scope = self._require_scope(context, dataset_id=dataset_id, document_id=document_id) + self._check_dataset_model(scope.dataset) + self._check_embedding_model(scope.dataset) + args = self._validate_values(values, scope.document.doc_form) + self.mutations.update_segment(scope, actor_id=context.account_id, segment_id=segment_id, args=args) + return self._detail(scope.document.ref.segment(segment_id)) + + def _detail(self, segment_ref: SegmentRef) -> SegmentDetail: + result = self._store.get_segment(segment_ref) + if result is None: + raise SegmentNotFoundError("Segment not found") + return result + + def delete_segment( + self, + context: RequestContext, + *, + dataset_id: str, + document_id: str, + segment_id: str, + ) -> None: + scope = self._require_scope(context, dataset_id=dataset_id, document_id=document_id) + self._check_dataset_model(scope.dataset) + return self.mutations.delete_segment(scope.document.ref.segment(segment_id)) + + def start_batch_import( + self, + context: RequestContext, + *, + dataset_id: str, + document_id: str, + upload_file_id: str, + ) -> SegmentBatchImport: + scope = self._require_scope( + context, + dataset_id=dataset_id, + document_id=document_id, + enforce_permission=False, + ) + filename = self._uploads.get_file_name( + workspace_id=scope.dataset.workspace_id, + upload_file_id=upload_file_id, + ) + if filename is None: + raise SegmentUploadFileNotFoundError("UploadFile not found") + if not filename.lower().endswith(".csv"): + raise SegmentInvalidFileTypeError("Invalid file type. Only CSV files are allowed") + + job_id = self._job_id_factory() + try: + self._indexing_state.set_batch_waiting(job_id) + self._batch_dispatcher.dispatch( + job_id=job_id, + upload_file_id=upload_file_id, + dataset_id=scope.dataset.id, + document_id=scope.document.id, + workspace_id=scope.dataset.workspace_id, + actor_id=context.account_id, + ) + except Exception as error: + raise SegmentBatchImportDispatchError(str(error)) from error + return SegmentBatchImport(job_id=job_id, job_status="waiting") + + def get_batch_import_status(self, job_id: str) -> SegmentBatchImport: + status = self._indexing_state.get_batch_status(job_id) + if status is None: + raise SegmentBatchImportNotFoundError("The job does not exist.") + return SegmentBatchImport(job_id=job_id, job_status=status) + + def create_child_chunk( + self, + context: RequestContext, + *, + dataset_id: str, + document_id: str, + segment_id: str, + content: str, + ) -> ChildChunkRecord: + scope = self._require_scope(context, dataset_id=dataset_id, document_id=document_id) + self._check_embedding_model(scope.dataset) + return self.mutations.create_child_chunk( + scope.document.ref.segment(segment_id), content=content, actor_id=context.account_id + ) + + def list_child_chunks( + self, + *, + workspace_id: str, + dataset_id: str, + document_id: str, + segment_id: str, + query: ChildChunkListFilter, + ) -> ChildChunkPage: + scope = self._require_scope_without_actor( + workspace_id=workspace_id, + dataset_id=dataset_id, + document_id=document_id, + ) + self._check_dataset_model(scope.dataset) + result = self._store.list_child_chunks(scope.document.ref.segment(segment_id), query) + if result is None: + raise SegmentNotFoundError("Segment not found") + return result + + def update_child_chunks( + self, + context: RequestContext, + *, + dataset_id: str, + document_id: str, + segment_id: str, + chunks: Sequence[ChildChunkUpdateArgs], + ) -> tuple[ChildChunkRecord, ...]: + scope = self._require_scope(context, dataset_id=dataset_id, document_id=document_id) + self._check_dataset_model(scope.dataset) + return self.mutations.update_child_chunks( + scope.document.ref.segment(segment_id), chunks=chunks, actor_id=context.account_id + ) + + def delete_child_chunk( + self, + context: RequestContext, + *, + dataset_id: str, + document_id: str, + segment_id: str, + child_chunk_id: str, + ) -> None: + scope = self._require_scope(context, dataset_id=dataset_id, document_id=document_id) + self._check_dataset_model(scope.dataset) + return self.mutations.delete_child_chunk(scope.document.ref.segment(segment_id), child_chunk_id=child_chunk_id) + + def update_child_chunk( + self, + context: RequestContext, + *, + dataset_id: str, + document_id: str, + segment_id: str, + child_chunk_id: str, + content: str, + ) -> ChildChunkRecord: + scope = self._require_scope(context, dataset_id=dataset_id, document_id=document_id) + self._check_dataset_model(scope.dataset) + return self.mutations.update_child_chunk( + scope.document.ref.segment(segment_id), + child_chunk_id=child_chunk_id, + content=content, + actor_id=context.account_id, + ) + + def _validate_values(self, values: Mapping[str, object], doc_form: str) -> SegmentUpdateArgs: + return validate_segment_values( + values, doc_form=doc_form, attachment_limit=self._limits.SINGLE_CHUNK_ATTACHMENT_LIMIT + ) + + def _require_scope( + self, + context: RequestContext, + *, + dataset_id: str, + document_id: str, + enforce_permission: bool = True, + ) -> SegmentMutationScope: + scope = self._scopes.get_segment_scope( + dataset_ref=DatasetRef(context.active_workspace_id, dataset_id), + document_id=document_id, + actor_id=context.account_id if enforce_permission else None, + ) + dataset = scope.dataset + if dataset is None: + raise SegmentDatasetNotFoundError("Dataset not found") + if enforce_permission: + if scope.access_snapshot is None: + raise SegmentPermissionDeniedError("Dataset access state is missing") + try: + self._dataset_access.check_access(context, scope.access_snapshot) + except DatasetAccessDeniedError as error: + raise SegmentPermissionDeniedError(str(error)) from error + document = scope.document + if document is None: + raise SegmentDocumentNotFoundError("Document not found") + return SegmentMutationScope(dataset=dataset, document=document) + + def _require_scope_without_actor( + self, + *, + workspace_id: str, + dataset_id: str, + document_id: str, + ) -> SegmentMutationScope: + scope = self._scopes.get_segment_scope( + dataset_ref=DatasetRef(workspace_id, dataset_id), + document_id=document_id, + actor_id=None, + ) + if scope.dataset is None: + raise SegmentDatasetNotFoundError("Dataset not found") + if scope.document is None: + raise SegmentDocumentNotFoundError("Document not found") + return SegmentMutationScope(dataset=scope.dataset, document=scope.document) + + def _check_dataset_model(self, dataset: SegmentDatasetRecord) -> None: + if dataset.indexing_technique != _HIGH_QUALITY: + return + try: + self._model_guard.check(dataset) + except SegmentModelProviderError as error: + if error.kind == "token": + message = f"The dataset is unavailable, due to: {error.description}" + else: + message = ( + "No Embedding Model available. Please configure a valid provider in the Settings -> Model Provider." + ) + raise SegmentDatasetModelUnavailableError(message) from error + + def _check_embedding_model(self, dataset: SegmentDatasetRecord) -> None: + if dataset.indexing_technique != _HIGH_QUALITY: + return + try: + self._model_guard.check(dataset) + except SegmentModelProviderError as error: + if error.kind == "bad_request": + message = ( + "No Embedding Model available. Please configure a valid provider in the Settings -> Model Provider." + ) + else: + message = error.description + raise SegmentEmbeddingModelUnavailableError(message) from error + + +class SegmentMutationService: + """Shared writes for Console and Service API after their own admission checks. + + Callers supply an owner chain reconstructed from authorized database rows. + Persistence ports recheck that chain and own their transactions. Partial + update arguments are accepted here; transport validation stays at admission. + """ + + def __init__( + self, + *, + store: SegmentStore, + index: SegmentIndex, + text_hash: Callable[[str], str], + indexing_state: SegmentIndexingState, + ) -> None: + self._store = store + self._index = index + self._text_hash = text_hash + self._indexing_state = indexing_state + + def index_segments( + self, + document_ref: DocumentRef, + *, + segment_ids: Sequence[str], + keywords_list: Sequence[Sequence[str] | None] | None = None, + ) -> None: + """Index a committed batch and persist failures through the scoped store. + + Callers must finish their write transaction before invoking this method. + The index adapter reloads the complete owner chain before external I/O. + """ + if keywords_list is not None and len(keywords_list) != len(segment_ids): + raise ValueError("Keywords must correspond to the segment batch") + if not segment_ids: + return + try: + self._index.create_many(document_ref, segment_ids=segment_ids, keywords_list=keywords_list) + except Exception as error: + logger.exception("Batch indexing failed for document %s", document_ref.document_id) + self._store.set_segments( + document_ref, + segment_ids, + {"enabled": False, "disabled_at": naive_utc_now(), "status": "error", "error": str(error)}, + ) + raise + + def create_segment( + self, + scope: SegmentMutationScope, + *, + actor_id: str, + args: SegmentUpdateArgs, + ) -> SegmentRef: + now = naive_utc_now() + segment_ref = scope.document.ref.segment(str(uuid4())) + saved = False + try: + with self._indexing_state.lock(f"add_segment_lock_document_id_{scope.document.id}", timeout=600): + self._store.save_segment( + segment_ref, + { + "content": args.content, + "answer": args.answer if scope.document.doc_form == "qa_model" else None, + "word_count": len(args.content or "") + + (len(args.answer or "") if scope.document.doc_form == "qa_model" else 0), + "tokens": self._index.count_tokens(scope.dataset, args.content or ""), + "index_node_id": str(uuid4()), + "index_node_hash": self._text_hash(args.content or ""), + "created_by": actor_id, + "status": "completed", + "indexing_at": now, + "completed_at": now, + }, + create=True, + ) + saved = True + except Exception as error: + if saved: + self._mark_error(segment_ref, error) + raise + try: + self._index.create(segment_ref, keywords=args.keywords, attachment_ids=args.attachment_ids or ()) + except Exception as error: + self._mark_error(segment_ref, error) + return segment_ref + + def update_segment( + self, + scope: SegmentMutationScope, + *, + actor_id: str, + segment_id: str, + args: SegmentUpdateArgs, + ) -> None: + segment_ref = scope.document.ref.segment(segment_id) + previous = self._store.get_segment_update_state(segment_ref) + if previous is None: + raise SegmentNotFoundError("Segment not found") + if self._indexing_state.is_segment_indexing(segment_id): + raise ValueError("Segment is indexing, please try again later") + if args.enabled is False and previous.enabled: + self._store.save_segment( + segment_ref, {"enabled": False, "disabled_at": naive_utc_now(), "disabled_by": actor_id} + ) + self._indexing_state.mark_segment_indexing(segment_id) + self._index.change_status(scope.document.ref, [segment_id], "disable") + return + if not previous.enabled and args.enabled is not True: + raise ValueError("Can't update disabled segment") + content = args.content or previous.content + content_changed = content != previous.content + keyword_changed = bool(args.keywords) and Counter(previous.keywords or ()) != Counter(args.keywords or ()) + changes: dict[str, object] = { + "content": content, + "word_count": len(content) + (len(args.answer or "") if scope.document.doc_form == "qa_model" else 0), + "enabled": True, + "disabled_at": None, + "disabled_by": None, + } + if scope.document.doc_form == "qa_model": + changes["answer"] = args.answer + if keyword_changed and not content_changed: + changes["keywords"] = args.keywords + try: + if content_changed: + now = naive_utc_now() + changes.update( + index_node_hash=self._text_hash(content), + tokens=self._index.count_tokens( + scope.dataset, + content + (args.answer or "") if scope.document.doc_form == "qa_model" else content, + ), + status="completed", + indexing_at=now, + completed_at=now, + updated_at=now, + updated_by=actor_id, + ) + self._store.save_segment(segment_ref, changes) + if scope.document.doc_form == "hierarchical_model": + if args.regenerate_child_chunks: + self._index.update(segment_ref, keywords=args.keywords, regenerate_children=True) + elif content_changed or args.enabled or keyword_changed: + self._index.update(segment_ref, keywords=args.keywords, regenerate_children=False) + if scope.dataset.indexing_technique == _HIGH_QUALITY: + try: + self._index.update_summary(segment_ref, summary=args.summary, content_changed=content_changed) + except Exception: + logger.exception("Failed to update summary for segment %s", segment_id) + if args.attachment_ids is not None: + self._index.update_attachments(segment_ref, args.attachment_ids) + except Exception as error: + self._mark_error(segment_ref, error) + + def delete_segment(self, segment_ref: SegmentRef) -> None: + segments = self._store.get_segments(segment_ref.document, (segment_ref.segment_id,)) + if not segments: + raise SegmentNotFoundError("Segment not found") + segment = segments[0] + if self._indexing_state.is_segment_indexing(segment_ref.segment_id, deleting=True): + raise ValueError("Segment is deleting.") + if segment.enabled: + self._indexing_state.mark_segment_indexing(segment_ref.segment_id, deleting=True) + self._index.delete( + segment_ref.document, + (segment,), + self._children(segment_ref), + ) + self._store.delete_segments(segment_ref.document, (segment_ref.segment_id,)) + + def create_child_chunk(self, segment_ref: SegmentRef, *, content: str, actor_id: str) -> ChildChunkRecord: + with self._indexing_state.lock(f"add_child_lock_{segment_ref.segment_id}", timeout=20): + children = self._children(segment_ref) + child = self._new_child( + segment_ref, content, max((c.data.position for c in children), default=0) + 1, actor_id + ) + self._write_children(segment_ref, added=(child,)) + return child.data + + def update_child_chunks( + self, segment_ref: SegmentRef, *, chunks: Sequence[ChildChunkUpdateArgs], actor_id: str + ) -> tuple[ChildChunkRecord, ...]: + existing = self._children(segment_ref) + remaining = {child.data.id: child for child in existing} + added: list[ChildChunkState] = [] + updated: list[ChildChunkState] = [] + for chunk in chunks: + if chunk.id: + previous = remaining.pop(chunk.id, None) + if previous is not None and previous.data.content != chunk.content: + updated.append(self._updated_child(previous, chunk.content, actor_id)) + else: + added.append(self._new_child(segment_ref, chunk.content, len(existing) + len(added) + 1, actor_id)) + self._write_children(segment_ref, added=added, updated=updated, deleted=tuple(remaining.values())) + return tuple(child.data for child in sorted(added + updated, key=lambda child: child.data.position)) + + def delete_child_chunk(self, segment_ref: SegmentRef, *, child_chunk_id: str) -> None: + child = self._child(segment_ref, child_chunk_id) + try: + self._index.update_children(segment_ref, deleted=(child,)) + except Exception as error: + raise ChildChunkDeleteIndexApplicationError(str(error)) from error + self._store.save_children(segment_ref, deleted=(child,)) + + def update_child_chunk( + self, segment_ref: SegmentRef, *, child_chunk_id: str, content: str, actor_id: str + ) -> ChildChunkRecord: + child = self._updated_child(self._child(segment_ref, child_chunk_id), content, actor_id) + self._write_children(segment_ref, updated=(child,)) + return child.data + + def _mark_error(self, segment_ref: SegmentRef, error: Exception) -> None: + logger.exception("Segment indexing failed: %s", segment_ref.segment_id) + self._store.save_segment( + segment_ref, {"enabled": False, "disabled_at": naive_utc_now(), "status": "error", "error": str(error)} + ) + + def _children(self, segment_ref: SegmentRef) -> tuple[ChildChunkState, ...]: + children = self._store.get_children(segment_ref) + if children is None: + raise SegmentNotFoundError("Segment not found") + return children + + def _child(self, segment_ref: SegmentRef, child_id: str) -> ChildChunkState: + for child in self._children(segment_ref): + if child.data.id == child_id: + return child + raise ChildChunkNotFoundError("Child chunk not found") + + def _new_child(self, segment_ref: SegmentRef, content: str, position: int, actor_id: str) -> ChildChunkState: + now = naive_utc_now() + return ChildChunkState( + data=ChildChunkRecord( + id=str(uuid4()), + segment_id=segment_ref.segment_id, + content=content, + position=position, + word_count=len(content), + type="customized", + created_at=now, + updated_at=now, + ), + index_node_id=str(uuid4()), + index_node_hash=self._text_hash(content), + created_by=actor_id, + ) + + def _updated_child(self, child: ChildChunkState, content: str, actor_id: str) -> ChildChunkState: + return replace( + child, + data=child.data.model_copy( + update={ + "content": content, + "word_count": len(content), + "type": "customized", + "updated_at": naive_utc_now(), + } + ), + updated_by=actor_id, + ) + + def _write_children( + self, + segment_ref: SegmentRef, + *, + added: Sequence[ChildChunkState] = (), + updated: Sequence[ChildChunkState] = (), + deleted: Sequence[ChildChunkState] = (), + ) -> None: + try: + self._index.update_children(segment_ref, added=added, updated=updated, deleted=deleted) + self._store.save_children(segment_ref, added=added, updated=updated, deleted=deleted) + except Exception as error: + raise ChildChunkIndexingApplicationError(str(error)) from error + + +@dataclass(frozen=True, slots=True) +class SegmentMutationScope: + """Authorized dataset and document values for shared mutation policy.""" + + dataset: SegmentDatasetRecord + document: SegmentDocumentRecord + + def __post_init__(self) -> None: + if self.document.dataset_id != self.dataset.id or self.document.workspace_id != self.dataset.workspace_id: + raise ValueError("Document does not belong to the dataset") + + +def validate_segment_values(values: Mapping[str, object], *, doc_form: str, attachment_limit: int) -> SegmentUpdateArgs: + answer = values.get("answer") + if doc_form == "qa_model": + if not answer: + raise ValueError("Answer is required") + if not isinstance(answer, str) or not answer.strip(): + raise ValueError("Answer is empty") + content = values.get("content") + if not isinstance(content, str) or not content.strip(): + raise ValueError("Content is empty") + attachments = values.get("attachment_ids") + if attachments: + if not isinstance(attachments, list): + raise ValueError("Attachment IDs is invalid") + if len(attachments) > attachment_limit: + raise ValueError(f"Exceeded maximum attachment limit of {attachment_limit}") + return SegmentUpdateArgs.model_validate(values) diff --git a/api/services/knowledge/segments/indexing.py b/api/services/knowledge/segments/indexing.py new file mode 100644 index 00000000000..462e19f5ece --- /dev/null +++ b/api/services/knowledge/segments/indexing.py @@ -0,0 +1,356 @@ +"""Indexing adapters that materialize database inputs before external I/O.""" + +import logging +from collections.abc import Callable, Mapping, Sequence +from uuid import uuid4 + +from sqlalchemy.orm import Session + +from core.credit_usage import CreditUsageCreatedBy +from core.model_context import with_credit_usage_created_by +from core.model_manager import ModelManager +from core.rag.datasource.keyword.jieba.jieba import Jieba +from core.rag.datasource.vdb.vector_factory import Vector +from core.rag.entities import Rule +from core.rag.index_processor.constant.doc_type import DocType +from core.rag.index_processor.processor.paragraph_index_processor import ParagraphIndexProcessor +from core.rag.index_processor.processor.parent_child_index_processor import ParentChildIndexProcessor +from core.rag.models.document import Document +from graphon.model_runtime.entities.model_entities import ModelType +from libs.datetime_utils import naive_utc_now +from models.enums import SegmentType +from repositories.knowledge.segment_repository import SegmentIndexingSnapshot, SQLAlchemySegmentRepository +from repositories.knowledge.upload_file_repository import SQLAlchemyKnowledgeUploadRepository +from services.knowledge.entities.segments import ChildChunkRecord +from services.knowledge.resource_scope import DocumentRef, SegmentRef +from services.knowledge.segments.adapters import RedisSegmentClient +from services.knowledge.segments.application import ChildChunkState, SegmentDatasetRecord, SegmentIndexTarget +from services.knowledge.summaries.adapters import SummaryIndexAdapter + +logger = logging.getLogger(__name__) + + +class SegmentIndexingGateway: + def __init__( + self, + *, + segments: SQLAlchemySegmentRepository, + uploads: SQLAlchemyKnowledgeUploadRepository, + redis: RedisSegmentClient, + new_session: Callable[[], Session], + delete_task: Callable[..., object], + enable_task: Callable[..., object], + disable_task: Callable[..., object], + ) -> None: + self._segments = segments + self._uploads = uploads + self._redis = redis + self._new_session = new_session + self._delete_task = delete_task + self._enable_task = enable_task + self._disable_task = disable_task + + def count_tokens(self, dataset: SegmentDatasetRecord, text: str) -> int: + if dataset.indexing_technique != "high_quality": + return 0 + model = ModelManager.for_tenant(tenant_id=dataset.workspace_id).get_model_instance( + tenant_id=dataset.workspace_id, + provider=dataset.embedding_model_provider or "", + model_type=ModelType.TEXT_EMBEDDING, + model=dataset.embedding_model or "", + ) + return model.get_text_embedding_num_tokens(texts=[text])[0] + + @staticmethod + def _vector(snapshot: SegmentIndexingSnapshot) -> Vector: + return Vector(snapshot.dataset, session=None, vector_type=snapshot.vector_type) + + @staticmethod + def _document(snapshot: SegmentIndexingSnapshot) -> Document: + segment = snapshot.segment + return Document( + page_content=segment.content, + metadata={ + "doc_id": segment.index_node_id, + "doc_hash": segment.index_node_hash, + "document_id": segment.document_id, + "dataset_id": segment.dataset_id, + "doc_type": DocType.TEXT, + }, + ) + + @with_credit_usage_created_by(CreditUsageCreatedBy.KNOWLEDGE_INDEXING) + def create_many( + self, + document_ref: DocumentRef, + *, + segment_ids: Sequence[str], + keywords_list: Sequence[Sequence[str] | None] | None, + ) -> None: + """Materialize the whole scoped batch before indexing any of its records.""" + snapshots = self._segments.get_indexing_snapshots(document_ref, segment_ids) + if not snapshots: + return + first = snapshots[0] + documents = [self._document(snapshot) for snapshot in snapshots] + if first.document.doc_form == "hierarchical_model": + for snapshot in snapshots: + self._regenerate_children(document_ref.segment(snapshot.segment.id), snapshot, replace_existing=False) + elif first.dataset.indexing_technique == "high_quality": + self._vector(first).create(documents) + elif first.document.doc_form != "qa_model": + self._write_keywords(document_ref, first, documents, keywords_list, replace_existing=False) + if first.dataset.is_multimodal: + files = self._uploads.get_files( + workspace_id=document_ref.dataset.tenant_id, + file_ids=tuple(attachment_id for snapshot in snapshots for attachment_id in snapshot.attachment_ids), + ) + attachments = [ + document + for snapshot in snapshots + for document in self._attachment_documents( + snapshot, [files[file_id] for file_id in snapshot.attachment_ids if file_id in files] + ) + ] + if attachments and first.dataset.indexing_technique == "high_quality": + self._vector(first).create_multimodal(attachments, upload_files=files) + + @with_credit_usage_created_by(CreditUsageCreatedBy.KNOWLEDGE_INDEXING) + def create(self, segment_ref: SegmentRef, *, keywords: Sequence[str] | None, attachment_ids: Sequence[str]) -> None: + snapshot = self._segments.get_indexing_snapshot(segment_ref) + files = self._uploads.get_files(workspace_id=segment_ref.document.dataset.tenant_id, file_ids=attachment_ids) + self._segments.replace_attachments(segment_ref, tuple(files)) + if snapshot.document.doc_form == "hierarchical_model": + self._regenerate_children(segment_ref, snapshot, replace_existing=False) + elif snapshot.dataset.indexing_technique == "high_quality": + self._vector(snapshot).create([self._document(snapshot)]) + else: + self._update_keywords(segment_ref, snapshot, keywords, replace_existing=False) + if files and snapshot.dataset.is_multimodal: + self._vector(snapshot).create_multimodal( + self._attachment_documents(snapshot, files.values()), upload_files=files + ) + + @with_credit_usage_created_by(CreditUsageCreatedBy.KNOWLEDGE_INDEXING) + def update(self, segment_ref: SegmentRef, *, keywords: Sequence[str] | None, regenerate_children: bool) -> None: + snapshot = self._segments.get_indexing_snapshot(segment_ref) + if regenerate_children: + self._regenerate_children(segment_ref, snapshot, replace_existing=True) + elif snapshot.dataset.indexing_technique == "high_quality": + vector = self._vector(snapshot) + assert snapshot.segment.index_node_id is not None + vector.delete_by_ids([snapshot.segment.index_node_id]) + vector.add_texts([self._document(snapshot)], duplicate_check=True) + else: + self._update_keywords(segment_ref, snapshot, keywords, replace_existing=True) + + def _regenerate_children( + self, segment_ref: SegmentRef, snapshot: SegmentIndexingSnapshot, *, replace_existing: bool + ) -> None: + if snapshot.dataset.indexing_technique != "high_quality": + raise ValueError("The knowledge base index technique is not high quality!") + rule = snapshot.process_rule + if rule is None: + if replace_existing: + return + raise ValueError("No processing rule found.") + manager = ModelManager.for_tenant(tenant_id=snapshot.dataset.tenant_id) + model = ( + manager.get_model_instance( + tenant_id=snapshot.dataset.tenant_id, + provider=snapshot.dataset.embedding_model_provider, + model_type=ModelType.TEXT_EMBEDDING, + model=snapshot.dataset.embedding_model, + ) + if snapshot.dataset.embedding_model_provider + else manager.get_default_model_instance( + tenant_id=snapshot.dataset.tenant_id, model_type=ModelType.TEXT_EMBEDDING + ) + ) + children = ParentChildIndexProcessor().split_child_nodes( + self._document(snapshot), Rule.model_validate(rule.to_dict()["rules"]), rule.mode, model + ) + previous = self._segments.get_children(segment_ref) or () + vector = self._vector(snapshot) + if replace_existing and previous: + vector.delete_by_ids([child.index_node_id for child in previous if child.index_node_id]) + if children: + vector.create([Document.model_validate(child.model_dump()) for child in children]) + now = naive_utc_now() + added = tuple( + ChildChunkState( + data=ChildChunkRecord( + id=str(uuid4()), + segment_id=segment_ref.segment_id, + content=child.page_content, + position=position, + word_count=len(child.page_content), + type=SegmentType.AUTOMATIC, + created_at=now, + updated_at=now, + ), + index_node_id=child.metadata["doc_id"], + index_node_hash=child.metadata["doc_hash"], + created_by=snapshot.document.created_by, + ) + for position, child in enumerate(children, start=1) + ) + self._segments.save_children(segment_ref, added=added, deleted=previous if replace_existing else ()) + + @with_credit_usage_created_by(CreditUsageCreatedBy.KNOWLEDGE_INDEXING) + def update_children( + self, + segment_ref: SegmentRef, + *, + added: Sequence[ChildChunkState] = (), + updated: Sequence[ChildChunkState] = (), + deleted: Sequence[ChildChunkState] = (), + ) -> None: + snapshot = self._segments.get_indexing_snapshot(segment_ref) + if snapshot.dataset.indexing_technique != "high_quality": + return + documents = [ + Document( + page_content=child.data.content, + metadata={ + "doc_id": child.index_node_id, + "doc_hash": child.index_node_hash, + "document_id": segment_ref.document.document_id, + "dataset_id": segment_ref.document.dataset.dataset_id, + }, + ) + for child in (*added, *updated) + ] + vector = self._vector(snapshot) + deleted_ids = [] + for child in (*updated, *deleted): + assert child.index_node_id is not None + deleted_ids.append(child.index_node_id) + if deleted_ids: + vector.delete_by_ids(deleted_ids) + if documents: + vector.add_texts(documents, duplicate_check=True) + + @staticmethod + def _attachment_documents(snapshot: SegmentIndexingSnapshot, files) -> list[Document]: + return [ + Document( + page_content=file.name, + metadata={ + "doc_id": file.id, + "doc_hash": "", + "dataset_id": snapshot.dataset.id, + "document_id": snapshot.document.id, + "doc_type": DocType.IMAGE, + }, + ) + for file in files + ] + + @with_credit_usage_created_by(CreditUsageCreatedBy.KNOWLEDGE_INDEXING) + def update_attachments(self, segment_ref: SegmentRef, attachment_ids: Sequence[str]) -> None: + snapshot = self._segments.get_indexing_snapshot(segment_ref) + if snapshot.dataset.indexing_technique != "high_quality" or set(snapshot.attachment_ids) == set(attachment_ids): + return + files = self._uploads.get_files(workspace_id=snapshot.dataset.tenant_id, file_ids=attachment_ids) + if snapshot.dataset.is_multimodal: + vector = self._vector(snapshot) + if snapshot.attachment_ids: + vector.delete_by_ids(list(snapshot.attachment_ids)) + if files: + vector.create_multimodal(self._attachment_documents(snapshot, files.values()), upload_files=files) + self._segments.replace_attachments(segment_ref, tuple(files)) + + def update_summary(self, segment_ref: SegmentRef, *, summary: str | None, content_changed: bool) -> None: + if summary is None and not content_changed: + return + snapshot = self._segments.get_indexing_snapshot(segment_ref) + if snapshot.dataset.indexing_technique != "high_quality" or snapshot.document.doc_form == "qa_model": + return + setting = snapshot.dataset.summary_index_setting + previous = snapshot.summary + if summary is None or (previous is not None and summary == previous.summary_content): + if not (content_changed and previous and setting and setting.get("enable") is True): + return + try: + self._segments.save_summary(segment_ref, previous.summary_content or "") + summary, _ = ParagraphIndexProcessor.generate_summary_from_inputs( + snapshot.dataset.tenant_id, + snapshot.segment.content, + setting, + document_language=snapshot.document.doc_language, + image_loader=lambda: self._segments.get_summary_images(segment_ref), + ) + if not summary.strip(): + raise ValueError("Generated summary is empty") + except Exception as error: + self._segments.save_summary(segment_ref, previous.summary_content or "", error=str(error)) + raise + if not summary.strip(): + if previous and previous.summary_index_node_id: + try: + self._vector(snapshot).delete_by_ids([previous.summary_index_node_id]) + except Exception: + logger.exception("Failed to delete summary vector for %s", segment_ref.segment_id) + self._segments.save_summary(segment_ref, None) + return + record = self._segments.save_summary(segment_ref, summary) + assert record is not None + SummaryIndexAdapter.vectorize_summary( + record, + snapshot.segment, + snapshot.dataset, + new_session=self._new_session, + vector_factory=lambda: self._vector(snapshot), + ) + + def _update_keywords( + self, + segment_ref: SegmentRef, + snapshot: SegmentIndexingSnapshot, + keywords: Sequence[str] | None, + *, + replace_existing: bool, + ) -> None: + assert snapshot.segment.index_node_id is not None + self._write_keywords( + segment_ref.document, snapshot, [self._document(snapshot)], [keywords], replace_existing=replace_existing + ) + + def _write_keywords( + self, + document_ref: DocumentRef, + snapshot: SegmentIndexingSnapshot, + documents: list[Document], + keywords_list: Sequence[Sequence[str] | None] | None, + *, + replace_existing: bool, + ) -> None: + dataset_ref = document_ref.dataset + + def persist(storage_type: str, data: str, selected: Mapping[str, Sequence[str]]) -> None: + self._segments.save_keyword_table(dataset_ref, storage_type=storage_type, data=data, keywords=selected) + + Jieba(snapshot.dataset).update_texts( + documents, + read=lambda: self._segments.get_keyword_table(dataset_ref), + write=persist, + lock=self._redis.lock(f"keyword_indexing_lock_{snapshot.dataset.id}", timeout=600), + keywords_list=keywords_list, + replace_existing=replace_existing, + ) + + def delete( + self, document_ref: DocumentRef, segments: Sequence[SegmentIndexTarget], children: Sequence[ChildChunkState] + ) -> None: + self._delete_task( + [segment.index_node_id for segment in segments], + document_ref.dataset.dataset_id, + document_ref.document_id, + [segment.id for segment in segments], + [child.index_node_id for child in children if child.index_node_id], + ) + + def change_status(self, document_ref: DocumentRef, segment_ids: Sequence[str], action: str) -> None: + task = self._enable_task if action == "enable" else self._disable_task + task(list(segment_ids), document_ref.dataset.dataset_id, document_ref.document_id) diff --git a/api/services/knowledge/summaries/__init__.py b/api/services/knowledge/summaries/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/api/services/summary_index_service.py b/api/services/knowledge/summaries/adapters.py similarity index 63% rename from api/services/summary_index_service.py rename to api/services/knowledge/summaries/adapters.py index 2f5d832c0e1..f85bce12d77 100644 --- a/api/services/summary_index_service.py +++ b/api/services/knowledge/summaries/adapters.py @@ -1,13 +1,14 @@ -"""Summary index service for generating and managing document segment summaries.""" +"""Adapters for existing ORM-based summary callers and indexing infrastructure.""" import logging import time import uuid -from datetime import UTC, datetime +from collections.abc import Callable from typing import TypedDict, cast from sqlalchemy import select, update from sqlalchemy.orm import Session +from sqlalchemy.orm.attributes import set_committed_value from core.credit_usage import CreditUsageCreatedBy from core.db.session_factory import session_factory @@ -24,10 +25,21 @@ from libs import helper from models.dataset import Dataset, DocumentSegment, DocumentSegmentSummary from models.dataset import Document as DatasetDocument from models.enums import SummaryStatus +from repositories.knowledge.dataset_read_repository import get_segment_document +from repositories.knowledge.summary_repository import SQLAlchemySummaryRepository +from services.knowledge.resource_scope import DatasetRef +from services.knowledge.summaries.application import SummaryIndexService, SummaryVectorInput, SummaryVectorResult logger = logging.getLogger(__name__) +def _delete_summary_vectors(dataset: Dataset, node_ids: list[str], *, session: Session) -> None: + # Legacy callers own this transaction; committing here breaks session.begin(). + # Document cleanup tasks materialize their inputs before deleting external indexes. + vector_type = Vector.resolve_vector_type(dataset, session=session) + Vector(dataset, session=None, vector_type=vector_type).delete_by_ids(node_ids) + + class SummaryEntryDict(TypedDict): segment_id: str segment_position: int @@ -44,8 +56,8 @@ class DocumentSummaryStatusDetailDict(TypedDict): summaries: list[SummaryEntryDict] -class SummaryIndexService: - """Service for generating and managing summary indexes.""" +class SummaryIndexAdapter: + """Adapt existing indexing callers to summary generation, storage and vectors.""" @staticmethod @with_credit_usage_created_by(CreditUsageCreatedBy.KNOWLEDGE_INDEXING) @@ -65,7 +77,8 @@ class SummaryIndexService: summary_index_setting: Summary index configuration Keyword Args: - session: SQLAlchemy session used to load the segment's document. + session: Caller session, committed after loading generation inputs and + before invoking the model. Image reads use a separate bounded session. Returns: Tuple of (summary_content, llm_usage) where llm_usage is LLMUsage object @@ -73,24 +86,34 @@ class SummaryIndexService: Raises: ValueError: If summary_index_setting is invalid or generation fails """ - # Reuse the existing generate_summary method from ParagraphIndexProcessor - # Use lazy import to avoid circular import + # Lazy import avoids the index processor's circular dependency on this adapter. from core.rag.index_processor.processor.paragraph_index_processor import ParagraphIndexProcessor # Get document language to ensure summary is generated in the correct language # This is especially important for image-only chunks where text is empty or minimal document_language = None - document = segment.get_document(session=session) + document = get_segment_document(segment, session=session) if document and document.doc_language: document_language = document.doc_language - summary_content, usage = ParagraphIndexProcessor.generate_summary( - tenant_id=dataset.tenant_id, - text=segment.content, + tenant_id, text, segment_id = dataset.tenant_id, segment.content, segment.id + if segment.tenant_id != tenant_id or segment.dataset_id != dataset.id: + raise ValueError("Segment does not belong to the dataset") + session.commit() + + def load_images(): + with session_factory.create_session() as image_session: + images = ParagraphIndexProcessor._extract_images_from_segment_attachments( + tenant_id, segment_id, image_session + ) + return images or ParagraphIndexProcessor._extract_images_from_text(tenant_id, text, image_session) + + summary_content, usage = ParagraphIndexProcessor.generate_summary_from_inputs( + tenant_id=tenant_id, + text=text, summary_index_setting=summary_index_setting, - segment_id=segment.id, document_language=document_language, - session=session, + image_loader=load_images, ) if not summary_content: @@ -167,416 +190,79 @@ class SummaryIndexService: segment: DocumentSegment, dataset: Dataset, session: Session | None = None, + *, + new_session: Callable[[], Session] | None = None, + vector_factory: Callable[[], Vector] | None = None, ) -> None: - """ - Vectorize summary and store in vector database. + """Bridge already-authorized ORM callers to the detached summary use case. - Args: - summary_record: DocumentSegmentSummary record - segment: Original DocumentSegment - dataset: Dataset containing the segment - session: Optional SQLAlchemy session. If provided, uses this session instead of creating a new one. - If not provided, creates a new session and commits automatically. + Pending summary writes are committed before external I/O. Each outcome + attempt then owns a separate transaction, including retries and failures. """ if dataset.indexing_technique != IndexTechniqueType.HIGH_QUALITY: - logger.warning( - "Summary vectorization skipped for dataset %s: indexing_technique is not high_quality", - dataset.id, - ) return - - # Get summary_record_id for later session queries - summary_record_id = summary_record.id - # Save the original session parameter for use in error handling - original_session = session - - def create_vector() -> Vector: - if original_session is not None: - return Vector(dataset, session=original_session) - with session_factory.create_session() as vector_session: - return Vector(dataset, session=vector_session) - - logger.debug( - "Starting vectorization for segment %s, summary_record_id=%s, using_provided_session=%s", - segment.id, - summary_record_id, - original_session is not None, + ref = DatasetRef(dataset.tenant_id, dataset.id).document(segment.document_id).segment(segment.id) + if segment.dataset_id != dataset.id or segment.tenant_id != dataset.tenant_id: + raise ValueError("Segment does not belong to the dataset") + if ( + summary_record.dataset_id != dataset.id + or summary_record.document_id != segment.document_id + or summary_record.chunk_id != segment.id + ): + raise ValueError("Summary does not belong to the segment") + command = SummaryVectorInput( + ref, summary_record.id, summary_record.summary_content or "", summary_record.summary_index_node_id ) - - # Reuse existing index_node_id if available (like segment does), otherwise generate new one - old_summary_node_id = summary_record.summary_index_node_id - if old_summary_node_id: - # Reuse existing index_node_id (like segment behavior) - summary_index_node_id = old_summary_node_id - logger.debug("Reusing existing index_node_id %s for segment %s", summary_index_node_id, segment.id) - else: - # Generate new index node ID only for new summaries - summary_index_node_id = str(uuid.uuid4()) - logger.debug("Generated new index_node_id %s for segment %s", summary_index_node_id, segment.id) - - # Always regenerate hash (in case summary content changed) - summary_content = summary_record.summary_content - if not summary_content or not summary_content.strip(): + if not command.content.strip(): raise ValueError(f"Summary content is empty for segment {segment.id}, cannot vectorize") - summary_hash = helper.generate_text_hash(summary_content) - - # Delete old vector only if we're reusing the same index_node_id (to overwrite) - # If index_node_id changed, the old vector should have been deleted elsewhere - if old_summary_node_id and old_summary_node_id == summary_index_node_id: - try: - vector = create_vector() - vector.delete_by_ids([old_summary_node_id]) - except Exception as e: - logger.warning( - "Failed to delete old summary vector for segment %s: %s. Continuing with new vectorization.", - segment.id, - str(e), + if session is not None: + session.commit() + create_session = new_session or session_factory.create_session + with create_session() as read_session: + owned_dataset = read_session.scalar( + select(Dataset).where( + Dataset.id == ref.document.dataset.dataset_id, Dataset.tenant_id == ref.document.dataset.tenant_id ) - - # Calculate embedding tokens for summary (for logging and statistics) - embedding_tokens = 0 - try: - model_manager = ModelManager.for_tenant(tenant_id=dataset.tenant_id) - embedding_model = model_manager.get_model_instance( - tenant_id=dataset.tenant_id, - provider=dataset.embedding_model_provider, - model_type=ModelType.TEXT_EMBEDDING, - model=dataset.embedding_model, ) - if embedding_model: - tokens_list = embedding_model.get_text_embedding_num_tokens([summary_content]) - raw_embedding_tokens = tokens_list[0] if tokens_list else 0 - embedding_tokens = raw_embedding_tokens if isinstance(raw_embedding_tokens, int) else 0 - except Exception: - logger.warning("Failed to calculate embedding tokens for summary", exc_info=True) - - # Create document with summary content and metadata - summary_document = Document( - page_content=summary_content, - metadata={ - "doc_id": summary_index_node_id, - "doc_hash": summary_hash, - "dataset_id": dataset.id, - "document_id": segment.document_id, - "original_chunk_id": segment.id, # Key: link to original chunk - "doc_type": DocType.TEXT, - "is_summary": True, # Identifier for summary documents - }, + owned_segment = read_session.scalar( + select(DocumentSegment).where( + DocumentSegment.id == ref.segment_id, + DocumentSegment.tenant_id == ref.document.dataset.tenant_id, + DocumentSegment.dataset_id == ref.document.dataset.dataset_id, + DocumentSegment.document_id == ref.document.document_id, + ) + ) + owned_document = read_session.scalar( + select(DatasetDocument).where( + DatasetDocument.id == ref.document.document_id, + DatasetDocument.dataset_id == ref.document.dataset.dataset_id, + DatasetDocument.tenant_id == ref.document.dataset.tenant_id, + ) + ) + if owned_dataset is None or owned_segment is None or owned_document is None: + raise LookupError("Summary segment no longer exists in this dataset") + vector_type = ( + Vector.resolve_vector_type(owned_dataset, session=read_session) if vector_factory is None else None + ) + factory = vector_factory or (lambda: Vector(owned_dataset, session=None, vector_type=vector_type)) + service = SummaryIndexService( + store=SQLAlchemySummaryRepository(new_session=create_session), + vectors=SummaryVectorGateway(owned_dataset, factory), + new_id=lambda: str(uuid.uuid4()), + text_hash=helper.generate_text_hash, + sleep=time.sleep, ) - - # Vectorize and store with retry mechanism for connection errors - max_retries = 3 - retry_delay = 2.0 - - for attempt in range(max_retries): - try: - logger.debug( - "Attempting to vectorize summary for segment %s (attempt %s/%s)", - segment.id, - attempt + 1, - max_retries, - ) - vector = create_vector() - # Use duplicate_check=False to ensure re-vectorization even if old vector still exists - # The old vector should have been deleted above, but if deletion failed, - # we still want to re-vectorize (upsert will overwrite) - vector.add_texts([summary_document], duplicate_check=False) - logger.debug( - "Successfully added summary vector to database for segment %s (attempt %s/%s)", - segment.id, - attempt + 1, - max_retries, - ) - - # Log embedding token usage - if embedding_tokens > 0: - logger.info( - "Summary embedding for segment %s used %s tokens", - segment.id, - embedding_tokens, - ) - - # Success - update summary record with index node info - # Use provided session if available, otherwise create a new one - use_provided_session = session is not None - if not use_provided_session: - logger.debug("Creating new session for vectorization of segment %s", segment.id) - session_context = session_factory.create_session() - session = session_context.__enter__() - else: - logger.debug("Using provided session for vectorization of segment %s", segment.id) - session_context = None # Don't use context manager for provided session - - # At this point, session is guaranteed to be not None - # Type narrowing: session is definitely not None after the if/else above - if session is None: - raise RuntimeError("Session should not be None at this point") - - try: - # Declare summary_record_in_session variable - summary_record_in_session: DocumentSegmentSummary | None - - # If using provided session, merge the summary_record into it - if use_provided_session: - # Merge the summary_record into the provided session - logger.debug( - "Merging summary_record (id=%s) into provided session for segment %s", - summary_record_id, - segment.id, - ) - summary_record_in_session = session.merge(summary_record) - logger.debug( - "Successfully merged summary_record for segment %s, merged_id=%s", - segment.id, - summary_record_in_session.id, - ) - else: - # Query the summary record in the new session - logger.debug( - "Querying summary_record by id=%s for segment %s in new session", - summary_record_id, - segment.id, - ) - summary_record_in_session = session.scalar( - select(DocumentSegmentSummary) - .where(DocumentSegmentSummary.id == summary_record_id) - .limit(1) - ) - - if not summary_record_in_session: - # Record not found - try to find by chunk_id and dataset_id instead - logger.debug( - "Summary record not found by id=%s, trying chunk_id=%s and dataset_id=%s " - "for segment %s", - summary_record_id, - segment.id, - dataset.id, - segment.id, - ) - summary_record_in_session = session.scalar( - select(DocumentSegmentSummary) - .where( - DocumentSegmentSummary.chunk_id == segment.id, - DocumentSegmentSummary.dataset_id == dataset.id, - ) - .limit(1) - ) - - if not summary_record_in_session: - # Still not found - create a new one using the parameter data - logger.warning( - "Summary record not found in database for segment %s (id=%s), creating new one. " - "This may indicate a session isolation issue.", - segment.id, - summary_record_id, - ) - summary_record_in_session = DocumentSegmentSummary( - dataset_id=dataset.id, - document_id=segment.document_id, - chunk_id=segment.id, - summary_content=summary_content, - summary_index_node_id=summary_index_node_id, - summary_index_node_hash=summary_hash, - tokens=embedding_tokens, - status=SummaryStatus.COMPLETED, - enabled=True, - ) - if summary_record_in_session is None: - raise RuntimeError("summary_record_in_session should not be None at this point") - summary_record_in_session.id = summary_record_id - session.add(summary_record_in_session) - logger.info( - "Created new summary record (id=%s) for segment %s after vectorization", - summary_record_id, - segment.id, - ) - else: - # Found by chunk_id - update it - logger.info( - "Found summary record for segment %s by chunk_id " - "(id mismatch: expected %s, found %s). " - "This may indicate the record was created in a different session.", - segment.id, - summary_record_id, - summary_record_in_session.id, - ) - else: - logger.debug( - "Found summary_record (id=%s) for segment %s in new session", - summary_record_id, - segment.id, - ) - - # At this point, summary_record_in_session is guaranteed to be not None - if summary_record_in_session is None: - raise RuntimeError("summary_record_in_session should not be None at this point") - - # Update all fields including summary_content - # Always use the summary_content from the parameter (which is the latest from outer session) - # rather than relying on what's in the database, in case outer session hasn't committed yet - summary_record_in_session.summary_index_node_id = summary_index_node_id - summary_record_in_session.summary_index_node_hash = summary_hash - summary_record_in_session.tokens = embedding_tokens # Save embedding tokens - summary_record_in_session.status = SummaryStatus.COMPLETED - # Ensure summary_content is preserved (use the latest from summary_record parameter) - # This is critical: use the parameter value, not the database value - summary_record_in_session.summary_content = summary_content - # Explicitly update updated_at to ensure it's refreshed even if other fields haven't changed - summary_record_in_session.updated_at = datetime.now(UTC).replace(tzinfo=None) - session.add(summary_record_in_session) - - # Only commit if we created the session ourselves - if not use_provided_session: - logger.debug("Committing session for segment %s (self-created session)", segment.id) - session.commit() - logger.debug("Successfully committed session for segment %s", segment.id) - else: - # When using provided session, flush to ensure changes are written to database - # This prevents refresh() from overwriting our changes - logger.debug( - "Flushing session for segment %s (using provided session, caller will commit)", - segment.id, - ) - session.flush() - logger.debug("Successfully flushed session for segment %s", segment.id) - # If using provided session, let the caller handle commit - - logger.info( - "Successfully vectorized summary for segment %s, index_node_id=%s, index_node_hash=%s, " - "tokens=%s, summary_record_id=%s, use_provided_session=%s", - segment.id, - summary_index_node_id, - summary_hash, - embedding_tokens, - summary_record_in_session.id, - use_provided_session, - ) - # Update the original object for consistency - summary_record.summary_index_node_id = summary_index_node_id - summary_record.summary_index_node_hash = summary_hash - summary_record.tokens = embedding_tokens - summary_record.status = SummaryStatus.COMPLETED - summary_record.summary_content = summary_content - if summary_record_in_session.updated_at: - summary_record.updated_at = summary_record_in_session.updated_at - finally: - # Only close session if we created it ourselves - if not use_provided_session and session_context: - session_context.__exit__(None, None, None) - # Success, exit function - return - - except (ConnectionError, Exception) as e: - error_str = str(e).lower() - # Check if it's a connection-related error that might be transient - is_connection_error = any( - keyword in error_str - for keyword in [ - "connection", - "disconnected", - "timeout", - "network", - "could not connect", - "server disconnected", - "weaviate", - ] - ) - - if is_connection_error and attempt < max_retries - 1: - # Retry for connection errors - wait_time = retry_delay * (2**attempt) # Exponential backoff - logger.warning( - "Vectorization attempt %s/%s failed for segment %s (connection error): %s. " - "Retrying in %.1f seconds...", - attempt + 1, - max_retries, - segment.id, - str(e), - wait_time, - ) - time.sleep(wait_time) - continue - else: - # Final attempt failed or non-connection error - log and update status - logger.error( - "Failed to vectorize summary for segment %s after %s attempts: %s. " - "summary_record_id=%s, index_node_id=%s, use_provided_session=%s", - segment.id, - attempt + 1, - str(e), - summary_record_id, - summary_index_node_id, - session is not None, - exc_info=True, - ) - # Update error status in session - # Use the original_session saved at function start (the function parameter) - logger.debug( - "Updating error status for segment %s, summary_record_id=%s, has_original_session=%s", - segment.id, - summary_record_id, - original_session is not None, - ) - if original_session is not None: - # Keep the error update in the caller-owned transaction. Opening another - # session here can deadlock when the caller has already flushed this row. - summary_record.status = SummaryStatus.ERROR - summary_record.error = f"Vectorization failed: {str(e)}" - summary_record.updated_at = datetime.now(UTC).replace(tzinfo=None) - original_session.add(summary_record) - raise - - # Standalone callers still need this method to persist the error itself. - with session_factory.create_session() as error_session: - # Try to find the record by id first - # Note: Using assignment only (no type annotation) to avoid redeclaration error - summary_record_in_session = error_session.scalar( - select(DocumentSegmentSummary) - .where(DocumentSegmentSummary.id == summary_record_id) - .limit(1) - ) - if not summary_record_in_session: - # Try to find by chunk_id and dataset_id - logger.debug( - "Summary record not found by id=%s, trying chunk_id=%s and dataset_id=%s " - "for segment %s", - summary_record_id, - segment.id, - dataset.id, - segment.id, - ) - summary_record_in_session = error_session.scalar( - select(DocumentSegmentSummary) - .where( - DocumentSegmentSummary.chunk_id == segment.id, - DocumentSegmentSummary.dataset_id == dataset.id, - ) - .limit(1) - ) - - if summary_record_in_session: - summary_record_in_session.status = SummaryStatus.ERROR - summary_record_in_session.error = f"Vectorization failed: {str(e)}" - summary_record_in_session.updated_at = datetime.now(UTC).replace(tzinfo=None) - error_session.add(summary_record_in_session) - error_session.commit() - logger.info( - "Updated error status in new session for segment %s, record_id=%s", - segment.id, - summary_record_in_session.id, - ) - # Update the original object for consistency - summary_record.status = SummaryStatus.ERROR - summary_record.error = summary_record_in_session.error - summary_record.updated_at = summary_record_in_session.updated_at - else: - logger.warning( - "Could not update error status: summary record not found for segment %s (id=%s). " - "This may indicate a session isolation issue.", - segment.id, - summary_record_id, - ) - raise + try: + result = service.vectorize(command) + except Exception as error: + set_committed_value(summary_record, "status", SummaryStatus.ERROR) + set_committed_value(summary_record, "error", f"Vectorization failed: {error}") + raise + set_committed_value(summary_record, "summary_index_node_id", result.node_id) + set_committed_value(summary_record, "summary_index_node_hash", result.content_hash) + set_committed_value(summary_record, "tokens", result.tokens) + set_committed_value(summary_record, "status", SummaryStatus.COMPLETED) + set_committed_value(summary_record, "error", None) @staticmethod def batch_create_summary_records( @@ -726,7 +412,7 @@ class SummaryIndexService: session.commit() # Generate summary (returns summary_content and llm_usage) - summary_content, llm_usage = SummaryIndexService.generate_summary_for_segment( + summary_content, llm_usage = SummaryIndexAdapter.generate_summary_for_segment( segment, dataset, summary_index_setting, session=session ) @@ -746,7 +432,7 @@ class SummaryIndexService: llm_usage.completion_tokens, ) - SummaryIndexService.vectorize_summary(summary_record_in_session, segment, dataset, session=session) + SummaryIndexAdapter.vectorize_summary(summary_record_in_session, segment, dataset, session=session) # vectorize_summary mutates status and token fields; refresh before returning the ORM object. session.refresh(summary_record_in_session) session.commit() @@ -837,9 +523,10 @@ class SummaryIndexService: logger.info("No segments found for document %s", document.id) return [] + session.commit() # Batch create summary records with "not_started" status before processing # This ensures all records exist upfront, allowing status tracking - SummaryIndexService.batch_create_summary_records( + SummaryIndexAdapter.batch_create_summary_records( segments=segments, dataset=dataset, status=SummaryStatus.NOT_STARTED, @@ -860,14 +547,14 @@ class SummaryIndexService: pass try: - summary_record = SummaryIndexService.generate_and_vectorize_summary( + summary_record = SummaryIndexAdapter.generate_and_vectorize_summary( segment, dataset, summary_index_setting, session=session ) summary_records.append(summary_record) except Exception as e: logger.exception("Failed to generate summary for segment %s", segment.id) # Update summary record with error status - SummaryIndexService.update_summary_record_error( + SummaryIndexAdapter.update_summary_record_error( segment=segment, dataset=dataset, error=str(e), @@ -925,8 +612,7 @@ class SummaryIndexService: summary_node_ids = [s.summary_index_node_id for s in summaries if s.summary_index_node_id] if summary_node_ids: try: - vector = Vector(dataset, session=session) - vector.delete_by_ids(summary_node_ids) + _delete_summary_vectors(dataset, summary_node_ids, session=session) except Exception: logger.warning("Failed to remove summary vectors", exc_info=True) @@ -1006,7 +692,7 @@ class SummaryIndexService: try: # Re-vectorize summary (this will update status and tokens in its own session) # Pass the session to vectorize_summary to avoid session isolation issues - SummaryIndexService.vectorize_summary(summary, segment, dataset, session=session) + SummaryIndexAdapter.vectorize_summary(summary, segment, dataset, session=session) # Refresh the object from database to get the updated status and tokens from vectorize_summary session.refresh(summary) @@ -1050,14 +736,13 @@ class SummaryIndexService: if not summaries: return - # Delete from vector database + # Preserve the caller's transaction until it decides to commit or roll back. if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY: summary_node_ids = [s.summary_index_node_id for s in summaries if s.summary_index_node_id] if summary_node_ids: - vector = Vector(dataset, session=session) - vector.delete_by_ids(summary_node_ids) + _delete_summary_vectors(dataset, summary_node_ids, session=session) - # Delete summary records + # Delete summary records in the caller's existing transaction. for summary in summaries: session.delete(summary) @@ -1095,7 +780,7 @@ class SummaryIndexService: # Vectorization uses dataset.embedding_model, which doesn't require summary_index_setting # Skip qa_model documents - document = segment.get_document(session=session) + document = get_segment_document(segment, session=session) if document and document.doc_form == "qa_model": return None @@ -1117,8 +802,7 @@ class SummaryIndexService: old_summary_node_id = summary_record.summary_index_node_id if old_summary_node_id: try: - vector = Vector(dataset, session=session) - vector.delete_by_ids([old_summary_node_id]) + _delete_summary_vectors(dataset, [old_summary_node_id], session=session) except Exception as e: logger.warning( "Failed to delete old summary vector for segment %s: %s", @@ -1137,9 +821,6 @@ class SummaryIndexService: return None if summary_record: - # Update existing summary - old_summary_node_id = summary_record.summary_index_node_id - # Update summary content summary_record.summary_content = summary_content summary_record.status = SummaryStatus.GENERATING @@ -1148,20 +829,9 @@ class SummaryIndexService: # Flush to ensure summary_content is saved before vectorize_summary queries it session.flush() - # Delete old vector if exists (before vectorization) - if old_summary_node_id: - try: - vector = Vector(dataset, session=session) - vector.delete_by_ids([old_summary_node_id]) - except Exception as e: - logger.warning( - "Failed to delete old summary vector for segment %s: %s", - segment.id, - str(e), - ) else: # Create new summary record if doesn't exist - summary_record = SummaryIndexService.create_summary_record( + summary_record = SummaryIndexAdapter.create_summary_record( segment, dataset, summary_content, @@ -1171,7 +841,7 @@ class SummaryIndexService: try: # Vectorization must finish here so the manual summary is searchable immediately. - SummaryIndexService.vectorize_summary(summary_record, segment, dataset, session=session) + SummaryIndexAdapter.vectorize_summary(summary_record, segment, dataset, session=session) session.refresh(summary_record) session.commit() logger.info("Successfully updated and re-vectorized summary for segment %s", segment.id) @@ -1340,7 +1010,7 @@ class SummaryIndexService: return None # Get all summary records for these segments - summaries = SummaryIndexService.get_segments_summaries(segment_ids, dataset_id, session=session) + summaries = SummaryIndexAdapter.get_segments_summaries(segment_ids, dataset_id, session=session) summary_status_map = {chunk_id: summary.status for chunk_id, summary in summaries.items()} # Check if there are any "not_started" or "generating" status summaries @@ -1396,7 +1066,7 @@ class SummaryIndexService: # Get all summary records for these segments all_segment_ids = [seg.id for seg in segments] - summaries = SummaryIndexService.get_segments_summaries(all_segment_ids, dataset_id, session=session) + summaries = SummaryIndexAdapter.get_segments_summaries(all_segment_ids, dataset_id, session=session) summary_status_map = {chunk_id: summary.status for chunk_id, summary in summaries.items()} # Calculate summary_index_status for each document @@ -1433,6 +1103,8 @@ class SummaryIndexService: document_id: str, dataset_id: str, session: Session, + *, + tenant_id: str, ) -> DocumentSummaryStatusDetailDict: """ Get detailed summary status for a document. @@ -1440,6 +1112,7 @@ class SummaryIndexService: Args: document_id: Document ID dataset_id: Dataset ID + tenant_id: Owning tenant ID session: SQLAlchemy session used for segment lookup Returns: @@ -1453,10 +1126,11 @@ class SummaryIndexService: - timeout: Number of summaries that timed out - summaries: List of summary records with status and content preview """ - from services.dataset_service import SegmentService + from services.knowledge.dataset_service import SegmentService # Get all segments for this document segments = SegmentService.get_segments_by_document_and_dataset( + tenant_id=tenant_id, document_id=document_id, dataset_id=dataset_id, session=session, @@ -1470,7 +1144,7 @@ class SummaryIndexService: segment_ids = [segment.id for segment in segments] summaries = [] if segment_ids: - summaries = SummaryIndexService.get_document_summaries( + summaries = SummaryIndexAdapter.get_document_summaries( document_id=document_id, dataset_id=dataset_id, segment_ids=segment_ids, @@ -1528,3 +1202,37 @@ class SummaryIndexService: summary_status=cast(dict[str, int], status_counts), summaries=summary_list, ) + + +class SummaryVectorGateway: + def __init__(self, dataset: Dataset, factory: Callable[[], Vector]) -> None: + self._dataset, self._factory = dataset, factory + + def delete(self, node_id: str) -> None: + self._factory().delete_by_ids([node_id]) + + def count_tokens(self, content: str) -> int: + dataset = self._dataset + model = ModelManager.for_tenant(tenant_id=dataset.tenant_id).get_model_instance( + tenant_id=dataset.tenant_id, + provider=dataset.embedding_model_provider or "", + model_type=ModelType.TEXT_EMBEDDING, + model=dataset.embedding_model or "", + ) + tokens = model.get_text_embedding_num_tokens([content]) + return tokens[0] if tokens and isinstance(tokens[0], int) else 0 + + def write(self, command: SummaryVectorInput, result: SummaryVectorResult) -> None: + document = Document( + page_content=command.content, + metadata={ + "doc_id": result.node_id, + "doc_hash": result.content_hash, + "dataset_id": command.ref.document.dataset.dataset_id, + "document_id": command.ref.document.document_id, + "original_chunk_id": command.ref.segment_id, + "doc_type": DocType.TEXT, + "is_summary": True, + }, + ) + self._factory().add_texts([document], duplicate_check=False) diff --git a/api/services/knowledge/summaries/application.py b/api/services/knowledge/summaries/application.py new file mode 100644 index 00000000000..538bc465a3d --- /dev/null +++ b/api/services/knowledge/summaries/application.py @@ -0,0 +1,82 @@ +"""Summary vectorization policy independent of ORM, model providers, and vector stores.""" + +import logging +from collections.abc import Callable +from dataclasses import dataclass +from typing import Protocol + +from services.knowledge.resource_scope import SegmentRef + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class SummaryVectorInput: + ref: SegmentRef + summary_id: str + content: str + previous_node_id: str | None + + +@dataclass(frozen=True) +class SummaryVectorResult: + node_id: str + content_hash: str + tokens: int + + +class SummaryStore(Protocol): + def completed(self, command: SummaryVectorInput, result: SummaryVectorResult) -> None: ... + def failed(self, command: SummaryVectorInput, error: str) -> None: ... + + +class SummaryVectorBackend(Protocol): + def delete(self, node_id: str) -> None: ... + def count_tokens(self, content: str) -> int: ... + def write(self, command: SummaryVectorInput, result: SummaryVectorResult) -> None: ... + + +class SummaryIndexService: + def __init__( + self, + *, + store: SummaryStore, + vectors: SummaryVectorBackend, + new_id: Callable[[], str], + text_hash: Callable[[str], str], + sleep: Callable[[float], None], + ) -> None: + self._store, self._vectors = store, vectors + self._new_id, self._hash, self._sleep = new_id, text_hash, sleep + + def vectorize(self, command: SummaryVectorInput) -> SummaryVectorResult: + if not command.content.strip(): + raise ValueError(f"Summary content is empty for segment {command.ref.segment_id}, cannot vectorize") + node_id = command.previous_node_id or self._new_id() + if command.previous_node_id: + try: + self._vectors.delete(command.previous_node_id) + except Exception: + logger.warning("Failed to delete previous summary vector for %s", command.ref.segment_id, exc_info=True) + tokens = 0 + try: + tokens = self._vectors.count_tokens(command.content) + except Exception: + logger.warning("Failed to calculate summary embedding tokens", exc_info=True) + result = SummaryVectorResult(node_id, self._hash(command.content), tokens) + for attempt in range(3): + try: + self._vectors.write(command, result) + self._store.completed(command, result) + return result + except Exception as error: + transient = any( + word in str(error).lower() + for word in ("connection", "disconnected", "timeout", "network", "weaviate") + ) + if transient and attempt < 2: + self._sleep(2.0 * 2**attempt) + continue + self._store.failed(command, f"Vectorization failed: {error}") + raise + raise AssertionError("Unreachable vectorization retry state") diff --git a/api/services/metadata_service.py b/api/services/metadata_service.py deleted file mode 100644 index 317e371fc2a..00000000000 --- a/api/services/metadata_service.py +++ /dev/null @@ -1,399 +0,0 @@ -import copy -import logging - -from sqlalchemy import delete, func, select -from sqlalchemy.orm import Session - -from core.rag.index_processor.constant.built_in_field import BuiltInField, MetadataDataSource -from extensions.ext_redis import redis_client -from libs.datetime_utils import naive_utc_now -from libs.login import resolve_account_fallback -from models import Account -from models.dataset import Dataset, DatasetMetadata, DatasetMetadataBinding, Document -from models.enums import DatasetMetadataType -from services.dataset_ref_service import DatasetRefService -from services.dataset_service import DocumentService -from services.entities.knowledge_entities.knowledge_entities import ( - MetadataArgs, - MetadataOperationData, -) -from services.errors.metadata import MetadataResourceNotFoundError - -logger = logging.getLogger(__name__) - - -class MetadataService: - @staticmethod - def create_metadata( - dataset_id: str, - metadata_args: MetadataArgs, - current_user: Account | None = None, # TODO: the service_api is not migrated yet - current_tenant_id: str | None = None, - *, - session: Session, - ) -> DatasetMetadata: - # check if metadata name is too long - if len(metadata_args.name) > 255: - raise ValueError("Metadata name cannot exceed 255 characters.") - current_user, current_tenant_id = resolve_account_fallback(current_user, current_tenant_id) - # check if metadata name already exists - if session.scalar( - select(DatasetMetadata) - .where( - DatasetMetadata.tenant_id == current_tenant_id, - DatasetMetadata.dataset_id == dataset_id, - DatasetMetadata.name == metadata_args.name, - ) - .limit(1) - ): - raise ValueError("Metadata name already exists.") - for field in BuiltInField: - if field.value == metadata_args.name: - raise ValueError("Metadata name already exists in Built-in fields.") - metadata = DatasetMetadata( - tenant_id=current_tenant_id, - dataset_id=dataset_id, - type=metadata_args.type, - name=metadata_args.name, - created_by=current_user.id, - ) - session.add(metadata) - session.flush() - return metadata - - @staticmethod - def update_metadata_name( - dataset: Dataset, - metadata_id: str, - name: str, - current_user: Account, - *, - session: Session, - ) -> DatasetMetadata | None: - # check if metadata name is too long - if len(name) > 255: - raise ValueError("Metadata name cannot exceed 255 characters.") - - lock_key = f"dataset_metadata_lock_{dataset.id}" - # check if metadata name already exists - if session.scalar( - select(DatasetMetadata) - .where( - DatasetMetadata.tenant_id == dataset.tenant_id, - DatasetMetadata.dataset_id == dataset.id, - DatasetMetadata.name == name, - ) - .limit(1) - ): - raise ValueError("Metadata name already exists.") - for field in BuiltInField: - if field.value == name: - raise ValueError("Metadata name already exists in Built-in fields.") - try: - MetadataService.knowledge_base_metadata_lock_check(dataset.id, None) - metadata = session.scalar( - select(DatasetMetadata) - .where( - DatasetMetadata.id == metadata_id, - DatasetMetadata.tenant_id == dataset.tenant_id, - DatasetMetadata.dataset_id == dataset.id, - ) - .limit(1) - ) - if metadata is None: - raise ValueError("Metadata not found.") - old_name = metadata.name - metadata.name = name - metadata.updated_by = current_user.id - metadata.updated_at = naive_utc_now() - - # update related documents - dataset_metadata_bindings = session.scalars( - select(DatasetMetadataBinding).where( - DatasetMetadataBinding.metadata_id == metadata_id, - DatasetMetadataBinding.tenant_id == dataset.tenant_id, - DatasetMetadataBinding.dataset_id == dataset.id, - ) - ).all() - if dataset_metadata_bindings: - document_ids = [binding.document_id for binding in dataset_metadata_bindings] - documents = DocumentService.get_document_by_ids( - DatasetRefService.create_dataset_ref(dataset), document_ids, session - ) - for document in documents: - if not document.doc_metadata: - doc_metadata = {} - else: - doc_metadata = copy.deepcopy(document.doc_metadata) - value = doc_metadata.pop(old_name, None) - doc_metadata[name] = value - document.doc_metadata = doc_metadata - session.add(document) - session.commit() - return metadata - except Exception: - logger.exception("Update metadata name failed") - return None - finally: - redis_client.delete(lock_key) - - @staticmethod - def delete_metadata(dataset: Dataset, metadata_id: str, session: Session): - lock_key = f"dataset_metadata_lock_{dataset.id}" - try: - MetadataService.knowledge_base_metadata_lock_check(dataset.id, None) - metadata = session.scalar( - select(DatasetMetadata) - .where( - DatasetMetadata.id == metadata_id, - DatasetMetadata.tenant_id == dataset.tenant_id, - DatasetMetadata.dataset_id == dataset.id, - ) - .limit(1) - ) - if metadata is None: - raise ValueError("Metadata not found.") - session.delete(metadata) - - # deal related documents - dataset_metadata_bindings = session.scalars( - select(DatasetMetadataBinding).where( - DatasetMetadataBinding.metadata_id == metadata_id, - DatasetMetadataBinding.tenant_id == dataset.tenant_id, - DatasetMetadataBinding.dataset_id == dataset.id, - ) - ).all() - if dataset_metadata_bindings: - document_ids = [binding.document_id for binding in dataset_metadata_bindings] - documents = DocumentService.get_document_by_ids( - DatasetRefService.create_dataset_ref(dataset), document_ids, session - ) - for document in documents: - if not document.doc_metadata: - doc_metadata = {} - else: - doc_metadata = copy.deepcopy(document.doc_metadata) - doc_metadata.pop(metadata.name, None) - document.doc_metadata = doc_metadata - session.add(document) - session.commit() - return metadata - except Exception: - logger.exception("Delete metadata failed") - finally: - redis_client.delete(lock_key) - - @staticmethod - def get_built_in_fields(): - return [ - {"name": BuiltInField.document_name, "type": DatasetMetadataType.STRING}, - {"name": BuiltInField.uploader, "type": DatasetMetadataType.STRING}, - {"name": BuiltInField.upload_date, "type": DatasetMetadataType.TIME}, - {"name": BuiltInField.last_update_date, "type": DatasetMetadataType.TIME}, - {"name": BuiltInField.source, "type": DatasetMetadataType.STRING}, - ] - - @staticmethod - def enable_built_in_field(dataset: Dataset, session: Session): - if dataset.built_in_field_enabled: - return - lock_key = f"dataset_metadata_lock_{dataset.id}" - try: - MetadataService.knowledge_base_metadata_lock_check(dataset.id, None) - session.add(dataset) - documents = DocumentService.get_working_documents_by_dataset_id(dataset.id, session) - if documents: - for document in documents: - if not document.doc_metadata: - doc_metadata = {} - else: - doc_metadata = copy.deepcopy(document.doc_metadata) - doc_metadata[BuiltInField.document_name] = document.name - doc_metadata[BuiltInField.uploader] = document.get_uploader(session=session) - doc_metadata[BuiltInField.upload_date] = document.upload_date.timestamp() - doc_metadata[BuiltInField.last_update_date] = document.last_update_date.timestamp() - doc_metadata[BuiltInField.source] = MetadataDataSource[document.data_source_type] - document.doc_metadata = doc_metadata - session.add(document) - dataset.built_in_field_enabled = True - session.commit() - except Exception: - logger.exception("Enable built-in field failed") - finally: - redis_client.delete(lock_key) - - @staticmethod - def disable_built_in_field(dataset: Dataset, session: Session): - if not dataset.built_in_field_enabled: - return - lock_key = f"dataset_metadata_lock_{dataset.id}" - try: - MetadataService.knowledge_base_metadata_lock_check(dataset.id, None) - session.add(dataset) - documents = DocumentService.get_working_documents_by_dataset_id(dataset.id, session) - document_ids = [] - if documents: - for document in documents: - if not document.doc_metadata: - doc_metadata = {} - else: - doc_metadata = copy.deepcopy(document.doc_metadata) - doc_metadata.pop(BuiltInField.document_name, None) - doc_metadata.pop(BuiltInField.uploader, None) - doc_metadata.pop(BuiltInField.upload_date, None) - doc_metadata.pop(BuiltInField.last_update_date, None) - doc_metadata.pop(BuiltInField.source, None) - document.doc_metadata = doc_metadata - session.add(document) - document_ids.append(document.id) - dataset.built_in_field_enabled = False - session.commit() - except Exception: - logger.exception("Disable built-in field failed") - finally: - redis_client.delete(lock_key) - - @staticmethod - def update_documents_metadata( - dataset: Dataset, - metadata_args: MetadataOperationData, - current_user: Account, - *, - session: Session, - ): - metadata_ids = { - metadata_value.id - for operation in metadata_args.operation_data - for metadata_value in operation.metadata_list - } - metadatas = session.scalars( - select(DatasetMetadata).where( - DatasetMetadata.id.in_(metadata_ids), - DatasetMetadata.tenant_id == dataset.tenant_id, - DatasetMetadata.dataset_id == dataset.id, - ) - ).all() - metadata_by_id = {metadata.id: metadata for metadata in metadatas} - if metadata_ids != set(metadata_by_id): - raise MetadataResourceNotFoundError("Metadata not found.") - - document_ids = {operation.document_id for operation in metadata_args.operation_data} - owned_document_ids = set( - session.scalars( - select(Document.id).where( - Document.id.in_(document_ids), - Document.tenant_id == dataset.tenant_id, - Document.dataset_id == dataset.id, - ) - ).all() - ) - if document_ids != owned_document_ids: - raise MetadataResourceNotFoundError("Document not found.") - - for operation in metadata_args.operation_data: - lock_key = f"document_metadata_lock_{operation.document_id}" - try: - MetadataService.knowledge_base_metadata_lock_check(None, operation.document_id) - document = session.scalar( - select(Document) - .where( - Document.id == operation.document_id, - Document.tenant_id == dataset.tenant_id, - Document.dataset_id == dataset.id, - ) - .with_for_update() - .execution_options(populate_existing=True) - ) - if document is None: - raise MetadataResourceNotFoundError("Document not found.") - if operation.partial_update: - doc_metadata = copy.deepcopy(document.doc_metadata) if document.doc_metadata else {} - else: - doc_metadata = {} - for metadata_value in operation.metadata_list: - doc_metadata[metadata_by_id[metadata_value.id].name] = metadata_value.value - if dataset.built_in_field_enabled: - doc_metadata[BuiltInField.document_name] = document.name - doc_metadata[BuiltInField.uploader] = document.get_uploader(session=session) - doc_metadata[BuiltInField.upload_date] = document.upload_date.timestamp() - doc_metadata[BuiltInField.last_update_date] = document.last_update_date.timestamp() - doc_metadata[BuiltInField.source] = MetadataDataSource[document.data_source_type] - document.doc_metadata = doc_metadata - session.add(document) - - # deal metadata binding (in the same transaction as the doc_metadata update) - if not operation.partial_update: - session.execute( - delete(DatasetMetadataBinding).where( - DatasetMetadataBinding.tenant_id == dataset.tenant_id, - DatasetMetadataBinding.dataset_id == dataset.id, - DatasetMetadataBinding.document_id == document.id, - ) - ) - - for metadata_value in operation.metadata_list: - # check if binding already exists - if operation.partial_update: - existing_binding = session.scalar( - select(DatasetMetadataBinding) - .where( - DatasetMetadataBinding.tenant_id == dataset.tenant_id, - DatasetMetadataBinding.dataset_id == dataset.id, - DatasetMetadataBinding.document_id == document.id, - DatasetMetadataBinding.metadata_id == metadata_value.id, - ) - .limit(1) - ) - if existing_binding: - continue - - dataset_metadata_binding = DatasetMetadataBinding( - tenant_id=dataset.tenant_id, - dataset_id=dataset.id, - document_id=document.id, - metadata_id=metadata_value.id, - created_by=current_user.id, - ) - session.add(dataset_metadata_binding) - session.commit() - except Exception: - session.rollback() - logger.exception("Update documents metadata failed") - raise - finally: - redis_client.delete(lock_key) - - @staticmethod - def knowledge_base_metadata_lock_check(dataset_id: str | None, document_id: str | None): - if dataset_id: - lock_key = f"dataset_metadata_lock_{dataset_id}" - if redis_client.get(lock_key): - raise ValueError("Another knowledge base metadata operation is running, please wait a moment.") - redis_client.set(lock_key, 1, ex=3600) - if document_id: - lock_key = f"document_metadata_lock_{document_id}" - if redis_client.get(lock_key): - raise ValueError("Another document metadata operation is running, please wait a moment.") - redis_client.set(lock_key, 1, ex=3600) - - @staticmethod - def get_dataset_metadatas(dataset: Dataset, session: Session): - return { - "doc_metadata": [ - { - "id": item.get("id"), - "name": item.get("name"), - "type": item.get("type"), - "count": session.scalar( - select(func.count(DatasetMetadataBinding.id)).where( - DatasetMetadataBinding.metadata_id == item.get("id"), - DatasetMetadataBinding.dataset_id == dataset.id, - ) - ) - or 0, - } - for item in dataset.get_doc_metadata(session=session) - if item.get("id") != "built-in" - ], - "built_in_field_enabled": dataset.built_in_field_enabled, - } diff --git a/api/services/model_provider/__init__.py b/api/services/model_provider/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/api/services/model_provider_service.py b/api/services/model_provider/service.py similarity index 98% rename from api/services/model_provider_service.py rename to api/services/model_provider/service.py index 9aee8eb87a5..a1e218f8cca 100644 --- a/api/services/model_provider_service.py +++ b/api/services/model_provider/service.py @@ -7,6 +7,8 @@ from typing import TYPE_CHECKING, Any from sqlalchemy import and_, select from sqlalchemy.orm import Session +from services.credentials.query import CredentialQuery + if TYPE_CHECKING: from models.account import Account @@ -459,12 +461,15 @@ class ModelProviderService: for model in provider_configurations.get_models(provider=provider) ] - def get_provider_available_credentials(self, tenant_id: str, provider: str, user: "Account | None" = None): - return self._get_provider_manager(tenant_id).get_provider_available_credentials( - tenant_id=tenant_id, - provider_name=provider, - user=user, - ) + def get_provider_available_credentials( + self, tenant_id: str, provider: str, user: "Account | None" = None, *, credential_query: CredentialQuery + ) -> list[CredentialConfiguration]: + return [ + CredentialConfiguration(credential_id=record.id, credential_name=record.name) + for record in credential_query.list_models( + workspace_id=tenant_id, provider=provider, actor_id=user.id if user else None + ) + ] def get_provider_model_available_credentials( self, diff --git a/api/services/rag_pipeline/pipeline_generate_service.py b/api/services/rag_pipeline/pipeline_generate_service.py index 182e3fb283e..12e85611b04 100644 --- a/api/services/rag_pipeline/pipeline_generate_service.py +++ b/api/services/rag_pipeline/pipeline_generate_service.py @@ -10,7 +10,9 @@ from models.dataset import Pipeline from models.enums import IndexingStatus from models.model import Account, App, EndUser from models.workflow import Workflow -from services.dataset_ref_service import DatasetRefService, DocumentRef +from repositories.knowledge.dataset_read_repository import get_pipeline_dataset +from services.knowledge.dataset_service import DocumentService +from services.knowledge.resource_scope import DatasetRef, DocumentRef from services.rag_pipeline.rag_pipeline import RagPipelineService @@ -25,6 +27,7 @@ class PipelineGenerateService: streaming: bool = True, *, session: Session, + generator: PipelineGenerator, ): """ Pipeline Content Generate @@ -38,14 +41,14 @@ class PipelineGenerateService: try: workflow = cls._get_workflow(pipeline, invoke_from, session) if original_document_id := args.get("original_document_id"): - dataset = pipeline.retrieve_dataset(session) + dataset = get_pipeline_dataset(pipeline, session) if dataset is None or dataset.tenant_id != pipeline.tenant_id: raise ValueError("Pipeline dataset is required") - dataset_ref = DatasetRefService.create_dataset_ref(dataset) - document_ref = DatasetRefService.create_document_ref_from_id(dataset_ref, original_document_id) + dataset_ref = DatasetRef(tenant_id=dataset.tenant_id, dataset_id=dataset.id) + document_ref = dataset_ref.document(original_document_id) cls.update_document_status(document_ref, session=session) return PipelineGenerator.convert_to_event_stream( - PipelineGenerator().generate( + generator.generate( session=session, pipeline=pipeline, workflow=workflow, @@ -71,11 +74,19 @@ class PipelineGenerateService: @classmethod def generate_single_iteration( - cls, pipeline: Pipeline, user: Account, node_id: str, args: Any, session: Session, streaming: bool = True + cls, + pipeline: Pipeline, + user: Account, + node_id: str, + args: Any, + session: Session, + streaming: bool = True, + *, + generator: PipelineGenerator, ): workflow = cls._get_workflow(pipeline, InvokeFrom.DEBUGGER, session) return PipelineGenerator.convert_to_event_stream( - PipelineGenerator().single_iteration_generate( + generator.single_iteration_generate( pipeline=pipeline, workflow=workflow, node_id=node_id, @@ -88,11 +99,19 @@ class PipelineGenerateService: @classmethod def generate_single_loop( - cls, pipeline: Pipeline, user: Account, node_id: str, args: Any, session: Session, streaming: bool = True + cls, + pipeline: Pipeline, + user: Account, + node_id: str, + args: Any, + session: Session, + streaming: bool = True, + *, + generator: PipelineGenerator, ): workflow = cls._get_workflow(pipeline, InvokeFrom.DEBUGGER, session) return PipelineGenerator.convert_to_event_stream( - PipelineGenerator().single_loop_generate( + generator.single_loop_generate( pipeline=pipeline, workflow=workflow, node_id=node_id, @@ -130,7 +149,9 @@ class PipelineGenerateService: @classmethod def update_document_status(cls, document_ref: DocumentRef, *, session: Session) -> None: """Set a document in the owner-bound dataset to waiting.""" - document = DatasetRefService.get_document_by_ref(document_ref, session=session) + document = next( + iter(DocumentService.get_documents_by_ids(document_ref.dataset, [document_ref.document_id], session)), None + ) if document is None: raise ValueError("Pipeline document not found") document.indexing_status = IndexingStatus.WAITING diff --git a/api/services/rag_pipeline/pipeline_template/customized/customized_retrieval.py b/api/services/rag_pipeline/pipeline_template/customized/customized_retrieval.py index 2aa39cd5d4e..268dd51ab68 100644 --- a/api/services/rag_pipeline/pipeline_template/customized/customized_retrieval.py +++ b/api/services/rag_pipeline/pipeline_template/customized/customized_retrieval.py @@ -6,6 +6,7 @@ from sqlalchemy.orm import Session from libs.login import resolve_tenant_id_fallback from models.dataset import PipelineCustomizedTemplate +from repositories.knowledge.dataset_read_repository import get_pipeline_template_creator_name from services.rag_pipeline.pipeline_template.pipeline_template_base import PipelineTemplateRetrievalBase from services.rag_pipeline.pipeline_template.pipeline_template_type import PipelineTemplateType @@ -119,5 +120,5 @@ class CustomizedPipelineTemplateRetrieval(PipelineTemplateRetrievalBase): "chunk_structure": pipeline_template.chunk_structure, "export_data": pipeline_template.yaml_content, "graph": graph_data, - "created_by": pipeline_template.created_user_name(session=session), + "created_by": get_pipeline_template_creator_name(pipeline_template, session=session), } diff --git a/api/services/rag_pipeline/rag_pipeline.py b/api/services/rag_pipeline/rag_pipeline.py index c5f1e1935ca..35d973e24a6 100644 --- a/api/services/rag_pipeline/rag_pipeline.py +++ b/api/services/rag_pipeline/rag_pipeline.py @@ -15,6 +15,7 @@ import contexts from configs import dify_config from core.app.apps.pipeline.pipeline_generator import PipelineGenerator from core.app.entities.app_invoke_entities import InvokeFrom +from core.datasource.datasource_manager import DatasourceManager from core.datasource.entities.datasource_entities import ( DatasourceMessage, DatasourceProviderType, @@ -75,15 +76,18 @@ from models.workflow import ( WorkflowType, ) from repositories.factory import DifyAPIRepositoryFactory -from services.dataset_ref_service import DatasetRefService -from services.datasource_provider_service import DatasourceProviderService +from repositories.knowledge.dataset_read_repository import get_pipeline_dataset +from services.credentials.query import CredentialQuery +from services.data_source.provider_service import DatasourceProviderService from services.entities.knowledge_entities.rag_pipeline_entities import ( KnowledgeConfiguration, PipelineTemplateInfoEntity, ) from services.errors.app import IsDraftWorkflowError, WorkflowHashNotEqualError, WorkflowNotFoundError from services.errors.rag_pipeline import RagPipelineResourceNotFoundError +from services.knowledge.resource_scope import DatasetRef from services.rag_pipeline.pipeline_template.pipeline_template_factory import PipelineTemplateRetrievalFactory +from services.rag_pipeline.rag_pipeline_dsl_service import RagPipelineDslService from services.tools.builtin_tools_manage_service import BuiltinToolManageService from services.workflow_draft_variable_service import DraftVariableSaver, DraftVarLoader from services.workflow_node_execution_trace_service import ( @@ -491,7 +495,8 @@ class RagPipelineService: graph = workflow.graph_dict nodes = graph.get("nodes", []) - from services.dataset_service import DatasetService + # dataset_service imports RagPipelineService, so defer this reverse dependency. + from services.knowledge.dataset_service import DatasetService for node in nodes: if node.get("data", {}).get("type") == "knowledge-index": @@ -499,7 +504,7 @@ class RagPipelineService: knowledge_configuration = KnowledgeConfiguration.model_validate(knowledge_configuration) # update dataset - dataset = pipeline.retrieve_dataset(session=session) + dataset = get_pipeline_dataset(pipeline, session=session) if not dataset: raise ValueError("Dataset not found") DatasetService.update_rag_pipeline_dataset_settings( @@ -667,6 +672,8 @@ class RagPipelineService: datasource_type: str, is_published: bool, credential_id: str | None = None, + *, + datasource_providers: DatasourceProviderService, ) -> Generator[Mapping[str, Any], None, None]: """ Run published workflow datasource @@ -718,16 +725,13 @@ class RagPipelineService: # other type directly use original value variables_map[key] = param_value - from core.datasource.datasource_manager import DatasourceManager - datasource_runtime = DatasourceManager.get_datasource_runtime( provider_id=f"{datasource_node_data.get('plugin_id')}/{datasource_node_data.get('provider_name')}", datasource_name=datasource_node_data.get("datasource_name"), tenant_id=pipeline.tenant_id, datasource_type=DatasourceProviderType(datasource_type), ) - datasource_provider_service = DatasourceProviderService() - credentials = datasource_provider_service.get_datasource_credentials( + credentials = datasource_providers.get_datasource_credentials( tenant_id=pipeline.tenant_id, provider=datasource_node_data.get("provider_name"), plugin_id=datasource_node_data.get("plugin_id"), @@ -835,6 +839,8 @@ class RagPipelineService: datasource_type: str, is_published: bool, credential_id: str | None = None, + *, + datasource_providers: DatasourceProviderService, ) -> Mapping[str, Any]: """ Run published workflow datasource @@ -863,16 +869,13 @@ class RagPipelineService: if not user_inputs.get(key): user_inputs[key] = value["value"] - from core.datasource.datasource_manager import DatasourceManager - datasource_runtime = DatasourceManager.get_datasource_runtime( provider_id=f"{datasource_node_data.get('plugin_id')}/{datasource_node_data.get('provider_name')}", datasource_name=datasource_node_data.get("datasource_name"), tenant_id=pipeline.tenant_id, datasource_type=DatasourceProviderType(datasource_type), ) - datasource_provider_service = DatasourceProviderService() - credentials = datasource_provider_service.get_datasource_credentials( + credentials = datasource_providers.get_datasource_credentials( tenant_id=pipeline.tenant_id, provider=datasource_node_data.get("provider_name"), plugin_id=datasource_node_data.get("plugin_id"), @@ -1065,9 +1068,19 @@ class RagPipelineService: .limit(1) ) if dataset: - dataset_ref = DatasetRefService.create_dataset_ref(dataset) - document_ref = DatasetRefService.create_document_ref_from_id(dataset_ref, document_id.value) - document = DatasetRefService.get_document_by_ref(document_ref, session=self._session) + dataset_ref = DatasetRef(tenant_id=dataset.tenant_id, dataset_id=dataset.id) + document_ref = dataset_ref.document(document_id.value) + # dataset_service imports RagPipelineService; importing here avoids the cycle. + from services.knowledge.dataset_service import DocumentService + + document = next( + iter( + DocumentService.get_documents_by_ids( + document_ref.dataset, [document_ref.document_id], self._session + ) + ), + None, + ) if document: document.indexing_status = IndexingStatus.ERROR document.error = error @@ -1319,8 +1332,6 @@ class RagPipelineService: ) ) - from services.rag_pipeline.rag_pipeline_dsl_service import RagPipelineDslService - rag_pipeline_dsl_service = RagPipelineDslService(session) dsl = rag_pipeline_dsl_service.export_rag_pipeline_dsl(pipeline=pipeline, include_secret=True) pipeline_customized_template = PipelineCustomizedTemplate( @@ -1499,7 +1510,9 @@ class RagPipelineService: "uninstalled_recommended_plugins": uninstalled_plugin_list, } - def retry_error_document(self, dataset: Dataset, document: Document, user: Account | EndUser): + def retry_error_document( + self, dataset: Dataset, document: Document, user: Account | EndUser, *, generator: PipelineGenerator + ): """ Retry error document """ @@ -1515,7 +1528,7 @@ class RagPipelineService: workflow = self.get_published_workflow(pipeline) if not workflow: raise ValueError("Workflow not found") - PipelineGenerator().generate( + generator.generate( session=self._session, pipeline=pipeline, workflow=workflow, @@ -1534,7 +1547,15 @@ class RagPipelineService: is_retry=True, ) - def get_datasource_plugins(self, tenant_id: str, dataset_id: str, is_published: bool) -> list[dict]: + def get_datasource_plugins( + self, + tenant_id: str, + dataset_id: str, + is_published: bool, + *, + credential_query: CredentialQuery, + datasource_providers: DatasourceProviderService, + ) -> list[dict]: """ Get datasource plugins """ @@ -1601,12 +1622,11 @@ class RagPipelineService: user_input_variables.append(value) # get credentials - datasource_provider_service: DatasourceProviderService = DatasourceProviderService() - credentials: list[dict[Any, Any]] = datasource_provider_service.list_datasource_credentials( + credentials: list[dict[Any, Any]] = datasource_providers.list_datasource_credentials( tenant_id=tenant_id, provider=datasource_node_data.get("provider_name"), plugin_id=datasource_node_data.get("plugin_id"), - session=self._session, + credential_query=credential_query, ) credential_info_list: list[Any] = [] for credential in credentials: diff --git a/api/services/rag_pipeline/rag_pipeline_dsl_service.py b/api/services/rag_pipeline/rag_pipeline_dsl_service.py index 088cade8dfe..5090b89aca5 100644 --- a/api/services/rag_pipeline/rag_pipeline_dsl_service.py +++ b/api/services/rag_pipeline/rag_pipeline_dsl_service.py @@ -42,6 +42,7 @@ from models import Account from models.dataset import Dataset, DatasetCollectionBinding, Pipeline from models.enums import CollectionBindingType, DatasetRuntimeMode from models.workflow import Workflow, WorkflowType +from repositories.knowledge.dataset_read_repository import get_pipeline_dataset from services.dsl_content import DSL_MAX_SIZE, dsl_content_size from services.dsl_version import check_version_compatibility from services.entities.dsl_entities import CheckDependenciesResult, ImportMode, ImportStatus, PendingImportOwner @@ -217,7 +218,7 @@ class RagPipelineDslService: status=ImportStatus.FAILED, error="Pipeline not found", ) - dataset = pipeline.retrieve_dataset(session=self._session) + dataset = get_pipeline_dataset(pipeline, session=self._session) # If major version mismatch, store import info in Redis if status == ImportStatus.PENDING: @@ -414,7 +415,7 @@ class RagPipelineDslService: data=data, account=account, ) - dataset = pipeline.retrieve_dataset(session=self._session) + dataset = get_pipeline_dataset(pipeline, session=self._session) # create dataset name = pipeline.name @@ -656,7 +657,7 @@ class RagPipelineDslService: :param include_secret: Whether include secret variable :return: """ - dataset = pipeline.retrieve_dataset(session=self._session) + dataset = get_pipeline_dataset(pipeline, session=self._session) if not dataset: raise ValueError("Missing dataset for rag pipeline") icon_info = dataset.icon_info diff --git a/api/services/rag_pipeline/rag_pipeline_manage_service.py b/api/services/rag_pipeline/rag_pipeline_manage_service.py index c01140ff934..0278292abb3 100644 --- a/api/services/rag_pipeline/rag_pipeline_manage_service.py +++ b/api/services/rag_pipeline/rag_pipeline_manage_service.py @@ -2,7 +2,7 @@ import logging from core.plugin.entities.plugin_daemon import PluginDatasourceProviderEntity from core.plugin.impl.datasource import PluginDatasourceManager -from services.datasource_provider_service import DatasourceProviderService +from services.data_source.provider_service import DatasourceProviderService logger = logging.getLogger(__name__) @@ -11,6 +11,8 @@ class RagPipelineManageService: @staticmethod def list_rag_pipeline_datasources( tenant_id: str, + *, + datasource_providers: DatasourceProviderService, ) -> list[PluginDatasourceProviderEntity]: """ list rag pipeline datasources @@ -24,9 +26,8 @@ class RagPipelineManageService: # built-in providers that declare neither credentials nor OAuth never require authorization datasource.is_authorized = True continue - datasource_provider_service = DatasourceProviderService() try: - credentials = datasource_provider_service.get_datasource_credentials( + credentials = datasource_providers.get_datasource_credentials( tenant_id=tenant_id, provider=datasource.provider, plugin_id=datasource.plugin_id, diff --git a/api/services/rag_pipeline/rag_pipeline_task_proxy.py b/api/services/rag_pipeline/rag_pipeline_task_proxy.py index ab3b7338c08..3db7820d71f 100644 --- a/api/services/rag_pipeline/rag_pipeline_task_proxy.py +++ b/api/services/rag_pipeline/rag_pipeline_task_proxy.py @@ -1,8 +1,11 @@ import json import logging -from collections.abc import Callable, Sequence +from collections.abc import Sequence from functools import cached_property +from celery import signature +from celery.canvas import Signature + from configs import dify_config from core.app.entities.rag_pipeline_invoke_entities import RagPipelineInvokeEntity from core.rag.pipeline.queue import TenantIsolatedTaskQueue @@ -10,8 +13,6 @@ from enums import CloudPlan, DeploymentEdition from extensions.ext_database import db from services.feature_service import FeatureService from services.file_service import FileService -from tasks.rag_pipeline.priority_rag_pipeline_run_task import priority_rag_pipeline_run_task -from tasks.rag_pipeline.rag_pipeline_run_task import rag_pipeline_run_task logger = logging.getLogger(__name__) @@ -44,14 +45,14 @@ class RagPipelineTaskProxy: ) return upload_file.id - def _send_to_direct_queue(self, upload_file_id: str, task_func: Callable[[str, str], None]): + def _send_to_direct_queue(self, upload_file_id: str, task_func: Signature): logger.info("tenant %s send file %s to direct queue", self._dataset_tenant_id, upload_file_id) - task_func.delay( # type: ignore + task_func.delay( rag_pipeline_invoke_entities_file_id=upload_file_id, tenant_id=self._dataset_tenant_id, ) - def _send_to_tenant_queue(self, upload_file_id: str, task_func: Callable[[str, str], None]): + def _send_to_tenant_queue(self, upload_file_id: str, task_func: Signature): logger.info("tenant %s send file %s to tenant queue", self._dataset_tenant_id, upload_file_id) if self._tenant_isolated_task_queue.get_task_key(): # Add to waiting queue using List operations (lpush) @@ -60,20 +61,35 @@ class RagPipelineTaskProxy: else: # Set flag and execute task self._tenant_isolated_task_queue.set_task_waiting_time() - task_func.delay( # type: ignore + task_func.delay( rag_pipeline_invoke_entities_file_id=upload_file_id, tenant_id=self._dataset_tenant_id, ) logger.info("tenant %s init tasks: %s", self._dataset_tenant_id, upload_file_id) def _send_to_default_tenant_queue(self, upload_file_id: str): - self._send_to_tenant_queue(upload_file_id, rag_pipeline_run_task) + self._send_to_tenant_queue( + upload_file_id, + signature("tasks.rag_pipeline.rag_pipeline_run_task.rag_pipeline_run_task", queue="pipeline"), + ) def _send_to_priority_tenant_queue(self, upload_file_id: str): - self._send_to_tenant_queue(upload_file_id, priority_rag_pipeline_run_task) + self._send_to_tenant_queue( + upload_file_id, + signature( + "tasks.rag_pipeline.priority_rag_pipeline_run_task.priority_rag_pipeline_run_task", + queue="priority_pipeline", + ), + ) def _send_to_priority_direct_queue(self, upload_file_id: str): - self._send_to_direct_queue(upload_file_id, priority_rag_pipeline_run_task) + self._send_to_direct_queue( + upload_file_id, + signature( + "tasks.rag_pipeline.priority_rag_pipeline_run_task.priority_rag_pipeline_run_task", + queue="priority_pipeline", + ), + ) def _dispatch(self): upload_file_id = self._upload_invoke_entities() diff --git a/api/services/rag_pipeline/rag_pipeline_transform_service.py b/api/services/rag_pipeline/rag_pipeline_transform_service.py index 9304257de9a..fbf4dd526c3 100644 --- a/api/services/rag_pipeline/rag_pipeline_transform_service.py +++ b/api/services/rag_pipeline/rag_pipeline_transform_service.py @@ -19,6 +19,7 @@ from factories import variable_factory from models.dataset import Dataset, Document, DocumentPipelineExecutionLog, Pipeline from models.enums import DatasetRuntimeMode, DataSourceType from models.workflow import Workflow, WorkflowType +from repositories.knowledge.dataset_read_repository import get_dataset_doc_form from services.entities.knowledge_entities.rag_pipeline_entities import KnowledgeConfiguration, RetrievalSetting from services.errors.rag_pipeline import RagPipelineResourceNotFoundError from services.file_service import FileService @@ -56,7 +57,7 @@ class RagPipelineTransformService: if not datasource_type and not indexing_technique: return self._transform_to_empty_pipeline(dataset, account_id=account_id, session=session) - doc_form = dataset.get_doc_form(session=session) + doc_form = get_dataset_doc_form(dataset, session=session) if not doc_form: return self._transform_to_empty_pipeline(dataset, account_id=account_id, session=session) retrieval_model = RetrievalSetting.model_validate(dataset.retrieval_model) if dataset.retrieval_model else None diff --git a/api/services/rbac_resource_service.py b/api/services/rbac_resource_service.py index fa34638f0c6..65026008267 100644 --- a/api/services/rbac_resource_service.py +++ b/api/services/rbac_resource_service.py @@ -5,7 +5,7 @@ from typing import TYPE_CHECKING from sqlalchemy import select from sqlalchemy.orm import Session, scoped_session -from models.dataset import Dataset +from models.dataset import Dataset, Document from models.model import App if TYPE_CHECKING: @@ -32,6 +32,15 @@ class RBACResourceService: select(Dataset.maintainer).where(Dataset.id == dataset_id, Dataset.tenant_id == tenant_id) ) + @staticmethod + def get_dataset_id_by_document(session: Session | scoped_session, tenant_id: str, document_id: str) -> str | None: + dataset_id = session.scalar( + select(Dataset.id) + .join(Document, Document.dataset_id == Dataset.id) + .where(Dataset.tenant_id == tenant_id, Document.tenant_id == tenant_id, Document.id == document_id) + ) + return None if dataset_id is None else str(dataset_id) + @staticmethod def get_dataset_id_by_pipeline(session: Session | scoped_session, tenant_id: str, pipeline_id: str) -> str | None: dataset_id = session.scalar( diff --git a/api/services/tools/builtin_tools_manage_service.py b/api/services/tools/builtin_tools_manage_service.py index dc82f91d3c1..dd7fd0f2152 100644 --- a/api/services/tools/builtin_tools_manage_service.py +++ b/api/services/tools/builtin_tools_manage_service.py @@ -34,6 +34,7 @@ from models.account import Account from models.enums import PermissionEnum from models.provider_ids import ToolProviderID from models.tools import BuiltinToolProvider, ToolOAuthSystemClient, ToolOAuthTenantClient +from services.credentials.query import CredentialQuery, ToolCredentialRecord from services.tools.tools_transform_service import ToolTransformService logger = logging.getLogger(__name__) @@ -310,7 +311,7 @@ class BuiltinToolManageService: @staticmethod def create_tool_encrypter( tenant_id: str, - db_provider: BuiltinToolProvider, + db_provider: BuiltinToolProvider | ToolCredentialRecord, provider: str, provider_controller: BuiltinToolProviderController, ): @@ -346,9 +347,10 @@ class BuiltinToolManageService: def get_builtin_tool_provider_credentials( tenant_id: str, provider_name: str, - session: Session, user: Account | None = None, include_credential_ids: list[str] | None = None, + *, + credential_query: CredentialQuery, ) -> list[ToolProviderCredentialApiEntity]: """ get builtin tool provider credentials, filtered by visibility. @@ -364,84 +366,45 @@ class BuiltinToolManageService: rows are marked with ``from_other_member=True`` so the UI can render them as borrowed-from-teammate (selectable but not editable). """ - from models.credential_permission import CredentialType as CredPermType - from services.credential_permission_service import CredentialPermissionService - - with session.no_autoflush: - base_filter = ( - BuiltinToolProvider.tenant_id == tenant_id, - BuiltinToolProvider.provider == provider_name, + providers = credential_query.list_tools( + workspace_id=tenant_id, + provider=provider_name, + actor_id=user.id if user else None, + include_credential_ids=include_credential_ids or (), + ) + if not providers: + return [] + provider_controller = ToolManager.get_builtin_provider(providers[0].provider, tenant_id) + credentials: list[ToolProviderCredentialApiEntity] = [] + for provider in providers: + encrypter, _ = BuiltinToolManageService.create_tool_encrypter( + tenant_id, provider, provider.provider, provider_controller ) - order = (BuiltinToolProvider.is_default.desc(), BuiltinToolProvider.created_at.asc()) - visible_query = select(BuiltinToolProvider).where(*base_filter).order_by(*order) - if user is not None: - visible_query = CredentialPermissionService.apply_visibility_filter( - visible_query, - model_id_column=BuiltinToolProvider.id, - model_user_id_column=BuiltinToolProvider.user_id, - model_visibility_column=BuiltinToolProvider.visibility, - credential_type=CredPermType.BUILTIN_TOOL_PROVIDER, - user=user, + decrypted = encrypter.mask_plugin_credentials(encrypter.decrypt(provider.credentials)) + credentials.append( + ToolProviderCredentialApiEntity( + id=provider.id, + name=provider.name, + provider=provider.provider, + credential_type=provider.credential_type, + is_default=provider.is_default, + credentials=dict(decrypted), + visibility=provider.visibility, + created_by=provider.created_by, + partial_member_list=list(provider.partial_member_ids), + from_other_member=provider.from_other_member, ) - visible_providers = list(session.scalars(visible_query).all()) - - # Fetch any explicitly-included IDs that the visibility filter excluded. - borrowed_ids: set[str] = set() - borrowed_providers: list[BuiltinToolProvider] = [] - if include_credential_ids: - visible_id_set = {p.id for p in visible_providers} - wanted_ids = [cid for cid in include_credential_ids if cid and cid not in visible_id_set] - if wanted_ids: - borrowed_query = ( - select(BuiltinToolProvider) - .where(*base_filter, BuiltinToolProvider.id.in_(wanted_ids)) - .order_by(*order) - ) - borrowed_providers = list(session.scalars(borrowed_query).all()) - borrowed_ids = {p.id for p in borrowed_providers} - - providers = visible_providers + borrowed_providers - if not providers: - return [] - - # Only the first visible row should be flagged is_default in the response. - if visible_providers: - visible_providers[0].is_default = True - provider_controller = ToolManager.get_builtin_provider(providers[0].provider, tenant_id) - - credentials: list[ToolProviderCredentialApiEntity] = [] - for provider in providers: - encrypter, _ = BuiltinToolManageService.create_tool_encrypter( - tenant_id, provider, provider.provider, provider_controller - ) - decrypt_credential = encrypter.mask_plugin_credentials(encrypter.decrypt(provider.credentials)) - credential_entity = ToolTransformService.convert_builtin_provider_to_credential_entity( - provider=provider, - credentials=dict(decrypt_credential), - ) - # Attach visibility, creator, and partial member list to the response entity - vis = getattr(provider, "visibility", "all_team_members") - vis_str = vis.value if hasattr(vis, "value") else str(vis) - credential_entity.visibility = vis_str - credential_entity.created_by = getattr(provider, "user_id", "") or "" - if vis_str == "partial_members": - credential_entity.partial_member_list = list( - CredentialPermissionService.get_partial_member_list( - provider.id, CredPermType.BUILTIN_TOOL_PROVIDER, session=session - ) - ) - if provider.id in borrowed_ids: - credential_entity.from_other_member = True - credentials.append(credential_entity) - return credentials + ) + return credentials @staticmethod def get_builtin_tool_provider_credential_info( tenant_id: str, provider: str, - session: Session, user: Account | None = None, include_credential_ids: list[str] | None = None, + *, + credential_query: CredentialQuery, ) -> ToolProviderCredentialInfoApiEntity: """ get builtin tool provider credential info @@ -451,9 +414,9 @@ class BuiltinToolManageService: credentials = BuiltinToolManageService.get_builtin_tool_provider_credentials( tenant_id, provider, - session=session, user=user, include_credential_ids=include_credential_ids, + credential_query=credential_query, ) credential_info = ToolProviderCredentialInfoApiEntity( supported_credential_types=supported_credential_types, diff --git a/api/services/trigger/trigger_provider_service.py b/api/services/trigger/trigger_provider_service.py index 8506c523a61..58112bcc273 100644 --- a/api/services/trigger/trigger_provider_service.py +++ b/api/services/trigger/trigger_provider_service.py @@ -5,10 +5,12 @@ import uuid from collections.abc import Mapping from typing import TYPE_CHECKING, Any, TypedDict +from services.credentials.query import CredentialQuery + if TYPE_CHECKING: from models.account import Account -from sqlalchemy import delete, desc, func, select +from sqlalchemy import delete, func, select from sqlalchemy.orm import Session, sessionmaker from configs import dify_config @@ -39,7 +41,6 @@ from models.trigger import ( TriggerOAuthSystemClient, TriggerOAuthTenantClient, TriggerSubscription, - WorkflowPluginTrigger, ) logger = logging.getLogger(__name__) @@ -73,48 +74,28 @@ class TriggerProviderService: tenant_id: str, provider_id: TriggerProviderID, user: "Account | None" = None, + *, + credential_query: CredentialQuery, ) -> list[TriggerProviderSubscriptionApiEntity]: """List all trigger subscriptions for the current tenant, filtered by visibility.""" - from models.credential_permission import CredentialType as CredPermType - from services.credential_permission_service import CredentialPermissionService - - subscriptions: list[TriggerProviderSubscriptionApiEntity] = [] - workflows_in_use_map: dict[str, int] = {} - with Session(db.engine, expire_on_commit=False) as session: - # Get all subscriptions with visibility filtering - query = ( - select(TriggerSubscription) - .where( - TriggerSubscription.tenant_id == tenant_id, - TriggerSubscription.provider_id == str(provider_id), - ) - .order_by(desc(TriggerSubscription.created_at)) + subscriptions = [ + TriggerProviderSubscriptionApiEntity( + id=row.id, + name=row.name, + provider=row.provider, + credential_type=row.credential_type, + credentials=dict(row.credentials), + endpoint=generate_plugin_trigger_endpoint_url(row.endpoint_id), + parameters=dict(row.parameters), + properties=dict(row.properties), + workflows_in_use=row.workflows_in_use, ) - if user is not None: - query = CredentialPermissionService.apply_visibility_filter( - query, - model_id_column=TriggerSubscription.id, - model_user_id_column=TriggerSubscription.user_id, - model_visibility_column=TriggerSubscription.visibility, - credential_type=CredPermType.TRIGGER_SUBSCRIPTION, - user=user, - ) - subscriptions_db = session.scalars(query).all() - subscriptions = [subscription.to_api_entity() for subscription in subscriptions_db] - if not subscriptions: - return [] - usage_counts = session.execute( - select( - WorkflowPluginTrigger.subscription_id, - func.count(func.distinct(WorkflowPluginTrigger.app_id)).label("app_count"), - ) - .where( - WorkflowPluginTrigger.tenant_id == tenant_id, - WorkflowPluginTrigger.subscription_id.in_([s.id for s in subscriptions]), - ) - .group_by(WorkflowPluginTrigger.subscription_id) - ).all() - workflows_in_use_map = {str(row.subscription_id): int(row.app_count) for row in usage_counts} + for row in credential_query.list_trigger_subscriptions( + workspace_id=tenant_id, provider=str(provider_id), actor_id=user.id if user else None + ) + ] + if not subscriptions: + return [] provider_controller = TriggerManager.get_trigger_provider(tenant_id, provider_id) for subscription in subscriptions: @@ -135,8 +116,6 @@ class TriggerProviderService: properties_encrypter.mask_credentials(dict(properties_encrypter.decrypt(subscription.properties))) ) subscription.parameters = dict(subscription.parameters) - count = workflows_in_use_map.get(subscription.id) - subscription.workflows_in_use = count if count is not None else 0 return subscriptions diff --git a/api/services/vector_service.py b/api/services/vector_service.py deleted file mode 100644 index 9d0505ed78e..00000000000 --- a/api/services/vector_service.py +++ /dev/null @@ -1,379 +0,0 @@ -import logging - -from sqlalchemy import delete, select -from sqlalchemy.orm import Session - -from core.credit_usage import CreditUsageCreatedBy -from core.model_context import with_credit_usage_created_by -from core.model_manager import ModelInstance, ModelManager -from core.rag.datasource.keyword.keyword_factory import Keyword -from core.rag.datasource.vdb.vector_factory import Vector -from core.rag.entities import ParentMode -from core.rag.index_processor.constant.doc_type import DocType -from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType -from core.rag.index_processor.index_processor_base import BaseIndexProcessor -from core.rag.index_processor.index_processor_factory import IndexProcessorFactory -from core.rag.models.document import AttachmentDocument, Document -from graphon.model_runtime.entities.model_entities import ModelType -from models import UploadFile -from models.dataset import ChildChunk, Dataset, DatasetProcessRule, DocumentSegment, SegmentAttachmentBinding -from models.dataset import Document as DatasetDocument -from models.enums import SegmentType - -logger = logging.getLogger(__name__) - - -class VectorService: - @classmethod - @with_credit_usage_created_by(CreditUsageCreatedBy.KNOWLEDGE_INDEXING) - def create_segments_vector( - cls, - keywords_list: list[list[str]] | None, - segments: list[DocumentSegment], - dataset: Dataset, - doc_form: str, - session: Session, - ): - """Create vector records for document segments using the caller's active DB session.""" - documents: list[Document] = [] - multimodal_documents: list[AttachmentDocument] = [] - - for segment in segments: - if doc_form == IndexStructureType.PARENT_CHILD_INDEX: - dataset_document = session.get(DatasetDocument, segment.document_id) - if not dataset_document: - logger.warning( - "Expected DatasetDocument record to exist, but none was found, document_id=%s, segment_id=%s", - segment.document_id, - segment.id, - ) - continue - # get the process rule - processing_rule = session.get(DatasetProcessRule, dataset_document.dataset_process_rule_id) - if not processing_rule: - raise ValueError("No processing rule found.") - # get embedding model instance - if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY: - # check embedding model setting - model_manager = ModelManager.for_tenant(tenant_id=dataset.tenant_id) - - if dataset.embedding_model_provider: - embedding_model_instance = model_manager.get_model_instance( - tenant_id=dataset.tenant_id, - provider=dataset.embedding_model_provider, - model_type=ModelType.TEXT_EMBEDDING, - model=dataset.embedding_model, - ) - else: - embedding_model_instance = model_manager.get_default_model_instance( - tenant_id=dataset.tenant_id, - model_type=ModelType.TEXT_EMBEDDING, - ) - else: - raise ValueError("The knowledge base index technique is not high quality!") - cls.generate_child_chunks( - segment, - dataset_document, - dataset, - embedding_model_instance, - processing_rule, - False, - session=session, - ) - else: - rag_document = Document( - page_content=segment.content, - metadata={ - "doc_id": segment.index_node_id, - "doc_hash": segment.index_node_hash, - "document_id": segment.document_id, - "dataset_id": segment.dataset_id, - "doc_type": DocType.TEXT, - }, - ) - documents.append(rag_document) - if dataset.is_multimodal: - for attachment in segment.get_attachments(session=session): - multimodal_document: AttachmentDocument = AttachmentDocument( - page_content=attachment["name"], - metadata={ - "doc_id": attachment["id"], - "doc_hash": "", - "document_id": segment.document_id, - "dataset_id": segment.dataset_id, - "doc_type": DocType.IMAGE, - }, - ) - multimodal_documents.append(multimodal_document) - index_processor: BaseIndexProcessor = IndexProcessorFactory(doc_form).init_index_processor() - - if len(documents) > 0: - index_processor.load( - dataset, documents, None, with_keywords=True, keywords_list=keywords_list, session=session - ) - if len(multimodal_documents) > 0: - index_processor.load(dataset, [], multimodal_documents, with_keywords=False, session=session) - - @classmethod - @with_credit_usage_created_by(CreditUsageCreatedBy.KNOWLEDGE_INDEXING) - def update_segment_vector( - cls, keywords: list[str] | None, segment: DocumentSegment, dataset: Dataset, session: Session - ): - # update segment index task - - # format new index - document = Document( - page_content=segment.content, - metadata={ - "doc_id": segment.index_node_id, - "doc_hash": segment.index_node_hash, - "document_id": segment.document_id, - "dataset_id": segment.dataset_id, - }, - ) - assert segment.index_node_id - if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY: - # update vector index - vector = Vector(dataset=dataset, session=session) - vector.delete_by_ids([segment.index_node_id]) - vector.add_texts([document], duplicate_check=True) - else: - # update keyword index - keyword = Keyword(dataset) - keyword.delete_by_ids([segment.index_node_id], session) - - # save keyword index - if keywords and len(keywords) > 0: - keyword.add_texts([document], session, keywords_list=[keywords]) - else: - keyword.add_texts([document], session) - - @classmethod - @with_credit_usage_created_by(CreditUsageCreatedBy.KNOWLEDGE_INDEXING) - def generate_child_chunks( - cls, - segment: DocumentSegment, - dataset_document: DatasetDocument, - dataset: Dataset, - embedding_model_instance: ModelInstance, - processing_rule: DatasetProcessRule, - regenerate: bool = False, - *, - session: Session, - ): - """Generate child chunks and persist them with the caller's active DB session.""" - index_processor = IndexProcessorFactory(dataset.get_doc_form(session=session)).init_index_processor() - assert segment.index_node_id - if regenerate: - # delete child chunks - index_processor.clean( - dataset, - [segment.index_node_id], - with_keywords=True, - delete_child_chunks=True, - session=session, - ) - - # generate child chunks - document = Document( - page_content=segment.content, - metadata={ - "doc_id": segment.index_node_id, - "doc_hash": segment.index_node_hash, - "document_id": segment.document_id, - "dataset_id": segment.dataset_id, - "doc_type": DocType.TEXT, - }, - ) - # use full doc mode to generate segment's child chunk - processing_rule_dict = processing_rule.to_dict() - if processing_rule_dict["rules"] is not None: - processing_rule_dict["rules"]["parent_mode"] = ParentMode.FULL_DOC - documents = index_processor.transform( - [document], - embedding_model_instance=embedding_model_instance, - process_rule=processing_rule_dict, - tenant_id=dataset.tenant_id, - doc_language=dataset_document.doc_language, - session=session, - ) - # save child chunks - if documents and documents[0].children: - index_processor.load(dataset, documents, session=session) - - for position, child_chunk in enumerate(documents[0].children, start=1): - child_segment = ChildChunk( - tenant_id=dataset.tenant_id, - dataset_id=dataset.id, - document_id=dataset_document.id, - segment_id=segment.id, - position=position, - index_node_id=child_chunk.metadata["doc_id"], - index_node_hash=child_chunk.metadata["doc_hash"], - content=child_chunk.page_content, - word_count=len(child_chunk.page_content), - type=SegmentType.AUTOMATIC, - created_by=dataset_document.created_by, - ) - session.add(child_segment) - session.flush() - - @classmethod - @with_credit_usage_created_by(CreditUsageCreatedBy.KNOWLEDGE_INDEXING) - def create_child_chunk_vector(cls, child_segment: ChildChunk, dataset: Dataset, *, session: Session): - child_document = Document( - page_content=child_segment.content, - metadata={ - "doc_id": child_segment.index_node_id, - "doc_hash": child_segment.index_node_hash, - "document_id": child_segment.document_id, - "dataset_id": child_segment.dataset_id, - }, - ) - if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY: - # save vector index - vector = Vector(dataset=dataset, session=session) - vector.add_texts([child_document], duplicate_check=True) - - @classmethod - @with_credit_usage_created_by(CreditUsageCreatedBy.KNOWLEDGE_INDEXING) - def update_child_chunk_vector( - cls, - new_child_chunks: list[ChildChunk], - update_child_chunks: list[ChildChunk], - delete_child_chunks: list[ChildChunk], - dataset: Dataset, - *, - session: Session, - ): - documents = [] - delete_node_ids = [] - for new_child_chunk in new_child_chunks: - new_child_document = Document( - page_content=new_child_chunk.content, - metadata={ - "doc_id": new_child_chunk.index_node_id, - "doc_hash": new_child_chunk.index_node_hash, - "document_id": new_child_chunk.document_id, - "dataset_id": new_child_chunk.dataset_id, - }, - ) - documents.append(new_child_document) - for update_child_chunk in update_child_chunks: - assert update_child_chunk.index_node_id - child_document = Document( - page_content=update_child_chunk.content, - metadata={ - "doc_id": update_child_chunk.index_node_id, - "doc_hash": update_child_chunk.index_node_hash, - "document_id": update_child_chunk.document_id, - "dataset_id": update_child_chunk.dataset_id, - }, - ) - documents.append(child_document) - delete_node_ids.append(update_child_chunk.index_node_id) - for delete_child_chunk in delete_child_chunks: - assert delete_child_chunk.index_node_id - delete_node_ids.append(delete_child_chunk.index_node_id) - if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY: - # update vector index - vector = Vector(dataset=dataset, session=session) - if delete_node_ids: - vector.delete_by_ids(delete_node_ids) - if documents: - vector.add_texts(documents, duplicate_check=True) - - @classmethod - def delete_child_chunk_vector(cls, child_chunk: ChildChunk, dataset: Dataset, *, session: Session): - vector = Vector(dataset=dataset, session=session) - assert child_chunk.index_node_id - vector.delete_by_ids([child_chunk.index_node_id]) - - @classmethod - @with_credit_usage_created_by(CreditUsageCreatedBy.KNOWLEDGE_INDEXING) - def update_multimodel_vector( - cls, segment: DocumentSegment, attachment_ids: list[str], dataset: Dataset, session: Session - ): - if dataset.indexing_technique != IndexTechniqueType.HIGH_QUALITY: - return - - attachments = segment.get_attachments(session=session) - old_attachment_ids = [attachment["id"] for attachment in attachments] if attachments else [] - - # Check if there's any actual change needed - if set(attachment_ids) == set(old_attachment_ids): - return - - try: - vector = Vector(dataset=dataset, session=session) - if dataset.is_multimodal: - # Delete old vectors if they exist - if old_attachment_ids: - vector.delete_by_ids(old_attachment_ids) - - # Delete existing segment attachment bindings in one operation - session.execute(delete(SegmentAttachmentBinding).where(SegmentAttachmentBinding.segment_id == segment.id)) - - if not attachment_ids: - session.flush() - return - - # Bulk fetch upload files - only fetch needed fields - upload_file_list = session.scalars(select(UploadFile).where(UploadFile.id.in_(attachment_ids))).all() - - if not upload_file_list: - session.flush() - return - - # Create a mapping for quick lookup - upload_file_map = {upload_file.id: upload_file for upload_file in upload_file_list} - - # Prepare batch operations - bindings = [] - documents = [] - - # Create common metadata base to avoid repetition - base_metadata = { - "doc_hash": "", - "document_id": segment.document_id, - "dataset_id": segment.dataset_id, - "doc_type": DocType.IMAGE, - } - - # Process attachments in the order specified by attachment_ids - for attachment_id in attachment_ids: - upload_file = upload_file_map.get(attachment_id) - if not upload_file: - logger.warning("Upload file not found for attachment_id: %s", attachment_id) - continue - - # Create segment attachment binding - bindings.append( - SegmentAttachmentBinding( - tenant_id=segment.tenant_id, - dataset_id=segment.dataset_id, - document_id=segment.document_id, - segment_id=segment.id, - attachment_id=upload_file.id, - ) - ) - - # Create document for vector indexing - documents.append( - Document(page_content=upload_file.name, metadata={**base_metadata, "doc_id": upload_file.id}) - ) - - # Bulk insert all bindings at once - if bindings: - session.add_all(bindings) - - # Add documents to vector store if any - if documents and dataset.is_multimodal: - vector.create_multimodal(documents) - - session.flush() - - except Exception: - logger.exception("Failed to update multimodal vector for segment %s", segment.id) - session.rollback() - raise diff --git a/api/tasks/add_document_to_index_task.py b/api/tasks/add_document_to_index_task.py index 9bd08b7daa8..76486cb51e2 100644 --- a/api/tasks/add_document_to_index_task.py +++ b/api/tasks/add_document_to_index_task.py @@ -15,6 +15,12 @@ from libs.datetime_utils import naive_utc_now from models.dataset import DatasetAutoDisableLog, DocumentSegment from models.dataset import Document as DatasetDocument from models.enums import IndexingStatus, SegmentStatus +from repositories.knowledge.dataset_read_repository import ( + get_dataset_doc_form, + get_document_dataset, + get_segment_child_chunks, +) +from repositories.knowledge.segment_read_adapter import get_segment_attachments logger = logging.getLogger(__name__) @@ -44,7 +50,7 @@ def add_document_to_index_task(dataset_document_id: str): indexing_cache_key = f"document_{dataset_document.id}_indexing" try: - dataset = dataset_document.get_dataset(session=session) + dataset = get_document_dataset(dataset_document, session=session) if not dataset: raise Exception(f"Document {dataset_document.id} dataset {dataset_document.dataset_id} doesn't exist.") @@ -70,7 +76,7 @@ def add_document_to_index_task(dataset_document_id: str): }, ) if dataset_document.doc_form == IndexStructureType.PARENT_CHILD_INDEX: - child_chunks = segment.get_child_chunks(session=session) + child_chunks = get_segment_child_chunks(segment, session=session) if child_chunks: child_documents = [] for child_chunk in child_chunks: @@ -86,7 +92,7 @@ def add_document_to_index_task(dataset_document_id: str): child_documents.append(child_document) document.children = child_documents if dataset.is_multimodal: - for attachment in segment.get_attachments(session=session): + for attachment in get_segment_attachments(segment, session=session): multimodal_documents.append( AttachmentDocument( page_content=attachment["name"], @@ -101,7 +107,7 @@ def add_document_to_index_task(dataset_document_id: str): ) documents.append(document) - index_type = dataset.get_doc_form(session=session) + index_type = get_dataset_doc_form(dataset, session=session) index_processor = IndexProcessorFactory(index_type).init_index_processor() index_processor.load(dataset, documents, multimodal_documents=multimodal_documents, session=session) @@ -119,12 +125,12 @@ def add_document_to_index_task(dataset_document_id: str): session.commit() # Enable summary indexes for all segments in this document - from services.summary_index_service import SummaryIndexService + from services.knowledge.summaries.adapters import SummaryIndexAdapter segment_ids_list = [segment.id for segment in segments] if segment_ids_list: try: - SummaryIndexService.enable_summaries_for_segments( + SummaryIndexAdapter.enable_summaries_for_segments( dataset=dataset, segment_ids=segment_ids_list, ) diff --git a/api/tasks/annotation/add_annotation_to_index_task.py b/api/tasks/annotation/add_annotation_to_index_task.py index 9e0a7d83cc1..509b9a0355c 100644 --- a/api/tasks/annotation/add_annotation_to_index_task.py +++ b/api/tasks/annotation/add_annotation_to_index_task.py @@ -9,7 +9,7 @@ from core.rag.datasource.vdb.vector_factory import Vector from core.rag.index_processor.constant.index_type import IndexTechniqueType from core.rag.models.document import Document from models.dataset import Dataset -from services.dataset_service import DatasetCollectionBindingService +from services.knowledge.dataset_service import DatasetCollectionBindingService logger = logging.getLogger(__name__) diff --git a/api/tasks/annotation/batch_import_annotations_task.py b/api/tasks/annotation/batch_import_annotations_task.py index 12af1a13df8..41656e1ab11 100644 --- a/api/tasks/annotation/batch_import_annotations_task.py +++ b/api/tasks/annotation/batch_import_annotations_task.py @@ -13,7 +13,7 @@ from core.rag.models.document import Document from extensions.ext_redis import redis_client from models.dataset import Dataset from models.model import App, AppAnnotationSetting, MessageAnnotation -from services.dataset_service import DatasetCollectionBindingService +from services.knowledge.dataset_service import DatasetCollectionBindingService logger = logging.getLogger(__name__) diff --git a/api/tasks/annotation/delete_annotation_index_task.py b/api/tasks/annotation/delete_annotation_index_task.py index 4e467e2e0ba..f86d06e153f 100644 --- a/api/tasks/annotation/delete_annotation_index_task.py +++ b/api/tasks/annotation/delete_annotation_index_task.py @@ -8,7 +8,7 @@ from core.db.session_factory import session_factory from core.rag.datasource.vdb.vector_factory import Vector from core.rag.index_processor.constant.index_type import IndexTechniqueType from models.dataset import Dataset -from services.dataset_service import DatasetCollectionBindingService +from services.knowledge.dataset_service import DatasetCollectionBindingService logger = logging.getLogger(__name__) diff --git a/api/tasks/annotation/enable_annotation_reply_task.py b/api/tasks/annotation/enable_annotation_reply_task.py index 1247cdf32b7..bf4410f10aa 100644 --- a/api/tasks/annotation/enable_annotation_reply_task.py +++ b/api/tasks/annotation/enable_annotation_reply_task.py @@ -14,7 +14,7 @@ from libs.datetime_utils import naive_utc_now from models.dataset import Dataset from models.enums import CollectionBindingType from models.model import App, AppAnnotationSetting, MessageAnnotation -from services.dataset_service import DatasetCollectionBindingService +from services.knowledge.dataset_service import DatasetCollectionBindingService logger = logging.getLogger(__name__) diff --git a/api/tasks/annotation/update_annotation_to_index_task.py b/api/tasks/annotation/update_annotation_to_index_task.py index 1d4c90155e2..611934c5cb1 100644 --- a/api/tasks/annotation/update_annotation_to_index_task.py +++ b/api/tasks/annotation/update_annotation_to_index_task.py @@ -9,7 +9,7 @@ from core.rag.datasource.vdb.vector_factory import Vector from core.rag.index_processor.constant.index_type import IndexTechniqueType from core.rag.models.document import Document from models.dataset import Dataset -from services.dataset_service import DatasetCollectionBindingService +from services.knowledge.dataset_service import DatasetCollectionBindingService logger = logging.getLogger(__name__) diff --git a/api/tasks/batch_clean_document_task.py b/api/tasks/batch_clean_document_task.py index 040b9b4f349..e158be32df5 100644 --- a/api/tasks/batch_clean_document_task.py +++ b/api/tasks/batch_clean_document_task.py @@ -8,11 +8,11 @@ from sqlalchemy import delete, select from sqlalchemy.engine import CursorResult from core.db.session_factory import session_factory -from core.rag.index_processor.index_processor_factory import IndexProcessorFactory from core.tools.utils.web_reader_tool import get_image_upload_file_ids from extensions.ext_storage import storage -from models.dataset import Dataset, DatasetMetadataBinding, DocumentSegment, SegmentAttachmentBinding +from models.dataset import DatasetMetadataBinding, DocumentSegment, SegmentAttachmentBinding from models.model import UploadFile +from services.knowledge.indexing.adapters.cleanup import clean_document_indexes from tasks.refresh_billing_vector_space_task import schedule_billing_vector_space_refresh logger = logging.getLogger(__name__) @@ -88,32 +88,21 @@ def batch_clean_document_task( files = session.scalars(select(UploadFile).where(UploadFile.id.in_(file_ids))).all() storage_keys_to_delete.extend([f.key for f in files if f and f.key]) - # ============ Step 2: Clean vector index (external service, fresh session for dataset) ============ - if index_node_ids: - try: - # Fetch dataset in a fresh session to avoid DetachedInstanceError - with session_factory.create_session() as session, session.begin(): - dataset = session.scalar(select(Dataset).where(Dataset.id == dataset_id).limit(1)) - if not dataset: - logger.warning("Dataset not found for vector index cleanup, dataset_id: %s", dataset_id) - else: - index_processor = IndexProcessorFactory(doc_form).init_index_processor() - index_processor.clean( - dataset, - index_node_ids, - with_keywords=True, - delete_child_chunks=True, - delete_summaries=True, - session=session, - ) - dataset_tenant_id = dataset.tenant_id - except Exception: - logger.exception( - "Failed to clean vector index for dataset_id: %s, document_ids: %s, index_node_ids count: %d", - dataset_id, - document_ids, - len(index_node_ids), - ) + # A retry can still have summaries even after its segments were deleted. + try: + dataset_tenant_id = clean_document_indexes( + dataset_id=dataset_id, + document_ids=document_ids, + doc_form=doc_form, + new_session=session_factory.create_session, + ) + except Exception: + logger.exception( + "Failed to clean vector index for dataset_id: %s, document_ids: %s, index_node_ids count: %d", + dataset_id, + document_ids, + len(index_node_ids), + ) # ============ Step 3: Delete metadata binding (separate short transaction) ============ try: diff --git a/api/tasks/batch_create_segment_to_index_task.py b/api/tasks/batch_create_segment_to_index_task.py index a4a5feed649..1abdffc02d8 100644 --- a/api/tasks/batch_create_segment_to_index_task.py +++ b/api/tasks/batch_create_segment_to_index_task.py @@ -23,7 +23,7 @@ from libs.datetime_utils import naive_utc_now from models.dataset import Dataset, Document, DocumentSegment from models.enums import SegmentStatus from models.model import UploadFile -from services.vector_service import VectorService +from services.knowledge.resource_scope import DatasetRef logger = logging.getLogger(__name__) @@ -61,11 +61,15 @@ def batch_create_segment_to_index_task( with session_factory.create_session() as session: try: - dataset = session.get(Dataset, dataset_id) + dataset = session.scalar(select(Dataset).where(Dataset.id == dataset_id, Dataset.tenant_id == tenant_id)) if not dataset: raise ValueError("Dataset not exist.") - dataset_document = session.get(Document, document_id) + dataset_document = session.scalar( + select(Document).where( + Document.id == document_id, Document.dataset_id == dataset_id, Document.tenant_id == tenant_id + ) + ) if not dataset_document: raise ValueError("Document not exist.") @@ -76,7 +80,9 @@ def batch_create_segment_to_index_task( ): raise ValueError("Document is not available.") - upload_file = session.get(UploadFile, upload_file_id) + upload_file = session.scalar( + select(UploadFile).where(UploadFile.id == upload_file_id, UploadFile.tenant_id == tenant_id) + ) if not upload_file: raise ValueError("UploadFile not found.") @@ -123,7 +129,7 @@ def batch_create_segment_to_index_task( if len(content) == 0: raise ValueError("The CSV file is empty.") - document_segments = [] + segment_ids: list[str] = [] embedding_model = None if dataset_config["indexing_technique"] == IndexTechniqueType.HIGH_QUALITY: model_manager = ModelManager.for_tenant(tenant_id=dataset_config["tenant_id"]) @@ -143,7 +149,11 @@ def batch_create_segment_to_index_task( with session_factory.create_session() as session, session.begin(): max_position = ( session.scalar( - select(func.max(DocumentSegment.position)).where(DocumentSegment.document_id == document_config["id"]) + select(func.max(DocumentSegment.position)).where( + DocumentSegment.tenant_id == tenant_id, + DocumentSegment.dataset_id == dataset_id, + DocumentSegment.document_id == document_id, + ) ) or 0 ) @@ -172,21 +182,28 @@ def batch_create_segment_to_index_task( segment_document.word_count += len(segment["answer"]) word_count_change += segment_document.word_count session.add(segment_document) - document_segments.append(segment_document) + segment_ids.append(segment_document.id) - with session_factory.create_session() as session, session.begin(): - dataset_document = session.get(Document, document_id) - if dataset_document: - assert dataset_document.word_count is not None - dataset_document.word_count += word_count_change - session.add(dataset_document) - - with session_factory.create_session() as session, session.begin(): - dataset = session.get(Dataset, dataset_id) - if dataset: - VectorService.create_segments_vector( - None, document_segments, dataset, document_config["doc_form"], session=session + dataset_document = session.scalar( + select(Document).where( + Document.id == document_id, Document.dataset_id == dataset_id, Document.tenant_id == tenant_id ) + ) + if dataset_document is None: + raise ValueError("Document no longer exists.") + dataset_document.word_count = (dataset_document.word_count or 0) + word_count_change + session.add(dataset_document) + + # The registry imports this task while assembling the batch dispatcher. + from extensions.ext_application_services import application_services + + try: + application_services().knowledge.segments.mutations.index_segments( + DatasetRef(tenant_id, dataset_id).document(document_id), segment_ids=segment_ids + ) + except Exception: + redis_client.setex(indexing_cache_key, 600, "error") + raise redis_client.setex(indexing_cache_key, 600, "completed") end_at = time.perf_counter() diff --git a/api/tasks/clean_document_task.py b/api/tasks/clean_document_task.py index e09743a0018..42d0d108df3 100644 --- a/api/tasks/clean_document_task.py +++ b/api/tasks/clean_document_task.py @@ -6,11 +6,11 @@ from celery import shared_task from sqlalchemy import delete, select from core.db.session_factory import session_factory -from core.rag.index_processor.index_processor_factory import IndexProcessorFactory from core.tools.utils.web_reader_tool import get_image_upload_file_ids from extensions.ext_storage import storage from models.dataset import Dataset, DatasetMetadataBinding, DocumentSegment, SegmentAttachmentBinding from models.model import UploadFile +from services.knowledge.indexing.adapters.cleanup import clean_document_indexes from tasks.refresh_billing_vector_space_task import schedule_billing_vector_space_refresh logger = logging.getLogger(__name__) @@ -67,39 +67,27 @@ def clean_document_task( logger.exception("Cleaned document when document deleted failed") return - # check segment is exist - if index_node_ids: - # Wrap vector / keyword index cleanup in try/except so that a transient - # failure here (e.g. billing API hiccup propagated via FeatureService when - # ModelManager is initialized inside ``Vector(dataset)``) does not abort - # the entire task and leave document_segments / child_chunks / image_files - # / metadata bindings stranded in PG. Mirrors the pattern already used in - # ``clean_dataset_task`` so the document row's hard delete (already - # committed by the caller) does not produce orphan PG rows just because - # the vector backend or one of its transitive dependencies was unhappy. - try: - index_processor = IndexProcessorFactory(doc_form).init_index_processor() - with session_factory.create_session() as session, session.begin(): - dataset = session.scalar(select(Dataset).where(Dataset.id == dataset_id).limit(1)) - if dataset: - index_processor.clean( - dataset, - index_node_ids, - with_keywords=True, - delete_child_chunks=True, - delete_summaries=True, - session=session, - ) - vector_cleanup_succeeded = True - except Exception: - logger.exception( - "Failed to clean vector / keyword index in clean_document_task, " - "document_id=%s, dataset_id=%s, index_node_ids_count=%d. " - "Continuing with PG / storage cleanup; vector orphans can be reaped later.", - document_id, - dataset_id, - len(index_node_ids), + # A retry must also find summaries whose segments were already removed. + # Preserve relational cleanup if the external index is temporarily unavailable. + try: + vector_cleanup_succeeded = ( + clean_document_indexes( + dataset_id=dataset_id, + document_ids=[document_id], + doc_form=doc_form, + new_session=session_factory.create_session, ) + is not None + ) + except Exception: + logger.exception( + "Failed to clean vector / keyword index in clean_document_task, " + "document_id=%s, dataset_id=%s, index_node_ids_count=%d. " + "Continuing with PG / storage cleanup; vector orphans can be reaped later.", + document_id, + dataset_id, + len(index_node_ids), + ) total_image_files = [] with session_factory.create_session() as session, session.begin(): diff --git a/api/tasks/clean_notion_document_task.py b/api/tasks/clean_notion_document_task.py index e7500b9a8d7..7e7d1a7718c 100644 --- a/api/tasks/clean_notion_document_task.py +++ b/api/tasks/clean_notion_document_task.py @@ -6,8 +6,9 @@ from celery import shared_task from sqlalchemy import delete, select from core.db.session_factory import session_factory -from core.rag.index_processor.index_processor_factory import IndexProcessorFactory from models.dataset import Dataset, Document, DocumentSegment +from repositories.knowledge.dataset_read_repository import get_dataset_doc_form +from services.knowledge.indexing.adapters.cleanup import clean_document_indexes logger = logging.getLogger(__name__) @@ -30,8 +31,7 @@ def clean_notion_document_task(document_ids: list[str], dataset_id: str): if not dataset: raise Exception("Document has no dataset") - index_type = dataset.get_doc_form(session=session) - index_processor = IndexProcessorFactory(index_type).init_index_processor() + index_type = get_dataset_doc_form(dataset, session=session) document_delete_stmt = delete(Document).where(Document.id.in_(document_ids)) session.execute(document_delete_stmt) @@ -48,17 +48,12 @@ def clean_notion_document_task(document_ids: list[str], dataset_id: str): # exception escaping this task would produce orphans that no later request # can reference back. Mirrors the pattern in ``clean_dataset_task``. try: - with session_factory.create_session() as session, session.begin(): - dataset = session.scalar(select(Dataset).where(Dataset.id == dataset_id).limit(1)) - if dataset: - index_processor.clean( - dataset, - total_index_node_ids, - with_keywords=True, - delete_child_chunks=True, - delete_summaries=True, - session=session, - ) + clean_document_indexes( + dataset_id=dataset_id, + document_ids=document_ids, + doc_form=index_type, + new_session=session_factory.create_session, + ) except Exception: logger.exception( "Failed to clean vector / keyword index in clean_notion_document_task, " diff --git a/api/tasks/create_segment_to_index_task.py b/api/tasks/create_segment_to_index_task.py index dfa1fa7fd2b..eef7d0d2d50 100644 --- a/api/tasks/create_segment_to_index_task.py +++ b/api/tasks/create_segment_to_index_task.py @@ -12,6 +12,11 @@ from extensions.ext_redis import redis_client from libs.datetime_utils import naive_utc_now from models.dataset import DocumentSegment from models.enums import IndexingStatus, SegmentStatus +from repositories.knowledge.dataset_read_repository import ( + get_dataset_doc_form, + get_segment_dataset, + get_segment_document, +) logger = logging.getLogger(__name__) @@ -56,13 +61,13 @@ def create_segment_to_index_task(segment_id: str, keywords: list[str] | None = N }, ) - dataset = segment.get_dataset(session=session) + dataset = get_segment_dataset(segment, session=session) if not dataset: logger.info(click.style(f"Segment {segment.id} has no dataset, pass.", fg="cyan")) return - dataset_document = segment.get_document(session=session) + dataset_document = get_segment_document(segment, session=session) if not dataset_document: logger.info(click.style(f"Segment {segment.id} has no document, pass.", fg="cyan")) @@ -76,7 +81,7 @@ def create_segment_to_index_task(segment_id: str, keywords: list[str] | None = N logger.info(click.style(f"Segment {segment.id} document status is invalid, pass.", fg="cyan")) return - index_type = dataset.get_doc_form(session=session) + index_type = get_dataset_doc_form(dataset, session=session) index_processor = IndexProcessorFactory(index_type).init_index_processor() index_processor.load(dataset, [document], session=session) diff --git a/api/tasks/deal_dataset_index_update_task.py b/api/tasks/deal_dataset_index_update_task.py index dfbbe3de065..02eae8ab6cf 100644 --- a/api/tasks/deal_dataset_index_update_task.py +++ b/api/tasks/deal_dataset_index_update_task.py @@ -12,6 +12,8 @@ from core.rag.index_processor.index_processor_factory import IndexProcessorFacto from core.rag.models.document import AttachmentDocument, ChildDocument, Document from models.dataset import Dataset, DocumentSegment from models.dataset import Document as DatasetDocument +from repositories.knowledge.dataset_read_repository import get_dataset_doc_form, get_segment_child_chunks +from repositories.knowledge.segment_read_adapter import get_segment_attachments @shared_task(queue="dataset") @@ -31,7 +33,7 @@ def deal_dataset_index_update_task(dataset_id: str, action: str): if not dataset: raise Exception("Dataset not found") - index_type = dataset.get_doc_form(session=session) or IndexStructureType.PARAGRAPH_INDEX + index_type = get_dataset_doc_form(dataset, session=session) or IndexStructureType.PARAGRAPH_INDEX index_processor = IndexProcessorFactory(index_type).init_index_processor() if action == "upgrade": dataset_documents = session.scalars( @@ -150,7 +152,7 @@ def deal_dataset_index_update_task(dataset_id: str, action: str): }, ) if dataset_document.doc_form == IndexStructureType.PARENT_CHILD_INDEX: - child_chunks = segment.get_child_chunks(session=session) + child_chunks = get_segment_child_chunks(segment, session=session) if child_chunks: child_documents = [] for child_chunk in child_chunks: @@ -166,7 +168,7 @@ def deal_dataset_index_update_task(dataset_id: str, action: str): child_documents.append(child_document) document.children = child_documents if dataset.is_multimodal: - for attachment in segment.get_attachments(session=session): + for attachment in get_segment_attachments(segment, session=session): multimodal_documents.append( AttachmentDocument( page_content=attachment["name"], diff --git a/api/tasks/deal_dataset_vector_index_task.py b/api/tasks/deal_dataset_vector_index_task.py index 928134fb7b7..4cb8ac6b84c 100644 --- a/api/tasks/deal_dataset_vector_index_task.py +++ b/api/tasks/deal_dataset_vector_index_task.py @@ -12,6 +12,8 @@ from core.rag.index_processor.index_processor_factory import IndexProcessorFacto from core.rag.models.document import AttachmentDocument, ChildDocument, Document from models.dataset import Dataset, DocumentSegment from models.dataset import Document as DatasetDocument +from repositories.knowledge.dataset_read_repository import get_dataset_doc_form, get_segment_child_chunks +from repositories.knowledge.segment_read_adapter import get_segment_attachments logger = logging.getLogger(__name__) @@ -33,7 +35,7 @@ def deal_dataset_vector_index_task(dataset_id: str, action: str): if not dataset: raise Exception("Dataset not found") - index_type = dataset.get_doc_form(session=session) or IndexStructureType.PARAGRAPH_INDEX + index_type = get_dataset_doc_form(dataset, session=session) or IndexStructureType.PARAGRAPH_INDEX index_processor = IndexProcessorFactory(index_type).init_index_processor() if action == "remove": index_processor.clean(dataset, None, with_keywords=False, session=session) @@ -146,7 +148,7 @@ def deal_dataset_vector_index_task(dataset_id: str, action: str): }, ) if dataset_document.doc_form == IndexStructureType.PARENT_CHILD_INDEX: - child_chunks = segment.get_child_chunks(session=session) + child_chunks = get_segment_child_chunks(segment, session=session) if child_chunks: child_documents = [] for child_chunk in child_chunks: @@ -162,7 +164,7 @@ def deal_dataset_vector_index_task(dataset_id: str, action: str): child_documents.append(child_document) document.children = child_documents if dataset.is_multimodal: - for attachment in segment.get_attachments(session=session): + for attachment in get_segment_attachments(segment, session=session): multimodal_documents.append( AttachmentDocument( page_content=attachment["name"], diff --git a/api/tasks/disable_segment_from_index_task.py b/api/tasks/disable_segment_from_index_task.py index 6db3837e0cd..75c52b7d9f4 100644 --- a/api/tasks/disable_segment_from_index_task.py +++ b/api/tasks/disable_segment_from_index_task.py @@ -10,6 +10,7 @@ from core.rag.index_processor.index_processor_factory import IndexProcessorFacto from extensions.ext_redis import redis_client from models.dataset import DocumentSegment from models.enums import SegmentStatus +from repositories.knowledge.dataset_read_repository import get_segment_dataset, get_segment_document logger = logging.getLogger(__name__) @@ -38,13 +39,13 @@ def disable_segment_from_index_task(segment_id: str): indexing_cache_key = f"segment_{segment.id}_indexing" try: - dataset = segment.get_dataset(session=session) + dataset = get_segment_dataset(segment, session=session) if not dataset: logger.info(click.style(f"Segment {segment.id} has no dataset, pass.", fg="cyan")) return - dataset_document = segment.get_document(session=session) + dataset_document = get_segment_document(segment, session=session) if not dataset_document: logger.info(click.style(f"Segment {segment.id} has no document, pass.", fg="cyan")) @@ -65,10 +66,10 @@ def disable_segment_from_index_task(segment_id: str): session.commit() # Disable summary index for this segment - from services.summary_index_service import SummaryIndexService + from services.knowledge.summaries.adapters import SummaryIndexAdapter try: - SummaryIndexService.disable_summaries_for_segments( + SummaryIndexAdapter.disable_summaries_for_segments( dataset=dataset, segment_ids=[segment.id], disabled_by=segment.disabled_by, diff --git a/api/tasks/disable_segments_from_index_task.py b/api/tasks/disable_segments_from_index_task.py index da536e210e2..c6e5aa76db5 100644 --- a/api/tasks/disable_segments_from_index_task.py +++ b/api/tasks/disable_segments_from_index_task.py @@ -70,13 +70,13 @@ def disable_segments_from_index_task(segment_ids: list, dataset_id: str, documen session.commit() # Disable summary indexes for these segments - from services.summary_index_service import SummaryIndexService + from services.knowledge.summaries.adapters import SummaryIndexAdapter segment_ids_list = [segment.id for segment in segments] try: # Get disabled_by from first segment (they should all have the same disabled_by) disabled_by = segments[0].disabled_by if segments else None - SummaryIndexService.disable_summaries_for_segments( + SummaryIndexAdapter.disable_summaries_for_segments( dataset=dataset, segment_ids=segment_ids_list, disabled_by=disabled_by, diff --git a/api/tasks/document_indexing_sync_task.py b/api/tasks/document_indexing_sync_task.py index 7b12c045c50..f0e46c25bdf 100644 --- a/api/tasks/document_indexing_sync_task.py +++ b/api/tasks/document_indexing_sync_task.py @@ -6,14 +6,16 @@ import click from celery import shared_task from sqlalchemy import delete, select -from core.db.session_factory import session_factory -from core.indexing_runner import DocumentIsPausedError, IndexingRunner +from core.db.session_factory import get_session_maker, session_factory from core.rag.extractor.notion_extractor import NotionExtractor from core.rag.index_processor.index_processor_factory import IndexProcessorFactory +from extensions.application_services.data_sources import build_data_source_credentials from libs.datetime_utils import naive_utc_now from models.dataset import Dataset, Document, DocumentSegment from models.enums import IndexingStatus -from services.datasource_provider_service import DatasourceProviderService +from services.knowledge.indexing.adapters.execution import build_document_indexing_service +from services.knowledge.indexing.errors import DocumentIsPausedError +from services.knowledge.resource_scope import DatasetRef logger = logging.getLogger(__name__) @@ -72,7 +74,7 @@ def document_indexing_sync_task(dataset_id: str, document_id: str): index_node_ids = [segment.index_node_id for segment in segments if segment.index_node_id] # Get credentials from datasource provider - datasource_provider_service = DatasourceProviderService() + datasource_provider_service = build_data_source_credentials(database_client=get_session_maker()).providers credential = datasource_provider_service.get_datasource_credentials( tenant_id=tenant_id, credential_id=credential_id, @@ -97,6 +99,9 @@ def document_indexing_sync_task(dataset_id: str, document_id: str): return loader = NotionExtractor( + notion_token_loader=lambda: datasource_provider_service.get_stored_notion_access_token( + tenant_id=tenant_id, credential_id=credential_id + ), notion_workspace_id=workspace_id, notion_obj_id=page_id, notion_page_type=page_type, @@ -112,7 +117,7 @@ def document_indexing_sync_task(dataset_id: str, document_id: str): logger.info(click.style(f"Document {document_id} content changed, starting sync", fg="green")) try: - indexing_runner = IndexingRunner() + indexing_service = build_document_indexing_service(session_factory=session_factory.get_session_maker()) with session_factory.create_session() as session: document = session.scalar(select(Document).where(Document.id == document_id).limit(1)) if not document: @@ -151,13 +156,13 @@ def document_indexing_sync_task(dataset_id: str, document_id: str): segment_delete_stmt = delete(DocumentSegment).where(DocumentSegment.document_id == document_id) session.execute(segment_delete_stmt) + document_refs = [DatasetRef(document.tenant_id, document.dataset_id).document(document.id)] # Make the source update and segment deletion visible before extraction. session.commit() logger.info(click.style(f"Deleted segments for document {document_id}", fg="green")) - indexing_runner.run([document], session) - session.commit() + indexing_service.run(document_refs) end_at = time.perf_counter() logger.info(click.style(f"Sync completed for document {document_id} latency: {end_at - start_at}", fg="green")) except DocumentIsPausedError as ex: diff --git a/api/tasks/document_indexing_task.py b/api/tasks/document_indexing_task.py index 7d41c542bba..c9af37a1fcb 100644 --- a/api/tasks/document_indexing_task.py +++ b/api/tasks/document_indexing_task.py @@ -10,7 +10,6 @@ from sqlalchemy import select from configs import dify_config from core.db.session_factory import session_factory from core.entities.document_task import DocumentTask -from core.indexing_runner import DocumentIsPausedError, IndexingRunner from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType from core.rag.pipeline.queue import TenantIsolatedTaskQueue from enums import CloudPlan, DeploymentEdition @@ -18,6 +17,9 @@ from libs.datetime_utils import naive_utc_now from models.dataset import Dataset, Document from models.enums import IndexingStatus from services.feature_service import FeatureService +from services.knowledge.indexing.adapters.execution import build_document_indexing_service +from services.knowledge.indexing.errors import DocumentIsPausedError +from services.knowledge.resource_scope import DatasetRef from tasks.generate_summary_index_task import generate_summary_index_task logger = logging.getLogger(__name__) @@ -107,7 +109,9 @@ def _document_indexing(dataset_id: str, document_ids: Sequence[str]): # Phase 2: Execute indexing without holding locks from the parsing-status update. has_error = False try: - indexing_runner = IndexingRunner(enforce_vector_space_admission=True) + indexing_service = build_document_indexing_service( + session_factory=session_factory.get_session_maker(), enforce_vector_space_admission=True + ) with session_factory.create_session() as session: dataset = session.scalar(select(Dataset).where(Dataset.id == dataset_id).limit(1)) if not dataset: @@ -120,8 +124,9 @@ def _document_indexing(dataset_id: str, document_ids: Sequence[str]): ).all() ) - indexing_runner.run(documents, session) + document_refs = [DatasetRef(doc.tenant_id, doc.dataset_id).document(doc.id) for doc in documents] session.commit() + indexing_service.run(document_refs) end_at = time.perf_counter() logger.info(click.style(f"Processed dataset: {dataset_id} latency: {end_at - start_at}", fg="green")) diff --git a/api/tasks/document_indexing_update_task.py b/api/tasks/document_indexing_update_task.py index 18e2338e97a..73929fd54d0 100644 --- a/api/tasks/document_indexing_update_task.py +++ b/api/tasks/document_indexing_update_task.py @@ -6,7 +6,6 @@ from celery import shared_task from sqlalchemy import delete, select from core.db.session_factory import session_factory -from core.indexing_runner import DocumentIsPausedError, IndexingRunner from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType from core.rag.index_processor.index_processor_factory import IndexProcessorFactory from extensions.ext_storage import storage @@ -14,6 +13,9 @@ from libs.datetime_utils import naive_utc_now from models.dataset import Dataset, Document, DocumentSegment, SegmentAttachmentBinding from models.enums import IndexingStatus from models.model import UploadFile +from services.knowledge.indexing.adapters.execution import build_document_indexing_service +from services.knowledge.indexing.errors import DocumentIsPausedError +from services.knowledge.resource_scope import DatasetRef from tasks.generate_summary_index_task import generate_summary_index_task logger = logging.getLogger(__name__) @@ -164,9 +166,10 @@ def document_indexing_update_task(dataset_id: str, document_id: str): storage_key, ) - indexing_runner = IndexingRunner() - indexing_runner.run([document], session) + indexing_service = build_document_indexing_service(session_factory=session_factory.get_session_maker()) + document_refs = [DatasetRef(document.tenant_id, document.dataset_id).document(document.id)] session.commit() + indexing_service.run(document_refs) end_at = time.perf_counter() logger.info(click.style(f"update document: {document.id} latency: {end_at - start_at}", fg="green")) diff --git a/api/tasks/duplicate_document_indexing_task.py b/api/tasks/duplicate_document_indexing_task.py index c083dbdbbad..5afa6582811 100644 --- a/api/tasks/duplicate_document_indexing_task.py +++ b/api/tasks/duplicate_document_indexing_task.py @@ -9,7 +9,6 @@ from sqlalchemy import delete, select from configs import dify_config from core.db.session_factory import session_factory from core.entities.document_task import DocumentTask -from core.indexing_runner import DocumentIsPausedError, IndexingRunner from core.rag.index_processor.index_processor_factory import IndexProcessorFactory from core.rag.pipeline.queue import TenantIsolatedTaskQueue from enums import CloudPlan, DeploymentEdition @@ -17,6 +16,9 @@ from libs.datetime_utils import naive_utc_now from models.dataset import Dataset, Document, DocumentSegment from models.enums import IndexingStatus from services.feature_service import FeatureService +from services.knowledge.indexing.adapters.execution import build_document_indexing_service +from services.knowledge.indexing.errors import DocumentIsPausedError +from services.knowledge.resource_scope import DatasetRef logger = logging.getLogger(__name__) @@ -158,12 +160,11 @@ def _duplicate_document_indexing_task(dataset_id: str, document_ids: Sequence[st document.indexing_status = IndexingStatus.PARSING document.processing_started_at = naive_utc_now() session.add(document) + indexing_service = build_document_indexing_service(session_factory=session_factory.get_session_maker()) + document_refs = [DatasetRef(doc.tenant_id, doc.dataset_id).document(doc.id) for doc in documents] # Do not keep segment deletions or parsing status changes open during extraction. session.commit() - - indexing_runner = IndexingRunner() - indexing_runner.run(list(documents), session) - session.commit() + indexing_service.run(document_refs) end_at = time.perf_counter() logger.info(click.style(f"Processed dataset: {dataset_id} latency: {end_at - start_at}", fg="green")) except DocumentIsPausedError as ex: diff --git a/api/tasks/enable_segment_to_index_task.py b/api/tasks/enable_segment_to_index_task.py index 782dc82dde3..f4075b9df86 100644 --- a/api/tasks/enable_segment_to_index_task.py +++ b/api/tasks/enable_segment_to_index_task.py @@ -14,6 +14,12 @@ from extensions.ext_redis import redis_client from libs.datetime_utils import naive_utc_now from models.dataset import DocumentSegment from models.enums import IndexingStatus, SegmentStatus +from repositories.knowledge.dataset_read_repository import ( + get_segment_child_chunks, + get_segment_dataset, + get_segment_document, +) +from repositories.knowledge.segment_read_adapter import get_segment_attachments logger = logging.getLogger(__name__) @@ -52,13 +58,13 @@ def enable_segment_to_index_task(segment_id: str): }, ) - dataset = segment.get_dataset(session=session) + dataset = get_segment_dataset(segment, session=session) if not dataset: logger.info(click.style(f"Segment {segment.id} has no dataset, pass.", fg="cyan")) return - dataset_document = segment.get_document(session=session) + dataset_document = get_segment_document(segment, session=session) if not dataset_document: logger.info(click.style(f"Segment {segment.id} has no document, pass.", fg="cyan")) @@ -74,7 +80,7 @@ def enable_segment_to_index_task(segment_id: str): index_processor = IndexProcessorFactory(dataset_document.doc_form).init_index_processor() if dataset_document.doc_form == IndexStructureType.PARENT_CHILD_INDEX: - child_chunks = segment.get_child_chunks(session=session) + child_chunks = get_segment_child_chunks(segment, session=session) if child_chunks: child_documents = [] for child_chunk in child_chunks: @@ -91,7 +97,7 @@ def enable_segment_to_index_task(segment_id: str): document.children = child_documents multimodel_documents = [] if dataset.is_multimodal: - for attachment in segment.get_attachments(session=session): + for attachment in get_segment_attachments(segment, session=session): multimodel_documents.append( AttachmentDocument( page_content=attachment["name"], @@ -110,10 +116,10 @@ def enable_segment_to_index_task(segment_id: str): session.commit() # Enable summary index for this segment - from services.summary_index_service import SummaryIndexService + from services.knowledge.summaries.adapters import SummaryIndexAdapter try: - SummaryIndexService.enable_summaries_for_segments( + SummaryIndexAdapter.enable_summaries_for_segments( dataset=dataset, segment_ids=[segment.id], ) diff --git a/api/tasks/enable_segments_to_index_task.py b/api/tasks/enable_segments_to_index_task.py index 0de5effcd47..6a95d564728 100644 --- a/api/tasks/enable_segments_to_index_task.py +++ b/api/tasks/enable_segments_to_index_task.py @@ -15,6 +15,8 @@ from libs.datetime_utils import naive_utc_now from models.dataset import Dataset, DocumentSegment from models.dataset import Document as DatasetDocument from models.enums import SegmentStatus +from repositories.knowledge.dataset_read_repository import get_segment_child_chunks +from repositories.knowledge.segment_read_adapter import get_segment_attachments logger = logging.getLogger(__name__) @@ -73,7 +75,7 @@ def enable_segments_to_index_task(segment_ids: list, dataset_id: str, document_i ) if dataset_document.doc_form == IndexStructureType.PARENT_CHILD_INDEX: - child_chunks = segment.get_child_chunks(session=session) + child_chunks = get_segment_child_chunks(segment, session=session) if child_chunks: child_documents = [] for child_chunk in child_chunks: @@ -90,7 +92,7 @@ def enable_segments_to_index_task(segment_ids: list, dataset_id: str, document_i document.children = child_documents if dataset.is_multimodal: - for attachment in segment.get_attachments(session=session): + for attachment in get_segment_attachments(segment, session=session): multimodal_documents.append( AttachmentDocument( page_content=attachment["name"], @@ -110,11 +112,11 @@ def enable_segments_to_index_task(segment_ids: list, dataset_id: str, document_i session.commit() # Enable summary indexes for these segments - from services.summary_index_service import SummaryIndexService + from services.knowledge.summaries.adapters import SummaryIndexAdapter segment_ids_list = [segment.id for segment in segments] try: - SummaryIndexService.enable_summaries_for_segments( + SummaryIndexAdapter.enable_summaries_for_segments( dataset=dataset, segment_ids=segment_ids_list, ) diff --git a/api/tasks/generate_summary_index_task.py b/api/tasks/generate_summary_index_task.py index 9eda5716b8c..7beadb806f2 100644 --- a/api/tasks/generate_summary_index_task.py +++ b/api/tasks/generate_summary_index_task.py @@ -11,7 +11,7 @@ from core.db.session_factory import session_factory from core.rag.index_processor.constant.index_type import IndexTechniqueType from models.dataset import Dataset, DocumentSegment from models.dataset import Document as DatasetDocument -from services.summary_index_service import SummaryIndexService +from services.knowledge.summaries.adapters import SummaryIndexAdapter logger = logging.getLogger(__name__) @@ -45,7 +45,13 @@ def generate_summary_index_task(dataset_id: str, document_id: str, segment_ids: logger.error(click.style(f"Dataset not found: {dataset_id}", fg="red")) return - document = session.scalar(select(DatasetDocument).where(DatasetDocument.id == document_id).limit(1)) + document = session.scalar( + select(DatasetDocument).where( + DatasetDocument.id == document_id, + DatasetDocument.dataset_id == dataset.id, + DatasetDocument.tenant_id == dataset.tenant_id, + ) + ) if not document: logger.error(click.style(f"Document not found: {document_id}", fg="red")) return @@ -85,8 +91,9 @@ def generate_summary_index_task(dataset_id: str, document_id: str, segment_ids: # Determine if only parent chunks should be processed only_parent_chunks = dataset.chunk_structure == "parent_child_index" - # Generate summaries - summary_records = SummaryIndexService.generate_summaries_for_document( + # Release the task's read transaction before the adapter invokes models. + session.commit() + summary_records = SummaryIndexAdapter.generate_summaries_for_document( dataset=dataset, document=document, summary_index_setting=summary_index_setting, diff --git a/api/tasks/install_default_plugins_task.py b/api/tasks/install_default_plugins_task.py index b3ab5515d85..67b2a79b016 100644 --- a/api/tasks/install_default_plugins_task.py +++ b/api/tasks/install_default_plugins_task.py @@ -8,7 +8,7 @@ from configs import dify_config from core.helper import marketplace from core.plugin.entities.plugin_daemon import PluginInstallTaskStatus from core.plugin.plugin_service import PluginService -from services.model_provider_service import ModelProviderService +from services.model_provider.service import ModelProviderService logger = logging.getLogger(__name__) diff --git a/api/tasks/rag_pipeline/priority_rag_pipeline_run_task.py b/api/tasks/rag_pipeline/priority_rag_pipeline_run_task.py index c0e447f5afe..c76edab1b4f 100644 --- a/api/tasks/rag_pipeline/priority_rag_pipeline_run_task.py +++ b/api/tasks/rag_pipeline/priority_rag_pipeline_run_task.py @@ -14,15 +14,19 @@ from sqlalchemy import select from sqlalchemy.orm import Session, sessionmaker from configs import dify_config +from core.app.apps.pipeline.pipeline_generator import PipelineGenerator from core.app.entities.app_invoke_entities import InvokeFrom, RagPipelineGenerateEntity from core.app.entities.rag_pipeline_invoke_entities import RagPipelineInvokeEntity +from core.db.session_factory import get_session_maker from core.rag.pipeline.queue import TenantIsolatedTaskQueue from core.repositories.factory import DifyCoreRepositoryFactory +from extensions.application_services.data_sources import build_data_source_credentials from extensions.ext_database import db from models import Account, Tenant from models.dataset import Pipeline from models.enums import WorkflowRunTriggeredFrom from models.workflow import Workflow, WorkflowNodeExecutionTriggeredFrom +from repositories.knowledge.document_repository import SQLAlchemyDocumentRepository from services.file_service import FileService logger = logging.getLogger(__name__) @@ -170,9 +174,11 @@ def run_single_rag_pipeline_task(rag_pipeline_invoke_entity: Mapping[str, Any], # Direct execution without creating another thread # Since we're already in a thread pool, no need for nested threading - from core.app.apps.pipeline.pipeline_generator import PipelineGenerator - pipeline_generator = PipelineGenerator() + pipeline_generator = PipelineGenerator( + documents=SQLAlchemyDocumentRepository(session_factory=session_factory), + datasource_providers=build_data_source_credentials(database_client=get_session_maker()).providers, + ) # Using protected method intentionally for async execution with Session(db.engine, expire_on_commit=False) as session: pipeline_generator._generate( # type: ignore[attr-defined] diff --git a/api/tasks/rag_pipeline/rag_pipeline_run_task.py b/api/tasks/rag_pipeline/rag_pipeline_run_task.py index 1e8ab2b79f0..9774be64e9f 100644 --- a/api/tasks/rag_pipeline/rag_pipeline_run_task.py +++ b/api/tasks/rag_pipeline/rag_pipeline_run_task.py @@ -15,15 +15,19 @@ from sqlalchemy import select from sqlalchemy.orm import sessionmaker from configs import dify_config +from core.app.apps.pipeline.pipeline_generator import PipelineGenerator from core.app.entities.app_invoke_entities import InvokeFrom, RagPipelineGenerateEntity from core.app.entities.rag_pipeline_invoke_entities import RagPipelineInvokeEntity +from core.db.session_factory import get_session_maker from core.rag.pipeline.queue import TenantIsolatedTaskQueue from core.repositories.factory import DifyCoreRepositoryFactory +from extensions.application_services.data_sources import build_data_source_credentials from extensions.ext_database import db from models import Account, Tenant from models.dataset import Pipeline from models.enums import WorkflowRunTriggeredFrom from models.workflow import Workflow, WorkflowNodeExecutionTriggeredFrom +from repositories.knowledge.document_repository import SQLAlchemyDocumentRepository from services.file_service import FileService logger = logging.getLogger(__name__) @@ -184,9 +188,11 @@ def run_single_rag_pipeline_task(rag_pipeline_invoke_entity: Mapping[str, Any], # Direct execution without creating another thread # Since we're already in a thread pool, no need for nested threading - from core.app.apps.pipeline.pipeline_generator import PipelineGenerator - pipeline_generator = PipelineGenerator() + pipeline_generator = PipelineGenerator( + documents=SQLAlchemyDocumentRepository(session_factory=session_factory), + datasource_providers=build_data_source_credentials(database_client=get_session_maker()).providers, + ) # Using protected method intentionally for async execution pipeline_generator._generate( # type: ignore[attr-defined] session=session, diff --git a/api/tasks/recover_document_indexing_task.py b/api/tasks/recover_document_indexing_task.py index 4606981c44b..b081e1bdf8b 100644 --- a/api/tasks/recover_document_indexing_task.py +++ b/api/tasks/recover_document_indexing_task.py @@ -6,8 +6,10 @@ from celery import shared_task from sqlalchemy import select from core.db.session_factory import session_factory -from core.indexing_runner import DocumentIsPausedError, IndexingRunner from models.dataset import Document +from services.knowledge.indexing.adapters.execution import build_document_indexing_service +from services.knowledge.indexing.errors import DocumentIsPausedError +from services.knowledge.resource_scope import DatasetRef logger = logging.getLogger(__name__) @@ -34,14 +36,16 @@ def recover_document_indexing_task(dataset_id: str, document_id: str): return try: - indexing_runner = IndexingRunner() - if document.indexing_status in {"waiting", "parsing", "cleaning"}: - indexing_runner.run([document], session) - elif document.indexing_status == "splitting": - indexing_runner.run_in_splitting_status(document, session) - elif document.indexing_status == "indexing": - indexing_runner.run_in_indexing_status(document, session) + indexing_service = build_document_indexing_service(session_factory=session_factory.get_session_maker()) + document_ref = DatasetRef(document.tenant_id, document.dataset_id).document(document.id) + status = document.indexing_status session.commit() + if status in {"waiting", "parsing", "cleaning"}: + indexing_service.run([document_ref]) + elif status == "splitting": + indexing_service.run_in_splitting_status(document_ref) + elif status == "indexing": + indexing_service.run_in_indexing_status(document_ref) end_at = time.perf_counter() logger.info(click.style(f"Processed document: {document.id} latency: {end_at - start_at}", fg="green")) except DocumentIsPausedError as ex: diff --git a/api/tasks/regenerate_summary_index_task.py b/api/tasks/regenerate_summary_index_task.py index 69aeba389bf..2d89aca647a 100644 --- a/api/tasks/regenerate_summary_index_task.py +++ b/api/tasks/regenerate_summary_index_task.py @@ -15,7 +15,7 @@ from core.rag.index_processor.constant.index_type import IndexStructureType, Ind from models.dataset import Dataset, DocumentSegment, DocumentSegmentSummary from models.dataset import Document as DatasetDocument from models.enums import SummaryStatus -from services.summary_index_service import SummaryIndexService +from services.knowledge.summaries.adapters import SummaryIndexAdapter logger = logging.getLogger(__name__) @@ -150,22 +150,8 @@ def regenerate_summary_index_task( for segment, summary_record in segment_summary_pairs: try: - # Delete old vector - if summary_record.summary_index_node_id: - try: - from core.rag.datasource.vdb.vector_factory import Vector - - vector = Vector(dataset, session=session) - vector.delete_by_ids([summary_record.summary_index_node_id]) - except Exception as e: - logger.warning( - "Failed to delete old summary vector for segment %s: %s", - segment.id, - str(e), - ) - - # Re-vectorize with new embedding model - SummaryIndexService.vectorize_summary(summary_record, segment, dataset, session=session) + # The summary use case replaces the old vector after releasing this transaction. + SummaryIndexAdapter.vectorize_summary(summary_record, segment, dataset, session=session) session.commit() total_segments_processed += 1 @@ -261,7 +247,7 @@ def regenerate_summary_index_task( continue # Regenerate both summary content and vectors (for summary_model change) - SummaryIndexService.generate_and_vectorize_summary( + SummaryIndexAdapter.generate_and_vectorize_summary( segment, dataset, summary_index_setting, session=session ) total_segments_processed += 1 diff --git a/api/tasks/remove_document_from_index_task.py b/api/tasks/remove_document_from_index_task.py index ebffe41db0c..de50e196e8e 100644 --- a/api/tasks/remove_document_from_index_task.py +++ b/api/tasks/remove_document_from_index_task.py @@ -10,6 +10,7 @@ from core.rag.index_processor.index_processor_factory import IndexProcessorFacto from extensions.ext_redis import redis_client from libs.datetime_utils import naive_utc_now from models.dataset import Document, DocumentSegment +from repositories.knowledge.dataset_read_repository import get_document_dataset logger = logging.getLogger(__name__) @@ -38,7 +39,7 @@ def remove_document_from_index_task(document_id: str): indexing_cache_key = f"document_{document.id}_indexing" try: - dataset = document.get_dataset(session=session) + dataset = get_document_dataset(document, session=session) if not dataset: raise Exception("Document has no dataset") @@ -48,12 +49,12 @@ def remove_document_from_index_task(document_id: str): segments = session.scalars(select(DocumentSegment).where(DocumentSegment.document_id == document.id)).all() # Disable summary indexes for all segments in this document - from services.summary_index_service import SummaryIndexService + from services.knowledge.summaries.adapters import SummaryIndexAdapter segment_ids_list = [segment.id for segment in segments] if segment_ids_list: try: - SummaryIndexService.disable_summaries_for_segments( + SummaryIndexAdapter.disable_summaries_for_segments( dataset=dataset, segment_ids=segment_ids_list, disabled_by=document.disabled_by, diff --git a/api/tasks/retry_document_indexing_task.py b/api/tasks/retry_document_indexing_task.py index 08e0aaea5ab..744723eae44 100644 --- a/api/tasks/retry_document_indexing_task.py +++ b/api/tasks/retry_document_indexing_task.py @@ -6,16 +6,20 @@ from celery import shared_task from sqlalchemy import delete, select from configs import dify_config -from core.db.session_factory import session_factory -from core.indexing_runner import IndexingRunner +from core.app.apps.pipeline.pipeline_generator import PipelineGenerator +from core.db.session_factory import get_session_maker, session_factory from core.rag.index_processor.index_processor_factory import IndexProcessorFactory from enums import DeploymentEdition +from extensions.application_services.data_sources import build_data_source_credentials from extensions.ext_redis import redis_client from libs.datetime_utils import naive_utc_now from models import Account, Tenant from models.dataset import Dataset, Document, DocumentSegment from models.enums import IndexingStatus +from repositories.knowledge.document_repository import SQLAlchemyDocumentRepository from services.feature_service import FeatureService +from services.knowledge.indexing.adapters.execution import build_document_indexing_service +from services.knowledge.resource_scope import DatasetRef from services.rag_pipeline.rag_pipeline import RagPipelineService logger = logging.getLogger(__name__) @@ -109,17 +113,32 @@ def retry_document_indexing_task(dataset_id: str, document_ids: list[str], user_ document.indexing_status = IndexingStatus.PARSING document.processing_started_at = naive_utc_now() session.add(document) - # The runner performs slow extraction/indexing in a separate transaction phase. + document_ref = DatasetRef(document.tenant_id, document.dataset_id).document(document.id) + runtime_mode = dataset.runtime_mode + # Persist preparation before handing off to the selected execution service. session.commit() - if dataset.runtime_mode == "rag_pipeline": + if runtime_mode == "rag_pipeline": with session_factory.create_session() as rag_session: rag_pipeline_service = RagPipelineService(rag_session) - rag_pipeline_service.retry_error_document(dataset, document, user) + rag_pipeline_service.retry_error_document( + dataset, + document, + user, + generator=PipelineGenerator( + documents=SQLAlchemyDocumentRepository( + session_factory=session_factory.get_session_maker() + ), + datasource_providers=build_data_source_credentials( + database_client=get_session_maker() + ).providers, + ), + ) else: - indexing_runner = IndexingRunner(enforce_vector_space_admission=True) - indexing_runner.run([document], session) - session.commit() + indexing_service = build_document_indexing_service( + session_factory=session_factory.get_session_maker(), enforce_vector_space_admission=True + ) + indexing_service.run([document_ref]) redis_client.delete(retry_indexing_cache_key) except Exception as ex: session.rollback() diff --git a/api/tasks/sync_website_document_indexing_task.py b/api/tasks/sync_website_document_indexing_task.py index cceafa9623c..30dc205b44c 100644 --- a/api/tasks/sync_website_document_indexing_task.py +++ b/api/tasks/sync_website_document_indexing_task.py @@ -7,15 +7,16 @@ from sqlalchemy import delete, select from configs import dify_config from core.db.session_factory import session_factory -from core.indexing_runner import IndexingRunner from core.rag.index_processor.index_processor_factory import IndexProcessorFactory from enums import DeploymentEdition from extensions.ext_redis import redis_client from libs.datetime_utils import naive_utc_now from models.dataset import Dataset, DocumentSegment from models.enums import IndexingStatus -from services.dataset_ref_service import DatasetRefService +from repositories.knowledge.document_repository import _get_document from services.feature_service import FeatureService +from services.knowledge.indexing.adapters.execution import build_document_indexing_service +from services.knowledge.resource_scope import DatasetRef logger = logging.getLogger(__name__) @@ -36,9 +37,9 @@ def sync_website_document_indexing_task(dataset_id: str, document_id: str): if dataset is None: raise ValueError("Dataset not found") tenant_id = dataset.tenant_id - dataset_ref = DatasetRefService.create_dataset_ref(dataset) - document_ref = DatasetRefService.create_document_ref_from_id(dataset_ref, document_id) - document = DatasetRefService.get_document_by_ref(document_ref, session=session) + dataset_ref = DatasetRef(tenant_id=dataset.tenant_id, dataset_id=dataset.id) + document_ref = dataset_ref.document(document_id) + document = _get_document(session, document_ref) if document is None: logger.info(click.style(f"Document not found: {document_id}", fg="yellow")) return @@ -100,16 +101,14 @@ def sync_website_document_indexing_task(dataset_id: str, document_id: str): document.indexing_status = IndexingStatus.PARSING document.processing_started_at = naive_utc_now() session.add(document) + indexing_service = build_document_indexing_service(session_factory=session_factory.get_session_maker()) # Release document/segment locks before extraction starts. session.commit() - - indexing_runner = IndexingRunner() - indexing_runner.run([document], session) - session.commit() + indexing_service.run([document_ref]) redis_client.delete(sync_indexing_cache_key) except Exception as ex: session.rollback() - document = DatasetRefService.get_document_by_ref(document_ref, session=session) + document = _get_document(session, document_ref) if document: document.indexing_status = IndexingStatus.ERROR document.error = str(ex) diff --git a/api/tests/test_containers_integration_tests/controllers/console/auth/test_data_source_bearer_auth.py b/api/tests/test_containers_integration_tests/controllers/console/auth/test_data_source_bearer_auth.py index f51bf8f13e1..1723f00bcb2 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/auth/test_data_source_bearer_auth.py +++ b/api/tests/test_containers_integration_tests/controllers/console/auth/test_data_source_bearer_auth.py @@ -9,7 +9,7 @@ from sqlalchemy.orm import Session from models.source import DataSourceApiKeyAuthBinding from services.auth.errors import DataSourceApiKeyAuthCredentialValidationError -from services.entities.data_source_api_key_auth_entities import ( +from services.data_source.entities.api_key_auth import ( DataSourceApiKeyAuthBindingCreate, DataSourceApiKeyAuthCredentials, ) @@ -82,9 +82,7 @@ def test_create_binding_successful( "credentials": {"auth_type": "bearer", "config": {"api_key": "secret"}}, } - with patch( - "services.auth.data_source_api_key_auth_service.DataSourceApiKeyAuthService.create_binding" - ) as create_auth: + with patch("services.data_source.auth.api_key_service.DataSourceApiKeyAuthService.create_binding") as create_auth: response = test_client_with_containers.post( "/console/api/api-key-auth/data-source/binding", json=payload, @@ -111,7 +109,7 @@ def test_create_binding_failure( account, _tenant = create_console_account_and_tenant(db_session_with_containers) with patch( - "services.auth.data_source_api_key_auth_service.DataSourceApiKeyAuthService.create_binding", + "services.data_source.auth.api_key_service.DataSourceApiKeyAuthService.create_binding", side_effect=DataSourceApiKeyAuthCredentialValidationError("Credentials rejected"), ): response = test_client_with_containers.post( diff --git a/api/tests/test_containers_integration_tests/controllers/console/auth/test_data_source_oauth.py b/api/tests/test_containers_integration_tests/controllers/console/auth/test_data_source_oauth.py index c9e800235b7..307f80d62cd 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/auth/test_data_source_oauth.py +++ b/api/tests/test_containers_integration_tests/controllers/console/auth/test_data_source_oauth.py @@ -7,8 +7,8 @@ from sqlalchemy.orm import Session from extensions.ext_application_services import ApplicationServices from models.source import DataSourceOauthBinding -from services.data_source_oauth_service import InvalidDataSourceOAuthProviderError -from services.entities.data_source_oauth_entities import DataSourceOAuthCallback +from services.data_source.entities.oauth import DataSourceOAuthCallback +from services.data_source.oauth_service import InvalidDataSourceOAuthProviderError from tests.test_containers_integration_tests.controllers.console.helpers import ( authenticate_console_client, create_console_account_and_tenant, @@ -17,7 +17,7 @@ from tests.test_containers_integration_tests.controllers.console.helpers import def _services(service: MagicMock) -> MagicMock: services = create_autospec(ApplicationServices, instance=True, spec_set=True) - services.resolve_data_source_oauth.return_value = service + services.data_sources.resolve_oauth.return_value = service return services @@ -54,7 +54,7 @@ def test_get_oauth_url_invalid_provider( service = MagicMock() services = _services(service) - services.resolve_data_source_oauth.side_effect = InvalidDataSourceOAuthProviderError + services.data_sources.resolve_oauth.side_effect = InvalidDataSourceOAuthProviderError with patch( "controllers.console.auth.data_source_oauth.application_services", return_value=services, @@ -99,7 +99,7 @@ def test_oauth_callback_missing_code(test_client_with_containers: FlaskClient) - def test_oauth_callback_invalid_provider(test_client_with_containers: FlaskClient) -> None: service = MagicMock() services = _services(service) - services.resolve_data_source_oauth.side_effect = InvalidDataSourceOAuthProviderError + services.data_sources.resolve_oauth.side_effect = InvalidDataSourceOAuthProviderError with patch( "controllers.console.auth.data_source_oauth.application_services", return_value=services, 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 bc0a205ecb1..1a9389d299b 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 @@ -1,26 +1,51 @@ """Integration coverage for Notion page bindings backed by persisted documents.""" from inspect import unwrap +from types import SimpleNamespace from unittest.mock import MagicMock, patch from uuid import uuid4 from flask import Flask -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, sessionmaker from controllers.console.datasets.data_source import DataSourceNotionListApi, DataSourceNotionListQuery -from models import Account -from models.dataset import Document +from machinery.context import RequestContext +from models import Account, Tenant, TenantAccountJoin, TenantAccountRole +from models.dataset import Dataset, Document from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus +from repositories.knowledge.dataset_repository import SQLAlchemyDatasetRepository +from repositories.knowledge.document_repository import SQLAlchemyDocumentRepository +from repositories.workspace_member_query_repository import WorkspaceMemberQueryRepository +from services.data_source.credential_gateway import ActorDatasourceCredentialResolver +from services.data_source.notion_import_adapters import PluginNotionSourceGateway +from services.data_source.notion_import_application_service import NotionImportApplicationService +from services.knowledge.dataset_access import DatasetAccessService def test_notion_page_is_marked_bound_from_persisted_document( flask_app_with_containers: Flask, db_session_with_containers: Session, ) -> None: - tenant_id = str(uuid4()) + tenant = Tenant(name="Workspace") + tenant_id = tenant.id dataset_id = str(uuid4()) account = Account(name="Test User", email="user@example.com") account.id = str(uuid4()) + dataset = Dataset( + id=dataset_id, + tenant_id=tenant_id, + name="Notion dataset", + created_by=account.id, + data_source_type=DataSourceType.NOTION_IMPORT, + ) + db_session_with_containers.add_all( + [ + account, + tenant, + TenantAccountJoin(tenant_id=tenant_id, account_id=account.id, role=TenantAccountRole.OWNER, current=True), + dataset, + ] + ) document = Document( tenant_id=tenant_id, dataset_id=dataset_id, @@ -30,7 +55,7 @@ def test_notion_page_is_marked_bound_from_persisted_document( batch=f"batch-{uuid4()}", name="Notion Page", created_from=DocumentCreatedFrom.WEB, - created_by=str(uuid4()), + created_by=account.id, indexing_status=IndexingStatus.COMPLETED, enabled=True, ) @@ -62,28 +87,43 @@ def test_notion_page_is_marked_bound_from_persisted_document( datasource_provider_type=lambda: None, ) + sessions = sessionmaker(bind=db_session_with_containers.get_bind(), expire_on_commit=False) + datasets = SQLAlchemyDatasetRepository(session_factory=sessions) + credentials = MagicMock(spec=ActorDatasourceCredentialResolver) + credentials.resolve.return_value = {"token": "token"} + notion_imports = NotionImportApplicationService( + dataset_access=DatasetAccessService( + datasets=datasets, + workspace_roles=WorkspaceMemberQueryRepository(session_factory=sessions), + legacy_permissions_enabled=True, + ), + datasets=datasets, + documents=SQLAlchemyDocumentRepository(session_factory=sessions), + source=PluginNotionSourceGateway(credentials=credentials, runtime_loader=MagicMock(return_value=runtime)), + ) + context = RequestContext( + request_id="notion-binding-test", trace_id=None, account_id=account.id, active_workspace_id=tenant_id + ) + with ( flask_app_with_containers.test_request_context(f"/?credential_id=c1&dataset_id={dataset_id}"), patch( - "controllers.console.datasets.data_source.DatasourceProviderService.get_datasource_credentials", - return_value={"token": "token"}, - ), - patch( - "controllers.console.datasets.data_source.DatasetService.get_dataset", - return_value=MagicMock(data_source_type="notion_import"), - ), - patch( - "core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", - return_value=runtime, + "controllers.console.datasets.data_source.application_services", + return_value=SimpleNamespace(data_sources=SimpleNamespace(notion_imports=notion_imports)), ), ): - response, status = unwrap(DataSourceNotionListApi().get)( + response, status = unwrap(DataSourceNotionListApi.get)( DataSourceNotionListApi(), DataSourceNotionListQuery(credential_id="c1", dataset_id=dataset_id), - db_session_with_containers, - tenant_id, - account, + context, ) assert status == 200 assert response["notion_info"][0]["pages"][0]["is_bound"] is True + credentials.resolve.assert_called_once_with( + workspace_id=tenant_id, + actor_id=account.id, + credential_id="c1", + provider="notion_datasource", + plugin_id="langgenius/notion_datasource", + ) diff --git a/api/tests/test_containers_integration_tests/models/test_dataset_models.py b/api/tests/test_containers_integration_tests/models/test_dataset_models.py index 52d9b42c356..1d3df7aa58a 100644 --- a/api/tests/test_containers_integration_tests/models/test_dataset_models.py +++ b/api/tests/test_containers_integration_tests/models/test_dataset_models.py @@ -1,3 +1,9 @@ +from repositories.knowledge.dataset_read_repository import ( + get_dataset_available_segment_count, + get_document_hit_count, + get_document_segment_count, +) + """ Integration tests for Dataset and Document model properties using testcontainers. @@ -13,6 +19,15 @@ from sqlalchemy.orm import Session from models.dataset import Dataset, Document, DocumentSegment from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus, SegmentStatus +from repositories.knowledge.dataset_read_repository import ( + get_dataset_available_document_count, + get_dataset_document_count, + get_dataset_word_count, + get_next_segment, + get_previous_segment, + get_segment_dataset, + get_segment_document, +) class TestDatasetDocumentProperties: @@ -49,7 +64,7 @@ class TestDatasetDocumentProperties: db_session_with_containers.add(doc) db_session_with_containers.flush() - assert dataset.get_total_documents(session=db_session_with_containers) == 3 + assert get_dataset_document_count(dataset, session=db_session_with_containers) == 3 def test_dataset_available_documents_count(self, db_session_with_containers: Session) -> None: """Test dataset can count available documents.""" @@ -104,7 +119,7 @@ class TestDatasetDocumentProperties: db_session_with_containers.add_all([doc_available, doc_pending, doc_disabled]) db_session_with_containers.flush() - assert dataset.get_total_available_documents(session=db_session_with_containers) == 1 + assert get_dataset_available_document_count(dataset, session=db_session_with_containers) == 1 def test_dataset_word_count_aggregation(self, db_session_with_containers: Session) -> None: """Test dataset can aggregate word count from documents.""" @@ -132,7 +147,7 @@ class TestDatasetDocumentProperties: db_session_with_containers.add(doc) db_session_with_containers.flush() - assert dataset.get_word_count(session=db_session_with_containers) == 5000 + assert get_dataset_word_count(dataset, session=db_session_with_containers) == 5000 def test_dataset_available_segment_count(self, db_session_with_containers: Session) -> None: """Test Dataset.available_segment_count counts completed and enabled segments.""" @@ -188,7 +203,7 @@ class TestDatasetDocumentProperties: db_session_with_containers.add(seg_waiting) db_session_with_containers.flush() - assert dataset.get_available_segment_count(session=db_session_with_containers) == 2 + assert get_dataset_available_segment_count(dataset, session=db_session_with_containers) == 2 def test_document_segment_count_property(self, db_session_with_containers: Session) -> None: """Test document can count its segments.""" @@ -228,7 +243,7 @@ class TestDatasetDocumentProperties: db_session_with_containers.add(seg) db_session_with_containers.flush() - assert doc.get_segment_count(session=db_session_with_containers) == 3 + assert get_document_segment_count(doc, session=db_session_with_containers) == 3 def test_document_hit_count_aggregation(self, db_session_with_containers: Session) -> None: """Test document can aggregate hit count from segments.""" @@ -269,7 +284,7 @@ class TestDatasetDocumentProperties: db_session_with_containers.add(seg) db_session_with_containers.flush() - assert doc.get_hit_count(session=db_session_with_containers) == 25 + assert get_document_hit_count(doc, session=db_session_with_containers) == 25 class TestDocumentSegmentNavigationProperties: @@ -322,7 +337,7 @@ class TestDocumentSegmentNavigationProperties: db_session_with_containers.flush() # Act - related_dataset = segment.get_dataset(session=db_session_with_containers) + related_dataset = get_segment_dataset(segment, session=db_session_with_containers) # Assert assert related_dataset is not None @@ -369,7 +384,7 @@ class TestDocumentSegmentNavigationProperties: db_session_with_containers.flush() # Act - related_document = segment.get_document(session=db_session_with_containers) + related_document = get_segment_document(segment, session=db_session_with_containers) # Assert assert related_document is not None @@ -426,7 +441,7 @@ class TestDocumentSegmentNavigationProperties: db_session_with_containers.flush() # Act - prev_seg = segment.previous_segment(session=db_session_with_containers) + prev_seg = get_previous_segment(segment, session=db_session_with_containers) # Assert assert prev_seg is not None @@ -483,7 +498,7 @@ class TestDocumentSegmentNavigationProperties: db_session_with_containers.flush() # Act - next_seg = segment.next_segment(session=db_session_with_containers) + next_seg = get_next_segment(segment, session=db_session_with_containers) # Assert assert next_seg is not None diff --git a/api/tests/test_containers_integration_tests/services/auth/test_auth_integration.py b/api/tests/test_containers_integration_tests/services/auth/test_auth_integration.py index ec2c013459e..12ac010819f 100644 --- a/api/tests/test_containers_integration_tests/services/auth/test_auth_integration.py +++ b/api/tests/test_containers_integration_tests/services/auth/test_auth_integration.py @@ -12,14 +12,14 @@ from sqlalchemy.orm import Session, sessionmaker from machinery.context import RequestContext from models.source import DataSourceApiKeyAuthBinding -from repositories.data_source_api_key_auth_repository import SQLAlchemyDataSourceApiKeyAuthBindingRepository -from services.auth.data_source_api_key_auth_gateways import ( +from repositories.data_source.api_key_auth_repository import SQLAlchemyDataSourceApiKeyAuthBindingRepository +from services.auth.errors import DataSourceApiKeyAuthProviderUnavailableError +from services.data_source.auth.api_key_gateways import ( ProviderApiKeyAuthCredentialValidator, TenantApiKeyAuthCredentialEncryptor, ) -from services.auth.data_source_api_key_auth_service import DataSourceApiKeyAuthService -from services.auth.errors import DataSourceApiKeyAuthProviderUnavailableError -from services.entities.data_source_api_key_auth_entities import ( +from services.data_source.auth.api_key_service import DataSourceApiKeyAuthService +from services.data_source.entities.api_key_auth import ( DataSourceApiKeyAuthBindingCreate, DataSourceApiKeyAuthCredentials, ) @@ -67,8 +67,8 @@ class TestAuthIntegration: encryptor=TenantApiKeyAuthCredentialEncryptor(), ) - @patch("services.auth.firecrawl.firecrawl.httpx.post") - @patch("services.auth.data_source_api_key_auth_gateways.encrypter.encrypt_token") + @patch("services.data_source.auth.firecrawl.firecrawl.httpx.post") + @patch("services.data_source.auth.api_key_gateways.encrypter.encrypt_token") def test_end_to_end_auth_flow( self, mock_encrypt, @@ -96,9 +96,9 @@ class TestAuthIntegration: assert len(bindings) == 1 assert bindings[0].provider == _FIRECRAWL_PROVIDER - @patch("services.auth.data_source_api_key_auth_gateways.encrypter.encrypt_token") - @patch("services.auth.firecrawl.firecrawl.httpx.post") - @patch("services.auth.jina.jina._http_client.post") + @patch("services.data_source.auth.api_key_gateways.encrypter.encrypt_token") + @patch("services.data_source.auth.firecrawl.firecrawl.httpx.post") + @patch("services.data_source.auth.jina.jina._http_client.post") def test_multi_tenant_isolation( self, mock_jina_http, @@ -132,8 +132,8 @@ class TestAuthIntegration: assert len(result2) == 1 assert result2[0].provider == _JINA_PROVIDER - @patch("services.auth.firecrawl.firecrawl.httpx.post") - @patch("services.auth.data_source_api_key_auth_gateways.encrypter.encrypt_token", return_value="encrypted_key") + @patch("services.data_source.auth.firecrawl.firecrawl.httpx.post") + @patch("services.data_source.auth.api_key_gateways.encrypter.encrypt_token", return_value="encrypted_key") def test_concurrent_creation_safety( self, mock_encrypt, @@ -168,7 +168,7 @@ class TestAuthIntegration: assert len(results) == 5 assert len(exceptions) == 0 - @patch("services.auth.firecrawl.firecrawl.httpx.post") + @patch("services.data_source.auth.firecrawl.firecrawl.httpx.post") def test_network_failure_recovery( self, mock_http, diff --git a/api/tests/test_containers_integration_tests/services/dataset_collection_binding.py b/api/tests/test_containers_integration_tests/services/dataset_collection_binding.py index d180da9555e..706382b4b9f 100644 --- a/api/tests/test_containers_integration_tests/services/dataset_collection_binding.py +++ b/api/tests/test_containers_integration_tests/services/dataset_collection_binding.py @@ -13,7 +13,7 @@ from sqlalchemy.orm import Session from models.dataset import DatasetCollectionBinding from models.enums import CollectionBindingType -from services.dataset_service import DatasetCollectionBindingService +from services.knowledge.dataset_service import DatasetCollectionBindingService class DatasetCollectionBindingTestDataFactory: diff --git a/api/tests/test_containers_integration_tests/services/dataset_service_update_delete.py b/api/tests/test_containers_integration_tests/services/dataset_service_update_delete.py index 2d6545bc6af..587258f5864 100644 --- a/api/tests/test_containers_integration_tests/services/dataset_service_update_delete.py +++ b/api/tests/test_containers_integration_tests/services/dataset_service_update_delete.py @@ -17,9 +17,9 @@ from models import Account, AccountStatus, Tenant, TenantAccountJoin, TenantAcco from models.dataset import AppDatasetJoin, Dataset, DatasetPermissionEnum from models.enums import DataSourceType from models.model import App -from services.dataset_ref_service import DatasetRefService -from services.dataset_service import DatasetService from services.errors.account import NoPermissionError +from services.knowledge.dataset_service import DatasetService +from services.knowledge.resource_scope import DatasetRef class DatasetUpdateDeleteTestDataFactory: @@ -140,7 +140,7 @@ class TestDatasetServiceDeleteDataset: dataset = DatasetUpdateDeleteTestDataFactory.create_dataset(db_session_with_containers, tenant.id, owner.id) # Act - with patch("services.dataset_service.dataset_was_deleted") as mock_dataset_was_deleted: + with patch("services.knowledge.dataset_service.dataset_was_deleted") as mock_dataset_was_deleted: result = DatasetService.delete_dataset(dataset.id, owner, session=db_session_with_containers) # Assert @@ -228,7 +228,7 @@ class TestDatasetServiceDatasetUseCheck: dataset = DatasetUpdateDeleteTestDataFactory.create_dataset(db_session_with_containers, tenant.id, owner.id) app = DatasetUpdateDeleteTestDataFactory.create_app(db_session_with_containers, tenant.id, owner.id) DatasetUpdateDeleteTestDataFactory.create_app_dataset_join(db_session_with_containers, app.id, dataset.id) - dataset_ref = DatasetRefService.create_dataset_ref(dataset) + dataset_ref = DatasetRef(tenant_id=dataset.tenant_id, dataset_id=dataset.id) # Act result = DatasetService.dataset_use_check(dataset_ref, session=db_session_with_containers) @@ -253,7 +253,7 @@ class TestDatasetServiceDatasetUseCheck: db_session_with_containers, role=TenantAccountRole.OWNER ) dataset = DatasetUpdateDeleteTestDataFactory.create_dataset(db_session_with_containers, tenant.id, owner.id) - dataset_ref = DatasetRefService.create_dataset_ref(dataset) + dataset_ref = DatasetRef(tenant_id=dataset.tenant_id, dataset_id=dataset.id) # Act result = DatasetService.dataset_use_check(dataset_ref, session=db_session_with_containers) @@ -290,7 +290,7 @@ class TestDatasetServiceUpdateDatasetApiStatus: current_time = datetime.datetime(2023, 1, 1, 12, 0, 0) # Act - with patch("services.dataset_service.naive_utc_now", return_value=current_time): + with patch("services.knowledge.dataset_service.naive_utc_now", return_value=current_time): DatasetService.update_dataset_api_status(dataset, True, owner, session=db_session_with_containers) # Assert @@ -322,7 +322,7 @@ class TestDatasetServiceUpdateDatasetApiStatus: current_time = datetime.datetime(2023, 1, 1, 12, 0, 0) # Act - with patch("services.dataset_service.naive_utc_now", return_value=current_time): + with patch("services.knowledge.dataset_service.naive_utc_now", return_value=current_time): DatasetService.update_dataset_api_status(dataset, False, owner, session=db_session_with_containers) # Assert diff --git a/api/tests/test_containers_integration_tests/services/document_service_status.py b/api/tests/test_containers_integration_tests/services/document_service_status.py index 6ffaf49d947..67427ef50e7 100644 --- a/api/tests/test_containers_integration_tests/services/document_service_status.py +++ b/api/tests/test_containers_integration_tests/services/document_service_status.py @@ -20,8 +20,8 @@ from models import Account from models.dataset import Dataset, Document from models.enums import CreatorUserRole, DataSourceType, DocumentCreatedFrom, IndexingStatus from models.model import UploadFile -from services.dataset_service import DocumentService from services.errors.document import DocumentIndexingError +from services.knowledge.dataset_service import DocumentService FIXED_TIME = datetime.datetime(2023, 1, 1, 12, 0, 0) @@ -256,10 +256,10 @@ class TestDocumentServicePauseDocument: """ with ( patch( - "services.dataset_service.current_user", create_autospec(Account, instance=True) + "services.knowledge.dataset_service.current_user", create_autospec(Account, instance=True) ) as mock_current_user, - patch("services.dataset_service.redis_client") as mock_redis, - patch("services.dataset_service.naive_utc_now") as mock_naive_utc_now, + patch("services.knowledge.dataset_service.redis_client") as mock_redis, + patch("services.knowledge.dataset_service.naive_utc_now") as mock_naive_utc_now, ): current_time = datetime.datetime(2023, 1, 1, 12, 0, 0) user_id = str(uuid4()) @@ -301,7 +301,9 @@ class TestDocumentServicePauseDocument: ) # Act - DocumentService.pause_document(document, session=db_session_with_containers) + DocumentService.pause_document( + document, session=db_session_with_containers, actor_id=mock_document_service_dependencies["user_id"] + ) # Assert db_session_with_containers.refresh(document) @@ -336,7 +338,9 @@ class TestDocumentServicePauseDocument: ) # Act - DocumentService.pause_document(document, session=db_session_with_containers) + DocumentService.pause_document( + document, session=db_session_with_containers, actor_id=mock_document_service_dependencies["user_id"] + ) # Assert db_session_with_containers.refresh(document) @@ -366,7 +370,9 @@ class TestDocumentServicePauseDocument: ) # Act - DocumentService.pause_document(document, session=db_session_with_containers) + DocumentService.pause_document( + document, session=db_session_with_containers, actor_id=mock_document_service_dependencies["user_id"] + ) # Assert db_session_with_containers.refresh(document) @@ -398,7 +404,9 @@ class TestDocumentServicePauseDocument: # Act & Assert with pytest.raises(DocumentIndexingError): - DocumentService.pause_document(document, session=db_session_with_containers) + DocumentService.pause_document( + document, session=db_session_with_containers, actor_id=mock_document_service_dependencies["user_id"] + ) db_session_with_containers.refresh(document) assert document.is_paused is False @@ -429,7 +437,9 @@ class TestDocumentServicePauseDocument: # Act & Assert with pytest.raises(DocumentIndexingError): - DocumentService.pause_document(document, session=db_session_with_containers) + DocumentService.pause_document( + document, session=db_session_with_containers, actor_id=mock_document_service_dependencies["user_id"] + ) db_session_with_containers.refresh(document) assert document.is_paused is False @@ -468,8 +478,8 @@ class TestDocumentServiceRecoverDocument: - Recovery task """ with ( - patch("services.dataset_service.redis_client") as mock_redis, - patch("services.dataset_service.recover_document_indexing_task") as mock_task, + patch("services.knowledge.dataset_service.redis_client") as mock_redis, + patch("services.knowledge.dataset_service.recover_document_indexing_task") as mock_task, ): yield { "redis_client": mock_redis, @@ -588,10 +598,10 @@ class TestDocumentServiceRetryDocument: """ with ( patch( - "services.dataset_service.current_user", create_autospec(Account, instance=True) + "services.knowledge.dataset_service.current_user", create_autospec(Account, instance=True) ) as mock_current_user, - patch("services.dataset_service.redis_client") as mock_redis, - patch("services.dataset_service.retry_document_indexing_task") as mock_task, + patch("services.knowledge.dataset_service.redis_client") as mock_redis, + patch("services.knowledge.dataset_service.retry_document_indexing_task") as mock_task, ): user_id = str(uuid4()) mock_current_user.id = user_id @@ -640,7 +650,12 @@ class TestDocumentServiceRetryDocument: ) # Act - DocumentService.retry_document(dataset.id, [document], session=db_session_with_containers) + DocumentService.retry_document( + dataset.id, + [document], + session=db_session_with_containers, + actor_id=mock_document_service_dependencies["user_id"], + ) # Assert db_session_with_containers.refresh(document) @@ -689,7 +704,12 @@ class TestDocumentServiceRetryDocument: ) # Act - DocumentService.retry_document(dataset.id, [document1, document2], session=db_session_with_containers) + DocumentService.retry_document( + dataset.id, + [document1, document2], + session=db_session_with_containers, + actor_id=mock_document_service_dependencies["user_id"], + ) # Assert db_session_with_containers.refresh(document1) @@ -732,7 +752,12 @@ class TestDocumentServiceRetryDocument: # Act & Assert with pytest.raises(ValueError, match="Document is being retried, please try again later"): - DocumentService.retry_document(dataset.id, [document], session=db_session_with_containers) + DocumentService.retry_document( + dataset.id, + [document], + session=db_session_with_containers, + actor_id=mock_document_service_dependencies["user_id"], + ) db_session_with_containers.refresh(document) assert document.indexing_status == IndexingStatus.ERROR @@ -763,7 +788,12 @@ class TestDocumentServiceRetryDocument: mock_document_service_dependencies["redis_client"].lock.side_effect = [first_retry_lock, second_retry_lock] with pytest.raises(ValueError, match="Document is being retried, please try again later"): - DocumentService.retry_document(dataset.id, [document1, document2], session=db_session_with_containers) + DocumentService.retry_document( + dataset.id, + [document1, document2], + session=db_session_with_containers, + actor_id=mock_document_service_dependencies["user_id"], + ) db_session_with_containers.refresh(document1) db_session_with_containers.refresh(document2) @@ -773,13 +803,13 @@ class TestDocumentServiceRetryDocument: second_retry_lock.release.assert_not_called() mock_document_service_dependencies["retry_task"].delay.assert_not_called() - def test_retry_document_missing_current_user_error( + def test_retry_document_missing_actor_id_error( self, db_session_with_containers: Session, mock_document_service_dependencies ): """ - Test error when current_user is missing. + Test error when the explicit actor ID is missing. - Verifies that when current_user is None or has no ID, a ValueError + Verifies that when the actor ID is empty, a ValueError is raised. This test ensures: @@ -797,11 +827,9 @@ class TestDocumentServiceRetryDocument: indexing_status=IndexingStatus.ERROR, ) - mock_document_service_dependencies["current_user"].id = None - # Act & Assert with pytest.raises(ValueError, match="Current user or current user id not found"): - DocumentService.retry_document(dataset.id, [document], session=db_session_with_containers) + DocumentService.retry_document(dataset.id, [document], session=db_session_with_containers, actor_id="") db_session_with_containers.refresh(document) assert document.indexing_status == IndexingStatus.ERROR @@ -846,10 +874,10 @@ class TestDocumentServiceBatchUpdateDocumentStatus: - Async tasks """ with ( - patch("services.dataset_service.redis_client") as mock_redis, - patch("services.dataset_service.add_document_to_index_task") as mock_add_task, - patch("services.dataset_service.remove_document_from_index_task") as mock_remove_task, - patch("services.dataset_service.naive_utc_now") as mock_naive_utc_now, + patch("services.knowledge.dataset_service.redis_client") as mock_redis, + patch("services.knowledge.dataset_service.add_document_to_index_task") as mock_add_task, + patch("services.knowledge.dataset_service.remove_document_from_index_task") as mock_remove_task, + patch("services.knowledge.dataset_service.naive_utc_now") as mock_naive_utc_now, ): current_time = datetime.datetime(2023, 1, 1, 12, 0, 0) mock_naive_utc_now.return_value = current_time @@ -1146,7 +1174,7 @@ class TestDocumentServiceRenameDocument: - Database session """ with patch( - "services.dataset_service.current_user", create_autospec(Account, instance=True) + "services.knowledge.dataset_service.current_user", create_autospec(Account, instance=True) ) as mock_current_user: mock_current_user.current_tenant_id = str(uuid4()) diff --git a/api/tests/test_containers_integration_tests/services/test_dataset_permission_service.py b/api/tests/test_containers_integration_tests/services/test_dataset_permission_service.py index 3f7db13c41e..4285108c054 100644 --- a/api/tests/test_containers_integration_tests/services/test_dataset_permission_service.py +++ b/api/tests/test_containers_integration_tests/services/test_dataset_permission_service.py @@ -19,8 +19,8 @@ from models.dataset import ( DatasetPermissionEnum, ) from models.enums import DataSourceType -from services.dataset_service import DatasetPermissionService, DatasetService from services.errors.account import NoPermissionError +from services.knowledge.dataset_service import DatasetPermissionService, DatasetService class DatasetPermissionTestDataFactory: diff --git a/api/tests/test_containers_integration_tests/services/test_dataset_service.py b/api/tests/test_containers_integration_tests/services/test_dataset_service.py index 912e00b0b7d..81199f29705 100644 --- a/api/tests/test_containers_integration_tests/services/test_dataset_service.py +++ b/api/tests/test_containers_integration_tests/services/test_dataset_service.py @@ -17,10 +17,10 @@ from graphon.model_runtime.entities.model_entities import ModelType from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole from models.dataset import Dataset, DatasetPermissionEnum, Document, ExternalKnowledgeBindings, Pipeline from models.enums import DatasetRuntimeMode, DataSourceType, DocumentCreatedFrom, IndexingStatus -from services.dataset_service import DatasetService -from services.entities.knowledge_entities.knowledge_entities import RerankingModel, RetrievalModel from services.entities.knowledge_entities.rag_pipeline_entities import IconInfo, RagPipelineDatasetCreateEntity from services.errors.dataset import DatasetNameDuplicateError +from services.knowledge.dataset_service import DatasetService +from services.knowledge.entities.knowledge_entities import RerankingModel, RetrievalModel class DatasetServiceIntegrationDataFactory: @@ -176,7 +176,7 @@ class TestDatasetServiceCreateDataset: embedding_model = DatasetServiceIntegrationDataFactory.create_embedding_model() # Act - with patch("services.dataset_service.ModelManager.for_tenant") as mock_model_manager: + with patch("services.knowledge.dataset_service.ModelManager.for_tenant") as mock_model_manager: mock_model_manager.return_value.get_default_model_instance.return_value = embedding_model result = DatasetService.create_empty_dataset( @@ -229,7 +229,9 @@ class TestDatasetServiceCreateDataset: external_knowledge_id = "knowledge-123" # Act - with patch("services.dataset_service.ExternalDatasetService.get_external_knowledge_api") as mock_get_api: + with patch( + "services.knowledge.dataset_service.ExternalDatasetService.get_external_knowledge_api" + ) as mock_get_api: mock_get_api.return_value = Mock(id=external_knowledge_api_id) result = DatasetService.create_empty_dataset( tenant_id=tenant.id, @@ -269,8 +271,10 @@ class TestDatasetServiceCreateDataset: # Act with ( - patch("services.dataset_service.ModelManager.for_tenant") as mock_model_manager, - patch("services.dataset_service.DatasetService.check_reranking_model_setting") as mock_check_reranking, + patch("services.knowledge.dataset_service.ModelManager.for_tenant") as mock_model_manager, + patch( + "services.knowledge.dataset_service.DatasetService.check_reranking_model_setting" + ) as mock_check_reranking, ): mock_model_manager.return_value.get_default_model_instance.return_value = embedding_model @@ -303,8 +307,10 @@ class TestDatasetServiceCreateDataset: # Act with ( - patch("services.dataset_service.ModelManager.for_tenant") as mock_model_manager, - patch("services.dataset_service.DatasetService.check_embedding_model_setting") as mock_check_embedding, + patch("services.knowledge.dataset_service.ModelManager.for_tenant") as mock_model_manager, + patch( + "services.knowledge.dataset_service.DatasetService.check_embedding_model_setting" + ) as mock_check_embedding, ): mock_model_manager.return_value.get_model_instance.return_value = embedding_model @@ -386,7 +392,9 @@ class TestDatasetServiceCreateDataset: external_knowledge_api_id = str(uuid4()) # Act / Assert - with patch("services.dataset_service.ExternalDatasetService.get_external_knowledge_api") as mock_get_api: + with patch( + "services.knowledge.dataset_service.ExternalDatasetService.get_external_knowledge_api" + ) as mock_get_api: mock_get_api.return_value = None with pytest.raises(ValueError, match=r"External API template not found\.?"): DatasetService.create_empty_dataset( @@ -408,7 +416,9 @@ class TestDatasetServiceCreateDataset: external_knowledge_api_id = str(uuid4()) # Act / Assert - with patch("services.dataset_service.ExternalDatasetService.get_external_knowledge_api") as mock_get_api: + with patch( + "services.knowledge.dataset_service.ExternalDatasetService.get_external_knowledge_api" + ) as mock_get_api: mock_get_api.return_value = Mock(id=external_knowledge_api_id) with pytest.raises(ValueError, match="external_knowledge_id is required"): DatasetService.create_empty_dataset( @@ -440,7 +450,7 @@ class TestDatasetServiceCreateRagPipelineDataset: ) # Act - with patch("services.dataset_service.current_user", account): + with patch("services.knowledge.dataset_service.current_user", account): result = DatasetService.create_empty_rag_pipeline_dataset( tenant_id=tenant.id, rag_pipeline_dataset_create_entity=entity, @@ -475,8 +485,8 @@ class TestDatasetServiceCreateRagPipelineDataset: # Act with ( - patch("services.dataset_service.current_user", account), - patch("services.dataset_service.generate_incremental_name") as mock_generate_name, + patch("services.knowledge.dataset_service.current_user", account), + patch("services.knowledge.dataset_service.generate_incremental_name") as mock_generate_name, ): mock_generate_name.return_value = generated_name result = DatasetService.create_empty_rag_pipeline_dataset( @@ -516,7 +526,7 @@ class TestDatasetServiceCreateRagPipelineDataset: # Act / Assert with ( - patch("services.dataset_service.current_user", account), + patch("services.knowledge.dataset_service.current_user", account), pytest.raises(DatasetNameDuplicateError, match=f"Dataset with name {duplicate_name} already exists"), ): DatasetService.create_empty_rag_pipeline_dataset( @@ -538,7 +548,7 @@ class TestDatasetServiceCreateRagPipelineDataset: ) # Act - with patch("services.dataset_service.current_user", account): + with patch("services.knowledge.dataset_service.current_user", account): result = DatasetService.create_empty_rag_pipeline_dataset( tenant_id=tenant.id, rag_pipeline_dataset_create_entity=entity, @@ -567,7 +577,7 @@ class TestDatasetServiceCreateRagPipelineDataset: ) # Act - with patch("services.dataset_service.current_user", account): + with patch("services.knowledge.dataset_service.current_user", account): result = DatasetService.create_empty_rag_pipeline_dataset( tenant_id=tenant.id, rag_pipeline_dataset_create_entity=entity, @@ -621,7 +631,7 @@ class TestDatasetServiceUpdateAndDeleteDataset: ) # Act - with patch("services.dataset_service.dataset_was_deleted") as dataset_deleted_signal: + with patch("services.knowledge.dataset_service.dataset_was_deleted") as dataset_deleted_signal: result = DatasetService.delete_dataset(dataset.id, account, session=db_session_with_containers) # Assert @@ -642,7 +652,7 @@ class TestDatasetServiceUpdateAndDeleteDataset: ) # Act - with patch("services.dataset_service.dataset_was_deleted") as dataset_deleted_signal: + with patch("services.knowledge.dataset_service.dataset_was_deleted") as dataset_deleted_signal: result = DatasetService.delete_dataset(dataset.id, account, session=db_session_with_containers) # Assert @@ -663,7 +673,7 @@ class TestDatasetServiceUpdateAndDeleteDataset: ) # Act - with patch("services.dataset_service.dataset_was_deleted") as dataset_deleted_signal: + with patch("services.knowledge.dataset_service.dataset_was_deleted") as dataset_deleted_signal: result = DatasetService.delete_dataset(dataset.id, account, session=db_session_with_containers) # Assert @@ -747,13 +757,13 @@ class TestDocumentServicePauseRecoverRetry: def test_pause_document_success(self, db_session_with_containers: Session): from extensions.ext_redis import redis_client - from services.dataset_service import DocumentService + from services.knowledge.dataset_service import DocumentService doc, account = self._create_indexing_document(db_session_with_containers, indexing_status="indexing") - with patch("services.dataset_service.current_user") as mock_user: + with patch("services.knowledge.dataset_service.current_user") as mock_user: mock_user.id = account.id - DocumentService.pause_document(doc, session=db_session_with_containers) + DocumentService.pause_document(doc, session=db_session_with_containers, actor_id=account.id) db_session_with_containers.refresh(doc) assert doc.is_paused is True @@ -765,29 +775,29 @@ class TestDocumentServicePauseRecoverRetry: redis_client.delete(cache_key) def test_pause_document_invalid_status_error(self, db_session_with_containers: Session): - from services.dataset_service import DocumentService from services.errors.document import DocumentIndexingError + from services.knowledge.dataset_service import DocumentService doc, account = self._create_indexing_document(db_session_with_containers, indexing_status="completed") - with patch("services.dataset_service.current_user") as mock_user: + with patch("services.knowledge.dataset_service.current_user") as mock_user: mock_user.id = account.id with pytest.raises(DocumentIndexingError): - DocumentService.pause_document(doc, session=db_session_with_containers) + DocumentService.pause_document(doc, session=db_session_with_containers, actor_id=account.id) def test_recover_document_success(self, db_session_with_containers: Session): from extensions.ext_redis import redis_client - from services.dataset_service import DocumentService + from services.knowledge.dataset_service import DocumentService doc, account = self._create_indexing_document(db_session_with_containers, indexing_status="indexing") # Pause first - with patch("services.dataset_service.current_user") as mock_user: + with patch("services.knowledge.dataset_service.current_user") as mock_user: mock_user.id = account.id - DocumentService.pause_document(doc, session=db_session_with_containers) + DocumentService.pause_document(doc, session=db_session_with_containers, actor_id=account.id) # Recover - with patch("services.dataset_service.recover_document_indexing_task") as recover_task: + with patch("services.knowledge.dataset_service.recover_document_indexing_task") as recover_task: DocumentService.recover_document(doc, session=db_session_with_containers) db_session_with_containers.refresh(doc) @@ -801,7 +811,7 @@ class TestDocumentServicePauseRecoverRetry: def test_retry_document_indexing_success(self, db_session_with_containers: Session): from extensions.ext_redis import redis_client - from services.dataset_service import DocumentService + from services.knowledge.dataset_service import DocumentService factory = DatasetServiceIntegrationDataFactory account, tenant = factory.create_account_with_tenant(db_session_with_containers) @@ -814,11 +824,13 @@ class TestDocumentServicePauseRecoverRetry: db_session_with_containers.commit() with ( - patch("services.dataset_service.current_user") as mock_user, - patch("services.dataset_service.retry_document_indexing_task") as retry_task, + patch("services.knowledge.dataset_service.current_user") as mock_user, + patch("services.knowledge.dataset_service.retry_document_indexing_task") as retry_task, ): mock_user.id = account.id - DocumentService.retry_document(dataset.id, [doc1, doc2], session=db_session_with_containers) + DocumentService.retry_document( + dataset.id, [doc1, doc2], session=db_session_with_containers, actor_id=account.id + ) db_session_with_containers.refresh(doc1) db_session_with_containers.refresh(doc2) diff --git a/api/tests/test_containers_integration_tests/services/test_dataset_service_batch_update_document_status.py b/api/tests/test_containers_integration_tests/services/test_dataset_service_batch_update_document_status.py index 2f1022a47ca..698a4220be5 100644 --- a/api/tests/test_containers_integration_tests/services/test_dataset_service_batch_update_document_status.py +++ b/api/tests/test_containers_integration_tests/services/test_dataset_service_batch_update_document_status.py @@ -16,8 +16,8 @@ from sqlalchemy.orm import Session from core.rag.index_processor.constant.index_type import IndexStructureType from models.dataset import Dataset, Document from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus -from services.dataset_service import DocumentService from services.errors.document import DocumentIndexingError +from services.knowledge.dataset_service import DocumentService FIXED_TIME = datetime.datetime(2023, 1, 1, 12, 0, 0) @@ -140,10 +140,10 @@ class TestDatasetServiceBatchUpdateDocumentStatus: def patched_dependencies(self): """Patch non-DB collaborators only.""" with ( - patch("services.dataset_service.redis_client") as redis_client, - patch("services.dataset_service.add_document_to_index_task") as add_task, - patch("services.dataset_service.remove_document_from_index_task") as remove_task, - patch("services.dataset_service.naive_utc_now") as naive_utc_now, + patch("services.knowledge.dataset_service.redis_client") as redis_client, + patch("services.knowledge.dataset_service.add_document_to_index_task") as add_task, + patch("services.knowledge.dataset_service.remove_document_from_index_task") as remove_task, + patch("services.knowledge.dataset_service.naive_utc_now") as naive_utc_now, ): naive_utc_now.return_value = FIXED_TIME redis_client.get.return_value = None diff --git a/api/tests/test_containers_integration_tests/services/test_dataset_service_create_dataset.py b/api/tests/test_containers_integration_tests/services/test_dataset_service_create_dataset.py index 292ae46190a..623a9bea386 100644 --- a/api/tests/test_containers_integration_tests/services/test_dataset_service_create_dataset.py +++ b/api/tests/test_containers_integration_tests/services/test_dataset_service_create_dataset.py @@ -9,8 +9,8 @@ import pytest from sqlalchemy.orm import Session from models.account import Account, Tenant, TenantAccountJoin -from services.dataset_service import DatasetService from services.entities.knowledge_entities.rag_pipeline_entities import IconInfo, RagPipelineDatasetCreateEntity +from services.knowledge.dataset_service import DatasetService class TestDatasetServiceCreateRagPipelineDataset: @@ -53,7 +53,7 @@ class TestDatasetServiceCreateRagPipelineDataset: tenant, _ = self._create_tenant_and_account(db_session_with_containers) mock_user = Mock(id=None) - with patch("services.dataset_service.current_user", mock_user): + with patch("services.knowledge.dataset_service.current_user", mock_user): with pytest.raises(ValueError, match="Current user or current user id not found"): DatasetService.create_empty_rag_pipeline_dataset( tenant_id=tenant.id, diff --git a/api/tests/test_containers_integration_tests/services/test_dataset_service_delete_dataset.py b/api/tests/test_containers_integration_tests/services/test_dataset_service_delete_dataset.py index b3ee73721d4..dc57992a999 100644 --- a/api/tests/test_containers_integration_tests/services/test_dataset_service_delete_dataset.py +++ b/api/tests/test_containers_integration_tests/services/test_dataset_service_delete_dataset.py @@ -9,7 +9,8 @@ from core.rag.index_processor.constant.index_type import IndexStructureType, Ind from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole from models.dataset import Dataset, Document from models.enums import DataSourceType, DocumentCreatedFrom -from services.dataset_service import DatasetService +from repositories.knowledge.dataset_read_repository import get_dataset_doc_form +from services.knowledge.dataset_service import DatasetService class DatasetDeleteIntegrationDataFactory: @@ -142,7 +143,7 @@ class TestDatasetServiceDeleteDataset: dataset.indexing_technique, dataset.index_struct, dataset.collection_binding_id, - dataset.get_doc_form(session=db_session_with_containers), + get_dataset_doc_form(dataset, session=db_session_with_containers), dataset.pipeline_id, ) diff --git a/api/tests/test_containers_integration_tests/services/test_dataset_service_document.py b/api/tests/test_containers_integration_tests/services/test_dataset_service_document.py index da7b0416a3d..28dc5d2f583 100644 --- a/api/tests/test_containers_integration_tests/services/test_dataset_service_document.py +++ b/api/tests/test_containers_integration_tests/services/test_dataset_service_document.py @@ -16,9 +16,10 @@ from models.account import TenantAccountRole from models.dataset import Dataset, Document from models.enums import CreatorUserRole, DataSourceType, DocumentCreatedFrom, IndexingStatus from models.model import UploadFile -from services.dataset_ref_service import DatasetRefService -from services.dataset_service import DocumentService +from repositories.knowledge.dataset_read_repository import get_dataset_doc_form from services.errors.account import NoPermissionError +from services.knowledge.dataset_service import DocumentService +from services.knowledge.resource_scope import DatasetRef FIXED_UPLOAD_CREATED_AT = datetime.datetime(2024, 1, 1, 0, 0, 0) @@ -117,7 +118,9 @@ class DocumentServiceIntegrationFactory: @pytest.fixture def current_user_mock(): - with patch("services.dataset_service.current_user", create_autospec(Account, instance=True)) as current_user: + with patch( + "services.knowledge.dataset_service.current_user", create_autospec(Account, instance=True) + ) as current_user: current_user.id = str(uuid4()) current_user.current_tenant_id = str(uuid4()) current_user.current_role = TenantAccountRole.EDITOR @@ -144,7 +147,7 @@ def test_get_documents_by_ids_returns_empty_for_empty_input(db_session_with_cont dataset = DocumentServiceIntegrationFactory.create_dataset(db_session_with_containers) result = DocumentService.get_documents_by_ids( - DatasetRefService.create_dataset_ref(dataset), [], session=db_session_with_containers + DatasetRef(tenant_id=dataset.tenant_id, dataset_id=dataset.id), [], session=db_session_with_containers ) assert result == [] @@ -161,7 +164,7 @@ def test_get_documents_by_ids_uses_single_batch_query(db_session_with_containers ) result = DocumentService.get_documents_by_ids( - DatasetRefService.create_dataset_ref(dataset), [doc_a.id, doc_b.id], db_session_with_containers + DatasetRef(tenant_id=dataset.tenant_id, dataset_id=dataset.id), [doc_a.id, doc_b.id], db_session_with_containers ) assert {document.id for document in result} == {doc_a.id, doc_b.id} @@ -218,7 +221,9 @@ def test_get_document_download_url_uses_signed_url_helper(db_session_with_contai data_source_info={"upload_file_id": upload_file.id}, ) - with patch("services.dataset_service.file_helpers.get_signed_file_url", return_value="signed-url") as get_url: + with patch( + "services.knowledge.dataset_service.file_helpers.get_signed_file_url", return_value="signed-url" + ) as get_url: result = DocumentService.get_document_download_url(document, session=db_session_with_containers) assert result == "signed-url" @@ -287,7 +292,7 @@ def test_get_upload_file_for_upload_file_document_raises_when_file_service_retur data_source_info={"upload_file_id": "missing-file"}, ) - with patch("services.dataset_service.FileService.get_upload_files_by_ids", return_value={}): + with patch("services.knowledge.dataset_service.FileService.get_upload_files_by_ids", return_value={}): with pytest.raises(NotFound, match="Uploaded file not found"): DocumentService._get_upload_file_for_upload_file_document(document, session=db_session_with_containers) @@ -432,7 +437,7 @@ def test_prepare_document_batch_download_zip_translates_permission_error_to_forb ) with patch( - "services.dataset_service.DatasetService.check_dataset_permission", + "services.knowledge.dataset_service.DatasetService.check_dataset_permission", side_effect=NoPermissionError("denied"), ): with pytest.raises(Forbidden, match="denied"): @@ -552,7 +557,7 @@ def test_get_error_documents_by_dataset_ref_returns_error_and_paused_documents(d indexing_status=IndexingStatus.COMPLETED, ) - dataset_ref = DatasetRefService.create_dataset_ref(dataset) + dataset_ref = DatasetRef(tenant_id=dataset.tenant_id, dataset_id=dataset.id) result = DocumentService.get_error_documents_by_dataset_ref(dataset_ref, session=db_session_with_containers) assert {document.id for document in result} == {error_document.id, paused_document.id} @@ -574,7 +579,9 @@ def test_get_batch_documents_filters_by_current_user_tenant(db_session_with_cont batch=batch, ) - with patch("services.dataset_service.current_user", create_autospec(Account, instance=True)) as current_user: + with patch( + "services.knowledge.dataset_service.current_user", create_autospec(Account, instance=True) + ) as current_user: current_user.current_tenant_id = dataset.tenant_id result = DocumentService.get_batch_documents(dataset.id, batch, session=db_session_with_containers) @@ -608,7 +615,7 @@ def test_delete_document_emits_signal_and_commits(db_session_with_containers: Se data_source_info={"upload_file_id": upload_file.id}, ) - with patch("services.dataset_service.document_was_deleted.send") as signal_send: + with patch("services.knowledge.dataset_service.document_was_deleted.send") as signal_send: DocumentService.delete_document(document, session=db_session_with_containers) assert db_session_with_containers.get(Document, document.id) is None @@ -622,13 +629,13 @@ def test_delete_document_emits_signal_and_commits(db_session_with_containers: Se def test_delete_documents_ignores_empty_input(db_session_with_containers: Session): dataset = DocumentServiceIntegrationFactory.create_dataset(db_session_with_containers) - dataset_ref = DatasetRefService.create_dataset_ref(dataset) + dataset_ref = DatasetRef(tenant_id=dataset.tenant_id, dataset_id=dataset.id) - with patch("services.dataset_service.batch_clean_document_task.delay") as delay: + with patch("services.knowledge.dataset_service.batch_clean_document_task.delay") as delay: DocumentService.delete_documents( dataset_ref, [], - dataset.get_doc_form(session=db_session_with_containers), + get_dataset_doc_form(dataset, session=db_session_with_containers), session=db_session_with_containers, ) @@ -662,13 +669,13 @@ def test_delete_documents_deletes_rows_and_dispatches_cleanup_task(db_session_wi position=2, data_source_info={"upload_file_id": upload_file_b.id}, ) - dataset_ref = DatasetRefService.create_dataset_ref(dataset) + dataset_ref = DatasetRef(tenant_id=dataset.tenant_id, dataset_id=dataset.id) - with patch("services.dataset_service.batch_clean_document_task.delay") as delay: + with patch("services.knowledge.dataset_service.batch_clean_document_task.delay") as delay: DocumentService.delete_documents( dataset_ref, [document_a.id, document_b.id], - dataset.get_doc_form(session=db_session_with_containers), + get_dataset_doc_form(dataset, session=db_session_with_containers), session=db_session_with_containers, ) diff --git a/api/tests/test_containers_integration_tests/services/test_dataset_service_get_segments.py b/api/tests/test_containers_integration_tests/services/test_dataset_service_get_segments.py index 84b9bcd7c5d..c4546648314 100644 --- a/api/tests/test_containers_integration_tests/services/test_dataset_service_get_segments.py +++ b/api/tests/test_containers_integration_tests/services/test_dataset_service_get_segments.py @@ -16,7 +16,7 @@ from core.rag.index_processor.constant.index_type import IndexTechniqueType from models import Account, AccountStatus, Tenant, TenantAccountJoin, TenantAccountRole, TenantStatus from models.dataset import Dataset, DatasetPermissionEnum, Document, DocumentSegment from models.enums import DataSourceType, DocumentCreatedFrom, SegmentStatus -from services.dataset_service import SegmentService +from services.knowledge.dataset_service import SegmentService class SegmentServiceTestDataFactory: diff --git a/api/tests/test_containers_integration_tests/services/test_dataset_service_permissions.py b/api/tests/test_containers_integration_tests/services/test_dataset_service_permissions.py index e0dd574a175..9c7e2d48269 100644 --- a/api/tests/test_containers_integration_tests/services/test_dataset_service_permissions.py +++ b/api/tests/test_containers_integration_tests/services/test_dataset_service_permissions.py @@ -21,9 +21,9 @@ from models.dataset import ( DatasetPermissionEnum, ) from models.enums import DataSourceType -from services.dataset_ref_service import DatasetRef, DatasetRefService -from services.dataset_service import DatasetCollectionBindingService, DatasetPermissionService, DatasetService from services.errors.account import NoPermissionError +from services.knowledge.dataset_service import DatasetCollectionBindingService, DatasetPermissionService, DatasetService +from services.knowledge.resource_scope import DatasetRef from tests.unit_tests.config_override import config_overrides_context @@ -196,7 +196,7 @@ class TestDatasetServicePermissionsAndLifecycle: created_by=owner.id, ) - with patch("services.dataset_service.dataset_was_deleted.send") as send_deleted_signal: + with patch("services.knowledge.dataset_service.dataset_was_deleted.send") as send_deleted_signal: result = DatasetService.delete_dataset(dataset.id, user=owner, session=db_session_with_containers) assert result is True @@ -215,7 +215,7 @@ class TestDatasetServicePermissionsAndLifecycle: dataset_id=dataset.id, ) - dataset_ref = DatasetRefService.create_dataset_ref(dataset) + dataset_ref = DatasetRef(tenant_id=dataset.tenant_id, dataset_id=dataset.id) assert DatasetService.dataset_use_check(dataset_ref, session=db_session_with_containers) is True @@ -227,7 +227,7 @@ class TestDatasetServicePermissionsAndLifecycle: created_by=owner.id, ) - dataset_ref = DatasetRefService.create_dataset_ref(dataset) + dataset_ref = DatasetRef(tenant_id=dataset.tenant_id, dataset_id=dataset.id) assert DatasetService.dataset_use_check(dataset_ref, session=db_session_with_containers) is False @@ -400,7 +400,7 @@ class TestDatasetServicePermissionsAndLifecycle: ) now = datetime(2026, 4, 14, 18, 0, 0) - with patch("services.dataset_service.naive_utc_now", return_value=now): + with patch("services.knowledge.dataset_service.naive_utc_now", return_value=now): DatasetService.update_dataset_api_status(dataset, True, owner, session=db_session_with_containers) db_session_with_containers.refresh(dataset) @@ -414,7 +414,7 @@ class TestDatasetServicePermissionsAndLifecycle: features = SimpleNamespace(billing=SimpleNamespace(subscription=SimpleNamespace(plan="professional"))) dataset_ref = DatasetRef(tenant_id=tenant.id, dataset_id=str(uuid4())) - with patch("services.dataset_service.FeatureService.get_features", return_value=features): + with patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=features): result = DatasetService.get_dataset_auto_disable_logs(dataset_ref, session=db_session_with_containers) assert result == {"document_ids": [], "count": 0} @@ -440,9 +440,9 @@ class TestDatasetServicePermissionsAndLifecycle: document_id=str(uuid4()), ) features = SimpleNamespace(billing=SimpleNamespace(subscription=SimpleNamespace(plan="professional"))) - dataset_ref = DatasetRefService.create_dataset_ref(dataset) + dataset_ref = DatasetRef(tenant_id=dataset.tenant_id, dataset_id=dataset.id) - with patch("services.dataset_service.FeatureService.get_features", return_value=features): + with patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=features): result = DatasetService.get_dataset_auto_disable_logs(dataset_ref, session=db_session_with_containers) assert result["count"] == 2 diff --git a/api/tests/test_containers_integration_tests/services/test_dataset_service_retrieval.py b/api/tests/test_containers_integration_tests/services/test_dataset_service_retrieval.py index 5643bfbeba7..2da05037f2a 100644 --- a/api/tests/test_containers_integration_tests/services/test_dataset_service_retrieval.py +++ b/api/tests/test_containers_integration_tests/services/test_dataset_service_retrieval.py @@ -29,7 +29,7 @@ from models.dataset import ( ) from models.enums import DatasetQuerySource, DataSourceType, ProcessRuleMode, TagType from models.model import Tag, TagBinding -from services.dataset_service import DatasetService, DocumentService +from services.knowledge.dataset_service import DatasetService, DocumentService class DatasetRetrievalTestDataFactory: diff --git a/api/tests/test_containers_integration_tests/services/test_dataset_service_update_dataset.py b/api/tests/test_containers_integration_tests/services/test_dataset_service_update_dataset.py index 8f4099fa484..cfa8b74d466 100644 --- a/api/tests/test_containers_integration_tests/services/test_dataset_service_update_dataset.py +++ b/api/tests/test_containers_integration_tests/services/test_dataset_service_update_dataset.py @@ -17,8 +17,8 @@ from models.account import ( ) from models.dataset import Dataset, ExternalKnowledgeApis, ExternalKnowledgeBindings from models.enums import DataSourceType -from services.dataset_service import DatasetService from services.errors.account import NoPermissionError +from services.knowledge.dataset_service import DatasetService class DatasetUpdateTestDataFactory: @@ -370,7 +370,7 @@ class TestDatasetServiceUpdateDataset: "retrieval_model": "new_model", } - with patch("services.dataset_service.deal_dataset_vector_index_task") as mock_task: + with patch("services.knowledge.dataset_service.deal_dataset_vector_index_task") as mock_task: result = DatasetService.update_dataset(dataset.id, update_data, user, session=db_session_with_containers) mock_task.delay.assert_called_once_with(dataset.id, "remove") @@ -408,12 +408,12 @@ class TestDatasetServiceUpdateDataset: } with ( - patch("services.dataset_service.current_user", user), - patch("services.dataset_service.ModelManager.for_tenant") as mock_model_manager, + patch("services.knowledge.dataset_service.current_user", user), + patch("services.knowledge.dataset_service.ModelManager.for_tenant") as mock_model_manager, patch( - "services.dataset_service.DatasetCollectionBindingService.get_dataset_collection_binding" + "services.knowledge.dataset_service.DatasetCollectionBindingService.get_dataset_collection_binding" ) as mock_get_binding, - patch("services.dataset_service.deal_dataset_vector_index_task") as mock_task, + patch("services.knowledge.dataset_service.deal_dataset_vector_index_task") as mock_task, ): mock_model_manager.return_value.get_model_instance.return_value = embedding_model mock_get_binding.return_value = binding @@ -503,13 +503,13 @@ class TestDatasetServiceUpdateDataset: } with ( - patch("services.dataset_service.current_user", user), - patch("services.dataset_service.ModelManager.for_tenant") as mock_model_manager, + patch("services.knowledge.dataset_service.current_user", user), + patch("services.knowledge.dataset_service.ModelManager.for_tenant") as mock_model_manager, patch( - "services.dataset_service.DatasetCollectionBindingService.get_dataset_collection_binding" + "services.knowledge.dataset_service.DatasetCollectionBindingService.get_dataset_collection_binding" ) as mock_get_binding, - patch("services.dataset_service.deal_dataset_vector_index_task") as mock_task, - patch("services.dataset_service.regenerate_summary_index_task") as mock_regenerate_task, + patch("services.knowledge.dataset_service.deal_dataset_vector_index_task") as mock_task, + patch("services.knowledge.dataset_service.regenerate_summary_index_task") as mock_regenerate_task, ): mock_model_manager.return_value.get_model_instance.return_value = embedding_model mock_get_binding.return_value = binding @@ -589,8 +589,8 @@ class TestDatasetServiceUpdateDataset: } with ( - patch("services.dataset_service.current_user", user), - patch("services.dataset_service.ModelManager.for_tenant") as mock_model_manager, + patch("services.knowledge.dataset_service.current_user", user), + patch("services.knowledge.dataset_service.ModelManager.for_tenant") as mock_model_manager, ): mock_model_manager.return_value.get_model_instance.side_effect = Exception("No Embedding Model available") diff --git a/api/tests/test_containers_integration_tests/services/test_document_service_display_status.py b/api/tests/test_containers_integration_tests/services/test_document_service_display_status.py index 383a5f6374f..46b32dbaca9 100644 --- a/api/tests/test_containers_integration_tests/services/test_document_service_display_status.py +++ b/api/tests/test_containers_integration_tests/services/test_document_service_display_status.py @@ -7,7 +7,7 @@ from sqlalchemy.orm import Session from core.rag.index_processor.constant.index_type import IndexStructureType from models.dataset import Dataset, Document from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus -from services.dataset_service import DocumentService +from services.knowledge.dataset_service import DocumentService def _create_dataset(db_session_with_containers) -> Dataset: diff --git a/api/tests/test_containers_integration_tests/services/test_document_service_rename_document.py b/api/tests/test_containers_integration_tests/services/test_document_service_rename_document.py index 6a0ea17560a..773452527d0 100644 --- a/api/tests/test_containers_integration_tests/services/test_document_service_rename_document.py +++ b/api/tests/test_containers_integration_tests/services/test_document_service_rename_document.py @@ -13,7 +13,7 @@ from models import Account from models.dataset import Dataset, Document from models.enums import CreatorUserRole, DataSourceType, DocumentCreatedFrom from models.model import UploadFile -from services.dataset_service import DocumentService +from services.knowledge.dataset_service import DocumentService FIXED_UPLOAD_CREATED_AT = datetime.datetime(2024, 1, 1, 0, 0, 0) @@ -21,7 +21,9 @@ FIXED_UPLOAD_CREATED_AT = datetime.datetime(2024, 1, 1, 0, 0, 0) @pytest.fixture def mock_env(): """Patch only non-SQL dependency used by rename_document: current_user context.""" - with patch("services.dataset_service.current_user", create_autospec(Account, instance=True)) as current_user: + with patch( + "services.knowledge.dataset_service.current_user", create_autospec(Account, instance=True) + ) as current_user: current_user.current_tenant_id = str(uuid4()) current_user.id = str(uuid4()) yield {"current_user": current_user} diff --git a/api/tests/test_containers_integration_tests/services/test_metadata_partial_update.py b/api/tests/test_containers_integration_tests/services/test_metadata_partial_update.py index b38bc6f4862..016c85e9dad 100644 --- a/api/tests/test_containers_integration_tests/services/test_metadata_partial_update.py +++ b/api/tests/test_containers_integration_tests/services/test_metadata_partial_update.py @@ -2,7 +2,7 @@ from __future__ import annotations from concurrent.futures import ThreadPoolExecutor from threading import Barrier -from unittest.mock import patch +from unittest.mock import create_autospec from uuid import uuid4 import pytest @@ -13,12 +13,24 @@ from sqlalchemy.orm import Session, sessionmaker from models import Account, Tenant from models.dataset import Dataset, DatasetMetadata, DatasetMetadataBinding, Document from models.enums import DatasetMetadataType, DataSourceType, DocumentCreatedFrom -from services.entities.knowledge_entities.knowledge_entities import ( +from repositories.knowledge.metadata_repository import SQLAlchemyMetadataRepository +from services.knowledge.dataset_access import DatasetAccess +from services.knowledge.entities.knowledge_entities import ( DocumentMetadataOperation, MetadataDetail, MetadataOperationData, ) -from services.metadata_service import MetadataService +from services.knowledge.metadata.application import MetadataService +from services.knowledge.resource_scope import DatasetRef + + +def _metadata_service(session: Session) -> MetadataService: + return MetadataService( + store=SQLAlchemyMetadataRepository( + session_factory=sessionmaker(bind=session.get_bind(), expire_on_commit=False) + ), + dataset_access=create_autospec(DatasetAccess, instance=True), + ) def _create_dataset(db_session: Session, *, tenant_id: str, built_in_field_enabled: bool = False) -> Dataset: @@ -112,8 +124,8 @@ class TestMetadataPartialUpdate: ) metadata_args = MetadataOperationData(operation_data=[operation]) - MetadataService.update_documents_metadata( - dataset, metadata_args, current_account, session=db_session_with_containers + _metadata_service(db_session_with_containers).update_documents_metadata( + DatasetRef(dataset.tenant_id, dataset.id), metadata_args, actor_id=current_account.id ) db_session_with_containers.expire_all() @@ -171,7 +183,7 @@ class TestMetadataPartialUpdate: _context: object, _executemany: bool, ) -> None: - if "FROM documents" in statement and "FOR UPDATE" in statement: + if "FROM datasets" in statement and "FOR UPDATE" in statement: locked_selects.append(None) locked_select_barrier.wait(timeout=10) @@ -179,15 +191,13 @@ class TestMetadataPartialUpdate: with session_factory() as session: owned_dataset = session.get(Dataset, dataset_id) assert owned_dataset is not None - MetadataService.update_documents_metadata( - owned_dataset, metadata_args, current_account, session=session + _metadata_service(session).update_documents_metadata( + DatasetRef(owned_dataset.tenant_id, owned_dataset.id), metadata_args, actor_id=current_account.id ) event.listen(engine, "before_cursor_execute", synchronize_locked_select) try: with ( - patch.object(MetadataService, "knowledge_base_metadata_lock_check"), - patch("services.metadata_service.redis_client.delete"), ThreadPoolExecutor(max_workers=2) as executor, ): futures = [executor.submit(update, operation) for operation in operations] @@ -227,8 +237,8 @@ class TestMetadataPartialUpdate: ) metadata_args = MetadataOperationData(operation_data=[operation]) - MetadataService.update_documents_metadata( - dataset, metadata_args, current_account, session=db_session_with_containers + _metadata_service(db_session_with_containers).update_documents_metadata( + DatasetRef(dataset.tenant_id, dataset.id), metadata_args, actor_id=current_account.id ) db_session_with_containers.expire_all() @@ -273,8 +283,8 @@ class TestMetadataPartialUpdate: ) metadata_args = MetadataOperationData(operation_data=[operation]) - MetadataService.update_documents_metadata( - dataset, metadata_args, current_account, session=db_session_with_containers + _metadata_service(db_session_with_containers).update_documents_metadata( + DatasetRef(dataset.tenant_id, dataset.id), metadata_args, actor_id=current_account.id ) db_session_with_containers.expire_all() @@ -311,8 +321,20 @@ class TestMetadataPartialUpdate: ) metadata_args = MetadataOperationData(operation_data=[operation]) - with patch.object(db_session_with_containers, "commit", side_effect=RuntimeError("database connection lost")): + service = _metadata_service(db_session_with_containers) + + def fail_commit(_session: Session) -> None: + raise RuntimeError("database connection lost") + + event.listen(Session, "before_commit", fail_commit) + try: with pytest.raises(RuntimeError, match="database connection lost"): - MetadataService.update_documents_metadata( - dataset, metadata_args, current_account, session=db_session_with_containers + service.update_documents_metadata( + DatasetRef(dataset.tenant_id, dataset.id), metadata_args, actor_id=current_account.id ) + finally: + event.remove(Session, "before_commit", fail_commit) + db_session_with_containers.expire_all() + unchanged = db_session_with_containers.get(Document, document.id) + assert unchanged is not None + assert unchanged.doc_metadata == {"existing_key": "existing_value"} diff --git a/api/tests/test_containers_integration_tests/services/test_metadata_service.py b/api/tests/test_containers_integration_tests/services/test_metadata_service.py index 6f7bf625563..6caa5fd757d 100644 --- a/api/tests/test_containers_integration_tests/services/test_metadata_service.py +++ b/api/tests/test_containers_integration_tests/services/test_metadata_service.py @@ -1,115 +1,68 @@ -from collections.abc import Generator -from typing import TypedDict -from unittest.mock import Mock, patch +from unittest.mock import create_autospec import pytest from faker import Faker -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, sessionmaker from core.rag.index_processor.constant.built_in_field import BuiltInField from core.rag.index_processor.constant.index_type import IndexStructureType from models import Account, AccountStatus, Tenant, TenantAccountJoin, TenantAccountRole, TenantStatus from models.dataset import Dataset, DatasetMetadata, DatasetMetadataBinding, Document from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus -from services.entities.knowledge_entities.knowledge_entities import ( +from repositories.knowledge.metadata_repository import SQLAlchemyMetadataRepository +from services.errors.metadata import MetadataResourceNotFoundError +from services.knowledge.dataset_access import DatasetAccess +from services.knowledge.entities.knowledge_entities import ( DocumentMetadataOperation, MetadataArgs, MetadataDetail, MetadataOperationData, ) -from services.errors.metadata import MetadataResourceNotFoundError -from services.metadata_service import MetadataService +from services.knowledge.metadata.application import MetadataService +from services.knowledge.resource_scope import DatasetRef -class MetadataServiceDeps(TypedDict): - redis_client: Mock - document_service: Mock +def _metadata_service(session: Session) -> MetadataService: + return MetadataService( + store=SQLAlchemyMetadataRepository( + session_factory=sessionmaker(bind=session.get_bind(), expire_on_commit=False) + ), + dataset_access=create_autospec(DatasetAccess, instance=True), + ) class TestMetadataService: """Integration tests for MetadataService using testcontainers.""" - @pytest.fixture - def mock_external_service_dependencies(self) -> Generator[MetadataServiceDeps, None, None]: - """Mock setup for external service dependencies.""" - with ( - patch("services.metadata_service.redis_client") as mock_redis_client, - patch("services.dataset_service.DocumentService") as mock_document_service, - ): - # Setup default mock returns - mock_redis_client.get.return_value = None - mock_redis_client.set.return_value = True - mock_redis_client.delete.return_value = 1 - - yield { - "redis_client": mock_redis_client, - "document_service": mock_document_service, - } - - def _create_test_account_and_tenant( - self, - db_session_with_containers: Session, - mock_external_service_dependencies: MetadataServiceDeps, - ) -> tuple[Account, Tenant]: + def _create_test_account_and_tenant(self, db_session_with_containers: Session) -> tuple[Account, Tenant]: """ Helper method to create a test account and tenant for testing. Args: db_session_with_containers: Database session from testcontainers infrastructure - mock_external_service_dependencies: Mock dependencies Returns: tuple: (account, tenant) - Created account and tenant instances """ fake = Faker() - - # Create account - account = Account( - email=fake.email(), - name=fake.name(), - interface_language="en-US", - status=AccountStatus.ACTIVE, - ) - + account = Account(email=fake.email(), name=fake.name(), interface_language="en-US", status=AccountStatus.ACTIVE) db_session_with_containers.add(account) db_session_with_containers.commit() - - # Create tenant for the account - tenant = Tenant( - name=fake.company(), - status=TenantStatus.NORMAL, - ) + tenant = Tenant(name=fake.company(), status=TenantStatus.NORMAL) db_session_with_containers.add(tenant) db_session_with_containers.commit() - - # Create tenant-account join - join = TenantAccountJoin( - tenant_id=tenant.id, - account_id=account.id, - role=TenantAccountRole.OWNER, - current=True, - ) + join = TenantAccountJoin(tenant_id=tenant.id, account_id=account.id, role=TenantAccountRole.OWNER, current=True) db_session_with_containers.add(join) db_session_with_containers.commit() - - # Set current tenant for account account.set_current_tenant_with_session(tenant, session=db_session_with_containers) - return account, tenant - def _create_test_dataset( - self, - db_session_with_containers: Session, - mock_external_service_dependencies: MetadataServiceDeps, - account: Account, - tenant: Tenant, - ) -> Dataset: + def _create_test_dataset(self, db_session_with_containers: Session, account: Account, tenant: Tenant) -> Dataset: """ Helper method to create a test dataset for testing. Args: db_session_with_containers: Database session from testcontainers infrastructure - mock_external_service_dependencies: Mock dependencies account: Account instance tenant: Tenant instance @@ -117,7 +70,6 @@ class TestMetadataService: Dataset: Created dataset instance """ fake = Faker() - dataset = Dataset( tenant_id=tenant.id, name=fake.company(), @@ -126,25 +78,18 @@ class TestMetadataService: created_by=account.id, built_in_field_enabled=False, ) - db_session_with_containers.add(dataset) db_session_with_containers.commit() - return dataset def _create_test_document( - self, - db_session_with_containers: Session, - mock_external_service_dependencies: MetadataServiceDeps, - dataset: Dataset, - account: Account, + self, db_session_with_containers: Session, dataset: Dataset, account: Account ) -> Document: """ Helper method to create a test document for testing. Args: db_session_with_containers: Database session from testcontainers infrastructure - mock_external_service_dependencies: Mock dependencies dataset: Dataset instance account: Account instance @@ -152,7 +97,6 @@ class TestMetadataService: Document: Created document instance """ fake = Faker() - document = Document( tenant_id=dataset.tenant_id, dataset_id=dataset.id, @@ -166,357 +110,221 @@ class TestMetadataService: doc_form=IndexStructureType.PARAGRAPH_INDEX, doc_language="en", ) - db_session_with_containers.add(document) db_session_with_containers.commit() - return document - def test_create_metadata_success( - self, db_session_with_containers: Session, mock_external_service_dependencies: MetadataServiceDeps - ) -> None: + def test_create_metadata_success(self, db_session_with_containers: Session) -> None: """ Test successful metadata creation with valid parameters. """ - # Arrange: Create test data - account, tenant = self._create_test_account_and_tenant( - db_session_with_containers, mock_external_service_dependencies - ) - dataset = self._create_test_dataset( - db_session_with_containers, mock_external_service_dependencies, account, tenant - ) - + account, tenant = self._create_test_account_and_tenant(db_session_with_containers) + dataset = self._create_test_dataset(db_session_with_containers, account, tenant) metadata_args = MetadataArgs(type="string", name="test_metadata") - - # Act: Execute the method under test - result = MetadataService.create_metadata( - dataset.id, metadata_args, account, tenant.id, session=db_session_with_containers + result = _metadata_service(db_session_with_containers).create_metadata( + DatasetRef(tenant.id, dataset.id), metadata_args, actor_id=account.id ) - - # Assert: Verify the expected outcomes + db_session_with_containers.expire_all() assert result is not None assert result.name == "test_metadata" assert result.type == "string" assert result.dataset_id == dataset.id assert result.tenant_id == tenant.id assert result.created_by == account.id - - # Verify database state - - db_session_with_containers.refresh(result) + db_session_with_containers.expire_all() assert result.id is not None assert result.created_at is not None - def test_create_metadata_name_too_long( - self, db_session_with_containers: Session, mock_external_service_dependencies: MetadataServiceDeps - ) -> None: + def test_create_metadata_name_too_long(self, db_session_with_containers: Session) -> None: """ Test metadata creation fails when name exceeds 255 characters. """ - # Arrange: Create test data - account, tenant = self._create_test_account_and_tenant( - db_session_with_containers, mock_external_service_dependencies - ) - dataset = self._create_test_dataset( - db_session_with_containers, mock_external_service_dependencies, account, tenant - ) - - long_name = "a" * 256 # 256 characters, exceeding 255 limit + account, tenant = self._create_test_account_and_tenant(db_session_with_containers) + dataset = self._create_test_dataset(db_session_with_containers, account, tenant) + long_name = "a" * 256 metadata_args = MetadataArgs(type="string", name=long_name) - - # Act & Assert: Verify proper error handling with pytest.raises(ValueError, match="Metadata name cannot exceed 255 characters."): - MetadataService.create_metadata( - dataset.id, metadata_args, account, tenant.id, session=db_session_with_containers + _metadata_service(db_session_with_containers).create_metadata( + DatasetRef(tenant.id, dataset.id), metadata_args, actor_id=account.id ) + db_session_with_containers.expire_all() - def test_create_metadata_name_already_exists( - self, db_session_with_containers: Session, mock_external_service_dependencies: MetadataServiceDeps - ) -> None: + def test_create_metadata_name_already_exists(self, db_session_with_containers: Session) -> None: """ Test metadata creation fails when name already exists in the same dataset. """ - # Arrange: Create test data - account, tenant = self._create_test_account_and_tenant( - db_session_with_containers, mock_external_service_dependencies - ) - dataset = self._create_test_dataset( - db_session_with_containers, mock_external_service_dependencies, account, tenant - ) - - # Create first metadata + account, tenant = self._create_test_account_and_tenant(db_session_with_containers) + dataset = self._create_test_dataset(db_session_with_containers, account, tenant) first_metadata_args = MetadataArgs(type="string", name="duplicate_name") - MetadataService.create_metadata( - dataset.id, first_metadata_args, account, tenant.id, session=db_session_with_containers + _metadata_service(db_session_with_containers).create_metadata( + DatasetRef(tenant.id, dataset.id), first_metadata_args, actor_id=account.id ) - - # Try to create second metadata with same name + db_session_with_containers.expire_all() second_metadata_args = MetadataArgs(type="number", name="duplicate_name") - - # Act & Assert: Verify proper error handling with pytest.raises(ValueError, match="Metadata name already exists."): - MetadataService.create_metadata( - dataset.id, second_metadata_args, account, tenant.id, session=db_session_with_containers + _metadata_service(db_session_with_containers).create_metadata( + DatasetRef(tenant.id, dataset.id), second_metadata_args, actor_id=account.id ) + db_session_with_containers.expire_all() - def test_create_metadata_name_conflicts_with_built_in_field( - self, db_session_with_containers: Session, mock_external_service_dependencies: MetadataServiceDeps - ) -> None: + def test_create_metadata_name_conflicts_with_built_in_field(self, db_session_with_containers: Session) -> None: """ Test metadata creation fails when name conflicts with built-in field names. """ - # Arrange: Create test data - account, tenant = self._create_test_account_and_tenant( - db_session_with_containers, mock_external_service_dependencies - ) - dataset = self._create_test_dataset( - db_session_with_containers, mock_external_service_dependencies, account, tenant - ) - - # Try to create metadata with built-in field name + account, tenant = self._create_test_account_and_tenant(db_session_with_containers) + dataset = self._create_test_dataset(db_session_with_containers, account, tenant) built_in_field_name = BuiltInField.document_name metadata_args = MetadataArgs(type="string", name=built_in_field_name) - - # Act & Assert: Verify proper error handling with pytest.raises(ValueError, match="Metadata name already exists in Built-in fields."): - MetadataService.create_metadata( - dataset.id, metadata_args, account, tenant.id, session=db_session_with_containers + _metadata_service(db_session_with_containers).create_metadata( + DatasetRef(tenant.id, dataset.id), metadata_args, actor_id=account.id ) + db_session_with_containers.expire_all() - def test_update_metadata_name_success( - self, db_session_with_containers: Session, mock_external_service_dependencies: MetadataServiceDeps - ) -> None: + def test_update_metadata_name_success(self, db_session_with_containers: Session) -> None: """ Test successful metadata name update with valid parameters. """ - # Arrange: Create test data - account, tenant = self._create_test_account_and_tenant( - db_session_with_containers, mock_external_service_dependencies - ) - dataset = self._create_test_dataset( - db_session_with_containers, mock_external_service_dependencies, account, tenant - ) - - # Create metadata first + account, tenant = self._create_test_account_and_tenant(db_session_with_containers) + dataset = self._create_test_dataset(db_session_with_containers, account, tenant) metadata_args = MetadataArgs(type="string", name="old_name") - metadata = MetadataService.create_metadata( - dataset.id, metadata_args, account, tenant.id, session=db_session_with_containers + metadata = _metadata_service(db_session_with_containers).create_metadata( + DatasetRef(tenant.id, dataset.id), metadata_args, actor_id=account.id ) - - # Act: Execute the method under test + db_session_with_containers.expire_all() new_name = "new_name" - result = MetadataService.update_metadata_name( - dataset, metadata.id, new_name, account, session=db_session_with_containers + result = _metadata_service(db_session_with_containers).update_metadata_name( + DatasetRef(dataset.tenant_id, dataset.id), metadata.id, new_name, actor_id=account.id ) - - # Assert: Verify the expected outcomes + db_session_with_containers.expire_all() assert result is not None assert result.name == new_name assert result.updated_by == account.id assert result.updated_at is not None - - # Verify database state - - db_session_with_containers.refresh(result) + db_session_with_containers.expire_all() assert result.name == new_name - def test_update_metadata_name_too_long( - self, db_session_with_containers: Session, mock_external_service_dependencies: MetadataServiceDeps - ) -> None: + def test_update_metadata_name_too_long(self, db_session_with_containers: Session) -> None: """ Test metadata name update fails when new name exceeds 255 characters. """ - # Arrange: Create test data - account, tenant = self._create_test_account_and_tenant( - db_session_with_containers, mock_external_service_dependencies - ) - dataset = self._create_test_dataset( - db_session_with_containers, mock_external_service_dependencies, account, tenant - ) - - # Create metadata first + account, tenant = self._create_test_account_and_tenant(db_session_with_containers) + dataset = self._create_test_dataset(db_session_with_containers, account, tenant) metadata_args = MetadataArgs(type="string", name="old_name") - metadata = MetadataService.create_metadata( - dataset.id, metadata_args, account, tenant.id, session=db_session_with_containers + metadata = _metadata_service(db_session_with_containers).create_metadata( + DatasetRef(tenant.id, dataset.id), metadata_args, actor_id=account.id ) - - # Try to update with too long name - long_name = "a" * 256 # 256 characters, exceeding 255 limit - - # Act & Assert: Verify proper error handling + db_session_with_containers.expire_all() + long_name = "a" * 256 with pytest.raises(ValueError, match="Metadata name cannot exceed 255 characters."): - MetadataService.update_metadata_name( - dataset, metadata.id, long_name, account, session=db_session_with_containers + _metadata_service(db_session_with_containers).update_metadata_name( + DatasetRef(dataset.tenant_id, dataset.id), metadata.id, long_name, actor_id=account.id ) + db_session_with_containers.expire_all() - def test_update_metadata_name_already_exists( - self, db_session_with_containers: Session, mock_external_service_dependencies: MetadataServiceDeps - ) -> None: + def test_update_metadata_name_already_exists(self, db_session_with_containers: Session) -> None: """ Test metadata name update fails when new name already exists in the same dataset. """ - # Arrange: Create test data - account, tenant = self._create_test_account_and_tenant( - db_session_with_containers, mock_external_service_dependencies - ) - dataset = self._create_test_dataset( - db_session_with_containers, mock_external_service_dependencies, account, tenant - ) - - # Create two metadata entries + account, tenant = self._create_test_account_and_tenant(db_session_with_containers) + dataset = self._create_test_dataset(db_session_with_containers, account, tenant) first_metadata_args = MetadataArgs(type="string", name="first_metadata") - first_metadata = MetadataService.create_metadata( - dataset.id, first_metadata_args, account, tenant.id, session=db_session_with_containers + first_metadata = _metadata_service(db_session_with_containers).create_metadata( + DatasetRef(tenant.id, dataset.id), first_metadata_args, actor_id=account.id ) - + db_session_with_containers.expire_all() second_metadata_args = MetadataArgs(type="number", name="second_metadata") - second_metadata = MetadataService.create_metadata( - dataset.id, second_metadata_args, account, tenant.id, session=db_session_with_containers + second_metadata = _metadata_service(db_session_with_containers).create_metadata( + DatasetRef(tenant.id, dataset.id), second_metadata_args, actor_id=account.id ) - - # Try to update first metadata with second metadata's name + db_session_with_containers.expire_all() with pytest.raises(ValueError, match="Metadata name already exists."): - MetadataService.update_metadata_name( - dataset, first_metadata.id, "second_metadata", account, session=db_session_with_containers + _metadata_service(db_session_with_containers).update_metadata_name( + DatasetRef(dataset.tenant_id, dataset.id), first_metadata.id, "second_metadata", actor_id=account.id ) + db_session_with_containers.expire_all() - def test_update_metadata_name_conflicts_with_built_in_field( - self, db_session_with_containers: Session, mock_external_service_dependencies: MetadataServiceDeps - ) -> None: + def test_update_metadata_name_conflicts_with_built_in_field(self, db_session_with_containers: Session) -> None: """ Test metadata name update fails when new name conflicts with built-in field names. """ - # Arrange: Create test data - account, tenant = self._create_test_account_and_tenant( - db_session_with_containers, mock_external_service_dependencies - ) - dataset = self._create_test_dataset( - db_session_with_containers, mock_external_service_dependencies, account, tenant - ) - - # Create metadata first + account, tenant = self._create_test_account_and_tenant(db_session_with_containers) + dataset = self._create_test_dataset(db_session_with_containers, account, tenant) metadata_args = MetadataArgs(type="string", name="old_name") - metadata = MetadataService.create_metadata( - dataset.id, metadata_args, account, tenant.id, session=db_session_with_containers + metadata = _metadata_service(db_session_with_containers).create_metadata( + DatasetRef(tenant.id, dataset.id), metadata_args, actor_id=account.id ) - - # Try to update with built-in field name + db_session_with_containers.expire_all() built_in_field_name = BuiltInField.document_name - with pytest.raises(ValueError, match="Metadata name already exists in Built-in fields."): - MetadataService.update_metadata_name( - dataset, metadata.id, built_in_field_name, account, session=db_session_with_containers + _metadata_service(db_session_with_containers).update_metadata_name( + DatasetRef(dataset.tenant_id, dataset.id), metadata.id, built_in_field_name, actor_id=account.id ) + db_session_with_containers.expire_all() - def test_update_metadata_name_not_found( - self, db_session_with_containers: Session, mock_external_service_dependencies: MetadataServiceDeps - ) -> None: + def test_update_metadata_name_not_found(self, db_session_with_containers: Session) -> None: """ Test metadata name update fails when metadata ID does not exist. """ - # Arrange: Create test data - account, tenant = self._create_test_account_and_tenant( - db_session_with_containers, mock_external_service_dependencies - ) - dataset = self._create_test_dataset( - db_session_with_containers, mock_external_service_dependencies, account, tenant - ) - - # Try to update non-existent metadata + account, tenant = self._create_test_account_and_tenant(db_session_with_containers) + dataset = self._create_test_dataset(db_session_with_containers, account, tenant) import uuid - fake_metadata_id = str(uuid.uuid4()) # Use valid UUID format + fake_metadata_id = str(uuid.uuid4()) new_name = "new_name" + with pytest.raises(MetadataResourceNotFoundError, match="Metadata not found"): + _metadata_service(db_session_with_containers).update_metadata_name( + DatasetRef(dataset.tenant_id, dataset.id), fake_metadata_id, new_name, actor_id=account.id + ) + db_session_with_containers.expire_all() - # Act: Execute the method under test - result = MetadataService.update_metadata_name( - dataset, fake_metadata_id, new_name, account, session=db_session_with_containers - ) - - # Assert: Verify the method returns None when metadata is not found - assert result is None - - def test_delete_metadata_success( - self, db_session_with_containers: Session, mock_external_service_dependencies: MetadataServiceDeps - ) -> None: + def test_delete_metadata_success(self, db_session_with_containers: Session) -> None: """ Test successful metadata deletion with valid parameters. """ - # Arrange: Create test data - account, tenant = self._create_test_account_and_tenant( - db_session_with_containers, mock_external_service_dependencies - ) - dataset = self._create_test_dataset( - db_session_with_containers, mock_external_service_dependencies, account, tenant - ) - - # Create metadata first + account, tenant = self._create_test_account_and_tenant(db_session_with_containers) + dataset = self._create_test_dataset(db_session_with_containers, account, tenant) metadata_args = MetadataArgs(type="string", name="to_be_deleted") - metadata = MetadataService.create_metadata( - dataset.id, metadata_args, account, tenant.id, session=db_session_with_containers + metadata = _metadata_service(db_session_with_containers).create_metadata( + DatasetRef(tenant.id, dataset.id), metadata_args, actor_id=account.id ) - - # Act: Execute the method under test - result = MetadataService.delete_metadata(dataset, metadata.id, session=db_session_with_containers) - - # Assert: Verify the expected outcomes + db_session_with_containers.expire_all() + result = _metadata_service(db_session_with_containers).delete_metadata( + DatasetRef(dataset.tenant_id, dataset.id), metadata.id + ) + db_session_with_containers.expire_all() assert result is not None assert result.id == metadata.id - - # Verify metadata was deleted from database - deleted_metadata = db_session_with_containers.query(DatasetMetadata).filter_by(id=metadata.id).first() assert deleted_metadata is None - def test_delete_metadata_not_found( - self, db_session_with_containers: Session, mock_external_service_dependencies: MetadataServiceDeps - ) -> None: + def test_delete_metadata_not_found(self, db_session_with_containers: Session) -> None: """ Test metadata deletion fails when metadata ID does not exist. """ - # Arrange: Create test data - account, tenant = self._create_test_account_and_tenant( - db_session_with_containers, mock_external_service_dependencies - ) - dataset = self._create_test_dataset( - db_session_with_containers, mock_external_service_dependencies, account, tenant - ) - - # Try to delete non-existent metadata + account, tenant = self._create_test_account_and_tenant(db_session_with_containers) + dataset = self._create_test_dataset(db_session_with_containers, account, tenant) import uuid - fake_metadata_id = str(uuid.uuid4()) # Use valid UUID format + fake_metadata_id = str(uuid.uuid4()) + with pytest.raises(MetadataResourceNotFoundError, match="Metadata not found"): + _metadata_service(db_session_with_containers).delete_metadata( + DatasetRef(dataset.tenant_id, dataset.id), fake_metadata_id + ) + db_session_with_containers.expire_all() - # Act: Execute the method under test - result = MetadataService.delete_metadata(dataset, fake_metadata_id, session=db_session_with_containers) - - # Assert: Verify the method returns None when metadata is not found - assert result is None - - def test_delete_metadata_with_document_bindings( - self, db_session_with_containers: Session, mock_external_service_dependencies: MetadataServiceDeps - ) -> None: + def test_delete_metadata_with_document_bindings(self, db_session_with_containers: Session) -> None: """ Test metadata deletion successfully removes document metadata bindings. """ - # Arrange: Create test data - account, tenant = self._create_test_account_and_tenant( - db_session_with_containers, mock_external_service_dependencies - ) - dataset = self._create_test_dataset( - db_session_with_containers, mock_external_service_dependencies, account, tenant - ) - document = self._create_test_document( - db_session_with_containers, mock_external_service_dependencies, dataset, account - ) - - # Create metadata + account, tenant = self._create_test_account_and_tenant(db_session_with_containers) + dataset = self._create_test_dataset(db_session_with_containers, account, tenant) + document = self._create_test_document(db_session_with_containers, dataset, account) metadata_args = MetadataArgs(type="string", name="test_metadata") - metadata = MetadataService.create_metadata( - dataset.id, metadata_args, account, tenant.id, session=db_session_with_containers + metadata = _metadata_service(db_session_with_containers).create_metadata( + DatasetRef(tenant.id, dataset.id), metadata_args, actor_id=account.id ) - - # Create metadata binding + db_session_with_containers.expire_all() binding = DatasetMetadataBinding( tenant_id=tenant.id, dataset_id=dataset.id, @@ -524,70 +332,37 @@ class TestMetadataService: document_id=document.id, created_by=account.id, ) - db_session_with_containers.add(binding) db_session_with_containers.commit() - - # Set document metadata document.doc_metadata = {"test_metadata": "test_value"} db_session_with_containers.add(document) db_session_with_containers.commit() - - # Act: Execute the method under test - result = MetadataService.delete_metadata(dataset, metadata.id, session=db_session_with_containers) - - # Assert: Verify the expected outcomes + result = _metadata_service(db_session_with_containers).delete_metadata( + DatasetRef(dataset.tenant_id, dataset.id), metadata.id + ) + db_session_with_containers.expire_all() assert result is not None - - # Verify metadata was deleted from database deleted_metadata = db_session_with_containers.query(DatasetMetadata).filter_by(id=metadata.id).first() assert deleted_metadata is None - # Note: The service attempts to update document metadata but may not succeed - # due to mock configuration. The main functionality (metadata deletion) is verified. - @pytest.mark.parametrize("operation", ["rename", "delete"]) @pytest.mark.parametrize("binding_owner", ["metadata", "document"]) def test_metadata_changes_ignore_historical_foreign_document_binding( - self, - operation: str, - binding_owner: str, - db_session_with_containers: Session, - mock_external_service_dependencies: MetadataServiceDeps, + self, operation: str, binding_owner: str, db_session_with_containers: Session ) -> None: - account, tenant = self._create_test_account_and_tenant( - db_session_with_containers, mock_external_service_dependencies - ) - dataset = self._create_test_dataset( - db_session_with_containers, mock_external_service_dependencies, account, tenant - ) - foreign_account, foreign_tenant = self._create_test_account_and_tenant( - db_session_with_containers, mock_external_service_dependencies - ) - foreign_dataset = self._create_test_dataset( - db_session_with_containers, - mock_external_service_dependencies, - foreign_account, - foreign_tenant, - ) - foreign_document = self._create_test_document( - db_session_with_containers, - mock_external_service_dependencies, - foreign_dataset, - foreign_account, - ) + account, tenant = self._create_test_account_and_tenant(db_session_with_containers) + dataset = self._create_test_dataset(db_session_with_containers, account, tenant) + foreign_account, foreign_tenant = self._create_test_account_and_tenant(db_session_with_containers) + foreign_dataset = self._create_test_dataset(db_session_with_containers, foreign_account, foreign_tenant) + foreign_document = self._create_test_document(db_session_with_containers, foreign_dataset, foreign_account) foreign_document.enabled = True foreign_document.archived = False foreign_document.indexing_status = IndexingStatus.COMPLETED foreign_document.doc_metadata = {"old_name": "foreign-value"} - - metadata = MetadataService.create_metadata( - dataset.id, - MetadataArgs(type="string", name="old_name"), - account, - tenant.id, - session=db_session_with_containers, + metadata = _metadata_service(db_session_with_containers).create_metadata( + DatasetRef(tenant.id, dataset.id), MetadataArgs(type="string", name="old_name"), actor_id=account.id ) + db_session_with_containers.expire_all() db_session_with_containers.add( DatasetMetadataBinding( tenant_id=dataset.tenant_id if binding_owner == "metadata" else foreign_dataset.tenant_id, @@ -598,163 +373,83 @@ class TestMetadataService: ) ) db_session_with_containers.commit() - if operation == "rename": - MetadataService.update_metadata_name( - dataset, metadata.id, "new_name", account, session=db_session_with_containers + _metadata_service(db_session_with_containers).update_metadata_name( + DatasetRef(dataset.tenant_id, dataset.id), metadata.id, "new_name", actor_id=account.id ) else: - MetadataService.delete_metadata(dataset, metadata.id, session=db_session_with_containers) - + _metadata_service(db_session_with_containers).delete_metadata( + DatasetRef(dataset.tenant_id, dataset.id), metadata.id + ) + db_session_with_containers.expire_all() db_session_with_containers.refresh(foreign_document) assert foreign_document.doc_metadata == {"old_name": "foreign-value"} - def test_get_built_in_fields_success( - self, db_session_with_containers: Session, mock_external_service_dependencies: MetadataServiceDeps - ) -> None: + def test_get_built_in_fields_success(self, db_session_with_containers: Session) -> None: """ Test successful retrieval of built-in metadata fields. """ - # Act: Execute the method under test result = MetadataService.get_built_in_fields() - - # Assert: Verify the expected outcomes assert result is not None assert len(result) == 5 - - # Verify all expected built-in fields are present field_names = [field["name"] for field in result] field_types = [field["type"] for field in result] - assert BuiltInField.document_name in field_names assert BuiltInField.uploader in field_names assert BuiltInField.upload_date in field_names assert BuiltInField.last_update_date in field_names assert BuiltInField.source in field_names - - # Verify field types assert "string" in field_types assert "time" in field_types - def test_enable_built_in_field_success( - self, db_session_with_containers: Session, mock_external_service_dependencies: MetadataServiceDeps - ) -> None: + def test_enable_built_in_field_success(self, db_session_with_containers: Session) -> None: """ Test successful enabling of built-in fields for a dataset. """ - # Arrange: Create test data - account, tenant = self._create_test_account_and_tenant( - db_session_with_containers, mock_external_service_dependencies - ) - dataset = self._create_test_dataset( - db_session_with_containers, mock_external_service_dependencies, account, tenant - ) - document = self._create_test_document( - db_session_with_containers, mock_external_service_dependencies, dataset, account - ) - - # Mock DocumentService.get_working_documents_by_dataset_id - mock_external_service_dependencies["document_service"].get_working_documents_by_dataset_id.return_value = [ - document - ] - - # Verify dataset starts with built-in fields disabled + account, tenant = self._create_test_account_and_tenant(db_session_with_containers) + dataset = self._create_test_dataset(db_session_with_containers, account, tenant) + document = self._create_test_document(db_session_with_containers, dataset, account) assert dataset.built_in_field_enabled is False - - # Act: Execute the method under test - MetadataService.enable_built_in_field(dataset, session=db_session_with_containers) - - # Assert: Verify the expected outcomes - + _metadata_service(db_session_with_containers).enable_built_in_field(DatasetRef(dataset.tenant_id, dataset.id)) + db_session_with_containers.expire_all() db_session_with_containers.refresh(dataset) assert dataset.built_in_field_enabled - # Note: Document metadata update depends on DocumentService mock working correctly - # The main functionality (enabling built-in fields) is verified - - def test_enable_built_in_field_already_enabled( - self, db_session_with_containers: Session, mock_external_service_dependencies: MetadataServiceDeps - ) -> None: + def test_enable_built_in_field_already_enabled(self, db_session_with_containers: Session) -> None: """ Test enabling built-in fields when they are already enabled. """ - # Arrange: Create test data - account, tenant = self._create_test_account_and_tenant( - db_session_with_containers, mock_external_service_dependencies - ) - dataset = self._create_test_dataset( - db_session_with_containers, mock_external_service_dependencies, account, tenant - ) - - # Enable built-in fields first + account, tenant = self._create_test_account_and_tenant(db_session_with_containers) + dataset = self._create_test_dataset(db_session_with_containers, account, tenant) dataset.built_in_field_enabled = True - db_session_with_containers.add(dataset) db_session_with_containers.commit() - - # Mock DocumentService.get_working_documents_by_dataset_id - mock_external_service_dependencies["document_service"].get_working_documents_by_dataset_id.return_value = list[ - Document - ]() - - # Act: Execute the method under test - MetadataService.enable_built_in_field(dataset, session=db_session_with_containers) - - # Assert: Verify the method returns early without changes + _metadata_service(db_session_with_containers).enable_built_in_field(DatasetRef(dataset.tenant_id, dataset.id)) + db_session_with_containers.expire_all() db_session_with_containers.refresh(dataset) assert dataset.built_in_field_enabled is True - def test_enable_built_in_field_with_no_documents( - self, db_session_with_containers: Session, mock_external_service_dependencies: MetadataServiceDeps - ) -> None: + def test_enable_built_in_field_with_no_documents(self, db_session_with_containers: Session) -> None: """ Test enabling built-in fields for a dataset with no documents. """ - # Arrange: Create test data - account, tenant = self._create_test_account_and_tenant( - db_session_with_containers, mock_external_service_dependencies - ) - dataset = self._create_test_dataset( - db_session_with_containers, mock_external_service_dependencies, account, tenant - ) - - # Mock DocumentService.get_working_documents_by_dataset_id to return empty list - mock_external_service_dependencies["document_service"].get_working_documents_by_dataset_id.return_value = list[ - Document - ]() - - # Act: Execute the method under test - MetadataService.enable_built_in_field(dataset, session=db_session_with_containers) - - # Assert: Verify the expected outcomes - + account, tenant = self._create_test_account_and_tenant(db_session_with_containers) + dataset = self._create_test_dataset(db_session_with_containers, account, tenant) + _metadata_service(db_session_with_containers).enable_built_in_field(DatasetRef(dataset.tenant_id, dataset.id)) + db_session_with_containers.expire_all() db_session_with_containers.refresh(dataset) assert dataset.built_in_field_enabled - def test_disable_built_in_field_success( - self, db_session_with_containers: Session, mock_external_service_dependencies: MetadataServiceDeps - ) -> None: + def test_disable_built_in_field_success(self, db_session_with_containers: Session) -> None: """ Test successful disabling of built-in fields for a dataset. """ - # Arrange: Create test data - account, tenant = self._create_test_account_and_tenant( - db_session_with_containers, mock_external_service_dependencies - ) - dataset = self._create_test_dataset( - db_session_with_containers, mock_external_service_dependencies, account, tenant - ) - document = self._create_test_document( - db_session_with_containers, mock_external_service_dependencies, dataset, account - ) - - # Enable built-in fields first + account, tenant = self._create_test_account_and_tenant(db_session_with_containers) + dataset = self._create_test_dataset(db_session_with_containers, account, tenant) + document = self._create_test_document(db_session_with_containers, dataset, account) dataset.built_in_field_enabled = True - db_session_with_containers.add(dataset) db_session_with_containers.commit() - - # Set document metadata with built-in fields document.doc_metadata = { BuiltInField.document_name: document.name, BuiltInField.uploader: "test_uploader", @@ -764,128 +459,60 @@ class TestMetadataService: } db_session_with_containers.add(document) db_session_with_containers.commit() - - # Mock DocumentService.get_working_documents_by_dataset_id - mock_external_service_dependencies["document_service"].get_working_documents_by_dataset_id.return_value = [ - document - ] - - # Act: Execute the method under test - MetadataService.disable_built_in_field(dataset, session=db_session_with_containers) - - # Assert: Verify the expected outcomes + _metadata_service(db_session_with_containers).disable_built_in_field(DatasetRef(dataset.tenant_id, dataset.id)) + db_session_with_containers.expire_all() db_session_with_containers.refresh(dataset) assert dataset.built_in_field_enabled is False - # Note: Document metadata update depends on DocumentService mock working correctly - # The main functionality (disabling built-in fields) is verified - - def test_disable_built_in_field_already_disabled( - self, db_session_with_containers: Session, mock_external_service_dependencies: MetadataServiceDeps - ) -> None: + def test_disable_built_in_field_already_disabled(self, db_session_with_containers: Session) -> None: """ Test disabling built-in fields when they are already disabled. """ - # Arrange: Create test data - account, tenant = self._create_test_account_and_tenant( - db_session_with_containers, mock_external_service_dependencies - ) - dataset = self._create_test_dataset( - db_session_with_containers, mock_external_service_dependencies, account, tenant - ) - - # Verify dataset starts with built-in fields disabled + account, tenant = self._create_test_account_and_tenant(db_session_with_containers) + dataset = self._create_test_dataset(db_session_with_containers, account, tenant) assert dataset.built_in_field_enabled is False - - # Mock DocumentService.get_working_documents_by_dataset_id - mock_external_service_dependencies["document_service"].get_working_documents_by_dataset_id.return_value = list[ - Document - ]() - - # Act: Execute the method under test - MetadataService.disable_built_in_field(dataset, session=db_session_with_containers) - - # Assert: Verify the method returns early without changes - + _metadata_service(db_session_with_containers).disable_built_in_field(DatasetRef(dataset.tenant_id, dataset.id)) + db_session_with_containers.expire_all() db_session_with_containers.refresh(dataset) assert not dataset.built_in_field_enabled - def test_disable_built_in_field_with_no_documents( - self, db_session_with_containers: Session, mock_external_service_dependencies: MetadataServiceDeps - ) -> None: + def test_disable_built_in_field_with_no_documents(self, db_session_with_containers: Session) -> None: """ Test disabling built-in fields for a dataset with no documents. """ - # Arrange: Create test data - account, tenant = self._create_test_account_and_tenant( - db_session_with_containers, mock_external_service_dependencies - ) - dataset = self._create_test_dataset( - db_session_with_containers, mock_external_service_dependencies, account, tenant - ) - - # Enable built-in fields first + account, tenant = self._create_test_account_and_tenant(db_session_with_containers) + dataset = self._create_test_dataset(db_session_with_containers, account, tenant) dataset.built_in_field_enabled = True - db_session_with_containers.add(dataset) db_session_with_containers.commit() - - # Mock DocumentService.get_working_documents_by_dataset_id to return empty list - mock_external_service_dependencies["document_service"].get_working_documents_by_dataset_id.return_value = list[ - Document - ]() - - # Act: Execute the method under test - MetadataService.disable_built_in_field(dataset, session=db_session_with_containers) - - # Assert: Verify the expected outcomes + _metadata_service(db_session_with_containers).disable_built_in_field(DatasetRef(dataset.tenant_id, dataset.id)) + db_session_with_containers.expire_all() db_session_with_containers.refresh(dataset) assert dataset.built_in_field_enabled is False - def test_update_documents_metadata_success( - self, db_session_with_containers: Session, mock_external_service_dependencies: MetadataServiceDeps - ) -> None: + def test_update_documents_metadata_success(self, db_session_with_containers: Session) -> None: """ Test successful update of documents metadata. """ - # Arrange: Create test data - account, tenant = self._create_test_account_and_tenant( - db_session_with_containers, mock_external_service_dependencies - ) - dataset = self._create_test_dataset( - db_session_with_containers, mock_external_service_dependencies, account, tenant - ) - document = self._create_test_document( - db_session_with_containers, mock_external_service_dependencies, dataset, account - ) - - # Create metadata + account, tenant = self._create_test_account_and_tenant(db_session_with_containers) + dataset = self._create_test_dataset(db_session_with_containers, account, tenant) + document = self._create_test_document(db_session_with_containers, dataset, account) metadata_args = MetadataArgs(type="string", name="test_metadata") - metadata = MetadataService.create_metadata( - dataset.id, metadata_args, account, tenant.id, session=db_session_with_containers + metadata = _metadata_service(db_session_with_containers).create_metadata( + DatasetRef(tenant.id, dataset.id), metadata_args, actor_id=account.id ) - - # Mock DocumentService.get_document - mock_external_service_dependencies["document_service"].get_document.return_value = document - + db_session_with_containers.expire_all() metadata_detail = MetadataDetail(id=metadata.id, name=metadata.name, value="test_value") - operation = DocumentMetadataOperation(document_id=document.id, metadata_list=[metadata_detail]) - operation_data = MetadataOperationData(operation_data=[operation]) - - # Act: Execute the method under test - MetadataService.update_documents_metadata(dataset, operation_data, account, session=db_session_with_containers) - - # Assert: Verify the expected outcomes - - # Verify document metadata was updated + _metadata_service(db_session_with_containers).update_documents_metadata( + DatasetRef(dataset.tenant_id, dataset.id), operation_data, actor_id=account.id + ) + db_session_with_containers.expire_all() db_session_with_containers.refresh(document) assert document.doc_metadata is not None assert "test_metadata" in document.doc_metadata assert document.doc_metadata["test_metadata"] == "test_value" - - # Verify metadata binding was created binding = ( db_session_with_containers.query(DatasetMetadataBinding) .filter_by(metadata_id=metadata.id, document_id=document.id) @@ -897,50 +524,24 @@ class TestMetadataService: @pytest.mark.parametrize("foreign_resource", ["metadata", "document"]) def test_update_documents_metadata_rejects_foreign_owner_before_writes( - self, - foreign_resource: str, - db_session_with_containers: Session, - mock_external_service_dependencies: MetadataServiceDeps, + self, foreign_resource: str, db_session_with_containers: Session ) -> None: - account, tenant = self._create_test_account_and_tenant( - db_session_with_containers, mock_external_service_dependencies + account, tenant = self._create_test_account_and_tenant(db_session_with_containers) + dataset = self._create_test_dataset(db_session_with_containers, account, tenant) + document = self._create_test_document(db_session_with_containers, dataset, account) + metadata = _metadata_service(db_session_with_containers).create_metadata( + DatasetRef(tenant.id, dataset.id), MetadataArgs(type="string", name="owned"), actor_id=account.id ) - dataset = self._create_test_dataset( - db_session_with_containers, mock_external_service_dependencies, account, tenant - ) - document = self._create_test_document( - db_session_with_containers, mock_external_service_dependencies, dataset, account - ) - metadata = MetadataService.create_metadata( - dataset.id, - MetadataArgs(type="string", name="owned"), - account, - tenant.id, - session=db_session_with_containers, - ) - - foreign_account, foreign_tenant = self._create_test_account_and_tenant( - db_session_with_containers, mock_external_service_dependencies - ) - foreign_dataset = self._create_test_dataset( - db_session_with_containers, - mock_external_service_dependencies, - foreign_account, - foreign_tenant, - ) - foreign_document = self._create_test_document( - db_session_with_containers, - mock_external_service_dependencies, - foreign_dataset, - foreign_account, - ) - foreign_metadata = MetadataService.create_metadata( - foreign_dataset.id, + db_session_with_containers.expire_all() + foreign_account, foreign_tenant = self._create_test_account_and_tenant(db_session_with_containers) + foreign_dataset = self._create_test_dataset(db_session_with_containers, foreign_account, foreign_tenant) + foreign_document = self._create_test_document(db_session_with_containers, foreign_dataset, foreign_account) + foreign_metadata = _metadata_service(db_session_with_containers).create_metadata( + DatasetRef(foreign_tenant.id, foreign_dataset.id), MetadataArgs(type="string", name="foreign"), - foreign_account, - foreign_tenant.id, - session=db_session_with_containers, + actor_id=foreign_account.id, ) + db_session_with_containers.expire_all() operation = DocumentMetadataOperation( document_id=foreign_document.id if foreign_resource == "document" else document.id, metadata_list=[ @@ -951,209 +552,80 @@ class TestMetadataService: ) ], ) - with pytest.raises(MetadataResourceNotFoundError, match=f"{foreign_resource.capitalize()} not found"): - MetadataService.update_documents_metadata( - dataset, + _metadata_service(db_session_with_containers).update_documents_metadata( + DatasetRef(dataset.tenant_id, dataset.id), MetadataOperationData(operation_data=[operation]), - account, - session=db_session_with_containers, + actor_id=account.id, ) - + db_session_with_containers.expire_all() db_session_with_containers.refresh(document) db_session_with_containers.refresh(foreign_document) assert document.doc_metadata is None assert foreign_document.doc_metadata is None assert db_session_with_containers.query(DatasetMetadataBinding).count() == 0 - def test_update_documents_metadata_with_built_in_fields_enabled( - self, db_session_with_containers: Session, mock_external_service_dependencies: MetadataServiceDeps - ) -> None: + def test_update_documents_metadata_with_built_in_fields_enabled(self, db_session_with_containers: Session) -> None: """ Test update of documents metadata when built-in fields are enabled. """ - # Arrange: Create test data - account, tenant = self._create_test_account_and_tenant( - db_session_with_containers, mock_external_service_dependencies - ) - dataset = self._create_test_dataset( - db_session_with_containers, mock_external_service_dependencies, account, tenant - ) - document = self._create_test_document( - db_session_with_containers, mock_external_service_dependencies, dataset, account - ) - - # Enable built-in fields + account, tenant = self._create_test_account_and_tenant(db_session_with_containers) + dataset = self._create_test_dataset(db_session_with_containers, account, tenant) + document = self._create_test_document(db_session_with_containers, dataset, account) dataset.built_in_field_enabled = True - db_session_with_containers.add(dataset) db_session_with_containers.commit() - - # Create metadata metadata_args = MetadataArgs(type="string", name="test_metadata") - metadata = MetadataService.create_metadata( - dataset.id, metadata_args, account, tenant.id, session=db_session_with_containers + metadata = _metadata_service(db_session_with_containers).create_metadata( + DatasetRef(tenant.id, dataset.id), metadata_args, actor_id=account.id ) - - # Mock DocumentService.get_document - mock_external_service_dependencies["document_service"].get_document.return_value = document - + db_session_with_containers.expire_all() metadata_detail = MetadataDetail(id=metadata.id, name=metadata.name, value="test_value") - operation = DocumentMetadataOperation(document_id=document.id, metadata_list=[metadata_detail]) - operation_data = MetadataOperationData(operation_data=[operation]) - - # Act: Execute the method under test - MetadataService.update_documents_metadata(dataset, operation_data, account, session=db_session_with_containers) - - # Assert: Verify the expected outcomes - # Verify document metadata was updated with both custom and built-in fields + _metadata_service(db_session_with_containers).update_documents_metadata( + DatasetRef(dataset.tenant_id, dataset.id), operation_data, actor_id=account.id + ) + db_session_with_containers.expire_all() db_session_with_containers.refresh(document) assert document.doc_metadata is not None assert "test_metadata" in document.doc_metadata assert document.doc_metadata["test_metadata"] == "test_value" - # Note: Built-in fields would be added if DocumentService mock works correctly - # The main functionality (custom metadata update) is verified - - def test_update_documents_metadata_document_not_found( - self, db_session_with_containers: Session, mock_external_service_dependencies: MetadataServiceDeps - ) -> None: + def test_update_documents_metadata_document_not_found(self, db_session_with_containers: Session) -> None: """ Test update of documents metadata when document is not found. """ - # Arrange: Create test data - account, tenant = self._create_test_account_and_tenant( - db_session_with_containers, mock_external_service_dependencies - ) - dataset = self._create_test_dataset( - db_session_with_containers, mock_external_service_dependencies, account, tenant - ) - - # Create metadata + account, tenant = self._create_test_account_and_tenant(db_session_with_containers) + dataset = self._create_test_dataset(db_session_with_containers, account, tenant) metadata_args = MetadataArgs(type="string", name="test_metadata") - metadata = MetadataService.create_metadata( - dataset.id, metadata_args, account, tenant.id, session=db_session_with_containers + metadata = _metadata_service(db_session_with_containers).create_metadata( + DatasetRef(tenant.id, dataset.id), metadata_args, actor_id=account.id ) - + db_session_with_containers.expire_all() metadata_detail = MetadataDetail(id=metadata.id, name=metadata.name, value="test_value") - - # Use a valid UUID format that does not exist in the database operation = DocumentMetadataOperation( document_id="00000000-0000-0000-0000-000000000000", metadata_list=[metadata_detail] ) - operation_data = MetadataOperationData(operation_data=[operation]) - with pytest.raises(MetadataResourceNotFoundError, match="Document not found"): - MetadataService.update_documents_metadata( - dataset, operation_data, account, session=db_session_with_containers + _metadata_service(db_session_with_containers).update_documents_metadata( + DatasetRef(dataset.tenant_id, dataset.id), operation_data, actor_id=account.id ) + db_session_with_containers.expire_all() - def test_knowledge_base_metadata_lock_check_dataset_id( - self, db_session_with_containers: Session, mock_external_service_dependencies: MetadataServiceDeps - ) -> None: - """ - Test metadata lock check for dataset operations. - """ - # Arrange: Setup mocks - mock_external_service_dependencies["redis_client"].get.return_value = None - mock_external_service_dependencies["redis_client"].set.return_value = True - - dataset_id = "test-dataset-id" - - # Act: Execute the method under test - MetadataService.knowledge_base_metadata_lock_check(dataset_id, None) - - # Assert: Verify the expected outcomes - # Verify Redis lock was set - mock_external_service_dependencies["redis_client"].set.assert_called_once() - - # Verify lock key format - call_args = mock_external_service_dependencies["redis_client"].set.call_args - assert call_args[0][0] == f"dataset_metadata_lock_{dataset_id}" - - def test_knowledge_base_metadata_lock_check_document_id( - self, db_session_with_containers: Session, mock_external_service_dependencies: MetadataServiceDeps - ) -> None: - """ - Test metadata lock check for document operations. - """ - # Arrange: Setup mocks - mock_external_service_dependencies["redis_client"].get.return_value = None - mock_external_service_dependencies["redis_client"].set.return_value = True - - document_id = "test-document-id" - - # Act: Execute the method under test - MetadataService.knowledge_base_metadata_lock_check(None, document_id) - - # Assert: Verify the expected outcomes - # Verify Redis lock was set - mock_external_service_dependencies["redis_client"].set.assert_called_once() - - # Verify lock key format - call_args = mock_external_service_dependencies["redis_client"].set.call_args - assert call_args[0][0] == f"document_metadata_lock_{document_id}" - - def test_knowledge_base_metadata_lock_check_lock_exists( - self, db_session_with_containers: Session, mock_external_service_dependencies: MetadataServiceDeps - ) -> None: - """ - Test metadata lock check when lock already exists. - """ - # Arrange: Setup mocks to simulate existing lock - mock_external_service_dependencies["redis_client"].get.return_value = "1" # Lock exists - - dataset_id = "test-dataset-id" - - # Act & Assert: Verify proper error handling - with pytest.raises( - ValueError, match="Another knowledge base metadata operation is running, please wait a moment." - ): - MetadataService.knowledge_base_metadata_lock_check(dataset_id, None) - - def test_knowledge_base_metadata_lock_check_document_lock_exists( - self, db_session_with_containers: Session, mock_external_service_dependencies: MetadataServiceDeps - ) -> None: - """ - Test metadata lock check when document lock already exists. - """ - # Arrange: Setup mocks to simulate existing lock - mock_external_service_dependencies["redis_client"].get.return_value = "1" # Lock exists - - document_id = "test-document-id" - - # Act & Assert: Verify proper error handling - with pytest.raises(ValueError, match="Another document metadata operation is running, please wait a moment."): - MetadataService.knowledge_base_metadata_lock_check(None, document_id) - - def test_get_dataset_metadatas_success( - self, db_session_with_containers: Session, mock_external_service_dependencies: MetadataServiceDeps - ) -> None: + def test_get_dataset_metadatas_success(self, db_session_with_containers: Session) -> None: """ Test successful retrieval of dataset metadata information. """ - # Arrange: Create test data - account, tenant = self._create_test_account_and_tenant( - db_session_with_containers, mock_external_service_dependencies - ) - dataset = self._create_test_dataset( - db_session_with_containers, mock_external_service_dependencies, account, tenant - ) - - # Create metadata + account, tenant = self._create_test_account_and_tenant(db_session_with_containers) + dataset = self._create_test_dataset(db_session_with_containers, account, tenant) metadata_args = MetadataArgs(type="string", name="test_metadata") - metadata = MetadataService.create_metadata( - dataset.id, metadata_args, account, tenant.id, session=db_session_with_containers + metadata = _metadata_service(db_session_with_containers).create_metadata( + DatasetRef(tenant.id, dataset.id), metadata_args, actor_id=account.id ) - - # Create document and metadata binding - document = self._create_test_document( - db_session_with_containers, mock_external_service_dependencies, dataset, account - ) - + db_session_with_containers.expire_all() + document = self._create_test_document(db_session_with_containers, dataset, account) binding = DatasetMetadataBinding( tenant_id=tenant.id, dataset_id=dataset.id, @@ -1161,95 +633,61 @@ class TestMetadataService: document_id=document.id, created_by=account.id, ) - db_session_with_containers.add(binding) db_session_with_containers.commit() - - # Act: Execute the method under test - result = MetadataService.get_dataset_metadatas(dataset, session=db_session_with_containers) - - # Assert: Verify the expected outcomes + result = _metadata_service(db_session_with_containers).get_dataset_metadatas( + DatasetRef(dataset.tenant_id, dataset.id) + ) + db_session_with_containers.expire_all() assert result is not None assert "doc_metadata" in result assert "built_in_field_enabled" in result - - # Verify metadata information doc_metadata = result["doc_metadata"] assert len(doc_metadata) == 1 assert doc_metadata[0]["id"] == metadata.id assert doc_metadata[0]["name"] == metadata.name assert doc_metadata[0]["type"] == metadata.type - assert doc_metadata[0]["count"] == 1 # One document bound to this metadata - - # Verify built-in field status + assert doc_metadata[0]["count"] == 1 assert result["built_in_field_enabled"] is False - def test_get_dataset_metadatas_with_built_in_fields_enabled( - self, db_session_with_containers: Session, mock_external_service_dependencies: MetadataServiceDeps - ) -> None: + def test_get_dataset_metadatas_with_built_in_fields_enabled(self, db_session_with_containers: Session) -> None: """ Test retrieval of dataset metadata when built-in fields are enabled. """ - # Arrange: Create test data - account, tenant = self._create_test_account_and_tenant( - db_session_with_containers, mock_external_service_dependencies - ) - dataset = self._create_test_dataset( - db_session_with_containers, mock_external_service_dependencies, account, tenant - ) - - # Enable built-in fields + account, tenant = self._create_test_account_and_tenant(db_session_with_containers) + dataset = self._create_test_dataset(db_session_with_containers, account, tenant) dataset.built_in_field_enabled = True - db_session_with_containers.add(dataset) db_session_with_containers.commit() - - # Create metadata metadata_args = MetadataArgs(type="string", name="test_metadata") - metadata = MetadataService.create_metadata( - dataset.id, metadata_args, account, tenant.id, session=db_session_with_containers + metadata = _metadata_service(db_session_with_containers).create_metadata( + DatasetRef(tenant.id, dataset.id), metadata_args, actor_id=account.id ) - - # Act: Execute the method under test - result = MetadataService.get_dataset_metadatas(dataset, session=db_session_with_containers) - - # Assert: Verify the expected outcomes + db_session_with_containers.expire_all() + result = _metadata_service(db_session_with_containers).get_dataset_metadatas( + DatasetRef(dataset.tenant_id, dataset.id) + ) + db_session_with_containers.expire_all() assert result is not None assert "doc_metadata" in result assert "built_in_field_enabled" in result - - # Verify metadata information doc_metadata = result["doc_metadata"] - assert len(doc_metadata) == 1 # Only custom metadata, built-in fields are not included in this list - - # Verify built-in field status + assert len(doc_metadata) == 1 assert result["built_in_field_enabled"] is True - def test_get_dataset_metadatas_no_metadata( - self, db_session_with_containers: Session, mock_external_service_dependencies: MetadataServiceDeps - ) -> None: + def test_get_dataset_metadatas_no_metadata(self, db_session_with_containers: Session) -> None: """ Test retrieval of dataset metadata when no metadata exists. """ - # Arrange: Create test data - account, tenant = self._create_test_account_and_tenant( - db_session_with_containers, mock_external_service_dependencies + account, tenant = self._create_test_account_and_tenant(db_session_with_containers) + dataset = self._create_test_dataset(db_session_with_containers, account, tenant) + result = _metadata_service(db_session_with_containers).get_dataset_metadatas( + DatasetRef(dataset.tenant_id, dataset.id) ) - dataset = self._create_test_dataset( - db_session_with_containers, mock_external_service_dependencies, account, tenant - ) - - # Act: Execute the method under test - result = MetadataService.get_dataset_metadatas(dataset, session=db_session_with_containers) - - # Assert: Verify the expected outcomes + db_session_with_containers.expire_all() assert result is not None assert "doc_metadata" in result assert "built_in_field_enabled" in result - - # Verify metadata information doc_metadata = result["doc_metadata"] - assert len(doc_metadata) == 0 # No metadata exists - - # Verify built-in field status + assert len(doc_metadata) == 0 assert result["built_in_field_enabled"] is False diff --git a/api/tests/test_containers_integration_tests/services/test_model_provider_service.py b/api/tests/test_containers_integration_tests/services/test_model_provider_service.py index f250488bbd7..356bf77dbd2 100644 --- a/api/tests/test_containers_integration_tests/services/test_model_provider_service.py +++ b/api/tests/test_containers_integration_tests/services/test_model_provider_service.py @@ -8,7 +8,7 @@ from core.entities.model_entities import ModelStatus from graphon.model_runtime.entities.model_entities import FetchFrom, ModelType from models import Account, Tenant, TenantAccountJoin, TenantAccountRole from models.provider import Provider, ProviderModel, ProviderModelSetting, ProviderType -from services.model_provider_service import ModelProviderService +from services.model_provider.service import ModelProviderService class TestModelProviderService: @@ -19,10 +19,10 @@ class TestModelProviderService: """Mock setup for external service dependencies.""" with ( patch( - "services.model_provider_service.create_plugin_provider_manager", autospec=True + "services.model_provider.service.create_plugin_provider_manager", autospec=True ) as mock_provider_manager, patch( - "services.model_provider_service.create_plugin_model_provider_factory", autospec=True + "services.model_provider.service.create_plugin_model_provider_factory", autospec=True ) as mock_model_provider_factory, ): # Setup default mock returns diff --git a/api/tests/test_containers_integration_tests/tasks/test_add_document_to_index_task.py b/api/tests/test_containers_integration_tests/tasks/test_add_document_to_index_task.py index 94fd7602f5e..30857e22c0e 100644 --- a/api/tests/test_containers_integration_tests/tasks/test_add_document_to_index_task.py +++ b/api/tests/test_containers_integration_tests/tasks/test_add_document_to_index_task.py @@ -8,8 +8,8 @@ from sqlalchemy.orm import Session from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType from extensions.ext_redis import redis_client from models import Account, Tenant, TenantAccountJoin, TenantAccountRole -from models.dataset import Dataset, DatasetAutoDisableLog, Document, DocumentSegment -from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus, SegmentStatus +from models.dataset import ChildChunk, Dataset, DatasetAutoDisableLog, DatasetProcessRule, Document, DocumentSegment +from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus, ProcessRuleMode, SegmentStatus from tasks.add_document_to_index_task import add_document_to_index_task @@ -377,56 +377,75 @@ class TestAddDocumentToIndexTask: # Refresh dataset to ensure doc_form property reflects the updated document db_session_with_containers.refresh(dataset) - # Create segments with mock child chunks + # Persist each segment and its own child chunks segments = self._create_test_segments(db_session_with_containers, document, dataset) # Set up Redis cache key indexing_cache_key = f"document_{document.id}_indexing" redis_client.set(indexing_cache_key, "processing", ex=300) - # Mock the get_child_chunks method for each segment - with patch.object(DocumentSegment, "get_child_chunks", autospec=True) as mock_get_child_chunks: - # Setup mock to return child chunks for each segment - mock_child_chunks = [] - for i in range(2): # Each segment has 2 child chunks - mock_child = MagicMock() - mock_child.content = f"child_content_{i}" - mock_child.index_node_id = f"child_node_{i}" - mock_child.index_node_hash = f"child_hash_{i}" - mock_child_chunks.append(mock_child) - - mock_get_child_chunks.return_value = mock_child_chunks - - # Act: Execute the task - add_document_to_index_task(document.id) - - # Assert: Verify parent-child index processing - mock_external_service_dependencies["index_processor_factory"].assert_called_once_with( - IndexStructureType.PARENT_CHILD_INDEX + rule = DatasetProcessRule( + dataset_id=dataset.id, + mode=ProcessRuleMode.HIERARCHICAL, + rules='{"parent_mode": "paragraph"}', + created_by=document.created_by, + ) + db_session_with_containers.add(rule) + db_session_with_containers.flush() + document.dataset_process_rule_id = rule.id + for segment in segments: + db_session_with_containers.add_all( + [ + ChildChunk( + tenant_id=dataset.tenant_id, + dataset_id=dataset.id, + document_id=document.id, + segment_id=segment.id, + position=i, + content=f"child_content_{i}", + word_count=1, + created_by=document.created_by, + index_node_id=f"{segment.index_node_id}-child-{i}", + index_node_hash=f"{segment.index_node_id}-hash-{i}", + ) + for i in range(2) + ] ) - mock_external_service_dependencies["index_processor"].load.assert_called_once() + db_session_with_containers.commit() - # Verify the load method was called with correct parameters - call_args = mock_external_service_dependencies["index_processor"].load.call_args - assert call_args is not None - documents = call_args[0][1] # Second argument should be documents list - assert len(documents) == 3 # 3 segments + # Act: Execute the task + add_document_to_index_task(document.id) - # Verify each document has children - for doc in documents: - assert hasattr(doc, "children") - assert len(doc.children) == 2 # Each document has 2 children + # Assert: Verify parent-child index processing + mock_external_service_dependencies["index_processor_factory"].assert_called_once_with( + IndexStructureType.PARENT_CHILD_INDEX + ) + mock_external_service_dependencies["index_processor"].load.assert_called_once() - # Verify database state changes - db_session_with_containers.refresh(document) - for segment in segments: - db_session_with_containers.refresh(segment) - assert segment.enabled is True - assert segment.disabled_at is None - assert segment.disabled_by is None + # Verify the load method was called with correct parameters + call_args = mock_external_service_dependencies["index_processor"].load.call_args + assert call_args is not None + documents = call_args[0][1] # Second argument should be documents list + assert len(documents) == 3 # 3 segments - # Verify redis cache was cleared - assert redis_client.exists(indexing_cache_key) == 0 + # Each parent receives only its own persisted children, in position order. + for doc in documents: + assert doc.children is not None + assert [child.page_content for child in doc.children] == ["child_content_0", "child_content_1"] + assert [child.metadata["doc_id"] for child in doc.children] == [ + f"{doc.metadata['doc_id']}-child-{i}" for i in range(2) + ] + + # Verify database state changes + db_session_with_containers.refresh(document) + for segment in segments: + db_session_with_containers.refresh(segment) + assert segment.enabled is True + assert segment.disabled_at is None + assert segment.disabled_by is None + + # Verify redis cache was cleared + assert redis_client.exists(indexing_cache_key) == 0 def test_add_document_to_index_with_already_enabled_segments( self, db_session_with_containers: Session, mock_external_service_dependencies diff --git a/api/tests/test_containers_integration_tests/tasks/test_batch_clean_document_task.py b/api/tests/test_containers_integration_tests/tasks/test_batch_clean_document_task.py index adac71a2edf..574ba58b742 100644 --- a/api/tests/test_containers_integration_tests/tasks/test_batch_clean_document_task.py +++ b/api/tests/test_containers_integration_tests/tasks/test_batch_clean_document_task.py @@ -7,7 +7,7 @@ and storage containers to ensure proper cleanup of documents, segments, and file import json import uuid -from unittest.mock import Mock, patch +from unittest.mock import patch import pytest from faker import Faker @@ -21,6 +21,7 @@ from models import Account, Tenant, TenantAccountJoin, TenantAccountRole from models.dataset import Dataset, Document, DocumentSegment from models.enums import CreatorUserRole, DataSourceType, DocumentCreatedFrom, IndexingStatus, SegmentStatus from models.model import UploadFile +from repositories.knowledge.dataset_read_repository import get_dataset_doc_form from tasks.batch_clean_document_task import batch_clean_document_task @@ -31,25 +32,20 @@ class TestBatchCleanDocumentTask: def mock_external_service_dependencies(self): """Mock setup for external service dependencies.""" with ( - patch("extensions.ext_storage.storage") as mock_storage, - patch("core.rag.index_processor.index_processor_factory.IndexProcessorFactory") as mock_index_factory, - patch("core.tools.utils.web_reader_tool.get_image_upload_file_ids") as mock_get_image_ids, + patch("tasks.batch_clean_document_task.storage") as mock_storage, + patch("services.knowledge.indexing.adapters.cleanup.Vector", autospec=True) as mock_vector, + patch("tasks.batch_clean_document_task.get_image_upload_file_ids") as mock_get_image_ids, + patch("tasks.batch_clean_document_task.schedule_billing_vector_space_refresh"), ): # Setup default mock returns mock_storage.delete.return_value = None - # Mock index processor - mock_index_processor = Mock() - mock_index_processor.clean.return_value = None - mock_index_factory.return_value.init_index_processor.return_value = mock_index_processor - # Mock image file ID extraction mock_get_image_ids.return_value = [] yield { "storage": mock_storage, - "index_factory": mock_index_factory, - "index_processor": mock_index_processor, + "vector": mock_vector.return_value, "get_image_ids": mock_get_image_ids, } @@ -121,6 +117,7 @@ class TestBatchCleanDocumentTask: created_by=account.id, embedding_model="text-embedding-ada-002", embedding_model_provider="openai", + indexing_technique="high_quality", ) db_session_with_containers.add(dataset) @@ -251,18 +248,18 @@ class TestBatchCleanDocumentTask: # Store original IDs for verification document_id = document.id segment_id = segment.id + index_node_id = segment.index_node_id file_id = upload_file.id # Execute the task batch_clean_document_task( document_ids=[document_id], dataset_id=dataset.id, - doc_form=dataset.get_doc_form(session=db_session_with_containers), + doc_form=get_dataset_doc_form(dataset, session=db_session_with_containers), file_ids=[file_id], ) - # Verify that the task completed successfully - # The task should have processed the segment and cleaned up the database + mock_external_service_dependencies["vector"].delete_by_ids.assert_called_once_with([index_node_id]) # Verify database cleanup db_session_with_containers.commit() # Ensure all changes are committed @@ -316,7 +313,7 @@ class TestBatchCleanDocumentTask: batch_clean_document_task( document_ids=[document_id], dataset_id=dataset.id, - doc_form=dataset.get_doc_form(session=db_session_with_containers), + doc_form=get_dataset_doc_form(dataset, session=db_session_with_containers), file_ids=[], ) @@ -359,7 +356,7 @@ class TestBatchCleanDocumentTask: batch_clean_document_task( document_ids=[document_id], dataset_id=dataset.id, - doc_form=dataset.get_doc_form(session=db_session_with_containers), + doc_form=get_dataset_doc_form(dataset, session=db_session_with_containers), file_ids=[file_id], ) @@ -411,7 +408,7 @@ class TestBatchCleanDocumentTask: ) # Verify that no index processing occurred - mock_external_service_dependencies["index_processor"].clean.assert_not_called() + mock_external_service_dependencies["vector"].delete_by_ids.assert_not_called() # Verify that no storage operations occurred mock_external_service_dependencies["storage"].delete.assert_not_called() @@ -457,7 +454,7 @@ class TestBatchCleanDocumentTask: batch_clean_document_task( document_ids=[document_id], dataset_id=dataset.id, - doc_form=dataset.get_doc_form(session=db_session_with_containers), + doc_form=get_dataset_doc_form(dataset, session=db_session_with_containers), file_ids=[file_id], ) @@ -518,7 +515,7 @@ class TestBatchCleanDocumentTask: batch_clean_document_task( document_ids=document_ids, dataset_id=dataset.id, - doc_form=dataset.get_doc_form(session=db_session_with_containers), + doc_form=get_dataset_doc_form(dataset, session=db_session_with_containers), file_ids=file_ids, ) @@ -658,7 +655,7 @@ class TestBatchCleanDocumentTask: batch_clean_document_task( document_ids=document_ids, dataset_id=dataset.id, - doc_form=dataset.get_doc_form(session=db_session_with_containers), + doc_form=get_dataset_doc_form(dataset, session=db_session_with_containers), file_ids=file_ids, ) @@ -754,7 +751,7 @@ class TestBatchCleanDocumentTask: batch_clean_document_task( document_ids=[document_id], dataset_id=dataset.id, - doc_form=dataset.get_doc_form(session=db_session_with_containers), + doc_form=get_dataset_doc_form(dataset, session=db_session_with_containers), file_ids=[file_id], ) diff --git a/api/tests/test_containers_integration_tests/tasks/test_batch_create_segment_to_index_task.py b/api/tests/test_containers_integration_tests/tasks/test_batch_create_segment_to_index_task.py index b0f57dba084..a99a7301b12 100644 --- a/api/tests/test_containers_integration_tests/tasks/test_batch_create_segment_to_index_task.py +++ b/api/tests/test_containers_integration_tests/tasks/test_batch_create_segment_to_index_task.py @@ -59,7 +59,7 @@ class TestBatchCreateSegmentToIndexTask: "tasks.batch_create_segment_to_index_task.ModelManager.for_tenant", autospec=True, ) as mock_model_manager, - patch("tasks.batch_create_segment_to_index_task.VectorService", autospec=True) as mock_vector_service, + patch("extensions.ext_application_services.application_services") as mock_services, ): # Setup default mock returns mock_storage.download.return_value = None @@ -71,13 +71,13 @@ class TestBatchCreateSegmentToIndexTask: mock_model_manager_instance.get_model_instance.return_value = mock_embedding_model mock_model_manager.return_value = mock_model_manager_instance - # Mock vector service - mock_vector_service.create_segments_vector.return_value = None + mutations = mock_services.return_value.knowledge.segments.mutations + mutations.index_segments.return_value = None yield { "storage": mock_storage, "model_manager": mock_model_manager, - "vector_service": mock_vector_service, + "segment_mutations": mutations, "embedding_model": mock_embedding_model, } @@ -314,8 +314,8 @@ class TestBatchCreateSegmentToIndexTask: assert document.word_count > 0 # Verify vector service was called - mock_vector_service = mock_external_service_dependencies["vector_service"] - mock_vector_service.create_segments_vector.assert_called_once() + mutations = mock_external_service_dependencies["segment_mutations"] + mutations.index_segments.assert_called_once() # Check Redis cache was set from extensions.ext_redis import redis_client @@ -750,8 +750,8 @@ class TestBatchCreateSegmentToIndexTask: assert document.word_count > 0 # Verify vector service was called - mock_vector_service = mock_external_service_dependencies["vector_service"] - mock_vector_service.create_segments_vector.assert_called_once() + mutations = mock_external_service_dependencies["segment_mutations"] + mutations.index_segments.assert_called_once() # Check Redis cache was set from extensions.ext_redis import redis_client diff --git a/api/tests/test_containers_integration_tests/tasks/test_clean_dataset_task.py b/api/tests/test_containers_integration_tests/tasks/test_clean_dataset_task.py index e801597e434..ee47693ee66 100644 --- a/api/tests/test_containers_integration_tests/tasks/test_clean_dataset_task.py +++ b/api/tests/test_containers_integration_tests/tasks/test_clean_dataset_task.py @@ -41,6 +41,7 @@ from models.enums import ( SegmentStatus, ) from models.model import UploadFile +from repositories.knowledge.dataset_read_repository import get_dataset_doc_form from tasks.clean_dataset_task import clean_dataset_task @@ -295,7 +296,7 @@ class TestCleanDatasetTask: indexing_technique=dataset.indexing_technique, index_struct=dataset.index_struct, collection_binding_id=dataset.collection_binding_id, - doc_form=dataset.get_doc_form(session=db_session_with_containers), + doc_form=get_dataset_doc_form(dataset, session=db_session_with_containers), ) # Verify results @@ -419,7 +420,7 @@ class TestCleanDatasetTask: indexing_technique=dataset.indexing_technique, index_struct=dataset.index_struct, collection_binding_id=dataset.collection_binding_id, - doc_form=dataset.get_doc_form(session=db_session_with_containers), + doc_form=get_dataset_doc_form(dataset, session=db_session_with_containers), ) # Verify results @@ -552,7 +553,7 @@ class TestCleanDatasetTask: indexing_technique=dataset.indexing_technique, index_struct=dataset.index_struct, collection_binding_id=dataset.collection_binding_id, - doc_form=dataset.get_doc_form(session=db_session_with_containers), + doc_form=get_dataset_doc_form(dataset, session=db_session_with_containers), ) # Verify results - even with vector cleanup failure, documents and segments should be deleted @@ -638,7 +639,7 @@ class TestCleanDatasetTask: indexing_technique=dataset.indexing_technique, index_struct=dataset.index_struct, collection_binding_id=dataset.collection_binding_id, - doc_form=dataset.get_doc_form(session=db_session_with_containers), + doc_form=get_dataset_doc_form(dataset, session=db_session_with_containers), ) # Verify results @@ -751,7 +752,7 @@ class TestCleanDatasetTask: indexing_technique=dataset.indexing_technique, index_struct=dataset.index_struct, collection_binding_id=dataset.collection_binding_id, - doc_form=dataset.get_doc_form(session=db_session_with_containers), + doc_form=get_dataset_doc_form(dataset, session=db_session_with_containers), ) end_time = time.time() @@ -845,7 +846,7 @@ class TestCleanDatasetTask: indexing_technique=dataset.indexing_technique, index_struct=dataset.index_struct, collection_binding_id=dataset.collection_binding_id, - doc_form=dataset.get_doc_form(session=db_session_with_containers), + doc_form=get_dataset_doc_form(dataset, session=db_session_with_containers), ) # Verify results @@ -999,7 +1000,7 @@ class TestCleanDatasetTask: indexing_technique=dataset.indexing_technique, index_struct=dataset.index_struct, collection_binding_id=dataset.collection_binding_id, - doc_form=dataset.get_doc_form(session=db_session_with_containers), + doc_form=get_dataset_doc_form(dataset, session=db_session_with_containers), ) # Verify results diff --git a/api/tests/test_containers_integration_tests/tasks/test_clean_notion_document_task.py b/api/tests/test_containers_integration_tests/tasks/test_clean_notion_document_task.py index b2007608e42..d9cdd9aef79 100644 --- a/api/tests/test_containers_integration_tests/tasks/test_clean_notion_document_task.py +++ b/api/tests/test_containers_integration_tests/tasks/test_clean_notion_document_task.py @@ -7,7 +7,7 @@ containers to ensure proper cleanup of Notion documents, segments, and vector in import json import uuid -from unittest.mock import Mock, patch +from unittest.mock import patch import pytest from faker import Faker @@ -47,31 +47,13 @@ class TestCleanNotionDocumentTask: } @pytest.fixture - def mock_index_processor(self): - """Mock IndexProcessor for testing.""" - mock_processor = Mock() - mock_processor.clean = Mock() - return mock_processor - - @pytest.fixture - def mock_index_processor_factory(self, mock_index_processor): - """Mock IndexProcessorFactory for testing.""" - # Mock the actual IndexProcessorFactory class - with patch("tasks.clean_notion_document_task.IndexProcessorFactory") as mock_factory: - # Create a mock instance that will be returned when IndexProcessorFactory() is called - mock_instance = Mock() - mock_instance.init_index_processor.return_value = mock_index_processor - - # Set the mock_factory to return our mock_instance when called - mock_factory.return_value = mock_instance - - # Ensure the mock_index_processor has the clean method properly set - mock_index_processor.clean = Mock() - - yield mock_factory + def mock_vector(self): + """Keep database cleanup real while isolating the external vector backend.""" + with patch("services.knowledge.indexing.adapters.cleanup.Vector", autospec=True) as vector: + yield vector.return_value def test_clean_notion_document_task_success( - self, db_session_with_containers: Session, mock_index_processor_factory, mock_external_service_dependencies + self, db_session_with_containers: Session, mock_vector, mock_external_service_dependencies ): """ Test successful cleanup of Notion documents with proper database operations. @@ -79,7 +61,7 @@ class TestCleanNotionDocumentTask: This test verifies that the task correctly: 1. Deletes Document records from database 2. Deletes DocumentSegment records from database - 3. Calls index processor to clean vector and keyword indices + 3. Deletes the documents' vector entries 4. Commits all changes to database """ fake = Faker() @@ -97,6 +79,7 @@ class TestCleanNotionDocumentTask: # Create dataset dataset = Dataset( + indexing_technique="high_quality", id=str(uuid.uuid4()), tenant_id=tenant.id, name=fake.company(), @@ -164,19 +147,19 @@ class TestCleanNotionDocumentTask: # Verify segments are deleted assert _count_segments(db_session_with_containers, DocumentSegment.document_id.in_(document_ids)) == 0 - # Verify index processor was called - mock_processor = mock_index_processor_factory.return_value.init_index_processor.return_value - mock_processor.clean.assert_called_once() + assert _count_documents(db_session_with_containers, Document.id.in_(document_ids)) == 0 + mock_vector.delete_by_ids.assert_called_once() + assert set(mock_vector.delete_by_ids.call_args.args[0]) == set(index_node_ids) # This test successfully verifies: # 1. Document records are properly deleted from the database # 2. DocumentSegment records are properly deleted from the database - # 3. The index processor's clean method is called + # 3. The external vector deletion receives the selected node IDs # 4. Database transaction handling works correctly # 5. The task completes without errors def test_clean_notion_document_task_dataset_not_found( - self, db_session_with_containers: Session, mock_index_processor_factory, mock_external_service_dependencies + self, db_session_with_containers: Session, mock_vector, mock_external_service_dependencies ): """ Test cleanup task behavior when dataset is not found. @@ -192,11 +175,10 @@ class TestCleanNotionDocumentTask: with pytest.raises(Exception, match="Document has no dataset"): clean_notion_document_task(document_ids, non_existent_dataset_id) - # Verify that the index processor factory was not used - mock_index_processor_factory.return_value.init_index_processor.assert_not_called() + mock_vector.delete_by_ids.assert_not_called() def test_clean_notion_document_task_empty_document_list( - self, db_session_with_containers: Session, mock_index_processor_factory, mock_external_service_dependencies + self, db_session_with_containers: Session, mock_vector, mock_external_service_dependencies ): """ Test cleanup task behavior with empty document list. @@ -219,6 +201,7 @@ class TestCleanNotionDocumentTask: # Create dataset dataset = Dataset( + indexing_technique="high_quality", id=str(uuid.uuid4()), tenant_id=tenant.id, name=fake.company(), @@ -232,23 +215,16 @@ class TestCleanNotionDocumentTask: # Execute cleanup task with empty document list clean_notion_document_task([], dataset.id) - # Verify that the index processor was called once with empty node list - mock_processor = mock_index_processor_factory.return_value.init_index_processor.return_value - assert mock_processor.clean.call_count == 1 - args, kwargs = mock_processor.clean.call_args - # args: (dataset, total_index_node_ids) - assert isinstance(args[0], Dataset) - assert args[1] == [] - assert kwargs["session"] is not None + # An empty document selection must never trigger a dataset-wide vector deletion. + mock_vector.delete_by_ids.assert_not_called() def test_clean_notion_document_task_with_different_index_types( - self, db_session_with_containers: Session, mock_index_processor_factory, mock_external_service_dependencies + self, db_session_with_containers: Session, mock_vector, mock_external_service_dependencies ): """ Test cleanup task with different dataset index types. - This test verifies that the task correctly initializes different types - of index processors based on the dataset's doc_form configuration. + This test verifies that cleanup uses the dataset's document form. """ fake = Faker() @@ -270,6 +246,7 @@ class TestCleanNotionDocumentTask: for index_type in index_types: # Create dataset (doc_form will be set via document creation) dataset = Dataset( + indexing_technique="high_quality", id=str(uuid.uuid4()), tenant_id=tenant.id, name=f"{fake.company()}_{index_type}", @@ -320,17 +297,15 @@ class TestCleanNotionDocumentTask: # Execute cleanup task clean_notion_document_task([document.id], dataset.id) - # Note: This test successfully verifies cleanup with different document types. - # The task properly handles various index types and document configurations. - # Verify segments are deleted assert _count_segments(db_session_with_containers, DocumentSegment.document_id == document.id) == 0 + mock_vector.delete_by_ids.assert_called_once_with(["test_node"]) # Reset mock for next iteration - mock_index_processor_factory.reset_mock() + mock_vector.reset_mock() def test_clean_notion_document_task_with_segments_no_index_node_ids( - self, db_session_with_containers: Session, mock_index_processor_factory, mock_external_service_dependencies + self, db_session_with_containers: Session, mock_vector, mock_external_service_dependencies ): """ Test cleanup task with segments that have no index_node_ids. @@ -353,6 +328,7 @@ class TestCleanNotionDocumentTask: # Create dataset dataset = Dataset( + indexing_technique="high_quality", id=str(uuid.uuid4()), tenant_id=tenant.id, name=fake.company(), @@ -409,11 +385,10 @@ class TestCleanNotionDocumentTask: # Verify segments are deleted assert _count_segments(db_session_with_containers, DocumentSegment.document_id == document.id) == 0 - # Note: This test successfully verifies that segments without index_node_ids - # are properly deleted from the database. + mock_vector.delete_by_ids.assert_not_called() def test_clean_notion_document_task_partial_document_cleanup( - self, db_session_with_containers: Session, mock_index_processor_factory, mock_external_service_dependencies + self, db_session_with_containers: Session, mock_vector, mock_external_service_dependencies ): """ Test cleanup task with partial document cleanup scenario. @@ -436,6 +411,7 @@ class TestCleanNotionDocumentTask: # Create dataset dataset = Dataset( + indexing_technique="high_quality", id=str(uuid.uuid4()), tenant_id=tenant.id, name=fake.company(), @@ -515,7 +491,7 @@ class TestCleanNotionDocumentTask: # The database operations work correctly, isolating only the specified documents. def test_clean_notion_document_task_with_mixed_segment_statuses( - self, db_session_with_containers: Session, mock_index_processor_factory, mock_external_service_dependencies + self, db_session_with_containers: Session, mock_vector, mock_external_service_dependencies ): """ Test cleanup task with segments in different statuses. @@ -538,6 +514,7 @@ class TestCleanNotionDocumentTask: # Create dataset dataset = Dataset( + indexing_technique="high_quality", id=str(uuid.uuid4()), tenant_id=tenant.id, name=fake.company(), @@ -601,21 +578,21 @@ class TestCleanNotionDocumentTask: # Verify all segments are deleted regardless of status assert _count_segments(db_session_with_containers, DocumentSegment.document_id == document.id) == 0 - # Note: This test successfully verifies database operations. - # IndexProcessor verification would require more sophisticated mocking. + mock_vector.delete_by_ids.assert_called_once() + assert set(mock_vector.delete_by_ids.call_args.args[0]) == set(index_node_ids) - def test_clean_notion_document_task_continues_when_index_processor_fails( - self, db_session_with_containers: Session, mock_index_processor_factory, mock_external_service_dependencies + def test_clean_notion_document_task_continues_when_vector_cleanup_fails( + self, db_session_with_containers: Session, mock_vector, mock_external_service_dependencies ): """ - Index processor failure (e.g. transient billing API error propagated via + Vector cleanup failure (e.g. transient billing API error propagated via ``FeatureService`` when ``Vector(dataset)`` lazily resolves the embedding model) must NOT abort the cleanup task. The Document rows have already been hard-deleted in the first session block before vector cleanup runs, so any uncaught exception escaping the task would strand ``DocumentSegment`` rows in PG with no parent ``Document``. - Contract: the task swallows the index_processor exception, logs it, and + Contract: the task swallows the vector cleanup exception, logs it, and proceeds to delete the segments — leaving PG consistent. (Vector orphans, if any, can be reaped later by an offline scanner.) @@ -639,6 +616,7 @@ class TestCleanNotionDocumentTask: # Create dataset dataset = Dataset( + indexing_technique="high_quality", id=str(uuid.uuid4()), tenant_id=tenant.id, name=fake.company(), @@ -685,23 +663,22 @@ class TestCleanNotionDocumentTask: db_session_with_containers.add(segment) db_session_with_containers.commit() - # Simulate the production failure mode: index_processor.clean() raises a + # Simulate the production failure mode: vector deletion raises a # ValueError mirroring ``BillingService._send_request`` returning non-200. - mock_index_processor = mock_index_processor_factory.return_value.init_index_processor.return_value - mock_index_processor.clean.side_effect = ValueError( + mock_vector.delete_by_ids.side_effect = ValueError( "Unable to retrieve billing information. Please try again later or contact support." ) - # Execute cleanup task — must NOT raise even though clean() raises. + # Execute cleanup task — must NOT raise even though vector deletion raises. # Before the safety-net wrapper this would have re-raised the ValueError, # aborting the task and leaving DocumentSegment stranded in PG. clean_notion_document_task([document.id], dataset.id) # Vector cleanup was attempted exactly once. - mock_index_processor.clean.assert_called_once() + mock_vector.delete_by_ids.assert_called_once_with(["test_node"]) - # The crucial assertion: despite the index processor failure, the - # final session block (line 51-52, ``DELETE FROM document_segments``) + # The crucial assertion: despite the vector cleanup failure, the + # final session block (``DELETE FROM document_segments``) # still ran and committed. This is what the wrapper buys us — without # it the production incident left tens of thousands of orphan segments # per affected tenant. Aligns with the assertion shape used by the @@ -709,7 +686,7 @@ class TestCleanNotionDocumentTask: assert _count_segments(db_session_with_containers, DocumentSegment.document_id == document.id) == 0 def test_clean_notion_document_task_with_large_number_of_documents( - self, db_session_with_containers: Session, mock_index_processor_factory, mock_external_service_dependencies + self, db_session_with_containers: Session, mock_vector, mock_external_service_dependencies ): """ Test cleanup task with a large number of documents and segments. @@ -732,6 +709,7 @@ class TestCleanNotionDocumentTask: # Create dataset dataset = Dataset( + indexing_technique="high_quality", id=str(uuid.uuid4()), tenant_id=tenant.id, name=fake.company(), @@ -808,7 +786,7 @@ class TestCleanNotionDocumentTask: # The database efficiently handles large-scale deletions. def test_clean_notion_document_task_with_documents_from_different_tenants( - self, db_session_with_containers: Session, mock_index_processor_factory, mock_external_service_dependencies + self, db_session_with_containers: Session, mock_vector, mock_external_service_dependencies ): """ Test cleanup task with documents from different tenants. @@ -840,6 +818,7 @@ class TestCleanNotionDocumentTask: # Create dataset for each tenant dataset = Dataset( + indexing_technique="high_quality", id=str(uuid.uuid4()), tenant_id=tenant.id, name=f"{fake.company()}_{i}", @@ -924,7 +903,7 @@ class TestCleanNotionDocumentTask: # Only documents from the target dataset are affected, maintaining tenant separation. def test_clean_notion_document_task_with_documents_in_different_states( - self, db_session_with_containers: Session, mock_index_processor_factory, mock_external_service_dependencies + self, db_session_with_containers: Session, mock_vector, mock_external_service_dependencies ): """ Test cleanup task with documents in different indexing states. @@ -947,6 +926,7 @@ class TestCleanNotionDocumentTask: # Create dataset dataset = Dataset( + indexing_technique="high_quality", id=str(uuid.uuid4()), tenant_id=tenant.id, name=fake.company(), @@ -1030,7 +1010,7 @@ class TestCleanNotionDocumentTask: # All documents are deleted regardless of their indexing status. def test_clean_notion_document_task_with_documents_having_metadata( - self, db_session_with_containers: Session, mock_index_processor_factory, mock_external_service_dependencies + self, db_session_with_containers: Session, mock_vector, mock_external_service_dependencies ): """ Test cleanup task with documents that have rich metadata. @@ -1053,6 +1033,7 @@ class TestCleanNotionDocumentTask: # Create dataset with built-in fields enabled dataset = Dataset( + indexing_technique="high_quality", id=str(uuid.uuid4()), tenant_id=tenant.id, name=fake.company(), diff --git a/api/tests/test_containers_integration_tests/tasks/test_create_segment_to_index_task.py b/api/tests/test_containers_integration_tests/tasks/test_create_segment_to_index_task.py index 40bffa766b1..aa4257bdb05 100644 --- a/api/tests/test_containers_integration_tests/tasks/test_create_segment_to_index_task.py +++ b/api/tests/test_containers_integration_tests/tasks/test_create_segment_to_index_task.py @@ -19,6 +19,7 @@ from extensions.ext_redis import redis_client from models import Account, AccountStatus, Tenant, TenantAccountJoin, TenantAccountRole, TenantStatus from models.dataset import Dataset, Document, DocumentSegment from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus, SegmentStatus +from repositories.knowledge.dataset_read_repository import get_dataset_doc_form from tasks.create_segment_to_index_task import create_segment_to_index_task @@ -227,7 +228,7 @@ class TestCreateSegmentToIndexTask: # Verify index processor was called mock_external_service_dependencies["index_processor_factory"].assert_called_once_with( - dataset.get_doc_form(session=db_session_with_containers) + get_dataset_doc_form(dataset, session=db_session_with_containers) ) mock_external_service_dependencies["index_processor"].load.assert_called_once() @@ -555,7 +556,7 @@ class TestCreateSegmentToIndexTask: # Verify index processor was called mock_external_service_dependencies["index_processor_factory"].assert_called_once_with( - dataset.get_doc_form(session=db_session_with_containers) + get_dataset_doc_form(dataset, session=db_session_with_containers) ) mock_external_service_dependencies["index_processor"].load.assert_called_once() @@ -988,7 +989,7 @@ class TestCreateSegmentToIndexTask: # Verify index processor was called mock_external_service_dependencies["index_processor_factory"].assert_called_once_with( - dataset.get_doc_form(session=db_session_with_containers) + get_dataset_doc_form(dataset, session=db_session_with_containers) ) mock_external_service_dependencies["index_processor"].load.assert_called_once() @@ -1064,7 +1065,7 @@ class TestCreateSegmentToIndexTask: # Verify index processor was called mock_external_service_dependencies["index_processor_factory"].assert_called_once_with( - dataset.get_doc_form(session=db_session_with_containers) + get_dataset_doc_form(dataset, session=db_session_with_containers) ) mock_external_service_dependencies["index_processor"].load.assert_called_once() diff --git a/api/tests/test_containers_integration_tests/tasks/test_dataset_indexing_task.py b/api/tests/test_containers_integration_tests/tasks/test_dataset_indexing_task.py index eeb54ec5c6f..0e18308b391 100644 --- a/api/tests/test_containers_integration_tests/tasks/test_dataset_indexing_task.py +++ b/api/tests/test_containers_integration_tests/tasks/test_dataset_indexing_task.py @@ -9,12 +9,13 @@ from faker import Faker from sqlalchemy import select from sqlalchemy.orm import Session -from core.indexing_runner import DocumentIsPausedError from core.rag.index_processor.constant.index_type import IndexTechniqueType from enums import CloudPlan, DeploymentEdition from models import Account, AccountStatus, Tenant, TenantAccountJoin, TenantAccountRole, TenantStatus from models.dataset import Dataset, Document from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus +from services.knowledge.indexing.errors import DocumentIsPausedError +from services.knowledge.resource_scope import DocumentRef from tasks.document_indexing_task import ( _document_indexing, _document_indexing_with_tenant_queue, @@ -84,7 +85,7 @@ def session_close_tracker(): def patched_external_dependencies(): """Patch non-DB collaborators while keeping database behavior real.""" with ( - patch("tasks.document_indexing_task.IndexingRunner", autospec=True) as mock_indexing_runner, + patch("tasks.document_indexing_task.build_document_indexing_service", autospec=True) as mock_indexing_runner, patch("tasks.document_indexing_task.FeatureService", autospec=True) as mock_feature_service, patch("tasks.document_indexing_task.generate_summary_index_task", autospec=True) as mock_summary_task, ): @@ -213,7 +214,7 @@ class TestDatasetIndexingTaskIntegration: assert len(opened) >= 2 assert opened_ids <= closed_ids - def _runner_documents_arg(self, patched_external_dependencies) -> Sequence[Document]: + def _runner_documents_arg(self, patched_external_dependencies) -> Sequence[DocumentRef]: """Return the document batch passed to the runner.""" return patched_external_dependencies["indexing_runner_instance"].run.call_args.args[0] diff --git a/api/tests/test_containers_integration_tests/tasks/test_document_indexing_sync_task.py b/api/tests/test_containers_integration_tests/tasks/test_document_indexing_sync_task.py index 77cd2598333..bd3393c510d 100644 --- a/api/tests/test_containers_integration_tests/tasks/test_document_indexing_sync_task.py +++ b/api/tests/test_containers_integration_tests/tasks/test_document_indexing_sync_task.py @@ -15,11 +15,13 @@ import pytest from sqlalchemy import delete, func, select, update from sqlalchemy.orm import Session -from core.indexing_runner import DocumentIsPausedError, IndexingRunner from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType from models import Account, AccountStatus, Tenant, TenantAccountJoin, TenantAccountRole, TenantStatus from models.dataset import Dataset, Document, DocumentSegment from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus, SegmentStatus +from services.knowledge.indexing.errors import DocumentIsPausedError +from services.knowledge.indexing.execution import DocumentIndexingService +from services.knowledge.resource_scope import DatasetRef from tasks.document_indexing_sync_task import document_indexing_sync_task @@ -133,14 +135,14 @@ class TestDocumentIndexingSyncTask: def mock_external_dependencies(self): """Patch only external collaborators; keep DB access real.""" with ( - patch("tasks.document_indexing_sync_task.DatasourceProviderService") as mock_datasource_service_class, + patch("tasks.document_indexing_sync_task.build_data_source_credentials") as mock_credentials, patch("tasks.document_indexing_sync_task.NotionExtractor") as mock_notion_extractor_class, patch("tasks.document_indexing_sync_task.IndexProcessorFactory") as mock_index_processor_factory, - patch("tasks.document_indexing_sync_task.IndexingRunner") as mock_indexing_runner_class, + patch("tasks.document_indexing_sync_task.build_document_indexing_service") as mock_indexing_runner_class, ): datasource_service = Mock() datasource_service.get_datasource_credentials.return_value = {"integration_secret": "test_token"} - mock_datasource_service_class.return_value = datasource_service + mock_credentials.return_value.providers = datasource_service notion_extractor = Mock() notion_extractor.get_notion_last_edited_time.return_value = "2024-01-02T00:00:00Z" @@ -150,7 +152,7 @@ class TestDocumentIndexingSyncTask: index_processor.clean = Mock() mock_index_processor_factory.return_value.init_index_processor.return_value = index_processor - indexing_runner = Mock(spec=IndexingRunner) + indexing_runner = Mock(spec=DocumentIndexingService) indexing_runner.run = Mock() mock_indexing_runner_class.return_value = indexing_runner @@ -339,16 +341,14 @@ class TestDocumentIndexingSyncTask: clean_call_args = mock_external_dependencies["index_processor"].clean.call_args assert clean_call_args is not None clean_args, clean_kwargs = clean_call_args - assert getattr(clean_args[0], "id", None) == context["dataset"].id + assert clean_args[0].id == context["dataset"].id assert set(clean_args[1]) == set(context["node_ids"]) assert clean_kwargs.get("with_keywords") is True assert clean_kwargs.get("delete_child_chunks") is True - run_call_args = mock_external_dependencies["indexing_runner"].run.call_args - assert run_call_args is not None - run_documents = run_call_args[0][0] - assert len(run_documents) == 1 - assert getattr(run_documents[0], "id", None) == context["document"].id + mock_external_dependencies["indexing_runner"].run.assert_called_once_with( + [DatasetRef(context["tenant"].id, context["dataset"].id).document(context["document"].id)] + ) def test_dataset_not_found_during_cleaning(self, db_session_with_containers: Session, mock_external_dependencies): """Test that task still updates document and reindexes if dataset vanishes before clean.""" @@ -458,7 +458,7 @@ class TestDocumentIndexingSyncTask: clean_call_args = mock_external_dependencies["index_processor"].clean.call_args assert clean_call_args is not None clean_args, clean_kwargs = clean_call_args - assert getattr(clean_args[0], "id", None) == context["dataset"].id + assert clean_args[0].id == context["dataset"].id assert set(clean_args[1]) == set(context["node_ids"]) assert clean_kwargs.get("with_keywords") is True assert clean_kwargs.get("delete_child_chunks") is True diff --git a/api/tests/test_containers_integration_tests/tasks/test_document_indexing_task.py b/api/tests/test_containers_integration_tests/tasks/test_document_indexing_task.py index 85b3993b1fd..613bb04e012 100644 --- a/api/tests/test_containers_integration_tests/tasks/test_document_indexing_task.py +++ b/api/tests/test_containers_integration_tests/tasks/test_document_indexing_task.py @@ -11,6 +11,7 @@ from enums import CloudPlan, DeploymentEdition from models import Account, AccountStatus, Tenant, TenantAccountJoin, TenantAccountRole, TenantStatus from models.dataset import Dataset, Document from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus +from services.knowledge.resource_scope import DocumentRef from tasks.document_indexing_task import ( _document_indexing, # Core function _document_indexing_with_tenant_queue, # Tenant queue wrapper function @@ -36,7 +37,9 @@ class TestDocumentIndexingTasks: def mock_external_service_dependencies(self): """Mock setup for external service dependencies.""" with ( - patch("tasks.document_indexing_task.IndexingRunner", autospec=True) as mock_indexing_runner, + patch( + "tasks.document_indexing_task.build_document_indexing_service", autospec=True + ) as mock_indexing_runner, patch("tasks.document_indexing_task.FeatureService", autospec=True) as mock_feature_service, ): # Setup mock indexing runner @@ -51,7 +54,7 @@ class TestDocumentIndexingTasks: "features": mock_features, } - def _runner_documents_arg(self, mock_external_service_dependencies) -> list[Document]: + def _runner_documents_arg(self, mock_external_service_dependencies) -> list[DocumentRef]: """Return the document batch passed to the runner.""" return mock_external_service_dependencies["indexing_runner_instance"].run.call_args.args[0] @@ -243,7 +246,7 @@ class TestDocumentIndexingTasks: This test verifies: - Proper dataset retrieval from database - Correct document processing and status updates - - IndexingRunner integration + - build_document_indexing_service integration - Database state updates """ # Arrange: Create test data @@ -350,10 +353,10 @@ class TestDocumentIndexingTasks: self, db_session_with_containers: Session, mock_external_service_dependencies ): """ - Test handling of IndexingRunner exceptions. + Test handling of build_document_indexing_service exceptions. This test verifies: - - Exceptions from IndexingRunner are properly caught + - Exceptions from build_document_indexing_service are properly caught - Task completes without raising exceptions - Database session is properly closed - Error logging occurs @@ -364,7 +367,7 @@ class TestDocumentIndexingTasks: ) document_ids = [doc.id for doc in documents] - # Mock IndexingRunner to raise an exception + # Mock build_document_indexing_service to raise an exception mock_external_service_dependencies["indexing_runner_instance"].run.side_effect = Exception( "Indexing runner failed" ) @@ -393,7 +396,7 @@ class TestDocumentIndexingTasks: - Documents with different initial states are handled correctly - Only valid documents are processed - Database state updates are consistent - - IndexingRunner receives correct documents + - build_document_indexing_service receives correct documents """ # Arrange: Create test data dataset, base_documents = self._create_test_dataset_and_documents( @@ -541,7 +544,7 @@ class TestDocumentIndexingTasks: - Processing continues normally outside Cloud - No billing validation occurs - Documents are processed successfully - - IndexingRunner is called correctly + - build_document_indexing_service is called correctly """ # Arrange: Create test data with billing disabled dataset, documents = self._create_test_dataset_with_billing_features( @@ -572,7 +575,7 @@ class TestDocumentIndexingTasks: self, db_session_with_containers: Session, mock_external_service_dependencies ): """ - Test handling of DocumentIsPausedError from IndexingRunner. + Test handling of DocumentIsPausedError from build_document_indexing_service. This test verifies: - DocumentIsPausedError is properly caught and handled @@ -586,8 +589,8 @@ class TestDocumentIndexingTasks: ) document_ids = [doc.id for doc in documents] - # Mock IndexingRunner to raise DocumentIsPausedError - from core.indexing_runner import DocumentIsPausedError + # Mock build_document_indexing_service to raise DocumentIsPausedError + from services.knowledge.indexing.errors import DocumentIsPausedError mock_external_service_dependencies["indexing_runner_instance"].run.side_effect = DocumentIsPausedError( "Document indexing is paused" @@ -809,7 +812,7 @@ class TestDocumentIndexingTasks: tenant_id = dataset.tenant_id dataset_id = dataset.id - # Mock IndexingRunner to raise an exception + # Mock build_document_indexing_service to raise an exception mock_external_service_dependencies["indexing_runner_instance"].run.side_effect = Exception("Test error") # Mock the task function diff --git a/api/tests/test_containers_integration_tests/tasks/test_document_indexing_update_task.py b/api/tests/test_containers_integration_tests/tasks/test_document_indexing_update_task.py index bda54389419..55aa5cacfa9 100644 --- a/api/tests/test_containers_integration_tests/tasks/test_document_indexing_update_task.py +++ b/api/tests/test_containers_integration_tests/tasks/test_document_indexing_update_task.py @@ -9,6 +9,7 @@ from core.rag.index_processor.constant.index_type import IndexStructureType, Ind from models import Account, AccountStatus, Tenant, TenantAccountJoin, TenantAccountRole, TenantStatus from models.dataset import Dataset, Document, DocumentSegment from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus, SegmentStatus +from services.knowledge.resource_scope import DatasetRef from tasks.document_indexing_update_task import document_indexing_update_task @@ -17,11 +18,11 @@ class TestDocumentIndexingUpdateTask: def mock_external_dependencies(self): """Patch external collaborators used by the update task. - IndexProcessorFactory.init_index_processor().clean(...) - - IndexingRunner.run([...]) + - build_document_indexing_service.run([...]) """ with ( patch("tasks.document_indexing_update_task.IndexProcessorFactory", autospec=True) as mock_factory, - patch("tasks.document_indexing_update_task.IndexingRunner", autospec=True) as mock_runner, + patch("tasks.document_indexing_update_task.build_document_indexing_service", autospec=True) as mock_runner, ): processor_instance = MagicMock() mock_factory.return_value.init_index_processor.return_value = processor_instance @@ -140,19 +141,16 @@ class TestDocumentIndexingUpdateTask: assert clean_call is not None args, kwargs = clean_call # args[0] is a Dataset instance (from another session), so validate by id. - assert getattr(args[0], "id", None) == dataset.id + assert args[0].id == dataset.id # args[1] should contain our node_ids. assert set(args[1]) == set(node_ids) assert kwargs.get("with_keywords") is True assert kwargs.get("delete_child_chunks") is True - # Assert indexing runner invoked with the updated document - run_call = mock_external_dependencies["runner_instance"].run.call_args - assert run_call is not None - run_docs = run_call[0][0] - assert len(run_docs) == 1 - first = run_docs[0] - assert getattr(first, "id", None) == document.id + # Reindexing receives the detached reference with its complete ownership chain. + mock_external_dependencies["runner_instance"].run.assert_called_once_with( + [DatasetRef(dataset.tenant_id, dataset.id).document(document.id)] + ) def test_clean_error_is_logged_and_indexing_continues( self, db_session_with_containers: Session, mock_external_dependencies diff --git a/api/tests/test_containers_integration_tests/tasks/test_duplicate_document_indexing_task.py b/api/tests/test_containers_integration_tests/tasks/test_duplicate_document_indexing_task.py index 1dacc2bba10..c2fe023cc38 100644 --- a/api/tests/test_containers_integration_tests/tasks/test_duplicate_document_indexing_task.py +++ b/api/tests/test_containers_integration_tests/tasks/test_duplicate_document_indexing_task.py @@ -5,12 +5,13 @@ from faker import Faker from sqlalchemy import select from sqlalchemy.orm import Session -from core.indexing_runner import DocumentIsPausedError from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType from enums import CloudPlan, DeploymentEdition from models import Account, Tenant, TenantAccountJoin, TenantAccountRole from models.dataset import Dataset, Document, DocumentSegment from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus, SegmentStatus +from services.knowledge.indexing.errors import DocumentIsPausedError +from services.knowledge.resource_scope import DocumentRef from tasks.duplicate_document_indexing_task import ( _duplicate_document_indexing_task, # Core function _duplicate_document_indexing_task_with_tenant_queue, # Tenant queue wrapper function @@ -37,7 +38,9 @@ class TestDuplicateDocumentIndexingTasks: def mock_external_service_dependencies(self): """Mock setup for external service dependencies.""" with ( - patch("tasks.duplicate_document_indexing_task.IndexingRunner", autospec=True) as mock_indexing_runner, + patch( + "tasks.duplicate_document_indexing_task.build_document_indexing_service", autospec=True + ) as mock_indexing_runner, patch("tasks.duplicate_document_indexing_task.FeatureService", autospec=True) as mock_feature_service, patch( "tasks.duplicate_document_indexing_task.IndexProcessorFactory", autospec=True @@ -62,7 +65,7 @@ class TestDuplicateDocumentIndexingTasks: "index_processor": mock_processor, } - def _runner_documents_arg(self, mock_external_service_dependencies) -> list[Document]: + def _runner_documents_arg(self, mock_external_service_dependencies) -> list[DocumentRef]: """Return the document batch passed to the runner.""" return mock_external_service_dependencies["indexing_runner_instance"].run.call_args.args[0] @@ -311,7 +314,7 @@ class TestDuplicateDocumentIndexingTasks: This test verifies: - Proper dataset retrieval from database - Correct document processing and status updates - - IndexingRunner integration + - build_document_indexing_service integration - Database state updates """ # Arrange: Create test data @@ -469,10 +472,10 @@ class TestDuplicateDocumentIndexingTasks: self, db_session_with_containers: Session, mock_external_service_dependencies ): """ - Test handling of IndexingRunner exceptions. + Test handling of build_document_indexing_service exceptions. This test verifies: - - Exceptions from IndexingRunner are properly caught + - Exceptions from build_document_indexing_service are properly caught - Task completes without raising exceptions - Database session is properly closed - Error logging occurs @@ -483,7 +486,7 @@ class TestDuplicateDocumentIndexingTasks: ) document_ids = [doc.id for doc in documents] - # Mock IndexingRunner to raise an exception + # Mock build_document_indexing_service to raise an exception mock_external_service_dependencies["indexing_runner_instance"].run.side_effect = Exception( "Indexing runner failed" ) @@ -630,7 +633,7 @@ class TestDuplicateDocumentIndexingTasks: # Act: Execute the task with empty document list _duplicate_document_indexing_task(dataset.id, document_ids) - # Assert: Verify IndexingRunner was called with empty list + # Assert: Verify build_document_indexing_service was called with empty list # Note: The actual implementation does call run([]) with an empty list. mock_external_service_dependencies["indexing_runner"].assert_called_once() mock_external_service_dependencies["indexing_runner_instance"].run.assert_called_once() @@ -840,7 +843,7 @@ class TestDuplicateDocumentIndexingTasks: def test_duplicate_document_indexing_runner_error( self, db_session_with_containers: Session, mock_external_service_dependencies ): - """Test duplicate document indexing when IndexingRunner raises an error.""" + """Test duplicate document indexing when build_document_indexing_service raises an error.""" self._test_duplicate_document_indexing_task_indexing_runner_exception( db_session_with_containers, mock_external_service_dependencies ) diff --git a/api/tests/test_containers_integration_tests/tasks/test_enable_segments_to_index_task.py b/api/tests/test_containers_integration_tests/tasks/test_enable_segments_to_index_task.py index 6d3b90d41c7..525deeca421 100644 --- a/api/tests/test_containers_integration_tests/tasks/test_enable_segments_to_index_task.py +++ b/api/tests/test_containers_integration_tests/tasks/test_enable_segments_to_index_task.py @@ -7,8 +7,8 @@ from sqlalchemy.orm import Session from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType from extensions.ext_redis import redis_client from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole -from models.dataset import Dataset, Document, DocumentSegment -from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus, SegmentStatus +from models.dataset import ChildChunk, Dataset, DatasetProcessRule, Document, DocumentSegment +from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus, ProcessRuleMode, SegmentStatus from tasks.enable_segments_to_index_task import enable_segments_to_index_task @@ -370,7 +370,7 @@ class TestEnableSegmentsToIndexTask: # Refresh dataset to ensure doc_form property reflects the updated document db_session_with_containers.refresh(dataset) - # Create segments with mock child chunks + # Persist each segment and its own child chunks segments = self._create_test_segments(db_session_with_containers, document, dataset) # Set up Redis cache keys @@ -379,43 +379,62 @@ class TestEnableSegmentsToIndexTask: indexing_cache_key = f"segment_{segment.id}_indexing" redis_client.set(indexing_cache_key, "processing", ex=300) - # Mock the get_child_chunks method for each segment - with patch.object(DocumentSegment, "get_child_chunks", autospec=True) as mock_get_child_chunks: - # Setup mock to return child chunks for each segment - mock_child_chunks = [] - for i in range(2): # Each segment has 2 child chunks - mock_child = MagicMock() - mock_child.content = f"child_content_{i}" - mock_child.index_node_id = f"child_node_{i}" - mock_child.index_node_hash = f"child_hash_{i}" - mock_child_chunks.append(mock_child) - - mock_get_child_chunks.return_value = mock_child_chunks - - # Act: Execute the task - enable_segments_to_index_task(segment_ids, dataset.id, document.id) - - # Assert: Verify parent-child index processing - mock_external_service_dependencies["index_processor_factory"].assert_called_once_with( - IndexStructureType.PARENT_CHILD_INDEX + rule = DatasetProcessRule( + dataset_id=dataset.id, + mode=ProcessRuleMode.HIERARCHICAL, + rules='{"parent_mode": "paragraph"}', + created_by=document.created_by, + ) + db_session_with_containers.add(rule) + db_session_with_containers.flush() + document.dataset_process_rule_id = rule.id + for segment in segments: + db_session_with_containers.add_all( + [ + ChildChunk( + tenant_id=dataset.tenant_id, + dataset_id=dataset.id, + document_id=document.id, + segment_id=segment.id, + position=i, + content=f"child_content_{i}", + word_count=1, + created_by=document.created_by, + index_node_id=f"{segment.index_node_id}-child-{i}", + index_node_hash=f"{segment.index_node_id}-hash-{i}", + ) + for i in range(2) + ] ) - mock_external_service_dependencies["index_processor"].load.assert_called_once() + db_session_with_containers.commit() - # Verify the load method was called with correct parameters - call_args = mock_external_service_dependencies["index_processor"].load.call_args - assert call_args is not None - documents = call_args[0][1] # Second argument should be documents list - assert len(documents) == 3 # 3 segments + # Act: Execute the task + enable_segments_to_index_task(segment_ids, dataset.id, document.id) - # Verify each document has children - for doc in documents: - assert hasattr(doc, "children") - assert len(doc.children) == 2 # Each document has 2 children + # Assert: Verify parent-child index processing + mock_external_service_dependencies["index_processor_factory"].assert_called_once_with( + IndexStructureType.PARENT_CHILD_INDEX + ) + mock_external_service_dependencies["index_processor"].load.assert_called_once() - # Verify Redis cache keys were deleted - for segment in segments: - indexing_cache_key = f"segment_{segment.id}_indexing" - assert redis_client.exists(indexing_cache_key) == 0 + # Verify the load method was called with correct parameters + call_args = mock_external_service_dependencies["index_processor"].load.call_args + assert call_args is not None + documents = call_args[0][1] # Second argument should be documents list + assert len(documents) == 3 # 3 segments + + # Each parent receives only its own persisted children, in position order. + for doc in documents: + assert doc.children is not None + assert [child.page_content for child in doc.children] == ["child_content_0", "child_content_1"] + assert [child.metadata["doc_id"] for child in doc.children] == [ + f"{doc.metadata['doc_id']}-child-{i}" for i in range(2) + ] + + # Verify Redis cache keys were deleted + for segment in segments: + indexing_cache_key = f"segment_{segment.id}_indexing" + assert redis_client.exists(indexing_cache_key) == 0 def test_enable_segments_to_index_general_exception_handling( self, db_session_with_containers: Session, mock_external_service_dependencies diff --git a/api/tests/unit_tests/controllers/conftest.py b/api/tests/unit_tests/controllers/conftest.py index c24a06ff520..d0cdf0fa25f 100644 --- a/api/tests/unit_tests/controllers/conftest.py +++ b/api/tests/unit_tests/controllers/conftest.py @@ -1,11 +1,15 @@ -"""Explicit App service composition for controller tests that use app queries.""" +"""Application dependencies shared by controller tests.""" +import sys from dataclasses import dataclass +from types import SimpleNamespace +from unittest.mock import create_autospec import pytest from flask import Flask from sqlalchemy.orm import Session, sessionmaker +from core.app.apps.pipeline.pipeline_generator import PipelineGenerator from extensions.application_services.app import AppServices from extensions.ext_application_services import ( _batch_get_enterprise_webapp_access_modes, @@ -18,10 +22,62 @@ from repositories.webapp_access_query_repository import WebAppAccessQueryReposit from services.app.console_service import ConsoleAppService from services.app.import_service import AppImportService from services.app.query_service import AppQueryService +from services.credentials.query import CredentialQuery +from services.data_source.provider_service import DatasourceProviderService from services.tag_application_service import TagApplicationService from services.webapp_access_query_service import WebAppAccessQueryService +@pytest.fixture(autouse=True) +def datasource_application_dependencies(monkeypatch: pytest.MonkeyPatch) -> None: + registry = SimpleNamespace( + data_sources=SimpleNamespace(providers=create_autospec(DatasourceProviderService, instance=True, spec_set=True)) + ) + for name in ( + "controllers.console.datasets.rag_pipeline.rag_pipeline_workflow", + "controllers.console.datasets.rag_pipeline.datasource_content_preview", + "controllers.console.datasets.website", + ): + module = sys.modules.get(name) + if module is not None: + monkeypatch.setattr(module, "application_services", lambda: registry) + + +@dataclass(frozen=True) +class PipelineKnowledgeStub: + pipeline_generator: PipelineGenerator + + +@dataclass(frozen=True) +class PipelineDataSourceStub: + providers: DatasourceProviderService + + +@dataclass(frozen=True) +class PipelineApplicationStub: + knowledge: PipelineKnowledgeStub + credential_queries: CredentialQuery + data_sources: PipelineDataSourceStub + + +@pytest.fixture +def pipeline_application(monkeypatch: pytest.MonkeyPatch) -> PipelineGenerator: + from controllers.console.datasets.rag_pipeline import rag_pipeline_workflow as console_workflow + from controllers.service_api.dataset.rag_pipeline import rag_pipeline_workflow as service_api_workflow + + generator = create_autospec(PipelineGenerator, instance=True, spec_set=True) + registry = PipelineApplicationStub( + knowledge=PipelineKnowledgeStub(pipeline_generator=generator), + credential_queries=create_autospec(CredentialQuery, instance=True, spec_set=True), + data_sources=PipelineDataSourceStub( + providers=create_autospec(DatasourceProviderService, instance=True, spec_set=True) + ), + ) + monkeypatch.setattr(console_workflow, "application_services", lambda: registry) + monkeypatch.setattr(service_api_workflow, "application_services", lambda: registry) + return generator + + @dataclass class AppQueryTestServices: console: ConsoleAppService diff --git a/api/tests/unit_tests/controllers/console/auth/test_data_source_bearer_auth.py b/api/tests/unit_tests/controllers/console/auth/test_data_source_bearer_auth.py index 47f25498f3a..b04f8d51451 100644 --- a/api/tests/unit_tests/controllers/console/auth/test_data_source_bearer_auth.py +++ b/api/tests/unit_tests/controllers/console/auth/test_data_source_bearer_auth.py @@ -1,8 +1,8 @@ from __future__ import annotations +from dataclasses import dataclass from datetime import UTC, datetime from inspect import unwrap -from types import SimpleNamespace from unittest.mock import create_autospec, patch from uuid import UUID @@ -23,14 +23,14 @@ from controllers.console.auth.error import ( ) from libs.exception import BaseHTTPException from machinery.context import RequestContext -from services.auth.data_source_api_key_auth_service import DataSourceApiKeyAuthService from services.auth.errors import ( DataSourceApiKeyAuthCredentialValidationError, DataSourceApiKeyAuthProviderUnavailableError, InvalidDataSourceApiKeyAuthCredentialsError, UnsupportedDataSourceApiKeyAuthProviderError, ) -from services.entities.data_source_api_key_auth_entities import ( +from services.data_source.auth.api_key_service import DataSourceApiKeyAuthService +from services.data_source.entities.api_key_auth import ( DataSourceApiKeyAuthBindingCreate, DataSourceApiKeyAuthBindingRecord, DataSourceApiKeyAuthCredentials, @@ -46,8 +46,18 @@ def _request_context() -> RequestContext: ) -def _application_services(service: DataSourceApiKeyAuthService) -> SimpleNamespace: - return SimpleNamespace(data_source_api_key_auth=service) +@dataclass(frozen=True) +class _DataSources: + api_key_auth: DataSourceApiKeyAuthService + + +@dataclass(frozen=True) +class _ApplicationServices: + data_sources: _DataSources + + +def _application_services(service: DataSourceApiKeyAuthService) -> _ApplicationServices: + return _ApplicationServices(data_sources=_DataSources(api_key_auth=service)) def test_list_data_source_auth_passes_context_and_serializes_application_result() -> None: diff --git a/api/tests/unit_tests/controllers/console/auth/test_data_source_oauth.py b/api/tests/unit_tests/controllers/console/auth/test_data_source_oauth.py index 72d28c0385b..e47876a7c2e 100644 --- a/api/tests/unit_tests/controllers/console/auth/test_data_source_oauth.py +++ b/api/tests/unit_tests/controllers/console/auth/test_data_source_oauth.py @@ -13,12 +13,12 @@ from controllers.console.auth.data_source_oauth import ( ) from extensions.ext_application_services import ApplicationServices from machinery.context import RequestContext -from services.data_source_oauth_service import ( +from services.data_source.entities.oauth import DataSourceOAuthCallback +from services.data_source.oauth_service import ( DataSourceOAuthError, InvalidDataSourceOAuthCodeError, InvalidDataSourceOAuthProviderError, ) -from services.entities.data_source_oauth_entities import DataSourceOAuthCallback from tests.unit_tests.config_override import config_overrides_context @@ -33,7 +33,7 @@ def _request_context() -> RequestContext: def _services(service: MagicMock) -> MagicMock: services = create_autospec(ApplicationServices, instance=True, spec_set=True) - services.resolve_data_source_oauth.return_value = service + services.data_sources.resolve_oauth.return_value = service return services @@ -55,7 +55,7 @@ def test_get_delegates_to_application_service_and_serializes_response() -> None: def test_get_maps_unknown_provider_to_bad_request() -> None: service = MagicMock() services = _services(service) - services.resolve_data_source_oauth.side_effect = InvalidDataSourceOAuthProviderError + services.data_sources.resolve_oauth.side_effect = InvalidDataSourceOAuthProviderError with patch( "controllers.console.auth.data_source_oauth.application_services", diff --git a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_datasource_auth.py b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_datasource_auth.py index c1cbbaa8682..4874145451d 100644 --- a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_datasource_auth.py +++ b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_datasource_auth.py @@ -1,5 +1,7 @@ import inspect from datetime import UTC, datetime +from importlib import import_module +from types import SimpleNamespace from unittest.mock import ANY, MagicMock, patch import pytest @@ -28,13 +30,26 @@ from controllers.console.datasets.rag_pipeline.datasource_auth import ( from core.plugin.impl.oauth import OAuthHandler from graphon.model_runtime.errors.validate import CredentialsValidateFailedError from models.account import Account -from services.datasource_provider_service import DatasourceProviderService +from services.credentials.query import CredentialQuery +from services.data_source.credential_gateway import DatasourceProviderCredentialStore +from services.data_source.provider_service import DatasourceProviderService from services.plugin.oauth_service import OAuthProxyService from tests.unit_tests.model_factories import make_account _PROVIDER_ID = "langgenius/notion_datasource/notion" +@pytest.fixture(autouse=True) +def credential_query_dependency(monkeypatch: pytest.MonkeyPatch) -> None: + query = MagicMock(spec=CredentialQuery) + providers = DatasourceProviderService(credentials=MagicMock(spec=DatasourceProviderCredentialStore)) + monkeypatch.setattr( + import_module("controllers.console.datasets.rag_pipeline.datasource_auth"), + "application_services", + lambda: SimpleNamespace(credential_queries=query, data_sources=SimpleNamespace(providers=providers)), + ) + + def _account() -> Account: return make_account(account_id="user-1", name="Datasource Auth Tester", email="datasource-auth@example.com") diff --git a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_datasource_catalog.py b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_datasource_catalog.py index 08fcc368967..cc2029a4a75 100644 --- a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_datasource_catalog.py +++ b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_datasource_catalog.py @@ -5,6 +5,7 @@ import pytest from flask import Flask from pydantic import ValidationError +from controllers.console.datasets.rag_pipeline import rag_pipeline_workflow as workflow_module from controllers.console.datasets.rag_pipeline.rag_pipeline_workflow import ( DatasourceListApi, RagPipelineDatasourceListResponse, @@ -14,7 +15,6 @@ from core.datasource.entities.datasource_entities import DatasourceParameter from core.plugin.entities import OAuthSchema from core.plugin.entities.plugin_daemon import PluginDatasourceProviderEntity from core.plugin.impl.datasource import PluginDatasourceManager -from services.datasource_provider_service import DatasourceProviderService from services.rag_pipeline.rag_pipeline_manage_service import RagPipelineManageService @@ -89,7 +89,9 @@ def test_catalog_serializes_real_plugin_and_local_file_providers( PluginDatasourceManager, "_request_with_plugin_daemon_response", return_value=[provider] ) as request, patch.object( - DatasourceProviderService, "get_datasource_credentials", return_value=credentials + workflow_module.application_services().data_sources.providers, + "get_datasource_credentials", + return_value=credentials, ) as get_credentials, ): result = unwrap(api.get)(api, "tenant-1") @@ -138,7 +140,11 @@ def test_catalog_keeps_oauth_only_provider_unauthorized_without_credentials(app: with ( app.test_request_context("/rag/pipelines/datasource-plugins"), patch.object(PluginDatasourceManager, "fetch_datasource_providers", return_value=[provider]), - patch.object(DatasourceProviderService, "get_datasource_credentials", return_value=None) as get_credentials, + patch.object( + workflow_module.application_services().data_sources.providers, + "get_datasource_credentials", + return_value=None, + ) as get_credentials, ): result = unwrap(api.get)(api, "tenant-1") @@ -154,7 +160,11 @@ def test_catalog_keeps_provider_visible_after_credential_failure(app: Flask) -> with ( app.test_request_context("/rag/pipelines/datasource-plugins"), patch.object(PluginDatasourceManager, "fetch_datasource_providers", return_value=[provider]), - patch.object(DatasourceProviderService, "get_datasource_credentials", side_effect=ValueError("decrypt failed")), + patch.object( + workflow_module.application_services().data_sources.providers, + "get_datasource_credentials", + side_effect=ValueError("decrypt failed"), + ), ): result = unwrap(api.get)(api, "tenant-1") diff --git a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_datasource_content_preview.py b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_datasource_content_preview.py index b3aed7be87c..daae29fd75c 100644 --- a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_datasource_content_preview.py +++ b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_datasource_content_preview.py @@ -1,5 +1,5 @@ from inspect import unwrap -from unittest.mock import MagicMock, patch +from unittest.mock import ANY, MagicMock, patch import pytest from flask import Flask @@ -61,6 +61,7 @@ class TestDataSourceContentPreviewApi: datasource_type=payload["datasource_type"], is_published=True, credential_id=payload["credential_id"], + datasource_providers=ANY, ) assert status == 200 assert response == preview_result diff --git a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline.py b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline.py index 53445379c5c..36032f3650c 100644 --- a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline.py +++ b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline.py @@ -375,7 +375,7 @@ class TestPublishCustomizedPipelineTemplateApi: with ( config_overrides_context(RBAC_ENABLED=True), - patch.object(Pipeline, "retrieve_dataset", return_value=dataset), + patch.object(module, "get_pipeline_dataset", return_value=dataset), patch.object(module.DatasetService, "check_dataset_permission") as legacy_acl, patch.object(module.RagPipelineService, "publish_customized_pipeline_template") as publish, ): @@ -408,7 +408,7 @@ class TestPublishCustomizedPipelineTemplateApi: with ( config_overrides_context(RBAC_ENABLED=True), - patch.object(Pipeline, "retrieve_dataset", return_value=dataset), + patch.object(module, "get_pipeline_dataset", return_value=dataset), patch.object(module.RagPipelineService, "publish_customized_pipeline_template") as publish, ): response, status = method(api, CustomizedPipelineTemplatePayload.model_validate(payload), account, pipeline) @@ -424,7 +424,7 @@ class TestPublishCustomizedPipelineTemplateApi: with ( config_overrides_context(RBAC_ENABLED=True), - patch.object(Pipeline, "retrieve_dataset", return_value=object()), + patch.object(module, "get_pipeline_dataset", return_value=object()), patch.object( module.RagPipelineService, "publish_customized_pipeline_template", @@ -445,7 +445,7 @@ class TestPublishCustomizedPipelineTemplateApi: with ( config_overrides_context(RBAC_ENABLED=False), - patch.object(Pipeline, "retrieve_dataset", return_value=dataset), + patch.object(module, "get_pipeline_dataset", return_value=dataset), patch.object(module.DatasetService, "check_dataset_permission") as check_permission, patch.object(module.RagPipelineService, "publish_customized_pipeline_template") as publish, ): @@ -463,7 +463,7 @@ class TestPublishCustomizedPipelineTemplateApi: with ( config_overrides_context(RBAC_ENABLED=False), - patch.object(Pipeline, "retrieve_dataset", return_value=object()), + patch.object(module, "get_pipeline_dataset", return_value=object()), patch.object(module.DatasetService, "check_dataset_permission") as check_permission, patch.object(module.RagPipelineService, "publish_customized_pipeline_template") as publish, pytest.raises(Forbidden), @@ -481,7 +481,7 @@ class TestPublishCustomizedPipelineTemplateApi: with ( config_overrides_context(RBAC_ENABLED=False), - patch.object(Pipeline, "retrieve_dataset", return_value=object()), + patch.object(module, "get_pipeline_dataset", return_value=object()), patch.object( module.DatasetService, "check_dataset_permission", @@ -502,7 +502,7 @@ class TestPublishCustomizedPipelineTemplateApi: with ( config_overrides_context(RBAC_ENABLED=False), - patch.object(Pipeline, "retrieve_dataset", return_value=None), + patch.object(module, "get_pipeline_dataset", return_value=None), patch.object(module.DatasetService, "check_dataset_permission") as check_permission, patch.object(module.RagPipelineService, "publish_customized_pipeline_template") as publish, pytest.raises(NotFound, match="Dataset not found"), diff --git a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow.py b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow.py index 294e0c04f35..8794c7d7b6a 100644 --- a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow.py +++ b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow.py @@ -6,7 +6,7 @@ import json from collections.abc import Iterator from datetime import datetime from inspect import unwrap as unwrap_all -from unittest.mock import MagicMock, patch +from unittest.mock import patch from uuid import UUID import pytest @@ -25,6 +25,7 @@ from controllers.console.datasets.rag_pipeline.rag_pipeline_workflow import ( WorkflowUpdatePayload, ) from controllers.web.error import InvokeRateLimitError as InvokeRateLimitHttpError +from core.app.apps.pipeline.pipeline_generator import PipelineGenerator from models.account import Account, TenantAccountRole from models.dataset import Dataset, Pipeline from models.engine import db @@ -108,9 +109,9 @@ def database_app() -> Iterator[Flask]: db.init_app(app) with app.app_context(): - Account.__table__.create(db.engine) - WorkflowToolProvider.__table__.create(db.engine) - Workflow.__table__.create(db.engine) + db.metadata.tables[Account.__tablename__].create(db.engine) + db.metadata.tables[WorkflowToolProvider.__tablename__].create(db.engine) + db.metadata.tables[Workflow.__tablename__].create(db.engine) db.session.add(_account()) db.session.commit() @@ -387,6 +388,7 @@ def test_rag_pipeline_transform_skips_legacy_acl_when_rbac_is_enabled(sqlite_eng def test_rag_pipeline_run_uses_sqlite_session( app: Flask, sqlite_engine: Engine, + pipeline_application: PipelineGenerator, api_type: type, payload: dict[str, object], ) -> None: @@ -398,7 +400,7 @@ def test_rag_pipeline_run_uses_sqlite_session( Session(sqlite_engine) as session, app.test_request_context("/", json=payload), patch.object(module, "load_rag_pipeline", return_value=pipeline) as load_pipeline, - patch.object(module.PipelineGenerateService, "generate", return_value=MagicMock()) as generate, + patch.object(module.PipelineGenerateService, "generate", return_value={"event": "done"}) as generate, patch.object(module.helper, "compact_generate_response", return_value={"ok": True}), ): req_data = ( @@ -411,6 +413,7 @@ def test_rag_pipeline_run_uses_sqlite_session( assert response == {"ok": True} load_pipeline.assert_called_once_with(session, pipeline.id) assert generate.call_args.kwargs["session"] is session + assert generate.call_args.kwargs["generator"] is pipeline_application assert session.get_bind() is sqlite_engine @@ -418,9 +421,10 @@ def test_rag_pipeline_run_uses_sqlite_session( def test_rag_pipeline_run_translates_rate_limit( app: Flask, sqlite_engine: Engine, + pipeline_application: PipelineGenerator, api_type: type, ) -> None: - payload = { + payload: dict[str, object] = { "inputs": {}, "datasource_type": "x", "datasource_info_list": [], @@ -439,7 +443,10 @@ def test_rag_pipeline_run_translates_rate_limit( Session(sqlite_engine) as session, app.test_request_context("/", json=payload), patch.object(module, "load_rag_pipeline", return_value=pipeline), - patch.object(module.PipelineGenerateService, "generate", side_effect=InvokeRateLimitError("limit")), + patch.object(module.PipelineGenerateService, "generate", side_effect=InvokeRateLimitError("limit")) as generate, pytest.raises(InvokeRateLimitHttpError), ): handler(api, req_data, session, _account(), pipeline.id) + + assert generate.call_args.kwargs["session"] is session + assert generate.call_args.kwargs["generator"] is pipeline_application diff --git a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow_apis.py b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow_apis.py index fd26d0ec33f..9226d33433b 100644 --- a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow_apis.py +++ b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow_apis.py @@ -54,6 +54,9 @@ from models.enums import CreatorUserRole from models.workflow import Workflow, WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom from services.errors.app import IsDraftWorkflowError, WorkflowHashNotEqualError, WorkflowNotFoundError +pytestmark = pytest.mark.usefixtures("pipeline_application") + + DEFAULT_WORKFLOW_TENANT_ID = "00000000-0000-0000-0000-000000000001" DEFAULT_WORKFLOW_APP_ID = "00000000-0000-0000-0000-000000000002" DEFAULT_WORKFLOW_CREATED_BY = "00000000-0000-0000-0000-000000000003" diff --git a/api/tests/unit_tests/controllers/console/datasets/test_application_boundaries.py b/api/tests/unit_tests/controllers/console/datasets/test_application_boundaries.py new file mode 100644 index 00000000000..9a459f032c8 --- /dev/null +++ b/api/tests/unit_tests/controllers/console/datasets/test_application_boundaries.py @@ -0,0 +1,62 @@ +"""Guard all migrated dataset endpoints against transport/persistence regressions.""" + +import ast +from pathlib import Path + +import pytest + + +@pytest.mark.parametrize( + ("filename", "count"), + [("data_source.py", 7), ("datasets.py", 20), ("datasets_document.py", 24), ("external.py", 8)], +) +def test_every_endpoint_declares_admission_context_and_application_service(filename: str, count: int) -> None: + root = Path(__file__).resolve().parents[5] + tree = ast.parse((root / "controllers/console/datasets" / filename).read_text()) + # This namespace package is not traversed by import-linter's package graph. + forbidden = ( + "sqlalchemy", + "repositories", + "controllers.common.session", + "libs.login", + "services.knowledge.dataset_service", + "services.knowledge.external.service", + "services.hit_testing_service", + "services.knowledge.dataset_read_service", + "services.knowledge.datasets.adapters", + "services.knowledge.documents.adapters", + "services.knowledge.external.adapters", + ) + for node in ast.walk(tree): + modules = ( + [node.module or ""] + if isinstance(node, ast.ImportFrom) + else [item.name for item in node.names] + if isinstance(node, ast.Import) + else [] + ) + for module in modules: + assert not any(module == name or module.startswith(name + ".") for name in forbidden), module + endpoints = [ + method + for resource in tree.body + if isinstance(resource, ast.ClassDef) + for method in resource.body + if isinstance(method, ast.FunctionDef) and method.name in {"get", "post", "patch", "put", "delete"} + ] + assert len(endpoints) == count + for endpoint in endpoints: + assert any( + isinstance(d, ast.Call) and ast.unparse(d.func) == "console_account_admission" + for d in endpoint.decorator_list + ), endpoint.name + assert any( + arg.arg == "request_context" + and arg.annotation is not None + and ast.unparse(arg.annotation) == "RequestContext" + for arg in endpoint.args.args + ), endpoint.name + assert any( + isinstance(n, ast.Call) and ast.unparse(n.func) == "application_services" for n in ast.walk(endpoint) + ), endpoint.name + assert not {"session", "current_user", "current_tenant_id"}.intersection(arg.arg for arg in endpoint.args.args) 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 009c48f1ad5..b5463712cbf 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 @@ -1,302 +1,597 @@ from __future__ import annotations -import inspect -from collections.abc import Callable, Iterator +from collections.abc import Callable +from dataclasses import dataclass, field from datetime import UTC, datetime -from typing import Literal, cast -from unittest.mock import MagicMock, PropertyMock, patch -from uuid import UUID +from inspect import unwrap +from typing import TYPE_CHECKING, cast import pytest from flask import Flask -from sqlalchemy import select -from sqlalchemy.orm import Session -from werkzeug.exceptions import NotFound +from flask.testing import FlaskClient +from flask_restx import Api -from controllers.console.datasets import data_source as module -from controllers.console.datasets.data_source import DataSourceApi, DataSourceNotionListApi, DataSourceNotionListQuery -from models import Account, DataSourceOauthBinding -from models.dataset import Dataset -from models.engine import db -from models.enums import DataSourceType +from controllers.common.errors import NotFoundError +from controllers.common.rbac import DatasetId +from controllers.console import console_ns +from controllers.console.app.error import ProviderNotInitializeError -ControllerMethod = Callable[..., tuple[dict[str, object], int]] - - -def unwrap(func: object) -> ControllerMethod: - return cast(ControllerMethod, inspect.unwrap(cast(Callable[..., object], func))) - - -@pytest.fixture -def flask_app() -> Iterator[Flask]: - app = Flask(__name__) - app.config["TESTING"] = True - app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///:memory:" - db.init_app(app) - - with app.app_context(): - DataSourceOauthBinding.__table__.create(db.engine) - yield app - - -@pytest.fixture -def current_user() -> Account: - account = Account(name="Test User", email="user-1@example.com") - account.id = "user-1" - return account +# Importing the module registers only the resources under test on console_ns. +from controllers.console.datasets import data_source as data_source_controller +from controllers.console.datasets.error import IndexingEstimateError +from core.entities.knowledge_entities import IndexingEstimate +from core.rbac import RBACPermission +from enums import DeploymentEdition +from libs.login import AccountWithTenant +from machinery.context import RequestContext +from services.data_source.binding_application_service import ( + BindingMutationResult, + DataSourceBindingApplicationService, +) +from services.data_source.entities.notion_import import ( + AuthorizedNotionPage, + NotionPageIcon, + NotionPageType, + NotionWorkspace, +) +from services.data_source.entities.oauth import DataSourceBindingSummary +from services.data_source.notion_import_application_service import NotionImportApplicationService +from services.knowledge.dataset_access import DatasetAccessRecord, DatasetAccessService, DatasetAccessSnapshot +from services.knowledge.document_sync import DocumentSyncApplicationService, SyncDocumentRecord +from services.knowledge.entities.indexing_estimate import ( + NewSourcesEstimateCommand, + NotionEstimateSource, +) +from services.knowledge.indexing.estimate import ( + DatasetEstimateRecord, + EstimateDocumentRecord, + EstimateSourceNotFoundError, + IndexingEstimateCredentialUnavailableError, + IndexingEstimateExecutionError, + IndexingEstimateProviderUnavailableError, +) +from services.knowledge.resource_scope import DatasetRef, DocumentRef +from tests.unit_tests.controllers.rbac_introspection import rbac_checks +if TYPE_CHECKING: + from models import Account TENANT_ID = "11111111-1111-1111-1111-111111111111" -BINDING_ID = "22222222-2222-2222-2222-222222222222" +ACCOUNT_ID = "22222222-2222-2222-2222-222222222222" +BINDING_ID = "33333333-3333-3333-3333-333333333333" +DATASET_ID = "44444444-4444-4444-4444-444444444444" +DOCUMENT_ID = "55555555-5555-5555-5555-555555555555" +PAGE_ID = "66666666-6666-6666-6666-666666666666" +FOREIGN_DATASET_ID = "77777777-7777-7777-7777-777777777777" -def _add_binding(session: Session, *, disabled: bool) -> DataSourceOauthBinding: - binding = DataSourceOauthBinding( - tenant_id=TENANT_ID, - access_token="token", - provider="notion", - source_info={ - "workspace_name": "Workspace", - "workspace_id": "workspace-1", - "workspace_icon": None, - "total": 1, - "pages": [ - { - "page_id": "page-1", - "page_name": "Page", - "page_icon": {"type": "emoji", "emoji": "P", "url": None}, - "parent_id": "parent-1", - "type": "page", - } - ], - }, - disabled=disabled, +@dataclass +class RecordingBindingStore: + bindings: tuple[DataSourceBindingSummary, ...] = () + mutations: list[tuple[str, str, bool]] = field(default_factory=list) + mutation_result: BindingMutationResult = "updated" + + def list_enabled_bindings(self, *, workspace_id: str) -> tuple[DataSourceBindingSummary, ...]: + assert workspace_id == TENANT_ID + return self.bindings + + def change_disabled_state( + self, + *, + workspace_id: str, + binding_id: str, + disabled: bool, + ) -> BindingMutationResult: + self.mutations.append((workspace_id, binding_id, disabled)) + return self.mutation_result + + +@dataclass +class DatasetCatalog: + record: DatasetAccessRecord + + def get_access_snapshot( + self, + *, + workspace_id: str, + dataset_id: str, + actor_id: str, + ) -> DatasetAccessSnapshot | None: + assert actor_id == ACCOUNT_ID + if workspace_id != self.record.workspace_id or dataset_id != self.record.id: + return None + return DatasetAccessSnapshot(dataset=self.record, actor_has_partial_access=False) + + def get_estimate_record(self, dataset_ref: DatasetRef) -> DatasetEstimateRecord | None: + if dataset_ref != DatasetRef(self.record.workspace_id, self.record.id): + return None + return DatasetEstimateRecord( + id=self.record.id, + workspace_id=self.record.workspace_id, + indexing_technique="economy", + ) + + def is_notion_dataset(self, dataset_ref: DatasetRef) -> bool: + return dataset_ref == DatasetRef(self.record.workspace_id, self.record.id) + + +class OwnerWorkspaceRoles: + def get_legacy_role(self, *, workspace_id: str, account_id: str) -> str | None: + assert workspace_id == TENANT_ID + assert account_id == ACCOUNT_ID + return "owner" + + +@dataclass +class KnowledgeDocuments: + bound_page_ids: frozenset[str] = frozenset() + active_refs: tuple[DocumentRef, ...] = () + + def list_bound_notion_page_ids(self, dataset_ref: DatasetRef) -> frozenset[str]: + assert dataset_ref == DatasetRef(TENANT_ID, DATASET_ID) + return self.bound_page_ids + + def list_active_notion_refs(self, dataset_ref: DatasetRef) -> tuple[DocumentRef, ...]: + assert dataset_ref == DatasetRef(TENANT_ID, DATASET_ID) + return self.active_refs + + def get_estimate_document(self, document_ref: DocumentRef) -> EstimateDocumentRecord | None: + if document_ref.document_id != DOCUMENT_ID: + return None + return EstimateDocumentRecord( + id=DOCUMENT_ID, + workspace_id=TENANT_ID, + dataset_id=DATASET_ID, + data_source_type="notion_import", + data_source_info=None, + indexing_status="waiting", + doc_form="text_model", + doc_language="English", + dataset_process_rule_id=None, + ) + + def list_estimate_documents_by_batch( + self, dataset_ref: DatasetRef, batch: str + ) -> tuple[EstimateDocumentRecord, ...]: + document = self.get_estimate_document(dataset_ref.document(DOCUMENT_ID)) + if document is None or batch != "batch": + return () + return (document,) + + def get_sync_document(self, document_ref: DocumentRef) -> SyncDocumentRecord | None: + if document_ref.document_id != DOCUMENT_ID: + return None + return SyncDocumentRecord(id=DOCUMENT_ID, data_source_type="notion_import") + + +@dataclass +class RecordingNotionSource: + workspaces: tuple[NotionWorkspace, ...] + list_calls: list[tuple[str, str, str]] = field(default_factory=list) + preview_calls: list[tuple[str, str, str, str, str]] = field(default_factory=list) + + def list_authorized_pages( + self, + *, + workspace_id: str, + actor_id: str, + credential_id: str, + ) -> tuple[NotionWorkspace, ...]: + self.list_calls.append((workspace_id, actor_id, credential_id)) + return self.workspaces + + def preview_page( + self, + *, + workspace_id: str, + actor_id: str, + credential_id: str, + page_id: str, + page_type: str, + ) -> str: + self.preview_calls.append((workspace_id, actor_id, credential_id, page_id, page_type)) + return "preview text" + + +@dataclass +class RecordingEstimateService: + calls: list[tuple[RequestContext, NewSourcesEstimateCommand]] = field(default_factory=list) + error: Exception | None = None + + def estimate_new_sources(self, context: RequestContext, command: NewSourcesEstimateCommand) -> IndexingEstimate: + self.calls.append((context, command)) + if self.error is not None: + raise self.error + return IndexingEstimate(total_segments=3, preview=[]) + + +@dataclass +class RecordingDispatcher: + calls: list[tuple[str, str]] = field(default_factory=list) + + def dispatch(self, *, dataset_id: str, document_id: str) -> None: + self.calls.append((dataset_id, document_id)) + + +@dataclass(frozen=True) +class DataSourceApplicationServices: + bindings: DataSourceBindingApplicationService + notion_imports: NotionImportApplicationService + + +@dataclass(frozen=True) +class KnowledgeApplicationServices: + document_sync: DocumentSyncApplicationService + indexing_estimates: RecordingEstimateService + + +@dataclass(frozen=True) +class ApplicationServiceRegistry: + data_sources: DataSourceApplicationServices + knowledge: KnowledgeApplicationServices + + +@dataclass(frozen=True) +class ControllerCollaborators: + client: FlaskClient + account: AdmittedAccount + bindings: RecordingBindingStore + notion_source: RecordingNotionSource + estimates: RecordingEstimateService + dispatcher: RecordingDispatcher + + +@dataclass +class AdmittedAccount: + id: str + role: str + status: str + + +@pytest.fixture +def collaborators( + config_overrides: Callable[..., None], + monkeypatch: pytest.MonkeyPatch, +) -> ControllerCollaborators: + config_overrides( + DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, + LOGIN_DISABLED=True, + RBAC_ENABLED=False, + CONSOLE_API_URL="https://console-api.example.com", + ) + app = Flask(__name__) + app.config.update(TESTING=False, PROPAGATE_EXCEPTIONS=False) + Api(app).add_namespace(console_ns, path="/console/api") + + account = AdmittedAccount(id=ACCOUNT_ID, role="owner", status="active") + account_with_tenant = AccountWithTenant(account=cast("Account", account), tenant_id=TENANT_ID) + monkeypatch.setattr( + "controllers.console.flask_admission.current_account_with_tenant", + lambda: account_with_tenant, + ) + monkeypatch.setattr( + "controllers.console.wraps.current_account_with_tenant", + lambda: account_with_tenant, + ) + + dataset = DatasetAccessRecord( + id=DATASET_ID, + workspace_id=TENANT_ID, + maintainer_id=ACCOUNT_ID, + permission="only_me", + ) + dataset_catalog = DatasetCatalog(dataset) + dataset_access = DatasetAccessService( + datasets=dataset_catalog, + workspace_roles=OwnerWorkspaceRoles(), + legacy_permissions_enabled=True, + ) + documents = KnowledgeDocuments( + bound_page_ids=frozenset({PAGE_ID}), + active_refs=(DatasetRef(TENANT_ID, DATASET_ID).document(DOCUMENT_ID),), + ) + bindings = RecordingBindingStore() + notion_source = RecordingNotionSource( + workspaces=( + NotionWorkspace( + workspace_id="notion-workspace-1", + workspace_name="One", + workspace_icon=None, + pages=( + AuthorizedNotionPage( + page_id=PAGE_ID, + page_name="Page one", + page_icon=NotionPageIcon(type="emoji", emoji="📄"), + parent_id=None, + page_type=NotionPageType.PAGE, + ), + ), + ), + NotionWorkspace( + workspace_id="notion-workspace-2", + workspace_name="Two", + workspace_icon=None, + pages=(), + ), + ) + ) + estimates = RecordingEstimateService() + dispatcher = RecordingDispatcher() + app.extensions["application_services"] = ApplicationServiceRegistry( + data_sources=DataSourceApplicationServices( + bindings=DataSourceBindingApplicationService(bindings=bindings), + notion_imports=NotionImportApplicationService( + dataset_access=dataset_access, + datasets=dataset_catalog, + documents=documents, + source=notion_source, + ), + ), + knowledge=KnowledgeApplicationServices( + document_sync=DocumentSyncApplicationService( + dataset_access=dataset_access, + documents=documents, + dispatcher=dispatcher, + ), + indexing_estimates=estimates, + ), + ) + return ControllerCollaborators( + client=app.test_client(), + account=account, + bindings=bindings, + notion_source=notion_source, + estimates=estimates, + dispatcher=dispatcher, ) - binding.id = BINDING_ID - binding.created_at = datetime(2026, 5, 25, 1, 2, 3, tzinfo=UTC) - session.add(binding) - session.commit() - return binding -def test_get_data_source_integrates_serializes_orm_binding( - flask_app: Flask, -) -> None: - binding = _add_binding(db.session, disabled=False) - expected_created_at = int(binding.created_at.timestamp()) - - with flask_app.test_request_context("/"): - response, status = unwrap(DataSourceApi().get)(DataSourceApi(), TENANT_ID) - - assert status == 200 - assert response == { - "data": [ - { - "id": BINDING_ID, - "provider": "notion", - "created_at": expected_created_at, - "is_bound": True, - "disabled": False, - "source_info": { - "workspace_name": "Workspace", - "workspace_id": "workspace-1", - "workspace_icon": None, - "pages": [ - { - "page_name": "Page", - "page_id": "page-1", - "page_icon": {"type": "emoji", "url": None, "emoji": "P"}, - "parent_id": "parent-1", - "type": "page", - } - ], - "total": 1, - }, - "link": "http://localhost/console/api/oauth/data-source/notion", - } - ] - } - - -def test_get_data_source_integrates_preserves_empty_list_when_no_binding( - flask_app: Flask, -) -> None: - with flask_app.test_request_context("/"): - response, status = unwrap(DataSourceApi().get)(DataSourceApi(), TENANT_ID) - - assert status == 200 - assert response == {"data": []} +def test_integration_collection_only_exposes_get(collaborators: ControllerCollaborators) -> None: + client = collaborators.client + assert client.get("/console/api/data-source/integrates").status_code == 200 + assert client.patch("/console/api/data-source/integrates").status_code == 405 @pytest.mark.parametrize( - ("disabled", "action", "expected_disabled"), - [(True, "enable", False), (False, "disable", True)], + ("method", "path"), + [ + ("get", "/console/api/data-source/integrates"), + ("patch", f"/console/api/data-source/integrates/{BINDING_ID}/enable"), + ], ) -@pytest.mark.parametrize("sqlite_session", [(DataSourceOauthBinding,)], indirect=True) -def test_patch_data_source_binding_updates_state( - flask_app: Flask, - sqlite_session: Session, +def test_integration_management_rejects_non_admin_role( + collaborators: ControllerCollaborators, + method: str, + path: str, +) -> None: + collaborators.account.role = "editor" + + assert collaborators.client.open(path, method=method.upper()).status_code == 403 + + +def test_integration_management_enforces_workspace_credential_permission( + collaborators: ControllerCollaborators, + config_overrides: Callable[..., None], + monkeypatch: pytest.MonkeyPatch, +) -> None: + config_overrides(RBAC_ENABLED=True) + calls: list[dict[str, object]] = [] + + def record_access(tenant_id: str, account_id: str, **requirement: object) -> bool: + calls.append({"tenant_id": tenant_id, "account_id": account_id, **requirement}) + return True + + monkeypatch.setattr("controllers.common.rbac.checks.RBACService.CheckAccess.check", record_access) + + list_response = collaborators.client.get("/console/api/data-source/integrates") + mutation_response = collaborators.client.patch(f"/console/api/data-source/integrates/{BINDING_ID}/enable") + + assert list_response.status_code == 200 + assert mutation_response.status_code == 200 + assert calls == [ + { + "tenant_id": TENANT_ID, + "account_id": ACCOUNT_ID, + "scene": RBACPermission.CREDENTIAL_MANAGE, + "resource_type": None, + "resource_id": None, + }, + { + "tenant_id": TENANT_ID, + "account_id": ACCOUNT_ID, + "scene": RBACPermission.CREDENTIAL_MANAGE, + "resource_type": None, + "resource_id": None, + }, + ] + + +@pytest.mark.parametrize(("action", "disabled"), [("enable", False), ("disable", True)]) +def test_integration_item_only_exposes_patch_and_rejects_unknown_action( + collaborators: ControllerCollaborators, + action: str, disabled: bool, - action: Literal["enable", "disable"], - expected_disabled: bool, ) -> None: - _add_binding(sqlite_session, disabled=disabled) - sqlite_session.expunge_all() - - with flask_app.test_request_context("/"): - response, status = unwrap(DataSourceApi().patch)( - DataSourceApi(), sqlite_session, TENANT_ID, UUID(BINDING_ID), action - ) - - sqlite_session.flush() - sqlite_session.expire_all() - binding = sqlite_session.scalar(select(DataSourceOauthBinding).where(DataSourceOauthBinding.id == BINDING_ID)) - assert status == 200 - assert response == {"result": "success"} - assert binding is not None - assert binding.disabled is expected_disabled + client = collaborators.client + item_path = f"/console/api/data-source/integrates/{BINDING_ID}/{action}" + assert client.get(item_path).status_code == 405 + response = client.patch(item_path) + assert response.status_code == 200 + assert response.get_json() == {"result": "success"} + assert collaborators.bindings.mutations == [(TENANT_ID, BINDING_ID, disabled)] + invalid = client.patch(f"/console/api/data-source/integrates/{BINDING_ID}/remove") + assert invalid.status_code == 400 + assert invalid.get_json()["code"] == "invalid_action" + assert len(collaborators.bindings.mutations) == 1 -@pytest.mark.parametrize("sqlite_session", [(DataSourceOauthBinding,)], indirect=True) -def test_patch_data_source_binding_rejects_unknown_binding( - flask_app: Flask, - sqlite_session: Session, +def test_integration_list_serializes_application_result_and_http_link( + collaborators: ControllerCollaborators, ) -> None: - with flask_app.test_request_context("/"), pytest.raises(NotFound, match="Data source binding not found"): - unwrap(DataSourceApi().patch)(DataSourceApi(), sqlite_session, TENANT_ID, UUID(BINDING_ID), "enable") - - -@pytest.mark.parametrize(("disabled", "action"), [(False, "enable"), (True, "disable")]) -@pytest.mark.parametrize("sqlite_session", [(DataSourceOauthBinding,)], indirect=True) -def test_patch_data_source_binding_rejects_current_state( - flask_app: Flask, - sqlite_session: Session, - disabled: bool, - action: Literal["enable", "disable"], -) -> None: - _add_binding(sqlite_session, disabled=disabled) - sqlite_session.expunge_all() - - with flask_app.test_request_context("/"), pytest.raises(ValueError): - unwrap(DataSourceApi().patch)(DataSourceApi(), sqlite_session, TENANT_ID, UUID(BINDING_ID), action) - - -@pytest.mark.parametrize("sqlite_session", [()], indirect=True) -def test_notion_pre_import_pages_serializes_frontend_list_shape( - flask_app: Flask, - current_user: Account, - sqlite_session: Session, -) -> None: - page = MagicMock( - page_id="page-1", - page_name="Page", - type="page", - parent_id="parent-1", - page_icon={"type": "emoji", "emoji": "P", "url": None}, - ) - online_document_message = MagicMock( - result=[ - MagicMock( - workspace_id="workspace-1", - workspace_name="Workspace", - workspace_icon=None, - pages=[page], - ) - ] - ) - runtime = MagicMock( - get_online_document_pages=MagicMock(return_value=iter([online_document_message])), - datasource_provider_type=MagicMock(return_value="online_document"), - ) - with ( - flask_app.test_request_context("/?credential_id=credential-1"), - patch.object( - module.DatasourceProviderService, - "get_datasource_credentials", - return_value={"token": "token"}, - ), - patch.object(type(module.db), "engine", new_callable=PropertyMock, return_value=MagicMock()), - patch("core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", return_value=runtime), - ): - response, status = unwrap(DataSourceNotionListApi().get)( - DataSourceNotionListApi(), - DataSourceNotionListQuery(credential_id="credential-1"), - sqlite_session, - "tenant-1", - current_user, - ) - - assert status == 200 - assert response == { - "notion_info": [ - { + collaborators.bindings.bindings = ( + DataSourceBindingSummary( + id=BINDING_ID, + provider="notion", + created_at=datetime(2026, 5, 25, 1, 2, 3, tzinfo=UTC), + disabled=False, + source_info={ "workspace_name": "Workspace", "workspace_id": "workspace-1", "workspace_icon": None, - "pages": [ - { - "page_name": "Page", - "page_id": "page-1", - "page_icon": {"type": "emoji", "url": None, "emoji": "P"}, - "parent_id": "parent-1", - "type": "page", - "is_bound": False, - } - ], - } - ] - } - runtime.get_online_document_pages.assert_called_once() - assert runtime.get_online_document_pages.call_args.kwargs["datasource_parameters"] == {} - - -@pytest.mark.parametrize("sqlite_session", [()], indirect=True) -def test_notion_pre_import_pages_rejects_missing_credential( - flask_app: Flask, - current_user: Account, - sqlite_session: Session, -) -> None: - with ( - flask_app.test_request_context("/?credential_id=credential-1"), - patch.object(module.DatasourceProviderService, "get_datasource_credentials", return_value=None), - pytest.raises(NotFound, match="Credential not found"), - ): - unwrap(DataSourceNotionListApi().get)( - DataSourceNotionListApi(), - DataSourceNotionListQuery(credential_id="credential-1"), - sqlite_session, - TENANT_ID, - current_user, - ) - - -@pytest.mark.parametrize("sqlite_session", [()], indirect=True) -def test_notion_pre_import_pages_rejects_non_notion_dataset( - flask_app: Flask, - current_user: Account, - sqlite_session: Session, -) -> None: - dataset = Dataset( - tenant_id=TENANT_ID, - name="Non-Notion dataset", - created_by=current_user.id, - data_source_type=DataSourceType.WEBSITE_CRAWL, + "pages": [], + "total": 0, + }, + ), ) - with ( - flask_app.test_request_context("/?credential_id=credential-1&dataset_id=dataset-1"), - patch.object( - module.DatasourceProviderService, - "get_datasource_credentials", - return_value={"token": "token"}, - ), - patch.object(module.DatasetService, "get_dataset", return_value=dataset), - pytest.raises(ValueError, match="Dataset is not notion type"), - ): - unwrap(DataSourceNotionListApi().get)( - DataSourceNotionListApi(), - DataSourceNotionListQuery(credential_id="credential-1", dataset_id="dataset-1"), - sqlite_session, - TENANT_ID, - current_user, - ) + response = collaborators.client.get("/console/api/data-source/integrates") + + assert response.status_code == 200 + assert response.get_json()["data"][0] == { + "id": BINDING_ID, + "provider": "notion", + "created_at": 1779670923, + "is_bound": True, + "disabled": False, + "source_info": { + "workspace_name": "Workspace", + "workspace_id": "workspace-1", + "workspace_icon": None, + "pages": [], + "total": 0, + }, + "link": "https://console-api.example.com/console/api/oauth/data-source/notion", + } + + +def test_notion_list_preserves_workspace_groups_and_bound_state(collaborators: ControllerCollaborators) -> None: + response = collaborators.client.get( + "/console/api/notion/pre-import/pages", + query_string={"credential_id": "credential-1", "dataset_id": DATASET_ID}, + ) + + assert response.status_code == 200 + payload = response.get_json()["notion_info"] + assert [workspace["workspace_id"] for workspace in payload] == [ + "notion-workspace-1", + "notion-workspace-2", + ] + assert payload[0]["pages"][0]["is_bound"] is True + assert payload[0]["pages"][0]["type"] == "page" + assert "page_type" not in payload[0]["pages"][0] + assert collaborators.notion_source.list_calls == [(TENANT_ID, ACCOUNT_ID, "credential-1")] + + +@pytest.mark.parametrize("page_type", ["page", "database"]) +def test_notion_preview_validates_page_type_and_delegates( + collaborators: ControllerCollaborators, page_type: str +) -> None: + response = collaborators.client.get( + f"/console/api/notion/pages/{PAGE_ID}/{page_type}/preview", + query_string={"credential_id": "credential-1"}, + ) + invalid = collaborators.client.get( + f"/console/api/notion/pages/{PAGE_ID}/collection/preview", + query_string={"credential_id": "credential-1"}, + ) + + assert response.status_code == 200 + assert response.get_json() == {"content": "preview text"} + assert invalid.status_code == 400 + assert invalid.get_json()["code"] == "invalid_param" + assert collaborators.notion_source.preview_calls == [(TENANT_ID, ACCOUNT_ID, "credential-1", PAGE_ID, page_type)] + + +def test_notion_estimate_maps_nested_payload_to_shared_command(collaborators: ControllerCollaborators) -> None: + response = collaborators.client.post( + "/console/api/datasets/notion-indexing-estimate", + json={ + "notion_info_list": [ + { + "workspace_id": "notion-workspace-1", + "credential_id": "credential-1", + "pages": [{"page_id": PAGE_ID, "type": "page"}], + } + ], + "process_rule": {"mode": "automatic"}, + "doc_form": "text_model", + "doc_language": "English", + }, + ) + + assert response.status_code == 200 + assert response.get_json()["total_segments"] == 3 + context, command = collaborators.estimates.calls[0] + assert context.active_workspace_id == TENANT_ID + assert context.account_id == ACCOUNT_ID + assert command.process_rule == {"mode": "automatic"} + source = command.sources[0] + assert isinstance(source, NotionEstimateSource) + assert source.credential_id == "credential-1" + assert source.page_id == PAGE_ID + + +@pytest.mark.parametrize( + ("error", "expected_http_error"), + [ + (IndexingEstimateCredentialUnavailableError(), NotFoundError), + (EstimateSourceNotFoundError("page-1"), NotFoundError), + (IndexingEstimateProviderUnavailableError(), ProviderNotInitializeError), + (IndexingEstimateExecutionError(), IndexingEstimateError), + ], +) +def test_notion_estimate_maps_application_errors( + error: Exception, + expected_http_error: type[Exception], + collaborators: ControllerCollaborators, +) -> None: + collaborators.estimates.error = error + api = data_source_controller.DataSourceNotionIndexingEstimateApi() + method = unwrap(api.post) + payload = data_source_controller.NotionEstimatePayload.model_validate( + { + "notion_info_list": [ + { + "workspace_id": "notion-workspace", + "credential_id": "credential-1", + "pages": [{"page_id": "page-1", "type": "page"}], + } + ], + "process_rule": {"mode": "automatic"}, + } + ) + context = RequestContext("request-1", None, "account-1", "workspace-1") + + with collaborators.client.application.app_context(), pytest.raises(expected_http_error): + method(api, payload, context) + + +def test_sync_routes_dispatch_only_application_selected_documents(collaborators: ControllerCollaborators) -> None: + dataset_response = collaborators.client.get(f"/console/api/datasets/{DATASET_ID}/notion/sync") + document_response = collaborators.client.get( + f"/console/api/datasets/{DATASET_ID}/documents/{DOCUMENT_ID}/notion/sync" + ) + + assert dataset_response.status_code == 200 + assert dataset_response.get_json() == {"result": "success"} + assert document_response.status_code == 200 + assert document_response.get_json() == {"result": "success"} + assert collaborators.dispatcher.calls == [(DATASET_ID, DOCUMENT_ID), (DATASET_ID, DOCUMENT_ID)] + + +def test_sync_route_does_not_dispatch_cross_workspace_dataset(collaborators: ControllerCollaborators) -> None: + response = collaborators.client.get(f"/console/api/datasets/{FOREIGN_DATASET_ID}/notion/sync") + + assert response.status_code == 404 + assert response.get_json()["code"] == "not_found" + assert collaborators.dispatcher.calls == [] + + +@pytest.mark.parametrize( + "resource", + [data_source_controller.DataSourceNotionDatasetSyncApi, data_source_controller.DataSourceNotionDocumentSyncApi], +) +def test_notion_sync_requires_dataset_edit_permission(resource): + [check] = rbac_checks(resource.get) + assert check.scene is RBACPermission.DATASET_EDIT + assert isinstance(check.locator, DatasetId) 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 deleted file mode 100644 index 8bbea58ba8e..00000000000 --- a/api/tests/unit_tests/controllers/console/datasets/test_data_source_notion_apis.py +++ /dev/null @@ -1,204 +0,0 @@ -"""Unit tests for controllers.console.datasets.data_source Notion endpoints.""" - -from __future__ import annotations - -import inspect -from unittest.mock import MagicMock, patch - -import pytest -from flask import Flask -from sqlalchemy.orm import Session -from werkzeug.exceptions import NotFound - -from controllers.common.rbac import DatasetId -from controllers.console.datasets.data_source import ( - DataSourceNotionDatasetSyncApi, - DataSourceNotionDocumentSyncApi, - DataSourceNotionIndexingEstimateApi, - DataSourceNotionPreviewApi, - DataSourceNotionPreviewQuery, - NotionEstimatePayload, -) -from controllers.console.wraps import RBACPermission -from core.rag.index_processor.constant.index_type import IndexStructureType -from models import Account, Dataset, Document -from models.dataset import DataSourceType -from tests.unit_tests.controllers.rbac_introspection import rbac_checks -from tests.unit_tests.model_factories import make_dataset, make_document - - -def _dataset() -> Dataset: - return make_dataset(dataset_id="ds-1", created_by="u1") - - -def _document() -> Document: - return make_document( - document_id="d1", - dataset_id="ds-1", - data_source_type=DataSourceType.NOTION_IMPORT, - name="Notion page", - created_by="u1", - ) - - -@pytest.fixture -def current_user() -> Account: - account = Account(name="Test User", email="u1@example.com") - account.id = "u1" - return account - - -class TestDataSourceNotionPreviewApi: - def test_get_preview_success(self, app: Flask) -> None: - api = DataSourceNotionPreviewApi() - method = inspect.unwrap(api.get) - - extractor = MagicMock(extract=lambda: [MagicMock(page_content="hello")]) - - with ( - app.test_request_context("/?credential_id=c1"), - patch( - "controllers.console.datasets.data_source.DatasourceProviderService.get_datasource_credentials", - return_value={"integration_secret": "t"}, - ), - patch( - "controllers.console.datasets.data_source.NotionExtractor", - return_value=extractor, - ), - ): - response, status = method(api, DataSourceNotionPreviewQuery(credential_id="c1"), "tenant-1", "p1", "page") - - assert status == 200 - - -class TestDataSourceNotionIndexingEstimateApi: - @pytest.mark.parametrize("sqlite_session", [()], indirect=True) - def test_post_indexing_estimate_success(self, app: Flask, sqlite_session: Session) -> None: - api = DataSourceNotionIndexingEstimateApi() - method = inspect.unwrap(api.post) - - empty_rules: dict[str, object] = {} - payload: dict[str, object] = { - "notion_info_list": [ - { - "workspace_id": "w1", - "credential_id": "c1", - "pages": [{"page_id": "p1", "type": "page"}], - } - ], - "process_rule": {"rules": empty_rules}, - "doc_form": IndexStructureType.PARAGRAPH_INDEX, - "doc_language": "English", - } - - with ( - app.test_request_context("/", method="POST", json=payload, headers={"Content-Type": "application/json"}), - patch( - "controllers.console.datasets.data_source.DocumentService.estimate_args_validate", - ), - patch( - "controllers.console.datasets.data_source.IndexingRunner.indexing_estimate", - return_value=MagicMock(model_dump=lambda: {"total_pages": 1}), - ), - ): - response, status = method(api, NotionEstimatePayload.model_validate(payload), sqlite_session, "tenant-1") - - assert status == 200 - - -class TestDataSourceNotionDatasetSyncApi: - def test_get_requires_dataset_edit_permission(self) -> None: - [check] = rbac_checks(DataSourceNotionDatasetSyncApi.get) - - assert check.scene is RBACPermission.DATASET_EDIT - assert isinstance(check.locator, DatasetId) - - @pytest.mark.parametrize("sqlite_session", [()], indirect=True) - def test_get_success(self, app: Flask, sqlite_session: Session) -> None: - api = DataSourceNotionDatasetSyncApi() - method = inspect.unwrap(api.get) - - with ( - app.test_request_context("/"), - patch( - "controllers.console.datasets.data_source.DatasetService.get_dataset", - return_value=_dataset(), - ), - patch( - "controllers.console.datasets.data_source.DocumentService.get_document_by_dataset_id", - return_value=[_document()], - ), - patch( - "controllers.console.datasets.data_source.document_indexing_sync_task.delay", - return_value=None, - ), - ): - response, status = method(api, sqlite_session, "ds-1") - - assert status == 200 - - @pytest.mark.parametrize("sqlite_session", [()], indirect=True) - def test_get_dataset_not_found(self, app: Flask, sqlite_session: Session) -> None: - api = DataSourceNotionDatasetSyncApi() - method = inspect.unwrap(api.get) - - with ( - app.test_request_context("/"), - patch( - "controllers.console.datasets.data_source.DatasetService.get_dataset", - return_value=None, - ), - ): - with pytest.raises(NotFound): - method(api, sqlite_session, "ds-1") - - -class TestDataSourceNotionDocumentSyncApi: - def test_get_requires_dataset_edit_permission(self) -> None: - [check] = rbac_checks(DataSourceNotionDocumentSyncApi.get) - - assert check.scene is RBACPermission.DATASET_EDIT - assert isinstance(check.locator, DatasetId) - - @pytest.mark.parametrize("sqlite_session", [()], indirect=True) - def test_get_success(self, app: Flask, sqlite_session: Session) -> None: - api = DataSourceNotionDocumentSyncApi() - method = inspect.unwrap(api.get) - - with ( - app.test_request_context("/"), - patch( - "controllers.console.datasets.data_source.DatasetService.get_dataset", - return_value=_dataset(), - ), - patch( - "controllers.console.datasets.data_source.DocumentService.get_document", - return_value=_document(), - ), - patch( - "controllers.console.datasets.data_source.document_indexing_sync_task.delay", - return_value=None, - ), - ): - response, status = method(api, sqlite_session, "ds-1", "doc-1") - - assert status == 200 - - @pytest.mark.parametrize("sqlite_session", [()], indirect=True) - def test_get_document_not_found(self, app: Flask, sqlite_session: Session) -> None: - api = DataSourceNotionDocumentSyncApi() - method = inspect.unwrap(api.get) - - with ( - app.test_request_context("/"), - patch( - "controllers.console.datasets.data_source.DatasetService.get_dataset", - return_value=_dataset(), - ), - patch( - "controllers.console.datasets.data_source.DocumentService.get_document", - return_value=None, - ), - ): - with pytest.raises(NotFound): - method(api, sqlite_session, "ds-1", "doc-1") 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 81ea5da31d4..cf4a06e98e5 100644 --- a/api/tests/unit_tests/controllers/console/datasets/test_datasets.py +++ b/api/tests/unit_tests/controllers/console/datasets/test_datasets.py @@ -1,29 +1,20 @@ -import datetime -import json -from collections.abc import Callable -from contextlib import ExitStack from inspect import unwrap from types import SimpleNamespace -from unittest.mock import ANY, MagicMock, PropertyMock, patch +from unittest.mock import MagicMock, create_autospec, patch +from uuid import UUID import pytest -from flask import Flask -from sqlalchemy import event, select -from sqlalchemy.orm import Session, sessionmaker -from werkzeug.exceptions import BadRequest, Forbidden, NotFound +from werkzeug.exceptions import Forbidden, NotFound import services -from controllers.common.rbac import DatasetId, Workspace -from controllers.console import console_ns +from controllers.common.errors import InvalidArgumentError, NotFoundError +from controllers.common.rbac import DatasetId, RBACPermission, Workspace from controllers.console.app.error import ProviderNotInitializeError -from controllers.console.datasets import datasets as datasets_controller +from controllers.console.datasets import datasets as controller from controllers.console.datasets.datasets import ( DatasetApi, DatasetApiBaseUrlApi, - DatasetApiDeleteApi, - DatasetApiKeyApi, DatasetAutoDisableLogApi, - DatasetCreatePayload, DatasetEnableApiApi, DatasetErrorDocs, DatasetIndexingEstimateApi, @@ -37,119 +28,34 @@ from controllers.console.datasets.datasets import ( DatasetUpdatePayload, DatasetUseCheckApi, IndexingEstimatePayload, - _get_retrieval_methods_by_vector_type, + _new_estimate_sources, +) +from controllers.console.datasets.error import ( + DatasetAccessDeniedRequestError, + DatasetInUseError, + DatasetNameDuplicateError, + IndexingEstimateError, ) -from controllers.console.datasets.error import DatasetInUseError, DatasetNameDuplicateError, IndexingEstimateError -from controllers.console.wraps import RBACPermission -from core.entities.knowledge_entities import IndexingEstimate -from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError -from core.provider_manager import ProviderManager -from core.rag.datasource.vdb.vector_type import VectorType -from core.rag.index_processor.constant.index_type import IndexStructureType -from core.rag.retrieval.retrieval_methods import RetrievalMethod -from extensions.application_services.knowledge import build_dataset_api_key_service -from extensions.storage.storage_type import StorageType -from fields.dataset_fields import build_dataset_detail_prefetch from machinery.context import RequestContext -from models.account import Account, Tenant, TenantAccountRole -from models.dataset import AppDatasetJoin, Dataset, DatasetPermission, DatasetQuery, Document, DocumentSegment -from models.enums import CreatorUserRole, DataSourceType, DocumentCreatedFrom, IndexingStatus, SegmentStatus -from models.model import ApiToken, AppMode, DatasetApiTokenBinding, IconType, UploadFile -from services.dataset_ref_service import DatasetRef -from services.dataset_service import DatasetPermissionService, DatasetService -from services.enterprise import rbac_service as enterprise_rbac_service -from services.entities.app_entities import AppRecord +from services.data_source.entities.notion_import import NotionPageType +from services.knowledge.dataset_access import DatasetAccessDeniedError, DatasetNotFoundError +from services.knowledge.datasets.application import DatasetApplicationService, DatasetListFilter +from services.knowledge.entities.indexing_estimate import ( + NotionEstimateSource, + UploadFileEstimateSource, + WebsiteEstimateSource, +) +from services.knowledge.indexing.estimate import ( + EstimateSourceNotFoundError, + IndexingEstimateCredentialUnavailableError, + IndexingEstimateExecutionError, + IndexingEstimateProviderUnavailableError, + UnsupportedEstimateSourceError, +) from tests.unit_tests.controllers.rbac_introspection import rbac_checks -pytestmark = pytest.mark.usefixtures("app_query_services") - - -@pytest.fixture(autouse=True) -def dataset_model_property_defaults(): - getter_values: dict[str, object] = { - "get_app_count": 0, - "get_document_count": 0, - "get_word_count": 0, - "get_author_name": None, - "get_tags": [], - "get_doc_form": None, - "get_external_knowledge_info": None, - "get_doc_metadata": [], - "get_is_published": False, - "get_total_documents": 0, - "get_total_available_documents": 0, - } - getters = {} - with ExitStack() as stack: - for name, value in getter_values.items(): - getters[name] = stack.enter_context(patch.object(Dataset, name, autospec=True, return_value=value)) - stack.enter_context( - patch( - "controllers.console.datasets.datasets.enterprise_rbac_service.RBACService.MyPermissions.get", - return_value=enterprise_rbac_service.MyPermissionsResponse(), - ) - ) - stack.enter_context( - patch( - "controllers.console.datasets.datasets.enterprise_rbac_service.RBACService.DatasetPermissions.batch_get", - return_value={}, - ) - ) - yield getters - - -def make_dataset(**overrides) -> Dataset: - base = { - "id": "ds-1", - "tenant_id": "tenant-1", - "name": "Dataset", - "description": "desc", - "provider": "vendor", - "permission": "only_me", - "data_source_type": None, - "indexing_technique": "economy", - "created_by": "account-1", - "created_at": datetime.datetime(2024, 1, 1, 12, 0, 0, tzinfo=datetime.UTC), - "updated_by": None, - "updated_at": datetime.datetime(2024, 1, 1, 12, 0, 0, tzinfo=datetime.UTC), - "embedding_model": None, - "embedding_model_provider": None, - "retrieval_model": None, - "summary_index_setting": None, - "built_in_field_enabled": False, - "pipeline_id": None, - "runtime_mode": "general", - "chunk_structure": None, - "icon_info": None, - "enable_api": False, - "is_multimodal": False, - } - base.update(overrides) - return Dataset(**base) - - -def make_account(role: TenantAccountRole = TenantAccountRole.EDITOR) -> Account: - account = Account(name="Test User", email="user@example.com") - account.id = "account-1" - account.role = role - return account - - -def make_related_app(**overrides) -> AppRecord: - base = { - "id": "app-1", - "name": "App", - "description": "desc", - "mode_compatible_with_agent": AppMode.CHAT, - "icon_type": IconType.EMOJI, - "icon": "🤖", - "icon_background": "#fff", - "enable_site": False, - "enable_api": False, - "created_by": "account-1", - } - base.update(overrides) - return AppRecord(**base) +CONTEXT = RequestContext("request-1", None, "account-1", "tenant-1") +DATASET_ID = UUID(int=1) def test_dataset_delete_requires_dataset_delete_permission() -> None: @@ -159,1714 +65,262 @@ def test_dataset_delete_requires_dataset_delete_permission() -> None: assert isinstance(check.locator, DatasetId) -def make_document_status(**overrides) -> Document: - base = { - "id": "doc-1", - "tenant_id": "tenant-1", - "dataset_id": "dataset-1", - "position": 1, - "data_source_type": DataSourceType.UPLOAD_FILE, - "batch": "batch-1", - "name": "doc.txt", - "created_from": DocumentCreatedFrom.WEB, - "created_by": "account-1", - "indexing_status": IndexingStatus.COMPLETED, - "enabled": True, - "archived": False, - "processing_started_at": None, - "parsing_completed_at": None, - "cleaning_completed_at": None, - "splitting_completed_at": None, - "completed_at": None, - "paused_at": None, - "error": None, - "stopped_at": None, - } - base.update(overrides) - return Document(**base) - - -def make_document_segment( - *, - position: int, - completed: bool, - document_id: str = "doc-1", - tenant_id: str = "tenant-1", - dataset_id: str = "dataset-1", - status: SegmentStatus = SegmentStatus.WAITING, -) -> DocumentSegment: - return DocumentSegment( - tenant_id=tenant_id, - dataset_id=dataset_id, - document_id=document_id, - position=position, - content=f"segment {position}", - word_count=2, - tokens=2, - created_by="account-1", - status=status, - completed_at=datetime.datetime(2024, 1, 1, tzinfo=datetime.UTC) if completed else None, +@pytest.fixture +def datasets(monkeypatch): + service = create_autospec(DatasetApplicationService, instance=True, spec_set=True) + monkeypatch.setattr( + controller, "application_services", lambda: SimpleNamespace(knowledge=SimpleNamespace(datasets=service)) ) - - -class _UsesSQLiteSession: - session: Session - - @pytest.fixture(autouse=True) - def _inject_sqlite_session(self, sqlite_session: Session) -> None: - self.session = sqlite_session - - -class TestDatasetList(_UsesSQLiteSession): - def _mock_user(self): - user = make_account() - return user - - def test_get_success_basic(self, app: Flask): - api = DatasetListApi() - method = unwrap(api.get) - current_user = self._mock_user() - datasets = [make_dataset(icon_info={"icon": "📙", "icon_type": "emoji"})] - with app.test_request_context("/datasets"): - with ( - patch.object(DatasetService, "get_datasets", return_value=(datasets, 1)), - patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])), - ): - resp, status = method(api, self.session, "tenant-1", current_user) - assert status == 200 - assert resp["total"] == 1 - assert resp["data"][0]["embedding_available"] is True - assert resp["data"][0]["icon_info"] == { - "icon": "📙", - "icon_background": None, - "icon_type": "emoji", - "icon_url": None, - } - - def test_get_batch_loads_database_fields_with_caller_session(self, app: Flask, dataset_model_property_defaults): - api = DatasetListApi() - method = unwrap(api.get) - current_user = self._mock_user() - dataset = make_dataset() - session = self.session - with app.test_request_context("/datasets"): - with ( - patch.object(DatasetService, "get_datasets", return_value=([dataset], 1)), - patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])), - patch( - "controllers.console.datasets.datasets.build_dataset_detail_prefetch", - wraps=build_dataset_detail_prefetch, - ) as prefetch_mock, - ): - resp, status = method(api, session, "tenant-1", current_user) - - assert status == 200 - prefetch_mock.assert_called_once_with([dataset], session=session) - # The page is served from that single batch load, so no field may fall back - # to a per-dataset query. - for getter in dataset_model_property_defaults.values(): - getter.assert_not_called() - assert resp["data"][0]["document_count"] == 0 - assert resp["data"][0]["tags"] == [] - - def test_get_with_ids_filter(self, app: Flask): - api = DatasetListApi() - method = unwrap(api.get) - current_user = self._mock_user() - datasets = [make_dataset()] - with app.test_request_context("/datasets?ids=1&ids=2"): - with ( - patch.object(DatasetService, "get_datasets_by_ids", return_value=(datasets, 2)) as by_ids_mock, - patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])), - ): - resp, status = method(api, self.session, "tenant-1", current_user) - by_ids_mock.assert_called_once() - assert status == 200 - assert resp["total"] == 2 - - def test_get_attaches_current_user_permission_keys(self, app: Flask): - api = DatasetListApi() - method = unwrap(api.get) - current_user = self._mock_user() - dataset = make_dataset(id="dataset-1") - permissions = enterprise_rbac_service.MyPermissionsResponse( - dataset=enterprise_rbac_service.ResourcePermissionSnapshot( - default_permission_keys=["dataset.acl.readonly"], - overrides=[ - enterprise_rbac_service.ResourcePermissionKeys( - resource_id="dataset-1", permission_keys=["dataset.acl.readonly", "dataset.acl.edit"] - ) - ], - ) - ) - with app.test_request_context("/datasets"): - with ( - patch.object(DatasetService, "get_datasets", return_value=([dataset], 1)), - patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])), - patch( - "controllers.console.datasets.datasets.enterprise_rbac_service.RBACService.MyPermissions.get", - return_value=permissions, - ) as get_permissions, - ): - resp, status = method(api, self.session, "tenant-1", current_user) - get_permissions.assert_called_once_with("tenant-1", current_user.id, session=ANY) - assert status == 200 - assert resp["data"][0]["permission_keys"] == ["dataset.acl.readonly", "dataset.acl.edit"] - - def test_get_restricted_whitelist_blocks_own_dataset_fallback( - self, app: Flask, config_overrides: Callable[..., None] - ): - config_overrides(RBAC_ENABLED=True) - api = DatasetListApi() - method = unwrap(api.get) - current_user = self._mock_user() - permissions = enterprise_rbac_service.MyPermissionsResponse( - workspace=enterprise_rbac_service.WorkspacePermissionSnapshot( - permission_keys=["dataset.create_and_management"] - ) - ) - with app.test_request_context("/datasets"): - with ( - patch.object(DatasetService, "get_datasets", return_value=([], 0)) as get_datasets, - patch( - "controllers.console.datasets.datasets.enterprise_rbac_service.RBACService.MyPermissions.get", - return_value=permissions, - ), - patch( - "controllers.console.datasets.datasets.enterprise_rbac_service.RBACService.DatasetAccess.whitelist_resources", - return_value=SimpleNamespace(resource_ids=[]), - ), - patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])), - ): - method(api, self.session, "tenant-1", current_user) - assert get_datasets.call_args.kwargs["accessible_dataset_ids"] == [] - assert get_datasets.call_args.kwargs["include_own_datasets"] is False - - def test_get_default_read_is_unrestricted_when_whitelist_unrestricted( - self, app: Flask, config_overrides: Callable[..., None] - ): - config_overrides(RBAC_ENABLED=True) - api = DatasetListApi() - method = unwrap(api.get) - current_user = self._mock_user() - permissions = enterprise_rbac_service.MyPermissionsResponse( - dataset=enterprise_rbac_service.ResourcePermissionSnapshot(default_permission_keys=["dataset.preview"]) - ) - with app.test_request_context("/datasets"): - with ( - patch.object(DatasetService, "get_datasets", return_value=([], 0)) as get_datasets, - patch( - "controllers.console.datasets.datasets.enterprise_rbac_service.RBACService.MyPermissions.get", - return_value=permissions, - ), - patch( - "controllers.console.datasets.datasets.enterprise_rbac_service.RBACService.DatasetAccess.whitelist_resources", - return_value=SimpleNamespace(unrestricted=True, resource_ids=[]), - ), - patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])), - ): - method(api, self.session, "tenant-1", current_user) - assert get_datasets.call_args.kwargs["accessible_dataset_ids"] is None - - def test_get_restricted_whitelist_overrides_default_read_permission( - self, app: Flask, config_overrides: Callable[..., None] - ): - config_overrides(RBAC_ENABLED=True) - api = DatasetListApi() - method = unwrap(api.get) - current_user = self._mock_user() - permissions = enterprise_rbac_service.MyPermissionsResponse( - dataset=enterprise_rbac_service.ResourcePermissionSnapshot(default_permission_keys=["dataset.preview"]) - ) - with app.test_request_context("/datasets"): - with ( - patch.object(DatasetService, "get_datasets", return_value=([], 0)) as get_datasets, - patch( - "controllers.console.datasets.datasets.enterprise_rbac_service.RBACService.MyPermissions.get", - return_value=permissions, - ), - patch( - "controllers.console.datasets.datasets.enterprise_rbac_service.RBACService.DatasetAccess.whitelist_resources", - return_value=SimpleNamespace(resource_ids=["dataset-whitelist-only"]), - ), - patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])), - ): - method(api, MagicMock(), "tenant-1", current_user) - assert get_datasets.call_args.kwargs["accessible_dataset_ids"] == ["dataset-whitelist-only"] - assert get_datasets.call_args.kwargs["include_own_datasets"] is False - - def test_get_restricted_whitelist_ignores_dataset_read_overrides( - self, app: Flask, config_overrides: Callable[..., None] - ): - config_overrides(RBAC_ENABLED=True) - api = DatasetListApi() - method = unwrap(api.get) - current_user = self._mock_user() - permissions = enterprise_rbac_service.MyPermissionsResponse( - dataset=enterprise_rbac_service.ResourcePermissionSnapshot( - overrides=[ - enterprise_rbac_service.ResourcePermissionKeys( - resource_id="dataset-acl-shared", permission_keys=["dataset.acl.preview"] - ), - enterprise_rbac_service.ResourcePermissionKeys( - resource_id="dataset-full", permission_keys=["dataset.full_access"] - ), - enterprise_rbac_service.ResourcePermissionKeys( - resource_id="dataset-shared", permission_keys=["dataset.preview"] - ), - enterprise_rbac_service.ResourcePermissionKeys(resource_id="dataset-hidden", permission_keys=[]), - ] - ) - ) - with app.test_request_context("/datasets"): - with ( - patch.object(DatasetService, "get_datasets", return_value=([], 0)) as get_datasets, - patch( - "controllers.console.datasets.datasets.enterprise_rbac_service.RBACService.MyPermissions.get", - return_value=permissions, - ), - patch( - "controllers.console.datasets.datasets.enterprise_rbac_service.RBACService.DatasetAccess.whitelist_resources", - return_value=SimpleNamespace(resource_ids=["dataset-whitelist-only"]), - ), - patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])), - ): - method(api, self.session, "tenant-1", current_user) - assert get_datasets.call_args.kwargs["accessible_dataset_ids"] == [ - "dataset-whitelist-only", - ] - assert get_datasets.call_args.kwargs["include_own_datasets"] is False - - def test_get_with_ids_applies_dataset_visibility(self, app: Flask, config_overrides: Callable[..., None]): - config_overrides(RBAC_ENABLED=True) - api = DatasetListApi() - method = unwrap(api.get) - current_user = self._mock_user() - permissions = enterprise_rbac_service.MyPermissionsResponse() - with app.test_request_context("/datasets?ids=dataset-1"): - with ( - patch.object(DatasetService, "get_datasets_by_ids", return_value=([], 0)) as get_datasets_by_ids, - patch( - "controllers.console.datasets.datasets.enterprise_rbac_service.RBACService.MyPermissions.get", - return_value=permissions, - ), - patch( - "controllers.console.datasets.datasets.enterprise_rbac_service.RBACService.DatasetAccess.whitelist_resources", - return_value=SimpleNamespace(resource_ids=[]), - ), - patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])), - ): - method(api, self.session, "tenant-1", current_user) - session = get_datasets_by_ids.call_args.kwargs["session"] - assert session is self.session - assert get_datasets_by_ids.call_args.args == (["dataset-1"], "tenant-1") - assert get_datasets_by_ids.call_args.kwargs == { - "user": current_user, - "accessible_dataset_ids": [], - "include_own_datasets": False, - "session": session, - } - - def test_get_with_tag_ids(self, app: Flask): - api = DatasetListApi() - method = unwrap(api.get) - current_user = self._mock_user() - datasets = [make_dataset()] - with app.test_request_context("/datasets?tag_ids=tag1"): - with ( - patch.object(DatasetService, "get_datasets", return_value=(datasets, 1)), - patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])), - ): - resp, status = method(api, self.session, "tenant-1", current_user) - assert status == 200 - - def test_get_allows_legacy_weighted_score_without_weight_type(self, app: Flask): - api = DatasetListApi() - method = unwrap(api.get) - current_user = self._mock_user() - datasets = [ - make_dataset( - retrieval_model={ - "search_method": "hybrid_search", - "reranking_enable": True, - "reranking_mode": "weighted_score", - "reranking_model": None, - "weights": { - "vector_setting": { - "vector_weight": 0.7, - "embedding_model_name": "text-embedding", - "embedding_provider_name": "openai", - }, - "keyword_setting": {"keyword_weight": 0.3}, - }, - "top_k": 3, - "score_threshold_enabled": False, - "score_threshold": 0.0, - } - ) - ] - with app.test_request_context("/datasets"): - with ( - patch.object(DatasetService, "get_datasets", return_value=(datasets, 1)), - patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])), - ): - resp, status = method(api, self.session, "tenant-1", current_user) - assert status == 200 - assert resp["data"][0]["retrieval_model_dict"]["weights"]["weight_type"] is None - - def test_get_merges_partial_retrieval_model_defaults(self, app: Flask): - api = DatasetListApi() - method = unwrap(api.get) - current_user = self._mock_user() - datasets = [make_dataset(retrieval_model={"top_k": 4, "score_threshold_enabled": False})] - with app.test_request_context("/datasets"): - with ( - patch.object(DatasetService, "get_datasets", return_value=(datasets, 1)), - patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])), - ): - resp, status = method(api, self.session, "tenant-1", current_user) - assert status == 200 - retrieval_model = resp["data"][0]["retrieval_model_dict"] - assert retrieval_model["search_method"] == "semantic_search" - assert retrieval_model["reranking_enable"] is False - assert retrieval_model["top_k"] == 4 - assert retrieval_model["score_threshold_enabled"] is False - - def test_embedding_available_false(self, app: Flask): - api = DatasetListApi() - method = unwrap(api.get) - current_user = self._mock_user() - datasets = [ - make_dataset( - indexing_technique="high_quality", embedding_model="text-embed", embedding_model_provider="openai" - ) - ] - config = MagicMock() - config.get_models.return_value = [] - with app.test_request_context("/datasets"): - with ( - patch.object(DatasetService, "get_datasets", return_value=(datasets, 1)), - patch.object(ProviderManager, "get_configurations", return_value=config), - ): - resp, status = method(api, self.session, "tenant-1", current_user) - assert resp["data"][0]["embedding_available"] is False - - def test_partial_members_permission(self, app: Flask): - api = DatasetListApi() - method = unwrap(api.get) - current_user = self._mock_user() - datasets = [make_dataset(permission="partial_members")] - session = self.session - session.add(DatasetPermission(dataset_id="ds-1", account_id="u1", tenant_id="tenant-1")) - session.flush() - with app.test_request_context("/datasets"): - with ( - patch.object(DatasetService, "get_datasets", return_value=(datasets, 1)), - patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])), - ): - resp, status = method(api, session, "tenant-1", current_user) - assert resp["data"][0]["partial_member_list"] == ["u1"] - - -class TestDatasetListApiPost(_UsesSQLiteSession): - def test_post_success(self, app: Flask): - api = DatasetListApi() - method = unwrap(api.post) - payload = {"name": "My Dataset", "description": "desc", "indexing_technique": "economy", "provider": "vendor"} - user = make_account() - dataset = make_dataset(name=payload["name"], description=payload["description"]) - with ( - app.test_request_context("/datasets", json=payload), - patch.object(type(console_ns), "payload", payload), - patch.object(DatasetService, "create_empty_dataset", return_value=dataset), - ): - _, status = method(api, DatasetCreatePayload(**payload), self.session, "tenant-1", user) - assert status == 201 - - def test_post_forbidden(self, app: Flask): - api = DatasetListApi() - method = unwrap(api.post) - payload = {"name": "test"} - 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, DatasetCreatePayload(**payload), self.session, "tenant-1", user) - - def test_post_duplicate_name(self, app: Flask): - api = DatasetListApi() - method = unwrap(api.post) - payload = {"name": "duplicate"} - user = make_account() - with ( - app.test_request_context("/datasets", json=payload), - patch.object(type(console_ns), "payload", payload), - patch.object( - DatasetService, "create_empty_dataset", side_effect=services.errors.dataset.DatasetNameDuplicateError() - ), - ): - with pytest.raises(DatasetNameDuplicateError): - method(api, DatasetCreatePayload(**payload), self.session, "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, DatasetCreatePayload(), self.session, "tenant-1", make_account()) - - def test_post_invalid_indexing_technique(self, app: Flask): - api = DatasetListApi() - method = unwrap(api.post) - 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, DatasetCreatePayload(**payload), self.session, "tenant-1", make_account()) - - def test_post_invalid_provider(self, app: Flask): - api = DatasetListApi() - method = unwrap(api.post) - 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, DatasetCreatePayload(**payload), self.session, "tenant-1", make_account()) - - -class TestDatasetApiGet(_UsesSQLiteSession): - def test_get_success_basic(self, app: Flask): - api = DatasetApi() - method = unwrap(api.get) - dataset_id = "123e4567-e89b-12d3-a456-426614174000" - user = make_account() - tenant_id = "tenant-1" - dataset = make_dataset(id=dataset_id) - with ( - app.test_request_context(f"/datasets/{dataset_id}"), - patch.object(DatasetService, "get_dataset", return_value=dataset), - patch.object(DatasetService, "check_dataset_permission", return_value=None), - patch("controllers.console.datasets.datasets.create_plugin_provider_manager") as provider_manager_mock, - ): - provider_manager_mock.return_value.get_configurations.return_value.get_models.return_value = [] - data, status = method(api, self.session, tenant_id, user, dataset_id) - assert status == 200 - assert data["embedding_available"] is True - - def test_get_attaches_permission_keys_when_rbac_enabled(self, app: Flask, config_overrides: Callable[..., None]): - config_overrides(RBAC_ENABLED=True) - api = DatasetApi() - method = unwrap(api.get) - dataset_id = "123e4567-e89b-12d3-a456-426614174000" - user = make_account() - tenant_id = "tenant-1" - dataset = make_dataset(id=dataset_id) - with ( - app.test_request_context(f"/datasets/{dataset_id}"), - patch.object(DatasetService, "get_dataset", return_value=dataset), - patch.object(DatasetService, "check_dataset_permission", return_value=None), - patch( - "controllers.console.datasets.datasets.enterprise_rbac_service.RBACService.MyPermissions.get", - return_value=enterprise_rbac_service.MyPermissionsResponse( - dataset=enterprise_rbac_service.ResourcePermissionSnapshot( - overrides=[ - enterprise_rbac_service.ResourcePermissionKeys( - resource_id=dataset_id, permission_keys=["dataset.acl.readonly", "dataset.acl.edit"] - ) - ] - ) - ), - ) as get_permissions, - patch("controllers.console.datasets.datasets.create_plugin_provider_manager") as provider_manager_mock, - ): - provider_manager_mock.return_value.get_configurations.return_value.get_models.return_value = [] - data, status = method(api, self.session, tenant_id, user, dataset_id) - get_permissions.assert_called_once_with(tenant_id, user.id, dataset_id=dataset_id, session=ANY) - assert status == 200 - assert data["permission_keys"] == ["dataset.acl.readonly", "dataset.acl.edit"] - - def test_get_uses_default_external_retrieval_model(self, app: Flask): - api = DatasetApi() - method = unwrap(api.get) - dataset_id = "dataset-id" - dataset = make_dataset(id=dataset_id, retrieval_model=None) - with ( - app.test_request_context(f"/datasets/{dataset_id}"), - patch.object(DatasetService, "get_dataset", return_value=dataset), - patch.object(DatasetService, "check_dataset_permission", return_value=None), - patch("controllers.console.datasets.datasets.create_plugin_provider_manager") as provider_manager_mock, - ): - provider_manager_mock.return_value.get_configurations.return_value.get_models.return_value = [] - data, status = method(api, self.session, "tenant", make_account(), dataset_id) - assert status == 200 - assert data["external_retrieval_model"] == {"top_k": 2, "score_threshold": 0.0, "score_threshold_enabled": None} - - def test_get_dataset_not_found(self, app: Flask): - api = DatasetApi() - method = unwrap(api.get) - dataset_id = "missing-id" - with ( - app.test_request_context(f"/datasets/{dataset_id}"), - patch.object(DatasetService, "get_dataset", return_value=None), - ): - with pytest.raises(NotFound, match="Dataset not found"): - method(api, self.session, "tenant", make_account(), dataset_id) - - def test_get_permission_denied(self, app: Flask): - api = DatasetApi() - method = unwrap(api.get) - dataset_id = "dataset-id" - dataset = make_dataset(id=dataset_id) - with ( - app.test_request_context(f"/datasets/{dataset_id}"), - patch.object(DatasetService, "get_dataset", return_value=dataset), - patch.object( - DatasetService, - "check_dataset_permission", - side_effect=services.errors.account.NoPermissionError("no access"), - ), - ): - with pytest.raises(Forbidden, match="no access"): - method(api, self.session, "tenant", make_account(), dataset_id) - - def test_get_high_quality_embedding_unavailable(self, app: Flask): - api = DatasetApi() - method = unwrap(api.get) - dataset_id = "dataset-id" - user = make_account() - tenant_id = "tenant-1" - dataset = make_dataset( - id=dataset_id, - indexing_technique="high_quality", - embedding_model="text-embedding", - embedding_model_provider="openai", - ) - with ( - app.test_request_context(f"/datasets/{dataset_id}"), - patch.object(DatasetService, "get_dataset", return_value=dataset), - patch.object(DatasetService, "check_dataset_permission", return_value=None), - patch("controllers.console.datasets.datasets.create_plugin_provider_manager") as provider_manager_mock, - ): - provider_manager_mock.return_value.get_configurations.return_value.get_models.return_value = [] - data, _ = method(api, self.session, tenant_id, user, dataset_id) - assert data["embedding_available"] is False - - def test_get_partial_members_permission(self, app: Flask): - api = DatasetApi() - method = unwrap(api.get) - dataset_id = "dataset-id" - dataset = make_dataset(id=dataset_id, permission="partial_members") - partial_members = ["u1", "u2"] - with ( - app.test_request_context(f"/datasets/{dataset_id}"), - patch.object(DatasetService, "get_dataset", return_value=dataset), - patch.object(DatasetService, "check_dataset_permission", return_value=None), - patch.object(DatasetPermissionService, "get_dataset_partial_member_list", return_value=partial_members), - patch("controllers.console.datasets.datasets.create_plugin_provider_manager") as provider_manager_mock, - ): - provider_manager_mock.return_value.get_configurations.return_value.get_models.return_value = [] - data, _ = method(api, self.session, "tenant", make_account(), dataset_id) - assert data["partial_member_list"] == partial_members - - -class TestDatasetApiPatch(_UsesSQLiteSession): - def test_patch_success_basic(self, app: Flask): - api = DatasetApi() - method = unwrap(api.patch) - dataset_id = "dataset-id" - payload = {"name": "updated-name", "description": "updated description"} - user = make_account() - tenant_id = "tenant-1" - dataset = make_dataset(id=dataset_id, tenant_id=tenant_id) - with ( - app.test_request_context(f"/datasets/{dataset_id}"), - patch.object(type(console_ns), "payload", payload), - patch.object(DatasetService, "get_dataset", return_value=dataset), - patch.object(DatasetPermissionService, "check_permission", return_value=None), - patch.object(DatasetService, "update_dataset", return_value=dataset), - patch.object(DatasetPermissionService, "get_dataset_partial_member_list", return_value=[]), - ): - result, status = method(api, DatasetUpdatePayload(), self.session, tenant_id, user, dataset_id) - assert status == 200 - assert result["partial_member_list"] == [] - - def test_patch_dataset_not_found(self, app: Flask): - api = DatasetApi() - method = unwrap(api.patch) - with ( - app.test_request_context("/datasets/missing"), - patch.object(DatasetService, "get_dataset", return_value=None), - ): - with pytest.raises(NotFound, match="Dataset not found"): - method(api, DatasetUpdatePayload(), self.session, "tenant-1", make_account(), "missing") - - def test_patch_permission_denied(self, app: Flask): - api = DatasetApi() - method = unwrap(api.patch) - dataset_id = "dataset-id" - dataset = make_dataset(id=dataset_id) - payload = {"name": "x"} - with ( - app.test_request_context(f"/datasets/{dataset_id}"), - patch.object(type(console_ns), "payload", payload), - patch.object(DatasetService, "get_dataset", return_value=dataset), - patch.object(DatasetPermissionService, "check_permission", side_effect=Forbidden("no permission")), - ): - with pytest.raises(Forbidden): - method(api, DatasetUpdatePayload(), self.session, "tenant", make_account(), dataset_id) - - def test_patch_partial_members_update(self, app: Flask): - api = DatasetApi() - method = unwrap(api.patch) - dataset_id = "dataset-id" - payload = {"permission": "partial_members", "partial_member_list": [{"user_id": "u1"}, {"user_id": "u2"}]} - dataset = make_dataset(id=dataset_id, permission="partial_members") - with ( - app.test_request_context(f"/datasets/{dataset_id}"), - patch.object(type(console_ns), "payload", payload), - patch.object(DatasetService, "get_dataset", return_value=dataset), - patch.object(DatasetPermissionService, "check_permission", return_value=None), - patch.object(DatasetService, "update_dataset", return_value=dataset), - 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, DatasetUpdatePayload(), self.session, "tenant", make_account(), dataset_id) - assert result["partial_member_list"] == ["u1", "u2"] - - def test_patch_clear_partial_members(self, app: Flask): - api = DatasetApi() - method = unwrap(api.patch) - dataset_id = "dataset-id" - payload = {"permission": "only_me"} - dataset = make_dataset(id=dataset_id) - with ( - app.test_request_context(f"/datasets/{dataset_id}"), - patch.object(type(console_ns), "payload", payload), - patch.object(DatasetService, "get_dataset", return_value=dataset), - patch.object(DatasetPermissionService, "check_permission", return_value=None), - patch.object(DatasetService, "update_dataset", return_value=dataset), - patch.object(DatasetPermissionService, "clear_partial_member_list", return_value=None), - patch.object(DatasetPermissionService, "get_dataset_partial_member_list", return_value=[]), - ): - result, _ = method(api, DatasetUpdatePayload(), self.session, "tenant", make_account(), dataset_id) - assert result["partial_member_list"] == [] - - -class TestDatasetApiDelete(_UsesSQLiteSession): - def test_delete_success(self, app: Flask): - api = DatasetApi() - method = unwrap(api.delete) - dataset_id = "dataset-id" - user = make_account() - with ( - app.test_request_context(f"/datasets/{dataset_id}"), - patch.object(DatasetService, "delete_dataset", return_value=True), - patch.object(DatasetPermissionService, "clear_partial_member_list", return_value=None), - ): - result, status = method(api, self.session, user, dataset_id) - assert status == 204 - assert result == "" - - def test_delete_forbidden_no_permission(self, app: Flask): - api = DatasetApi() - method = unwrap(api.delete) - dataset_id = "dataset-id" - user = make_account(TenantAccountRole.NORMAL) - with app.test_request_context(f"/datasets/{dataset_id}"): - with pytest.raises(Forbidden): - method(api, self.session, user, dataset_id) - - def test_delete_dataset_not_found(self, app: Flask): - api = DatasetApi() - method = unwrap(api.delete) - dataset_id = "missing-dataset" - user = make_account() - with ( - app.test_request_context(f"/datasets/{dataset_id}"), - patch.object(DatasetService, "delete_dataset", return_value=False), - ): - with pytest.raises(NotFound, match="Dataset not found"): - method(api, self.session, user, dataset_id) - - def test_delete_dataset_in_use(self, app: Flask): - api = DatasetApi() - method = unwrap(api.delete) - dataset_id = "dataset-id" - user = make_account() - with ( - app.test_request_context(f"/datasets/{dataset_id}"), - patch.object(DatasetService, "delete_dataset", side_effect=services.errors.dataset.DatasetInUseError()), - ): - with pytest.raises(DatasetInUseError): - method(api, self.session, user, dataset_id) - - -class TestDatasetUseCheckApi(_UsesSQLiteSession): - @pytest.mark.parametrize("is_using", [True, False]) - def test_get_use_check(self, app: Flask, is_using: bool): - api = DatasetUseCheckApi() - method = unwrap(api.get) - dataset_id = "dataset-id" - dataset = make_dataset(id=dataset_id) - current_user = make_account() - session = self.session - with ( - app.test_request_context(f"/datasets/{dataset_id}/use-check"), - patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset) as get_dataset, - patch.object(DatasetService, "check_dataset_permission") as check_permission, - patch.object(DatasetService, "dataset_use_check", return_value=is_using) as dataset_use_check, - ): - result, status = method(api, session, "tenant-1", current_user, dataset_id) - assert status == 200 - assert result == {"is_using": is_using} - get_dataset.assert_called_once_with(dataset_id, "tenant-1", session=session) - check_permission.assert_called_once_with(dataset, current_user, session) - dataset_use_check.assert_called_once_with(DatasetRef("tenant-1", dataset_id), session) - - def test_get_use_check_delegates_permission_check_in_rbac_mode( - self, app: Flask, config_overrides: Callable[..., None] - ): - config_overrides(RBAC_ENABLED=True) - api = DatasetUseCheckApi() - method = unwrap(api.get) - dataset = make_dataset(id="dataset-id") - current_user = make_account() - session = self.session - with ( - app.test_request_context("/datasets/dataset-id/use-check"), - patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset), - patch.object(DatasetService, "check_dataset_permission") as check_permission, - patch.object(DatasetService, "dataset_use_check", return_value=False), - ): - _, status = method(api, session, "tenant-1", current_user, "dataset-id") - - assert status == 200 - check_permission.assert_called_once_with(dataset, current_user, session) + return service @pytest.mark.parametrize( - "api_cls", - [DatasetUseCheckApi, DatasetIndexingStatusApi, DatasetErrorDocs, DatasetAutoDisableLogApi], + ("resource", "operation"), + [ + (DatasetApi, lambda service: service.get_dataset), + (DatasetUseCheckApi, lambda service: service.is_in_use), + (DatasetQueryApi, lambda service: service.queries), + (DatasetRelatedAppListApi, lambda service: service.related_apps), + (DatasetIndexingStatusApi, lambda service: service.indexing_status), + (DatasetErrorDocs, lambda service: service.error_documents), + (DatasetPermissionUserListApi, lambda service: service.partial_members), + (DatasetAutoDisableLogApi, lambda service: service.auto_disable_logs), + ], ) -def test_dataset_scoped_read_permission_denied(app: Flask, api_cls, sqlite_session: Session): - api = api_cls() - method = unwrap(api.get) - dataset = make_dataset(id="dataset-1") - session = sqlite_session - with ( - app.test_request_context("/"), - patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset), - patch.object( - DatasetService, - "check_dataset_permission", - side_effect=services.errors.account.NoPermissionError("no permission"), - ), - ): - with pytest.raises(Forbidden, match="no permission"): - method(api, session, "tenant-1", make_account(), "dataset-1") +@pytest.mark.parametrize( + ("error", "http_error"), [(DatasetNotFoundError(), NotFound), (DatasetAccessDeniedError(), Forbidden)] +) +def test_scoped_reads_pass_context_and_map_access_errors(app, datasets, resource, operation, error, http_error): + method = operation(datasets) + method.side_effect = error + with app.test_request_context("/"), pytest.raises(http_error): + unwrap(resource.get)(resource(), CONTEXT, DATASET_ID) + assert method.call_args.args == (CONTEXT,) + assert method.call_args.kwargs["dataset_id"] == str(DATASET_ID) -class TestDatasetQueryApi(_UsesSQLiteSession): - def _query_record(self, index: int = 1) -> DatasetQuery: - query = DatasetQuery( - dataset_id="dataset-id", - content=json.dumps([{"content_type": "text_query", "content": f"question {index}", "file_info": None}]), - source="hit_testing", - source_app_id=None, - created_by_role=CreatorUserRole.ACCOUNT, - created_by=f"account-{index}", - ) - query.id = f"query-{index}" - query.created_at = datetime.datetime(2024, 1, index, 12, 0, 0, tzinfo=datetime.UTC) - return query +def test_list_parses_repeated_filters_and_serializes_page(app, datasets): + datasets.list_datasets.return_value = {"data": [], "page": 2, "limit": 3, "total": 4, "has_more": False} + with app.test_request_context("/?page=2&limit=3&ids=a&ids=b&tag_ids=x&tag_ids=y&include_all=true&keyword=term"): + result, status = unwrap(DatasetListApi.get)(DatasetListApi(), CONTEXT) + datasets.list_datasets.assert_called_once_with( + CONTEXT, + DatasetListFilter(page=2, limit=3, ids=["a", "b"], tag_ids=["x", "y"], include_all=True, keyword="term"), + ) + assert status == 200 + assert result == datasets.list_datasets.return_value - def test_get_queries_success(self, app: Flask): - api = DatasetQueryApi() - method = unwrap(api.get) - dataset_id = "dataset-id" - current_user = make_account() - dataset = make_dataset(id=dataset_id) - queries = [self._query_record(1), self._query_record(2)] - with ( - app.test_request_context("/datasets/queries?page=1&limit=20"), - patch.object(DatasetService, "get_dataset", return_value=dataset), - patch.object(DatasetService, "check_dataset_permission", return_value=None), - patch.object(DatasetService, "get_dataset_queries", return_value=(queries, 2)), - ): - response, status = method(api, self.session, current_user, dataset_id) - assert status == 200 - assert response["total"] == 2 - assert response["page"] == 1 - assert response["limit"] == 20 - assert response["has_more"] is False - assert len(response["data"]) == 2 - assert response["data"][0] == { - "id": "query-1", - "queries": [{"content_type": "text_query", "content": "question 1", "file_info": None}], - "source": "hit_testing", - "source_app_id": None, - "created_by_role": "account", - "created_by": "account-1", - "created_at": 1704110400, + +def test_patch_does_not_turn_omitted_fields_into_updates(datasets): + datasets.update_dataset.side_effect = DatasetNotFoundError() + with pytest.raises(NotFound): + unwrap(DatasetApi.patch)(DatasetApi(), DatasetUpdatePayload(name="Changed"), CONTEXT, DATASET_ID) + datasets.update_dataset.assert_called_once_with(CONTEXT, dataset_id=str(DATASET_ID), values={"name": "Changed"}) + + +@pytest.mark.parametrize( + ("error", "http_error"), + [ + (services.errors.dataset.DatasetNameDuplicateError(), DatasetNameDuplicateError), + (services.errors.dataset.DatasetInUseError(), DatasetInUseError), + ], +) +def test_domain_errors_are_mapped_at_transport(datasets, error, http_error): + datasets.delete_dataset.side_effect = error + with pytest.raises(http_error): + unwrap(DatasetApi.delete)(DatasetApi(), CONTEXT, DATASET_ID) + + +def test_delete_success_returns_empty_204(datasets): + assert unwrap(DatasetApi.delete)(DatasetApi(), CONTEXT, DATASET_ID) == ("", 204) + datasets.delete_dataset.assert_called_once_with(CONTEXT, dataset_id=str(DATASET_ID)) + + +def test_request_base_url_and_explicit_status_are_forwarded(app, datasets): + datasets.api_base_url.return_value = "https://api.example/v1" + with app.test_request_context("/", base_url="https://console.example/"): + assert unwrap(DatasetApiBaseUrlApi.get)(DatasetApiBaseUrlApi(), CONTEXT) == { + "api_base_url": "https://api.example/v1" } + datasets.api_base_url.assert_called_once_with(CONTEXT, request_base_url="https://console.example") + assert unwrap(DatasetEnableApiApi.post)(DatasetEnableApiApi(), CONTEXT, DATASET_ID, "disable") == ( + {"result": "success"}, + 200, + ) + datasets.set_api_enabled.assert_called_once_with(CONTEXT, dataset_id=str(DATASET_ID), status="disable") - def test_get_image_query_uses_caller_session(self, app: Flask): - api = DatasetQueryApi() - method = unwrap(api.get) - dataset = make_dataset(id="dataset-id") - query = self._query_record() - query.content = json.dumps([{"content_type": "image_query", "content": "file-1"}]) - upload_file = UploadFile( - tenant_id="tenant-1", - storage_type=StorageType.LOCAL, - key="image.png", - name="image.png", - size=10, - extension="png", - mime_type="image/png", - created_by_role=CreatorUserRole.ACCOUNT, - created_by="account-1", - created_at=datetime.datetime(2024, 1, 1, tzinfo=datetime.UTC), - used=False, - ) - upload_file.id = "file-1" - session = self.session - session.add(upload_file) - session.flush() - with ( - app.test_request_context("/datasets/queries"), - patch.object(DatasetService, "get_dataset", return_value=dataset), - patch.object(DatasetService, "check_dataset_permission", return_value=None), - patch.object(DatasetService, "get_dataset_queries", return_value=([query], 1)), - patch("models.dataset.sign_upload_file_preview_url", return_value="signed-url"), - ): - response, status = method(api, session, make_account(), "dataset-id") - assert status == 200 - assert response["data"][0]["queries"] == [ +def test_retrieval_settings_forward_mock_flag(datasets): + datasets.retrieval_settings.return_value = {"retrieval_method": ["semantic_search"]} + assert ( + unwrap(DatasetRetrievalSettingApi.get)(DatasetRetrievalSettingApi(), CONTEXT) + == datasets.retrieval_settings.return_value + ) + unwrap(DatasetRetrievalSettingMockApi.get)(DatasetRetrievalSettingMockApi(), CONTEXT, "milvus") + datasets.retrieval_settings.assert_called_with(CONTEXT, vector_type="milvus", is_mock=True) + + +def test_new_estimate_sources_maps_each_supported_transport_shape() -> None: + upload_sources = _new_estimate_sources( + {"data_source_type": "upload_file", "file_info_list": {"file_ids": ["file-1", "file-2"]}} + ) + notion_sources = _new_estimate_sources( + { + "data_source_type": "notion_import", + "notion_info_list": [ + { + "workspace_id": "notion-workspace", + "credential_id": "credential-1", + "pages": [{"page_id": "page-1", "type": "page"}], + } + ], + } + ) + website_sources = _new_estimate_sources( + { + "data_source_type": "website_crawl", + "website_info_list": { + "provider": "firecrawl", + "job_id": "job-1", + "urls": ["https://example.com/a", "https://example.com/b"], + "only_main_content": True, + }, + } + ) + + assert upload_sources == (UploadFileEstimateSource("file-1"), UploadFileEstimateSource("file-2")) + assert notion_sources == (NotionEstimateSource("notion-workspace", "page-1", NotionPageType.PAGE, "credential-1"),) + assert website_sources == ( + WebsiteEstimateSource("firecrawl", "job-1", "https://example.com/a", only_main_content=True), + WebsiteEstimateSource("firecrawl", "job-1", "https://example.com/b", only_main_content=True), + ) + + +@pytest.mark.parametrize(("value", "expected"), [("false", False), ("true", True), (False, False), (True, True)]) +def test_website_estimate_parses_boolean_values(value: str | bool, expected: bool) -> None: + sources = _new_estimate_sources( + { + "data_source_type": "website_crawl", + "website_info_list": { + "provider": "firecrawl", + "job_id": "job-1", + "urls": ["https://example.com"], + "only_main_content": value, + }, + } + ) + assert sources == (WebsiteEstimateSource("firecrawl", "job-1", "https://example.com", only_main_content=expected),) + + +@pytest.mark.parametrize("values", [{"only_main_content": "invalid"}, {"urls": [42]}]) +def test_website_estimate_rejects_invalid_field_types(values: dict[str, object]) -> None: + with pytest.raises(ValueError): + _new_estimate_sources( { - "content_type": "image_query", - "content": "file-1", - "file_info": { - "id": "file-1", - "name": "image.png", - "size": 10, - "extension": "png", - "mime_type": "image/png", - "source_url": "signed-url", + "data_source_type": "website_crawl", + "website_info_list": { + "provider": "firecrawl", + "job_id": "job-1", + "urls": ["https://example.com"], + **values, }, } - ] - assert session.get(UploadFile, "file-1") is upload_file - - def test_get_queries_dataset_not_found(self, app: Flask): - api = DatasetQueryApi() - method = unwrap(api.get) - dataset_id = "dataset-id" - current_user = make_account() - with ( - app.test_request_context("/datasets/queries"), - patch.object(DatasetService, "get_dataset", return_value=None), - ): - with pytest.raises(NotFound, match="Dataset not found"): - method(api, self.session, current_user, dataset_id) - - def test_get_queries_permission_denied(self, app: Flask): - api = DatasetQueryApi() - method = unwrap(api.get) - dataset_id = "dataset-id" - current_user = make_account() - dataset = make_dataset(id=dataset_id) - with ( - app.test_request_context("/datasets/queries"), - patch.object(DatasetService, "get_dataset", return_value=dataset), - patch.object( - DatasetService, - "check_dataset_permission", - side_effect=services.errors.account.NoPermissionError("no access"), - ), - ): - with pytest.raises(Forbidden): - method(api, self.session, current_user, dataset_id) - - def test_get_queries_pagination_has_more(self, app: Flask): - api = DatasetQueryApi() - method = unwrap(api.get) - dataset_id = "dataset-id" - current_user = make_account() - dataset = make_dataset(id=dataset_id) - queries = [self._query_record(index) for index in range(1, 21)] - with ( - app.test_request_context("/datasets/queries?page=1&limit=20"), - patch.object(DatasetService, "get_dataset", return_value=dataset), - patch.object(DatasetService, "check_dataset_permission", return_value=None), - patch.object(DatasetService, "get_dataset_queries", return_value=(queries, 40)), - ): - response, status = method(api, self.session, current_user, dataset_id) - assert status == 200 - assert response["has_more"] is True - assert len(response["data"]) == 20 - - -class TestDatasetIndexingEstimateApi(_UsesSQLiteSession): - def _upload_file(self, *, tenant_id: str = "tenant-1", file_id: str = "file-1") -> UploadFile: - upload_file = UploadFile( - tenant_id=tenant_id, - storage_type=StorageType.LOCAL, - key="key", - name="name.txt", - size=1, - extension="txt", - mime_type="text/plain", - created_by_role=CreatorUserRole.ACCOUNT, - created_by="user-1", - created_at=datetime.datetime.now(tz=datetime.UTC), - used=False, ) - upload_file.id = file_id - return upload_file - def _base_payload(self): - return { - "info_list": {"data_source_type": "upload_file", "file_info_list": {"file_ids": ["file-1"]}}, - "process_rule": {"chunk_size": 100}, - "indexing_technique": "high_quality", - "doc_form": IndexStructureType.PARAGRAPH_INDEX, - "doc_language": "English", - "dataset_id": None, - } - @pytest.mark.parametrize( - ("dataset_id", "expected_scene", "expected_locator_type", "expected_path_args"), - [ - ("dataset-1", RBACPermission.DATASET_USE, DatasetId, {"dataset_id": "dataset-1"}), - (None, RBACPermission.DATASET_CREATE_AND_MANAGEMENT, Workspace, None), - ], +def test_new_estimate_sources_deduplicates_upload_ids_without_reordering() -> None: + sources = _new_estimate_sources( + {"data_source_type": "upload_file", "file_info_list": {"file_ids": ["file-2", "file-1", "file-2"]}} ) - def test_post_authorizes_existing_and_new_dataset_workflows( - self, - app: Flask, - dataset_id: str | None, - expected_scene: RBACPermission, - expected_locator_type: type[DatasetId] | type[Workspace], - expected_path_args: dict[str, str] | None, - ) -> None: - api = DatasetIndexingEstimateApi() - method = unwrap(api.post) - payload = {**self._base_payload(), "dataset_id": dataset_id} - session = self.session - session.add(self._upload_file()) - session.flush() - current_user = make_account() - with ( - app.test_request_context("/"), - patch("controllers.console.datasets.datasets.DocumentService.estimate_args_validate", return_value=None), - patch( - "controllers.console.datasets.datasets.IndexingRunner.indexing_estimate", - return_value=IndexingEstimate(total_segments=1, preview=[]), - ), - patch("controllers.console.datasets.datasets.enforce_rbac_checks", create=True) as enforce_checks, - ): - method(api, IndexingEstimatePayload(**payload), session, "tenant-1", current_user) - - enforce_checks.assert_called_once() - kwargs = enforce_checks.call_args.kwargs - assert kwargs["tenant_id"] == "tenant-1" - assert kwargs["account_id"] == current_user.id - assert kwargs.get("path_args") == expected_path_args - [check] = kwargs["checks"] - assert check.scene is expected_scene - assert isinstance(check.locator, expected_locator_type) - - def test_post_success_upload_file(self, app: Flask): - api = DatasetIndexingEstimateApi() - method = unwrap(api.post) - payload = self._base_payload() - mock_file = self._upload_file() - session = self.session - session.add(mock_file) - session.flush() - - mock_response = IndexingEstimate(total_segments=100, preview=[]) - - with ( - app.test_request_context("/"), - patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload), - 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, - IndexingEstimatePayload(**payload), - session, - "tenant-1", - make_account(), - ) - assert status == 200 - assert response == { - "tokens": 0, - "total_price": 0, - "currency": "USD", - "total_segments": 100, - "preview": [], - } - - def test_post_file_not_found(self, app: Flask): - api = DatasetIndexingEstimateApi() - method = unwrap(api.post) - payload = self._base_payload() - session = self.session - with ( - app.test_request_context("/"), - patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload), - patch("controllers.console.datasets.datasets.DocumentService.estimate_args_validate", return_value=None), - ): - with pytest.raises(NotFound): - method( - api, - IndexingEstimatePayload(**payload), - session, - "tenant-1", - make_account(), - ) - - def test_post_llm_bad_request_error(self, app: Flask): - api = DatasetIndexingEstimateApi() - method = unwrap(api.post) - mock_file = self._upload_file() - payload = self._base_payload() - session = self.session - session.add(mock_file) - session.flush() - with ( - app.test_request_context("/"), - patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload), - patch("controllers.console.datasets.datasets.DocumentService.estimate_args_validate", return_value=None), - patch( - "controllers.console.datasets.datasets.IndexingRunner.indexing_estimate", - side_effect=LLMBadRequestError(), - ), - ): - with pytest.raises(ProviderNotInitializeError): - method( - api, - IndexingEstimatePayload(**payload), - session, - "tenant-1", - make_account(), - ) - - def test_post_provider_token_not_init(self, app: Flask): - api = DatasetIndexingEstimateApi() - method = unwrap(api.post) - mock_file = self._upload_file() - payload = self._base_payload() - session = self.session - session.add(mock_file) - session.flush() - with ( - app.test_request_context("/"), - patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload), - patch("controllers.console.datasets.datasets.DocumentService.estimate_args_validate", return_value=None), - patch( - "controllers.console.datasets.datasets.IndexingRunner.indexing_estimate", - side_effect=ProviderTokenNotInitError("token missing"), - ), - ): - with pytest.raises(ProviderNotInitializeError): - method( - api, - IndexingEstimatePayload(**payload), - session, - "tenant-1", - make_account(), - ) - - def test_post_generic_exception(self, app: Flask): - api = DatasetIndexingEstimateApi() - method = unwrap(api.post) - mock_file = self._upload_file() - payload = self._base_payload() - session = self.session - session.add(mock_file) - session.flush() - with ( - app.test_request_context("/"), - patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload), - patch("controllers.console.datasets.datasets.DocumentService.estimate_args_validate", return_value=None), - patch( - "controllers.console.datasets.datasets.IndexingRunner.indexing_estimate", side_effect=Exception("boom") - ), - ): - with pytest.raises(IndexingEstimateError): - method( - api, - IndexingEstimatePayload(**payload), - session, - "tenant-1", - make_account(), - ) + assert sources == (UploadFileEstimateSource("file-2"), UploadFileEstimateSource("file-1")) -class TestDatasetRelatedAppListApi(_UsesSQLiteSession): - def test_get_success(self, app: Flask): - api = DatasetRelatedAppListApi() - method = unwrap(api.get) - dataset = make_dataset(id="dataset-1") - app1 = make_related_app(id="app-1", name="App 1") - app2 = make_related_app(id="app-2", name="App 2") - join1 = AppDatasetJoin(app_id="app-1", dataset_id="dataset-1") - join2 = AppDatasetJoin(app_id="app-2", dataset_id="dataset-1") - session = self.session - with ( - app.test_request_context("/"), - patch("controllers.console.datasets.datasets.DatasetService.get_dataset", return_value=dataset), - patch("controllers.console.datasets.datasets.DatasetService.check_dataset_permission", return_value=None), - patch("controllers.console.datasets.datasets.DatasetService.get_related_apps", return_value=[join1, join2]), - patch( - "services.app.query_service.AppQueryService.related_apps", - return_value=[app1, app2], - ) as related_apps, - ): - response, status = method(api, session, make_account(), "dataset-1") - assert status == 200 - assert response["total"] == 2 - assert response["data"] == [ - { - "id": "app-1", - "name": "App 1", - "description": "desc", - "mode": "chat", - "icon_type": "emoji", - "icon": "🤖", - "icon_background": "#fff", - "icon_url": None, - }, - { - "id": "app-2", - "name": "App 2", - "description": "desc", - "mode": "chat", - "icon_type": "emoji", - "icon": "🤖", - "icon_background": "#fff", - "icon_url": None, - }, - ] - related_apps.assert_called_once_with(dataset.tenant_id, ["app-1", "app-2"]) - - def test_get_dataset_not_found(self, app: Flask): - api = DatasetRelatedAppListApi() - method = unwrap(api.get) - with ( - app.test_request_context("/"), - patch("controllers.console.datasets.datasets.DatasetService.get_dataset", return_value=None), - ): - with pytest.raises(NotFound): - method(api, self.session, make_account(), "dataset-1") - - def test_get_permission_denied(self, app: Flask): - api = DatasetRelatedAppListApi() - method = unwrap(api.get) - dataset = make_dataset(id="dataset-1") - with ( - app.test_request_context("/"), - patch("controllers.console.datasets.datasets.DatasetService.get_dataset", return_value=dataset), - patch( - "controllers.console.datasets.datasets.DatasetService.check_dataset_permission", - side_effect=services.errors.account.NoPermissionError("no permission"), - ), - ): - with pytest.raises(Forbidden): - method(api, self.session, make_account(), "dataset-1") - - def test_get_filters_none_apps(self, app: Flask): - api = DatasetRelatedAppListApi() - method = unwrap(api.get) - dataset = make_dataset(id="dataset-1") - app1 = make_related_app() - join1 = AppDatasetJoin(app_id="app-1", dataset_id="dataset-1") - join2 = AppDatasetJoin(app_id="app-2", dataset_id="dataset-1") - session = self.session - with ( - app.test_request_context("/"), - patch("controllers.console.datasets.datasets.DatasetService.get_dataset", return_value=dataset), - patch("controllers.console.datasets.datasets.DatasetService.check_dataset_permission", return_value=None), - patch("controllers.console.datasets.datasets.DatasetService.get_related_apps", return_value=[join1, join2]), - patch( - "services.app.query_service.AppQueryService.related_apps", - return_value=[app1], - ), - ): - response, status = method(api, session, make_account(), "dataset-1") - assert status == 200 - assert response["total"] == 1 - assert response["data"] == [ - { - "id": "app-1", - "name": "App", - "description": "desc", - "mode": "chat", - "icon_type": "emoji", - "icon": "🤖", - "icon_background": "#fff", - "icon_url": None, - } - ] +@pytest.mark.parametrize( + "info_list", + [ + {"data_source_type": "upload_file", "file_info_list": {}}, + {"data_source_type": "notion_import", "notion_info_list": [{"workspace_id": "workspace"}]}, + { + "data_source_type": "notion_import", + "notion_info_list": [ + { + "workspace_id": "workspace", + "credential_id": "credential", + "pages": [{"page_id": "page", "type": "unknown"}], + } + ], + }, + {"data_source_type": "unsupported"}, + ], +) +def test_new_estimate_sources_rejects_malformed_transport_shapes(info_list: dict[str, object]) -> None: + with pytest.raises(ValueError): + _new_estimate_sources(info_list) -class TestDatasetIndexingStatusApi(_UsesSQLiteSession): - def test_get_success_with_documents(self, app: Flask): - api = DatasetIndexingStatusApi() - method = unwrap(api.get) - dataset = make_dataset(id="dataset-1") - current_user = make_account() - document = make_document_status() - session = self.session - session.add(document) - session.add_all([make_document_segment(position=position, completed=True) for position in range(1, 4)]) - session.flush() - with ( - app.test_request_context("/"), - patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset) as get_dataset, - patch.object(DatasetService, "check_dataset_permission") as check_permission, - ): - response, status = method(api, session, "tenant-1", current_user, "dataset-1") - assert status == 200 - assert "data" in response - assert len(response["data"]) == 1 - item = response["data"][0] - assert item["completed_segments"] == 3 - assert item["total_segments"] == 3 - get_dataset.assert_called_once_with("dataset-1", "tenant-1", session=session) - check_permission.assert_called_once_with(dataset, current_user, session) - assert session.get(Document, "doc-1") is document +@pytest.mark.parametrize( + ("error", "expected_http_error"), + [ + (IndexingEstimateCredentialUnavailableError(), NotFoundError), + (EstimateSourceNotFoundError("source-1"), NotFoundError), + (DatasetNotFoundError(), NotFoundError), + (DatasetAccessDeniedError(), DatasetAccessDeniedRequestError), + (UnsupportedEstimateSourceError("unsupported"), InvalidArgumentError), + (IndexingEstimateProviderUnavailableError(), ProviderNotInitializeError), + (IndexingEstimateExecutionError(), IndexingEstimateError), + ], +) +def test_new_source_estimate_maps_application_errors( + error: Exception, + expected_http_error: type[Exception], +) -> None: + estimates = MagicMock() + estimates.estimate_new_sources.side_effect = error + registry = SimpleNamespace(knowledge=SimpleNamespace(indexing_estimates=estimates)) + api = DatasetIndexingEstimateApi() + method = unwrap(api.post) + payload = IndexingEstimatePayload( + info_list={"data_source_type": "upload_file", "file_info_list": {"file_ids": ["file-1"]}}, + process_rule={"mode": "automatic"}, + indexing_technique="economy", + ) + context = RequestContext("request-1", None, "account-1", "workspace-1") - def test_get_success_no_documents(self, app: Flask): - api = DatasetIndexingStatusApi() - method = unwrap(api.get) - dataset = make_dataset(id="dataset-1") - session = self.session - with ( - app.test_request_context("/"), - patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset), - patch.object(DatasetService, "check_dataset_permission"), - ): - response, status = method(api, session, "tenant-1", make_account(), "dataset-1") - assert status == 200 - assert response == {"data": []} - - def test_segment_counts_different_values(self, app: Flask): - api = DatasetIndexingStatusApi() - method = unwrap(api.get) - dataset = make_dataset(id="dataset-1") - document = make_document_status(indexing_status=IndexingStatus.INDEXING) - session = self.session - session.add(document) - session.add_all([make_document_segment(position=position, completed=position <= 2) for position in range(1, 6)]) - session.flush() - with ( - app.test_request_context("/"), - patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset), - patch.object(DatasetService, "check_dataset_permission"), - ): - response, status = method(api, session, "tenant-1", make_account(), "dataset-1") - assert status == 200 - item = response["data"][0] - assert item["completed_segments"] == 2 - assert item["total_segments"] == 5 - - def test_get_batch_segment_counts_uses_one_aggregate_query(self, app: Flask, sqlite_engine): - api = DatasetIndexingStatusApi() - method = unwrap(api.get) - dataset = make_dataset(id="dataset-1") - documents = [make_document_status(id=f"doc-{index}", position=index) for index in range(1, 6)] - session = self.session - session.add_all(documents) - session.add_all( - [ - make_document_segment( - position=position, - completed=position <= 2, - document_id=document.id, - ) - for document in documents - for position in range(1, 4) - ] - ) - session.add_all( - [ - make_document_segment( - position=4, - completed=True, - document_id=documents[0].id, - status=SegmentStatus.RE_SEGMENT, - ), - make_document_segment( - position=5, - completed=True, - document_id=documents[0].id, - tenant_id="other-tenant", - ), - ] - ) - session.flush() - - select_statements: list[str] = [] - - def record_select(_connection, _cursor, statement, _parameters, _context, _executemany): - if statement.lstrip().upper().startswith("SELECT"): - select_statements.append(statement) - - event.listen(sqlite_engine, "before_cursor_execute", record_select) - try: - with ( - app.test_request_context("/"), - patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset), - patch.object(DatasetService, "check_dataset_permission"), - ): - response, status = method(api, session, "tenant-1", make_account(), "dataset-1") - finally: - event.remove(sqlite_engine, "before_cursor_execute", record_select) - - assert status == 200 - assert len(response["data"]) == 5 - assert all(item["completed_segments"] == 2 for item in response["data"]) - assert all(item["total_segments"] == 3 for item in response["data"]) - assert len(select_statements) == 2 + with patch("controllers.console.datasets.datasets.application_services", return_value=registry): + with pytest.raises(expected_http_error): + method(api, payload, context) -class TestDatasetApiKeyApi(_UsesSQLiteSession): - context = RequestContext("request", None, "actor", "tenant-1") +@pytest.mark.parametrize( + ("dataset_id", "scene", "locator_type"), + [ + ("dataset-1", RBACPermission.DATASET_USE, DatasetId), + (None, RBACPermission.DATASET_CREATE_AND_MANAGEMENT, Workspace), + ], +) +def test_new_source_estimate_authorizes_before_execution(dataset_id, scene, locator_type): + estimates = MagicMock() + registry = SimpleNamespace(knowledge=SimpleNamespace(indexing_estimates=estimates)) + payload = IndexingEstimatePayload( + info_list={"data_source_type": "upload_file", "file_info_list": {"file_ids": ["file-1"]}}, + process_rule={"mode": "automatic"}, + indexing_technique="economy", + dataset_id=dataset_id, + ) + with ( + patch.object(controller, "application_services", return_value=registry), + patch.object(controller, "enforce_rbac_checks", side_effect=Forbidden) as enforce_checks, + pytest.raises(Forbidden), + ): + unwrap(DatasetIndexingEstimateApi.post)(DatasetIndexingEstimateApi(), payload, CONTEXT) - @pytest.fixture(autouse=True) - def key_services(self, sqlite_session: Session, sqlite_session_factory: sessionmaker[Session], monkeypatch): - tenant = Tenant(name="Workspace") - tenant.id = "tenant-1" - sqlite_session.add(tenant) - sqlite_session.commit() - service = build_dataset_api_key_service(database_client=sqlite_session_factory) - monkeypatch.setattr( - datasets_controller, - "application_services", - lambda: SimpleNamespace(dataset_api_keys=service), - ) - - def test_get_api_keys_success(self, app: Flask): - api = DatasetApiKeyApi() - self.session.add_all( - [ - ApiToken(id="key-1", tenant_id="tenant-1", type="dataset", token="dataset-aaaa1111bbbb"), - ApiToken(id="key-2", tenant_id="tenant-1", type="dataset", token="dataset-cccc2222dddd"), - ] - ) - self.session.commit() - with app.test_request_context("/"): - response = unwrap(api.get)(api, self.context) - assert [key["token"] for key in response["data"]] == ["datas...bbbb", "datas...dddd"] - - def test_post_create_api_key_success(self, app: Flask): - api = DatasetApiKeyApi() - with app.test_request_context("/"): - response, status = unwrap(api.post)(api, self.context) - assert status == 200 - assert response["token"].startswith("dataset-") - assert response["type"] == "dataset" - assert response["dataset_ids"] == [] - assert response["created_at"] is not None - stored = self.session.get(ApiToken, response["id"]) - assert stored is not None - assert stored.token == response["token"] - - def test_post_exceed_max_keys(self, app: Flask): - api = DatasetApiKeyApi() - self.session.add_all( - [ApiToken(id=f"key-{i}", tenant_id="tenant-1", type="dataset", token=f"ds-{i}") for i in range(10)] - ) - self.session.commit() - with app.test_request_context("/"), pytest.raises(BadRequest) as exc_info: - unwrap(api.post)(api, self.context) - assert vars(exc_info.value)["data"] == { - "message": "Cannot create more than 10 API keys for this resource type.", - "custom": "max_keys_exceeded", - } - - def test_scoped_key_persists_bindings_and_lists_masked_scope(self, app: Flask): - api = DatasetApiKeyApi() - self.session.add_all( - [ - make_dataset(id="ds-1", tenant_id="tenant-1"), - make_dataset(id="ds-2", tenant_id="tenant-1"), - ] - ) - self.session.commit() - with app.test_request_context("/", json={"dataset_ids": ["ds-1", "ds-2", "ds-1"]}): - response, status = unwrap(api.post)(api, self.context) - listed = unwrap(api.get)(api, self.context) - assert status == 200 - assert response["dataset_ids"] == ["ds-1", "ds-2"] - bindings = self.session.scalars( - select(DatasetApiTokenBinding).where(DatasetApiTokenBinding.api_token_id == response["id"]) - ).all() - assert {binding.dataset_id for binding in bindings} == {"ds-1", "ds-2"} - assert set(listed["data"][0]["dataset_ids"]) == {"ds-1", "ds-2"} - assert listed["data"][0]["token"] != response["token"] - - def test_post_rejects_dataset_ids_from_another_tenant(self, app: Flask): - api = DatasetApiKeyApi() - self.session.add(make_dataset(id="foreign-ds", tenant_id="foreign")) - self.session.commit() - with app.test_request_context("/", json={"dataset_ids": ["foreign-ds"]}), pytest.raises(BadRequest) as exc_info: - unwrap(api.post)(api, self.context) - assert "Unknown knowledge base id(s)" in vars(exc_info.value)["data"]["message"] - assert self.session.scalar(select(ApiToken)) is None - - @pytest.mark.parametrize("dataset_ids", ["not-a-list", [1]]) - def test_post_rejects_invalid_dataset_ids(self, app: Flask, dataset_ids): - api = DatasetApiKeyApi() - with app.test_request_context("/", json={"dataset_ids": dataset_ids}), pytest.raises(BadRequest) as exc_info: - unwrap(api.post)(api, self.context) - assert vars(exc_info.value)["data"]["message"] == "dataset_ids must be a list of strings." - - def test_delete_success(self, app: Flask): - api = DatasetApiDeleteApi() - self.session.add(ApiToken(id="api-key-id", tenant_id="tenant-1", type="dataset", token="dataset-secret")) - self.session.commit() - with app.test_request_context("/"), patch("services.api_token_service.ApiTokenCache.delete") as invalidate: - result = unwrap(api.delete)(api, self.context, "api-key-id") - assert result == ("", 204) - self.session.expire_all() - assert self.session.get(ApiToken, "api-key-id") is None - invalidate.assert_called_once_with("dataset-secret", "dataset") - - def test_delete_key_not_found(self, app: Flask): - api = DatasetApiDeleteApi() - with app.test_request_context("/"), pytest.raises(NotFound): - unwrap(api.delete)(api, self.context, "api-key-id") - - -class TestDatasetEnableApiApi(_UsesSQLiteSession): - @pytest.mark.parametrize(("status_value", "enabled"), [("enable", True), ("disable", False)]) - def test_update_api_status(self, app: Flask, status_value: str, enabled: bool): - api = DatasetEnableApiApi() - method = unwrap(api.post) - dataset = make_dataset(id="dataset-1") - current_user = make_account() - session = self.session - with ( - app.test_request_context("/"), - patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset) as get_dataset, - patch.object(DatasetService, "check_dataset_permission") as check_permission, - patch.object(DatasetService, "update_dataset_api_status") as update_status, - ): - response, status = method(api, session, "tenant-1", current_user, "dataset-1", status_value) - assert status == 200 - assert response["result"] == "success" - get_dataset.assert_called_once_with("dataset-1", "tenant-1", session=session) - check_permission.assert_called_once_with(dataset, current_user, session) - update_status.assert_called_once_with(dataset, enabled, current_user, session) - - def test_rejects_non_editor(self, app: Flask): - api = DatasetEnableApiApi() - method = unwrap(api.post) - dataset = make_dataset(id="dataset-1") - session = self.session - with ( - app.test_request_context("/"), - patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset), - patch.object(DatasetService, "check_dataset_permission"), - patch.object(DatasetService, "update_dataset_api_status") as update_status, - ): - with pytest.raises(Forbidden): - method( - api, - session, - "tenant-1", - make_account(TenantAccountRole.NORMAL), - "dataset-1", - "enable", - ) - update_status.assert_not_called() - - -class TestDatasetApiBaseUrlApi: - def test_get_api_base_url_from_config(self, app: Flask, config_overrides: Callable[..., None]): - config_overrides(SERVICE_API_URL="https://example.com") - api = DatasetApiBaseUrlApi() - method = unwrap(api.get) - with ( - app.test_request_context("/"), - ): - response = method(api) - assert response["api_base_url"] == "https://example.com/v1" - - def test_get_api_base_url_from_request(self, app: Flask, config_overrides: Callable[..., None]): - config_overrides(SERVICE_API_URL=None) - api = DatasetApiBaseUrlApi() - method = unwrap(api.get) - with ( - app.test_request_context("http://localhost:5000/"), - ): - response = method(api) - assert response["api_base_url"] == "http://localhost:5000/v1" - - def test_get_api_base_url_no_double_v1(self, app: Flask, config_overrides: Callable[..., None]): - config_overrides(SERVICE_API_URL="https://example.com/v1") - api = DatasetApiBaseUrlApi() - method = unwrap(api.get) - with ( - app.test_request_context("/"), - ): - response = method(api) - assert response["api_base_url"] == "https://example.com/v1" - - -class TestDatasetRetrievalSettingApi: - def test_get_success(self, app: Flask, config_overrides: Callable[..., None]): - config_overrides(VECTOR_STORE="qdrant") - api = DatasetRetrievalSettingApi() - method = unwrap(api.get) - with ( - app.test_request_context("/"), - patch( - "controllers.console.datasets.datasets._get_retrieval_methods_by_vector_type", - return_value={"retrieval_method": ["semantic", "hybrid"]}, - ), - ): - response = method(api) - assert "retrieval_method" in response - - def test_tidb_vector_returns_semantic_only_when_fulltext_disabled(self, config_overrides: Callable[..., None]): - config_overrides(TIDB_VECTOR_ENABLE_FULLTEXT_SEARCH=False) - response = _get_retrieval_methods_by_vector_type(VectorType.TIDB_VECTOR) - - assert response["retrieval_method"] == [RetrievalMethod.SEMANTIC_SEARCH.value] - - def test_tidb_vector_returns_full_methods_when_fulltext_enabled(self, config_overrides: Callable[..., None]): - config_overrides(TIDB_VECTOR_ENABLE_FULLTEXT_SEARCH=True) - response = _get_retrieval_methods_by_vector_type(VectorType.TIDB_VECTOR) - - assert response["retrieval_method"] == [ - RetrievalMethod.SEMANTIC_SEARCH.value, - RetrievalMethod.FULL_TEXT_SEARCH.value, - RetrievalMethod.HYBRID_SEARCH.value, - ] - - -class TestDatasetRetrievalSettingMockApi(_UsesSQLiteSession): - def test_get_success(self, app: Flask): - api = DatasetRetrievalSettingMockApi() - method = unwrap(api.get) - with ( - app.test_request_context("/"), - patch( - "controllers.console.datasets.datasets._get_retrieval_methods_by_vector_type", - return_value={"retrieval_method": ["semantic"]}, - ), - ): - response = method(api, "milvus") - assert response["retrieval_method"] == ["semantic"] - - -class TestDatasetErrorDocs(_UsesSQLiteSession): - def test_get_success(self, app: Flask): - api = DatasetErrorDocs() - method = unwrap(api.get) - dataset = make_dataset(id="dataset-1") - error_doc = make_document_status(id="error-doc", indexing_status=IndexingStatus.ERROR, error="failed") - current_user = make_account() - session = self.session - with ( - app.test_request_context("/"), - patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset) as get_dataset, - patch.object(DatasetService, "check_dataset_permission") as check_permission, - patch( - "controllers.console.datasets.datasets.DocumentService.get_error_documents_by_dataset_ref", - return_value=[error_doc], - ) as get_error_documents, - ): - response, status = method(api, session, "tenant-1", current_user, "dataset-1") - assert status == 200 - assert response["total"] == 1 - get_dataset.assert_called_once_with("dataset-1", "tenant-1", session=session) - check_permission.assert_called_once_with(dataset, current_user, session) - get_error_documents.assert_called_once_with(DatasetRef("tenant-1", "dataset-1"), session) - - def test_get_dataset_not_found(self, app: Flask): - api = DatasetErrorDocs() - method = unwrap(api.get) - session = self.session - with ( - app.test_request_context("/"), - patch.object(DatasetService, "get_dataset_for_tenant", return_value=None) as get_dataset, - ): - with pytest.raises(NotFound): - method(api, session, "tenant-1", make_account(), "dataset-1") - get_dataset.assert_called_once_with("dataset-1", "tenant-1", session=session) - - -class TestDatasetPermissionUserListApi(_UsesSQLiteSession): - def test_get_success(self, app: Flask): - api = DatasetPermissionUserListApi() - method = unwrap(api.get) - dataset = make_dataset(id="dataset-1") - users = ["u1", "u2"] - with ( - app.test_request_context("/"), - patch("controllers.console.datasets.datasets.DatasetService.get_dataset", return_value=dataset), - patch("controllers.console.datasets.datasets.DatasetService.check_dataset_permission", return_value=None), - patch( - "controllers.console.datasets.datasets.DatasetPermissionService.get_dataset_partial_member_list", - return_value=users, - ), - ): - response, status = method(api, self.session, make_account(), "dataset-1") - assert status == 200 - assert response["data"] == users - - def test_get_permission_denied(self, app: Flask): - api = DatasetPermissionUserListApi() - method = unwrap(api.get) - dataset = make_dataset(id="dataset-1") - with ( - app.test_request_context("/"), - patch("controllers.console.datasets.datasets.DatasetService.get_dataset", return_value=dataset), - patch( - "controllers.console.datasets.datasets.DatasetService.check_dataset_permission", - side_effect=services.errors.account.NoPermissionError("no permission"), - ), - ): - with pytest.raises(Forbidden): - method(api, self.session, make_account(), "dataset-1") - - -class TestDatasetAutoDisableLogApi(_UsesSQLiteSession): - def test_get_success(self, app: Flask): - api = DatasetAutoDisableLogApi() - method = unwrap(api.get) - dataset = make_dataset(id="dataset-1") - logs = {"document_ids": ["doc-1"], "count": 1} - current_user = make_account() - session = self.session - with ( - app.test_request_context("/"), - patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset) as get_dataset, - patch.object(DatasetService, "check_dataset_permission") as check_permission, - patch.object(DatasetService, "get_dataset_auto_disable_logs", return_value=logs) as get_logs, - ): - response, status = method(api, session, "tenant-1", current_user, "dataset-1") - assert status == 200 - assert response == logs - get_dataset.assert_called_once_with("dataset-1", "tenant-1", session=session) - check_permission.assert_called_once_with(dataset, current_user, session) - get_logs.assert_called_once_with(DatasetRef("tenant-1", "dataset-1"), session) - - def test_get_dataset_not_found(self, app: Flask): - api = DatasetAutoDisableLogApi() - method = unwrap(api.get) - session = self.session - with ( - app.test_request_context("/"), - patch.object(DatasetService, "get_dataset_for_tenant", return_value=None) as get_dataset, - ): - with pytest.raises(NotFound): - method(api, session, "tenant-1", make_account(), "dataset-1") - get_dataset.assert_called_once_with("dataset-1", "tenant-1", session=session) + enforce_checks.assert_called_once() + kwargs = enforce_checks.call_args.kwargs + assert kwargs["tenant_id"] == CONTEXT.active_workspace_id + assert kwargs["account_id"] == CONTEXT.account_id + assert kwargs["path_args"] == ({"dataset_id": dataset_id} if dataset_id else None) + [check] = kwargs["checks"] + assert check.scene is scene + assert isinstance(check.locator, locator_type) + estimates.estimate_new_sources.assert_not_called() diff --git a/api/tests/unit_tests/controllers/console/datasets/test_datasets_document.py b/api/tests/unit_tests/controllers/console/datasets/test_datasets_document.py index b4a737b8171..1286f1a3f90 100644 --- a/api/tests/unit_tests/controllers/console/datasets/test_datasets_document.py +++ b/api/tests/unit_tests/controllers/console/datasets/test_datasets_document.py @@ -1,17 +1,16 @@ import datetime -import json from inspect import unwrap -from unittest.mock import MagicMock, patch +from types import SimpleNamespace +from unittest.mock import MagicMock, create_autospec, patch +from uuid import UUID import pytest -from flask import Flask -from sqlalchemy import event -from sqlalchemy.orm import Session from werkzeug.exceptions import Forbidden, NotFound -import services -from controllers.common.rbac import DatasetId -from controllers.console import console_ns +from controllers.common.errors import InvalidArgumentError, NotFoundError +from controllers.common.rbac import DatasetByDocument, DatasetId, RBACPermission +from controllers.console.app.error import ProviderNotInitializeError +from controllers.console.datasets import datasets_document as controller from controllers.console.datasets.datasets_document import ( DatasetDocumentListApi, DatasetInitApi, @@ -31,7 +30,6 @@ from controllers.console.datasets.datasets_document import ( DocumentRecoverApi, DocumentRenameApi, DocumentRenamePayload, - DocumentResource, DocumentRetryApi, DocumentRetryPayload, DocumentStatusApi, @@ -41,190 +39,466 @@ from controllers.console.datasets.datasets_document import ( WebsiteDocumentSyncApi, ) from controllers.console.datasets.error import ( + DatasetAccessDeniedRequestError, DocumentAlreadyFinishedError, DocumentIndexingError, IndexingEstimateError, InvalidActionError, InvalidMetadataError, ) -from controllers.console.wraps import RBACPermission -from core.entities.knowledge_entities import IndexingEstimate -from core.rag.index_processor.constant.index_type import IndexStructureType -from extensions.storage.storage_type import StorageType -from models.account import Account, TenantAccountRole -from models.dataset import ( - Dataset, - DatasetPermissionEnum, - DatasetProcessRule, - DocumentPipelineExecutionLog, - DocumentSegment, +from machinery.context import RequestContext +from services.knowledge.dataset_access import DatasetAccessDeniedError, DatasetNotFoundError +from services.knowledge.documents.application import ( + DatasetDocumentApplicationService, + DocumentIndexingStateError, + DocumentInvalidActionError, + DocumentListFilter, + DocumentNotFoundError, + DocumentProviderError, ) -from models.dataset import Document as DatasetDocument -from models.enums import CreatorUserRole, DataSourceType, DocumentCreatedFrom, IndexingStatus -from models.model import UploadFile -from services.dataset_ref_service import DatasetRef, DocumentRef -from services.vector_space_admission_service import ( - VECTOR_SPACE_ADMISSION_ERROR_CODE, - format_vector_space_admission_error, +from services.knowledge.indexing.estimate import ( + EstimateDocumentAlreadyFinishedError, + EstimateDocumentNotFoundError, + EstimateSourceNotFoundError, + IndexingEstimateCredentialUnavailableError, + IndexingEstimateExecutionError, + IndexingEstimateProviderUnavailableError, + UnsupportedEstimateSourceError, ) -from tests.unit_tests.config_override import config_overrides_context from tests.unit_tests.controllers.rbac_introspection import rbac_checks - -def make_serializable_document(**overrides): - return make_document(data_source_info=json.dumps({"upload_file_id": "file-1"}), **overrides) +CONTEXT = RequestContext("request-1", None, "account-1", "tenant-1") +DS = UUID(int=1) +DOC = UUID(int=2) -def make_document_detail(**overrides): - doc_metadata_details = overrides.pop("doc_metadata_details", []) - return make_document( - data_source_info=json.dumps({"upload_file_id": "file-1"}), - created_at=datetime.datetime(2024, 1, 1, tzinfo=datetime.UTC), - tokens=10, - doc_type="others", - doc_metadata=doc_metadata_details, - doc_language="English", - **overrides, +@pytest.fixture +def documents(monkeypatch): + documents = create_autospec(DatasetDocumentApplicationService, instance=True, spec_set=True) + monkeypatch.setattr( + controller, "application_services", lambda: SimpleNamespace(knowledge=SimpleNamespace(documents=documents)) ) + monkeypatch.setattr(controller, "check_knowledge_rate_limit", lambda: None) + return documents -def make_dataset(**overrides): - attrs = { - "id": "ds-1", - "tenant_id": "tenant-1", - "name": "Dataset", - "indexing_technique": "economy", - "chunk_structure": IndexStructureType.PARAGRAPH_INDEX, - "created_by": "u1", - "summary_index_setting": {"enable": True}, +def test_process_rule_authorizes_document_dataset_before_read(app, documents): + with ( + app.test_request_context(f"/?document_id={DOC}"), + patch.object(controller, "enforce_rbac_checks", side_effect=Forbidden) as enforce_checks, + pytest.raises(Forbidden), + ): + unwrap(GetProcessRuleApi.get)(GetProcessRuleApi(), CONTEXT) + + enforce_checks.assert_called_once() + kwargs = enforce_checks.call_args.kwargs + assert kwargs["tenant_id"] == CONTEXT.active_workspace_id + assert kwargs["account_id"] == CONTEXT.account_id + assert kwargs["path_args"] == {"document_id": str(DOC)} + [check] = kwargs["checks"] + assert check.scene is RBACPermission.DATASET_READONLY + assert isinstance(check.locator, DatasetByDocument) + documents.get_process_rule.assert_not_called() + + +def test_default_process_rule_does_not_require_dataset_permission(app, documents): + documents.get_process_rule.return_value = {"mode": "automatic", "rules": {}, "limits": {}} + with app.test_request_context("/"), patch.object(controller, "enforce_rbac_checks") as enforce_checks: + unwrap(GetProcessRuleApi.get)(GetProcessRuleApi(), CONTEXT) + enforce_checks.assert_not_called() + documents.get_process_rule.assert_called_once_with(CONTEXT, document_id=None) + + +@pytest.mark.parametrize( + ("resource", "verb", "operation", "kwargs"), + [ + (GetProcessRuleApi, GetProcessRuleApi.get, lambda service: service.get_process_rule, {}), + ( + DatasetDocumentListApi, + DatasetDocumentListApi.get, + lambda service: service.list_documents, + {"dataset_id": DS}, + ), + ( + DatasetDocumentListApi, + DatasetDocumentListApi.post, + lambda service: service.create_documents, + {"dataset_id": DS}, + ), + ( + DatasetDocumentListApi, + DatasetDocumentListApi.delete, + lambda service: service.delete_documents, + {"dataset_id": DS}, + ), + (DatasetInitApi, DatasetInitApi.post, lambda service: service.initialize_dataset, {}), + ( + DocumentBatchIndexingStatusApi, + DocumentBatchIndexingStatusApi.get, + lambda service: service.get_batch_indexing_status, + {"dataset_id": DS, "batch": "batch-1"}, + ), + ( + DocumentIndexingStatusApi, + DocumentIndexingStatusApi.get, + lambda service: service.get_indexing_status, + {"dataset_id": DS, "document_id": DOC}, + ), + (DocumentApi, DocumentApi.get, lambda service: service.get_document, {"dataset_id": DS, "document_id": DOC}), + ( + DocumentApi, + DocumentApi.delete, + lambda service: service.delete_document, + {"dataset_id": DS, "document_id": DOC}, + ), + ( + DocumentDownloadApi, + DocumentDownloadApi.get, + lambda service: service.get_download_url, + {"dataset_id": DS, "document_id": DOC}, + ), + ( + DocumentBatchDownloadZipApi, + DocumentBatchDownloadZipApi.post, + lambda service: service.build_download_zip, + {"dataset_id": DS}, + ), + ( + DocumentProcessingApi, + DocumentProcessingApi.patch, + lambda service: service.update_processing, + {"dataset_id": DS, "document_id": DOC, "action": "pause"}, + ), + ( + DocumentMetadataApi, + DocumentMetadataApi.put, + lambda service: service.update_metadata, + { + "dataset_id": DS, + "document_id": DOC, + "req_data": DocumentMetadataUpdatePayload(doc_type="book", doc_metadata={"title": "Book"}), + }, + ), + ( + DocumentStatusApi, + DocumentStatusApi.patch, + lambda service: service.change_status, + {"dataset_id": DS, "action": "enable"}, + ), + ( + DocumentPauseApi, + DocumentPauseApi.patch, + lambda service: service.pause_document, + {"dataset_id": DS, "document_id": DOC}, + ), + ( + DocumentRecoverApi, + DocumentRecoverApi.patch, + lambda service: service.recover_document, + {"dataset_id": DS, "document_id": DOC}, + ), + ( + DocumentRetryApi, + DocumentRetryApi.post, + lambda service: service.retry_documents, + {"dataset_id": DS, "req_data": DocumentRetryPayload(document_ids=[str(DOC)])}, + ), + ( + DocumentRenameApi, + DocumentRenameApi.post, + lambda service: service.rename_document, + {"dataset_id": DS, "document_id": DOC, "req_data": DocumentRenamePayload(name="New name")}, + ), + ( + WebsiteDocumentSyncApi, + WebsiteDocumentSyncApi.get, + lambda service: service.sync_website, + {"dataset_id": DS, "document_id": DOC}, + ), + ( + DocumentPipelineExecutionLogApi, + DocumentPipelineExecutionLogApi.get, + lambda service: service.get_execution_log, + {"dataset_id": DS, "document_id": DOC}, + ), + ( + DocumentGenerateSummaryApi, + DocumentGenerateSummaryApi.post, + lambda service: service.generate_summary, + {"dataset_id": DS, "req_data": GenerateSummaryPayload(document_list=[str(DOC)])}, + ), + ( + DocumentSummaryStatusApi, + DocumentSummaryStatusApi.get, + lambda service: service.get_summary_status, + {"dataset_id": DS, "document_id": DOC}, + ), + ], +) +def test_endpoints_pass_request_context_and_map_application_errors(app, documents, resource, verb, operation, kwargs): + operation(documents).side_effect = DocumentNotFoundError("Document not found.") + with app.test_request_context( + "/?document_id=" + str(DOC), json={"indexing_technique": "economy", "document_ids": [str(DOC)]} + ): + with pytest.raises(NotFound): + unwrap(verb)(resource(), request_context=CONTEXT, **kwargs) + called = operation(documents) + assert called.call_count == 1 + assert called.call_args.args == (CONTEXT,) + if "dataset_id" in kwargs: + assert called.call_args.kwargs["dataset_id"] == str(DS) + if "document_id" in kwargs: + assert called.call_args.kwargs["document_id"] == str(DOC) + + +@pytest.mark.parametrize(("fetch", "expected"), [("true", True), ("YES", True), ("false", False), ("invalid", False)]) +def test_list_preserves_filters_pagination_and_fetch_parsing(app, documents, fetch, expected): + documents.list_documents.return_value = {"data": [], "has_more": False, "page": 2, "limit": 3, "total": 3} + with app.test_request_context(f"/?page=2&limit=3&keyword=needle&status=available&sort=hit_count&fetch={fetch}"): + response = unwrap(DatasetDocumentListApi.get)(DatasetDocumentListApi(), CONTEXT, DS) + documents.list_documents.assert_called_once_with( + CONTEXT, + dataset_id=str(DS), + query=DocumentListFilter( + page=2, limit=3, search="needle", status="available", sort="hit_count", fetch=expected + ), + ) + assert response == {"data": [], "has_more": False, "page": 2, "limit": 3, "total": 3} + + +@pytest.mark.parametrize("metadata", ["only", "without", "all"]) +def test_document_detail_metadata_projection_and_raw_source(app, documents, metadata): + documents.get_document.return_value = { + "id": str(DOC), + "doc_type": "others", + "doc_metadata": [], + "name": "Name", + "data_source_info": {"file_path": "/path/to/local/file"}, + "data_source_detail_dict": {}, } - attrs.update(overrides) - return Dataset(**attrs) + with app.test_request_context(f"/?metadata={metadata}"): + response, status = unwrap(DocumentApi.get)(DocumentApi(), CONTEXT, DS, DOC) + assert status == 200 + documents.get_document.assert_called_once_with( + CONTEXT, dataset_id=str(DS), document_id=str(DOC), metadata_only=metadata == "only" + ) + if metadata == "only": + assert response == {"id": str(DOC), "doc_type": "others", "doc_metadata": []} + else: + assert response["data_source_info"] == {"file_path": "/path/to/local/file"} + assert ("doc_metadata" in response) is (metadata == "all") -def make_document(**overrides): - attrs = { - "id": "doc-1", - "tenant_id": "tenant-1", - "dataset_id": "ds-1", - "position": 1, - "data_source_type": DataSourceType.UPLOAD_FILE, - "data_source_info": None, +def test_invalid_metadata_fails_before_application_call(app, documents): + with app.test_request_context("/?metadata=invalid"), pytest.raises(InvalidMetadataError): + unwrap(DocumentApi.get)(DocumentApi(), CONTEXT, DS, DOC) + documents.get_document.assert_not_called() + + +@pytest.mark.parametrize( + ("resource", "verb", "kwargs"), + [ + (DocumentApi, DocumentApi.delete, {"document_id": DOC}), + (DatasetDocumentListApi, DatasetDocumentListApi.delete, {}), + (DocumentPauseApi, DocumentPauseApi.patch, {"document_id": DOC}), + (DocumentRecoverApi, DocumentRecoverApi.patch, {"document_id": DOC}), + (DocumentRetryApi, DocumentRetryApi.post, {"req_data": DocumentRetryPayload(document_ids=[str(DOC)])}), + ], +) +@pytest.mark.usefixtures("documents") +def test_no_content_mutations_return_empty_204(app, resource, verb, kwargs): + with app.test_request_context("/?document_id=" + str(DOC)): + result = unwrap(verb)(resource(), request_context=CONTEXT, dataset_id=DS, **kwargs) + assert result == ("", 204) + + +@pytest.mark.parametrize( + ("resource", "verb", "operation", "kwargs"), + [ + ( + DocumentProcessingApi, + DocumentProcessingApi.patch, + lambda service: service.update_processing, + {"document_id": DOC, "action": "resume"}, + ), + (DocumentStatusApi, DocumentStatusApi.patch, lambda service: service.change_status, {"action": "disable"}), + ( + WebsiteDocumentSyncApi, + WebsiteDocumentSyncApi.get, + lambda service: service.sync_website, + {"document_id": DOC}, + ), + ( + DocumentGenerateSummaryApi, + DocumentGenerateSummaryApi.post, + lambda service: service.generate_summary, + {"req_data": GenerateSummaryPayload(document_list=[str(DOC)])}, + ), + ], +) +def test_success_mutations_return_success_response(app, documents, resource, verb, operation, kwargs): + with app.test_request_context("/?document_id=" + str(DOC)): + result = unwrap(verb)(resource(), request_context=CONTEXT, dataset_id=DS, **kwargs) + assert result == ({"result": "success"}, 200) + assert operation(documents).call_count == 1 + + +@pytest.mark.parametrize("resource", [DatasetInitApi, DatasetDocumentListApi]) +def test_creation_serializes_materialized_dataset_and_documents(app, documents, resource): + operation = documents.initialize_dataset if resource is DatasetInitApi else documents.create_documents + operation.return_value = { + "dataset": {"id": str(DS), "name": "Dataset", "created_at": datetime.datetime(2024, 1, 1, tzinfo=datetime.UTC)}, + "documents": [ + { + "id": str(DOC), + "name": "Document", + "data_source_detail_dict": {}, + "data_source_info_dict": {}, + "doc_metadata_details": [], + } + ], "batch": "batch-1", - "name": "Document", - "created_from": DocumentCreatedFrom.WEB, - "created_by": "u1", - "indexing_status": IndexingStatus.COMPLETED, - "enabled": True, - "archived": False, - "doc_metadata": None, - "doc_form": IndexStructureType.PARAGRAPH_INDEX, - "need_summary": False, } - attrs.update(overrides) - return DatasetDocument(**attrs) + with app.test_request_context("/", json={"indexing_technique": "economy"}): + response = unwrap(resource.post)( + resource(), request_context=CONTEXT, **({} if resource is DatasetInitApi else {"dataset_id": DS}) + ) + assert response["dataset"]["created_at"] == 1704067200 + assert response["documents"][0]["data_source_info"] == {} + assert response["documents"][0]["doc_metadata"] == [] + assert response["batch"] == "batch-1" -def make_account(role: TenantAccountRole = TenantAccountRole.EDITOR) -> Account: - account = Account(name="Test User", email="test@example.com") - account.id = "u1" - account.role = role - return account +@pytest.mark.parametrize("resource", [DocumentIndexingStatusApi, DocumentBatchIndexingStatusApi]) +def test_indexing_status_keeps_timestamps_counts_and_error_fields(app, documents, resource): + value = { + "id": str(DOC), + "indexing_status": "paused", + "completed_segments": 2, + "total_segments": 3, + **dict.fromkeys( + ( + "processing_started_at", + "parsing_completed_at", + "cleaning_completed_at", + "splitting_completed_at", + "completed_at", + "paused_at", + "error", + "stopped_at", + ) + ), + } + value["processing_started_at"] = datetime.datetime(2024, 1, 1, tzinfo=datetime.UTC) + batch = resource is DocumentBatchIndexingStatusApi + operation = documents.get_batch_indexing_status if batch else documents.get_indexing_status + operation.return_value = {"data": [value]} if batch else value + with app.test_request_context("/"): + response = unwrap(resource.get)( + resource(), + request_context=CONTEXT, + dataset_id=DS, + **({"batch": "batch-1"} if batch else {"document_id": DOC}), + ) + row = response["data"][0] if batch else response + assert row["indexing_status"] == "paused" + assert row["processing_started_at"] == 1704067200 + assert row["completed_segments"] == 2 + assert row["total_segments"] == 3 -def make_segment(*, position: int, completed: bool = True, document_id: str = "doc-1") -> DocumentSegment: - return DocumentSegment( - tenant_id="tenant-1", - dataset_id="ds-1", - document_id=document_id, - position=position, - content=f"segment {position}", - word_count=2, - tokens=2, - created_by="u1", - completed_at=datetime.datetime(2024, 1, 1, tzinfo=datetime.UTC) if completed else None, +@pytest.mark.parametrize( + ("error", "http_error"), + [ + (DatasetNotFoundError(), NotFound), + (DocumentNotFoundError("missing"), NotFound), + (DatasetAccessDeniedError(), Forbidden), + (DocumentIndexingStateError("indexing"), DocumentIndexingError), + (DocumentInvalidActionError("invalid"), InvalidActionError), + (DocumentProviderError("provider", kind="uninitialized"), ProviderNotInitializeError), + ], +) +def test_document_errors_are_mapped_at_transport_boundary(error, http_error): + with pytest.raises(http_error): + controller._raise_document_error(error) + + +def test_unexpected_error_is_not_hidden(): + error = RuntimeError("unexpected") + with pytest.raises(RuntimeError) as raised: + controller._raise_document_error(error) + assert raised.value is error + + +class TestIndexingEstimateExceptionMapping: + @staticmethod + def _registry(method_name: str, error: Exception) -> SimpleNamespace: + estimates = MagicMock() + method = estimates.estimate_document if method_name == "estimate_document" else estimates.estimate_batch + method.side_effect = error + return SimpleNamespace(knowledge=SimpleNamespace(indexing_estimates=estimates)) + + @pytest.mark.parametrize( + ("error", "expected_http_error"), + [ + (DatasetNotFoundError(), NotFoundError), + (EstimateDocumentNotFoundError(), NotFoundError), + (EstimateSourceNotFoundError("source-1"), NotFoundError), + (IndexingEstimateCredentialUnavailableError(), NotFoundError), + (DatasetAccessDeniedError(), DatasetAccessDeniedRequestError), + (EstimateDocumentAlreadyFinishedError(), DocumentAlreadyFinishedError), + (UnsupportedEstimateSourceError("unsupported"), InvalidArgumentError), + (IndexingEstimateProviderUnavailableError(), ProviderNotInitializeError), + (IndexingEstimateExecutionError(), IndexingEstimateError), + ], ) + def test_document_estimate_maps_application_errors( + self, + error: Exception, + expected_http_error: type[Exception], + ) -> None: + api = DocumentIndexingEstimateApi() + method = unwrap(api.get) + context = RequestContext("request-1", None, "account-1", "workspace-1") + registry = self._registry("estimate_document", error) + with patch( + "controllers.console.datasets.datasets_document.application_services", + return_value=registry, + ): + with pytest.raises(expected_http_error): + method(api, context, UUID(int=1), UUID(int=2)) -def make_upload_file() -> UploadFile: - upload_file = UploadFile( - tenant_id="tenant-1", - storage_type=StorageType.LOCAL, - key="document.txt", - name="document.txt", - size=10, - extension="txt", - mime_type="text/plain", - created_by_role=CreatorUserRole.ACCOUNT, - created_by="u1", - created_at=datetime.datetime(2024, 1, 1, tzinfo=datetime.UTC), - used=False, + @pytest.mark.parametrize( + ("error", "expected_http_error"), + [ + (DatasetNotFoundError(), NotFoundError), + (EstimateDocumentNotFoundError(), NotFoundError), + (EstimateSourceNotFoundError("source-1"), NotFoundError), + (IndexingEstimateCredentialUnavailableError(), NotFoundError), + (DatasetAccessDeniedError(), DatasetAccessDeniedRequestError), + (EstimateDocumentAlreadyFinishedError(), DocumentAlreadyFinishedError), + (UnsupportedEstimateSourceError("unsupported"), InvalidArgumentError), + (IndexingEstimateProviderUnavailableError(), ProviderNotInitializeError), + (IndexingEstimateExecutionError(), IndexingEstimateError), + ], ) - upload_file.id = "file-1" - return upload_file + def test_batch_estimate_maps_application_errors( + self, + error: Exception, + expected_http_error: type[Exception], + ) -> None: + api = DocumentBatchIndexingEstimateApi() + method = unwrap(api.get) + context = RequestContext("request-1", None, "account-1", "workspace-1") + registry = self._registry("estimate_batch", error) - -@pytest.fixture -def tenant_ctx(): - account = Account(name="Test User", email="test@example.com") - account.id = "u1" - account.role = TenantAccountRole.EDITOR - return (account, "tenant-1") - - -@pytest.fixture(autouse=True) -def bypass_knowledge_rate_limit(): - with patch("controllers.console.datasets.datasets_document.check_knowledge_rate_limit") as check: - yield check - - -@pytest.fixture -def patch_tenant(tenant_ctx): - return tenant_ctx - - -@pytest.fixture -def dataset(): - return make_dataset() - - -@pytest.fixture -def document(): - return make_document( - indexing_status=IndexingStatus.INDEXING, - is_paused=False, - data_source_info=json.dumps({"upload_file_id": "file-1"}), - ) - - -@pytest.fixture -def patch_dataset(dataset): - with patch("controllers.console.datasets.datasets_document.DatasetService.get_dataset", return_value=dataset): - yield - - -@pytest.fixture -def patch_scoped_dataset(dataset): - with patch( - "controllers.console.datasets.datasets_document.DatasetService.get_dataset_for_tenant", return_value=dataset - ): - yield - - -@pytest.fixture -def patch_permission(): - with patch( - "controllers.console.datasets.datasets_document.DatasetService.check_dataset_permission", return_value=None - ): - yield - - -class _UsesSQLiteSession: - session: Session - - @pytest.fixture(autouse=True) - def _inject_sqlite_session(self, sqlite_session: Session) -> None: - self.session = sqlite_session + with patch( + "controllers.console.datasets.datasets_document.application_services", + return_value=registry, + ): + with pytest.raises(expected_http_error): + method(api, context, UUID(int=1), "unknown-batch") @pytest.mark.parametrize( @@ -270,1508 +544,3 @@ def test_document_mutation_routes_require_operation_specific_permissions(method, assert check.scene is expected_scene assert isinstance(check.locator, DatasetId) - - -class TestGetProcessRuleApi(_UsesSQLiteSession): - def test_get_default_success(self, app: Flask, patch_tenant): - api = GetProcessRuleApi() - method = unwrap(api.get) - user, _ = patch_tenant - with app.test_request_context("/"): - response = method(api, self.session, user) - assert "rules" in response - - def test_get_with_document_preserves_legacy_segmentation_delimiter(self, app: Flask, patch_tenant): - api = GetProcessRuleApi() - method = unwrap(api.get) - user, _ = patch_tenant - - document = make_document() - session = self.session - dataset = make_dataset() - process_rule = DatasetProcessRule( - dataset_id="ds-1", - mode="custom", - rules=json.dumps({"segmentation": {"delimiter": "---", "max_tokens": 123}}), - created_by="u1", - ) - session.add(process_rule) - session.flush() - - with ( - app.test_request_context("/?document_id=doc-1"), - patch( - "controllers.console.datasets.datasets_document.DocumentService.get_document_by_id", - return_value=document, - ) as mock_get_document, - patch( - "controllers.console.datasets.datasets_document.DatasetService.get_dataset", - return_value=dataset, - ), - patch( - "controllers.console.datasets.datasets_document.DatasetService.check_dataset_permission", - return_value=None, - ), - ): - response = method(api, session, user) - - mock_get_document.assert_called_once_with("doc-1", session) - assert response["rules"]["segmentation"]["separator"] == "---" - assert response["rules"]["segmentation"]["max_tokens"] == 123 - assert "delimiter" not in response["rules"]["segmentation"] - - def test_get_with_document_preserves_null_rules(self, app: Flask, patch_tenant): - api = GetProcessRuleApi() - method = unwrap(api.get) - user, _ = patch_tenant - session = self.session - dataset = make_dataset() - process_rule = DatasetProcessRule(dataset_id="ds-1", mode="custom", rules=None, created_by="u1") - session.add(process_rule) - session.flush() - with ( - app.test_request_context("/?document_id=doc-1"), - patch( - "controllers.console.datasets.datasets_document.DocumentService.get_document_by_id", - return_value=make_document(), - ), - patch( - "controllers.console.datasets.datasets_document.DatasetService.get_dataset", - return_value=dataset, - ), - patch( - "controllers.console.datasets.datasets_document.DatasetService.check_dataset_permission", - return_value=None, - ), - ): - response = method(api, session, user) - - assert response["mode"] == "custom" - assert response["rules"] is None - - def test_get_with_document_requires_dataset_readonly_permission(self, app: Flask, patch_tenant): - api = GetProcessRuleApi() - method = unwrap(api.get) - user, tenant_id = patch_tenant - document = make_document() - dataset = make_dataset(tenant_id=tenant_id) - - with ( - app.test_request_context("/?document_id=doc-1"), - patch( - "controllers.console.datasets.datasets_document.DocumentService.get_document_by_id", - return_value=document, - ), - patch( - "controllers.console.datasets.datasets_document.DatasetService.get_dataset", - return_value=dataset, - ), - patch( - "controllers.console.datasets.datasets_document.DatasetService.check_dataset_permission", - return_value=None, - ), - patch("controllers.console.datasets.datasets_document.enforce_rbac_checks", create=True) as enforce_checks, - ): - method(api, self.session, user) - - enforce_checks.assert_called_once() - kwargs = enforce_checks.call_args.kwargs - assert kwargs["tenant_id"] == tenant_id - assert kwargs["account_id"] == user.id - assert kwargs["path_args"] == {"dataset_id": dataset.id} - [check] = kwargs["checks"] - assert check.scene is RBACPermission.DATASET_READONLY - assert isinstance(check.locator, DatasetId) - - def test_get_with_document_dataset_not_found(self, app: Flask, patch_tenant): - api = GetProcessRuleApi() - method = unwrap(api.get) - user, _ = patch_tenant - document = make_document() - session = self.session - with ( - app.test_request_context("/?document_id=doc-1"), - patch( - "controllers.console.datasets.datasets_document.DocumentService.get_document_by_id", - return_value=document, - ), - patch("controllers.console.datasets.datasets_document.DatasetService.get_dataset", return_value=None), - ): - with pytest.raises(NotFound): - method(api, session, user) - - -class TestDatasetDocumentListApi(_UsesSQLiteSession): - def test_get_with_fetch_true_counts_segments(self, app: Flask, patch_tenant, patch_dataset, patch_permission): - api = DatasetDocumentListApi() - method = unwrap(api.get) - user, tenant_id = patch_tenant - doc = make_serializable_document() - pagination = MagicMock(items=[doc], total=1) - session = self.session - session.add_all([make_segment(position=1), make_segment(position=2)]) - session.flush() - with ( - app.test_request_context("/?fetch=true"), - patch("controllers.console.datasets.datasets_document.paginate_query", return_value=pagination), - patch( - "controllers.console.datasets.datasets_document.DocumentService.enrich_documents_with_summary_index_status", - return_value=None, - ), - ): - resp = method(api, session, tenant_id, user, "ds-1") - assert resp["data"][0]["id"] == "doc-1" - assert resp["data"][0]["completed_segments"] == 2 - assert resp["data"][0]["total_segments"] == 2 - - def test_get_with_search_status_and_created_at_sort( - self, app: Flask, patch_tenant, patch_dataset, patch_permission - ): - api = DatasetDocumentListApi() - method = unwrap(api.get) - user, tenant_id = patch_tenant - pagination = MagicMock(items=[make_serializable_document()], total=1) - with ( - app.test_request_context("/?keyword=test&status=enabled&sort=created_at"), - patch("controllers.console.datasets.datasets_document.paginate_query", return_value=pagination), - patch( - "controllers.console.datasets.datasets_document.DocumentService.apply_display_status_filter", - side_effect=lambda q, s: q, - ), - patch( - "controllers.console.datasets.datasets_document.DocumentService.enrich_documents_with_summary_index_status", - return_value=None, - ), - ): - resp = method(api, self.session, tenant_id, user, "ds-1") - assert resp["total"] == 1 - - def test_get_success(self, app: Flask, patch_tenant, patch_dataset, patch_permission): - api = DatasetDocumentListApi() - method = unwrap(api.get) - user, tenant_id = patch_tenant - document = make_serializable_document() - pagination = MagicMock(items=[document], total=1) - session = self.session - with ( - app.test_request_context("/"), - patch("controllers.console.datasets.datasets_document.paginate_query", return_value=pagination), - patch( - "controllers.console.datasets.datasets_document.DocumentService.enrich_documents_with_summary_index_status", - return_value=None, - ), - ): - response = method(api, session, tenant_id, user, "ds-1") - assert response["total"] == 1 - assert response["data"][0]["id"] == "doc-1" - assert "completed_segments" not in response["data"][0] - assert "total_segments" not in response["data"][0] - assert response["data"][0]["data_source_info"] == {"upload_file_id": "file-1"} - assert response["data"][0]["doc_metadata"] == [] - - def test_post_success(self, app: Flask, patch_tenant, patch_dataset, patch_permission): - api = DatasetDocumentListApi() - method = unwrap(api.post) - user, _ = patch_tenant - payload = {"indexing_technique": "economy"} - created_dataset = make_dataset() - created_document = make_document() - session = self.session - with ( - app.test_request_context("/", json=payload), - patch.object(type(console_ns), "payload", payload), - patch( - "controllers.console.datasets.datasets_document.DatasetService.get_dataset", - return_value=created_dataset, - ), - patch( - "controllers.console.datasets.datasets_document.DocumentService.document_create_args_validate", - return_value=None, - ), - patch( - "controllers.console.datasets.datasets_document.DocumentService.save_document_with_dataset_id", - return_value=([created_document], "batch-1"), - ), - ): - response = method(api, session, user, "ds-1") - assert "documents" in response - assert response["dataset"]["id"] == "ds-1" - assert response["documents"][0]["id"] == "doc-1" - assert response["documents"][0]["data_source_info"] == {} - assert response["documents"][0]["doc_metadata"] == [] - assert "data_source_info_dict" not in response["documents"][0] - assert "doc_metadata_details" not in response["documents"][0] - - def test_post_forbidden(self, app: Flask): - api = DatasetDocumentListApi() - method = unwrap(api.post) - user = make_account(TenantAccountRole.NORMAL) - with ( - app.test_request_context("/", json={}), - patch.object(type(console_ns), "payload", {}), - patch("controllers.console.datasets.datasets_document.DatasetService.get_dataset", return_value=dataset), - ): - with pytest.raises(Forbidden): - method(api, self.session, user, "ds-1") - - def test_get_with_fetch_true_and_invalid_fetch(self, app: Flask, patch_tenant, patch_dataset, patch_permission): - api = DatasetDocumentListApi() - method = unwrap(api.get) - user, tenant_id = patch_tenant - pagination = MagicMock(items=[make_serializable_document()], total=1) - with ( - app.test_request_context("/?fetch=maybe"), - patch("controllers.console.datasets.datasets_document.paginate_query", return_value=pagination), - patch( - "controllers.console.datasets.datasets_document.DocumentService.enrich_documents_with_summary_index_status", - return_value=None, - ), - ): - response = method(api, self.session, tenant_id, user, "ds-1") - assert response["total"] == 1 - - def test_get_sort_hit_count(self, app: Flask, patch_tenant, patch_dataset, patch_permission): - api = DatasetDocumentListApi() - method = unwrap(api.get) - user, tenant_id = patch_tenant - pagination = MagicMock(items=[], total=0) - with ( - app.test_request_context("/?sort=hit_count"), - patch("controllers.console.datasets.datasets_document.paginate_query", return_value=pagination), - patch( - "controllers.console.datasets.datasets_document.DocumentService.enrich_documents_with_summary_index_status", - return_value=None, - ), - ): - response = method(api, self.session, tenant_id, user, "ds-1") - assert response["total"] == 0 - - -class TestDatasetInitApi(_UsesSQLiteSession): - def test_post_success_serializes_created_dataset_and_documents(self, app: Flask, patch_tenant): - api = DatasetInitApi() - method = unwrap(api.post) - user, tenant_id = patch_tenant - payload = {"indexing_technique": "economy"} - created_dataset = make_dataset() - created_document = make_document(id="doc-init") - session = self.session - with ( - app.test_request_context("/", json=payload), - patch.object(type(console_ns), "payload", payload), - config_overrides_context(RBAC_ENABLED=True), - patch( - "controllers.console.datasets.datasets_document.DocumentService.document_create_args_validate", - return_value=None, - ), - patch( - "controllers.console.datasets.datasets_document.DocumentService.save_document_without_dataset_id", - return_value=(created_dataset, [created_document], "batch-init"), - ), - patch( - "controllers.console.datasets.datasets_document.enterprise_rbac_service.RBACService.DatasetAccess.replace_whitelist" - ), - ): - response = method(api, session, tenant_id, user) - assert response["dataset"]["id"] == "ds-1" - assert response["documents"][0]["id"] == "doc-init" - assert response["documents"][0]["data_source_info"] == {} - assert response["documents"][0]["doc_metadata"] == [] - assert response["batch"] == "batch-init" - assert created_dataset.permission == DatasetPermissionEnum.ALL_TEAM - - def test_post_success_syncs_creator_access_on_upload_create(self, app: Flask, patch_tenant): - """RBAC: the console upload create must keep the creator bound to their own dataset. - - Regression guard for the #42430 flip to automatic_include_workspace_members=False: - unlike POST /console/api/datasets (DatasetService.create_empty_dataset syncs the - creator internally), this path bypasses that helper, so the controller has to sync - the creator explicitly or the creator gets 403 on their own dataset. - """ - api = DatasetInitApi() - method = unwrap(api.post) - user, tenant_id = patch_tenant - payload = {"indexing_technique": "economy"} - created_dataset = make_dataset() - created_document = make_document(id="doc-init") - session = self.session - captured = {} - with ( - app.test_request_context("/", json=payload), - patch.object(type(console_ns), "payload", payload), - config_overrides_context(RBAC_ENABLED=True), - patch( - "controllers.console.datasets.datasets_document.DocumentService.document_create_args_validate", - return_value=None, - ), - patch( - "controllers.console.datasets.datasets_document.DocumentService.save_document_without_dataset_id", - return_value=(created_dataset, [created_document], "batch-init"), - ), - patch( - "controllers.console.datasets.datasets_document.enterprise_rbac_service.RBACService.DatasetAccess.replace_whitelist" - ), - patch( - "controllers.console.datasets.datasets_document.enterprise_rbac_service.try_sync_creator_access_policy_member_bindings" - ) as sync_creator, - ): - response = method(api, session, tenant_id, user) - captured["sync_calls"] = sync_creator.call_args_list - assert sync_creator.called, "upload create must sync the creator's own access byte-through" - assert response["dataset"]["id"] == "ds-1" - - -class TestDocumentResource(_UsesSQLiteSession): - def test_get_document_resolves_owner_chain(self, dataset): - api = DocumentResource() - session = self.session - user = make_account() - document = make_document() - - with ( - patch( - "controllers.console.datasets.datasets_document.DatasetService.get_dataset_for_tenant", - return_value=dataset, - ) as get_dataset, - patch( - "controllers.console.datasets.datasets_document.DatasetService.check_dataset_permission" - ) as check_permission, - patch( - "controllers.console.datasets.datasets_document.DatasetRefService.get_document_by_ref", - return_value=document, - ) as get_document, - ): - assert api.get_document(session, "ds-1", "doc-1", user, "tenant-1") is document - - get_dataset.assert_called_once_with("ds-1", "tenant-1", session=session) - check_permission.assert_called_once_with(dataset, user, session) - get_document.assert_called_once_with( - DocumentRef(dataset=DatasetRef(tenant_id="tenant-1", dataset_id="ds-1"), document_id="doc-1"), - session=session, - ) - - def test_get_document_delegates_permission_check_in_rbac_mode(self, dataset): - api = DocumentResource() - session = self.session - user = make_account() - with ( - config_overrides_context(RBAC_ENABLED=True), - patch( - "controllers.console.datasets.datasets_document.DatasetService.get_dataset_for_tenant", - return_value=dataset, - ), - patch( - "controllers.console.datasets.datasets_document.DatasetService.check_dataset_permission" - ) as check_permission, - patch( - "controllers.console.datasets.datasets_document.DatasetRefService.get_document_by_ref", - return_value=make_document(), - ), - ): - api.get_document(session, "ds-1", "doc-1", user, "tenant-1") - - check_permission.assert_called_once_with(dataset, user, session) - - -class TestDocumentApi(_UsesSQLiteSession): - def test_get_success(self, app: Flask, patch_tenant): - api = DocumentApi() - method = unwrap(api.get) - user, tenant_id = patch_tenant - document = make_document_detail() - with ( - app.test_request_context("/"), - patch.object(api, "get_document", return_value=document), - patch("controllers.console.datasets.datasets_document.DatasetService.get_process_rules", return_value={}), - ): - response, status = method(api, self.session, tenant_id, user, "ds-1", "doc-1") - assert status == 200 - - def test_get_invalid_metadata(self, app: Flask, patch_tenant): - api = DocumentApi() - method = unwrap(api.get) - user, tenant_id = patch_tenant - with ( - app.test_request_context("/?metadata=wrong"), - patch.object(api, "get_document", return_value=make_document()), - ): - with pytest.raises(InvalidMetadataError): - method(api, self.session, tenant_id, user, "ds-1", "doc-1") - - def test_delete_success(self, app: Flask, patch_tenant, patch_dataset): - api = DocumentApi() - method = unwrap(api.delete) - user, tenant_id = patch_tenant - with ( - app.test_request_context("/"), - patch( - "controllers.console.datasets.datasets_document.DatasetService.check_dataset_model_setting", - return_value=None, - ), - patch.object(api, "get_document", return_value=make_document()), - patch("controllers.console.datasets.datasets_document.DocumentService.delete_document", return_value=None), - ): - response, status = method(api, self.session, tenant_id, user, "ds-1", "doc-1") - assert status == 204 - - def test_delete_indexing_error(self, app: Flask, patch_tenant, patch_dataset): - api = DocumentApi() - method = unwrap(api.delete) - user, tenant_id = patch_tenant - with ( - app.test_request_context("/"), - patch( - "controllers.console.datasets.datasets_document.DatasetService.check_dataset_model_setting", - return_value=None, - ), - patch.object(api, "get_document", return_value=make_document()), - patch( - "controllers.console.datasets.datasets_document.DocumentService.delete_document", - side_effect=services.errors.document.DocumentIndexingError(), - ), - ): - with pytest.raises(DocumentIndexingError): - method(api, self.session, tenant_id, user, "ds-1", "doc-1") - - -class TestDocumentDownloadApi(_UsesSQLiteSession): - def test_download_success(self, app: Flask, patch_tenant): - api = DocumentDownloadApi() - method = unwrap(api.get) - user, tenant_id = patch_tenant - document = make_document() - with ( - app.test_request_context("/"), - patch.object(api, "get_document", return_value=document), - patch( - "controllers.console.datasets.datasets_document.DocumentService.get_document_download_url", - return_value="url", - ), - ): - response = method(api, self.session, tenant_id, user, "ds-1", "doc-1") - assert response["url"] == "url" - - -class TestDocumentProcessingApi(_UsesSQLiteSession): - def test_processing_forbidden_when_not_editor(self, app: Flask): - api = DocumentProcessingApi() - method = unwrap(api.patch) - user = make_account(TenantAccountRole.NORMAL) - with app.test_request_context("/"), patch.object(api, "get_document", return_value=make_document()): - with pytest.raises(Forbidden): - method(api, self.session, "tenant-1", user, "ds-1", "doc-1", "pause") - - def test_resume_from_error_state(self, app: Flask, patch_tenant): - api = DocumentProcessingApi() - method = unwrap(api.patch) - user, tenant_id = patch_tenant - doc = make_document(indexing_status=IndexingStatus.ERROR, is_paused=True) - session = self.session - with app.test_request_context("/"), patch.object(api, "get_document", return_value=doc): - _, status = method(api, session, tenant_id, user, "ds-1", "doc-1", "resume") - assert status == 200 - - def test_resume_success(self, app: Flask, patch_tenant): - api = DocumentProcessingApi() - method = unwrap(api.patch) - user, tenant_id = patch_tenant - document = make_document(indexing_status=IndexingStatus.PAUSED, is_paused=True) - session = self.session - with app.test_request_context("/"), patch.object(api, "get_document", return_value=document): - response, status = method(api, session, tenant_id, user, "ds-1", "doc-1", "resume") - assert status == 200 - - def test_pause_success(self, app: Flask, patch_tenant): - api = DocumentProcessingApi() - method = unwrap(api.patch) - user, tenant_id = patch_tenant - document = make_document(indexing_status=IndexingStatus.INDEXING) - session = self.session - with app.test_request_context("/"), patch.object(api, "get_document", return_value=document): - response, status = method(api, session, tenant_id, user, "ds-1", "doc-1", "pause") - assert status == 200 - - def test_pause_invalid(self, app: Flask, patch_tenant): - api = DocumentProcessingApi() - method = unwrap(api.patch) - user, tenant_id = patch_tenant - document = make_document(indexing_status=IndexingStatus.COMPLETED) - with app.test_request_context("/"), patch.object(api, "get_document", return_value=document): - with pytest.raises(InvalidActionError): - method(api, self.session, tenant_id, user, "ds-1", "doc-1", "pause") - - -class TestDocumentMetadataApi(_UsesSQLiteSession): - def test_put_metadata_schema_filtering(self, app: Flask, patch_tenant): - api = DocumentMetadataApi() - method = unwrap(api.put) - user, tenant_id = patch_tenant - doc = make_document() - payload = {"doc_type": "invoice", "doc_metadata": {"amount": 10, "invalid": "x"}} - schema = {"amount": int} - session = self.session - req_data = DocumentMetadataUpdatePayload.model_validate(payload) - with ( - app.test_request_context("/"), - patch.object(api, "get_document", return_value=doc), - patch( - "controllers.console.datasets.datasets_document.DocumentService.DOCUMENT_METADATA_SCHEMA", - {"invoice": schema}, - ), - ): - method(api, req_data, session, tenant_id, user, "ds-1", "doc-1") - assert doc.doc_metadata == {"amount": 10} - - def test_put_success(self, app: Flask, patch_tenant): - api = DocumentMetadataApi() - method = unwrap(api.put) - user, tenant_id = patch_tenant - document = make_document() - payload = {"doc_type": "others", "doc_metadata": {"a": 1}} - session = self.session - req_data = DocumentMetadataUpdatePayload.model_validate(payload) - with ( - app.test_request_context("/"), - patch.object(api, "get_document", return_value=document), - patch( - "controllers.console.datasets.datasets_document.DocumentService.DOCUMENT_METADATA_SCHEMA", - {"others": {}}, - ), - ): - response, status = method(api, req_data, session, tenant_id, user, "ds-1", "doc-1") - assert status == 200 - - def test_put_invalid_payload(self, app: Flask, patch_tenant): - api = DocumentMetadataApi() - method = unwrap(api.put) - user, tenant_id = patch_tenant - req_data = DocumentMetadataUpdatePayload.model_validate({}) - with app.test_request_context("/"), patch.object(api, "get_document", return_value=make_document()): - with pytest.raises(ValueError): - method(api, req_data, self.session, tenant_id, user, "ds-1", "doc-1") - - def test_put_invalid_doc_type(self, app: Flask, patch_tenant): - api = DocumentMetadataApi() - method = unwrap(api.put) - user, tenant_id = patch_tenant - payload = {"doc_type": "invalid", "doc_metadata": {}} - req_data = DocumentMetadataUpdatePayload.model_validate(payload) - with ( - app.test_request_context("/"), - patch.object(api, "get_document", return_value=make_document()), - patch( - "controllers.console.datasets.datasets_document.DocumentService.DOCUMENT_METADATA_SCHEMA", - {"others": {}}, - ), - ): - with pytest.raises(ValueError): - method(api, req_data, self.session, tenant_id, user, "ds-1", "doc-1") - - -class TestDocumentStatusApi(_UsesSQLiteSession): - def test_patch_success(self, app: Flask, patch_tenant, patch_dataset): - api = DocumentStatusApi() - method = unwrap(api.patch) - user, _ = patch_tenant - with ( - app.test_request_context("/?document_id=doc-1"), - patch( - "controllers.console.datasets.datasets_document.DatasetService.check_dataset_model_setting", - return_value=None, - ), - patch( - "controllers.console.datasets.datasets_document.DatasetService.check_dataset_permission", - return_value=None, - ), - patch( - "controllers.console.datasets.datasets_document.DocumentService.batch_update_document_status", - return_value=None, - ), - ): - response, status = method(api, self.session, user, "ds-1", "enable") - assert status == 200 - - def test_patch_invalid_action(self, app: Flask, patch_tenant, patch_dataset): - api = DocumentStatusApi() - method = unwrap(api.patch) - user, _ = patch_tenant - with ( - app.test_request_context("/?document_id=doc-1"), - patch( - "controllers.console.datasets.datasets_document.DatasetService.check_dataset_model_setting", - return_value=None, - ), - patch( - "controllers.console.datasets.datasets_document.DatasetService.check_dataset_permission", - return_value=None, - ), - patch( - "controllers.console.datasets.datasets_document.DocumentService.batch_update_document_status", - side_effect=ValueError("x"), - ), - ): - with pytest.raises(InvalidActionError): - method(api, self.session, user, "ds-1", "enable") - - -class TestDocumentRetryApi(_UsesSQLiteSession): - def test_retry_archived_document_skipped(self, app: Flask, patch_tenant, patch_scoped_dataset, patch_permission): - api = DocumentRetryApi() - method = unwrap(api.post) - user, tenant_id = patch_tenant - payload = {"document_ids": ["doc-1"]} - req_data = DocumentRetryPayload.model_validate(payload) - doc = make_document(indexing_status=IndexingStatus.INDEXING) - session = self.session - session.add(doc) - session.flush() - with ( - app.test_request_context("/"), - patch("controllers.console.datasets.datasets_document.DocumentService.check_archived", return_value=True), - patch("controllers.console.datasets.datasets_document.DocumentService.retry_document") as retry_mock, - ): - resp, status = method(api, req_data, session, tenant_id, user, "ds-1") - assert status == 204 - retry_mock.assert_called_once_with("ds-1", [], session) - - def test_retry_success(self, app: Flask, patch_tenant, patch_scoped_dataset, patch_permission): - api = DocumentRetryApi() - method = unwrap(api.post) - user, tenant_id = patch_tenant - payload = {"document_ids": ["doc-1"]} - req_data = DocumentRetryPayload.model_validate(payload) - document = make_document(indexing_status=IndexingStatus.INDEXING) - session = self.session - session.add(document) - session.flush() - with ( - app.test_request_context("/"), - patch("controllers.console.datasets.datasets_document.DocumentService.check_archived", return_value=False), - patch( - "controllers.console.datasets.datasets_document.DocumentService.retry_document", return_value=None - ) as retry_mock, - ): - response, status = method(api, req_data, session, tenant_id, user, "ds-1") - assert status == 204 - retry_mock.assert_called_once_with("ds-1", [document], session) - - def test_retry_loads_selected_documents_in_one_scoped_query( - self, app: Flask, patch_tenant, patch_scoped_dataset, patch_permission - ): - api = DocumentRetryApi() - method = unwrap(api.post) - user, tenant_id = patch_tenant - payload = {"document_ids": ["doc-1", "doc-2"]} - req_data = DocumentRetryPayload.model_validate(payload) - first_document = make_document(id="doc-1", indexing_status=IndexingStatus.ERROR) - second_document = make_document(id="doc-2", position=2, indexing_status=IndexingStatus.ERROR) - decoy = make_document(id="doc-decoy", tenant_id="other-tenant", dataset_id="other-dataset") - session = self.session - session.add_all([first_document, second_document, decoy]) - session.flush() - - with ( - app.test_request_context("/"), - patch("controllers.console.datasets.datasets_document.DocumentService.check_archived", return_value=False), - patch( - "controllers.console.datasets.datasets_document.DocumentService.retry_document", return_value=None - ) as retry_mock, - ): - response, status = method(api, req_data, session, tenant_id, user, "ds-1") - - assert status == 204 - retry_mock.assert_called_once_with("ds-1", [first_document, second_document], session) - - def test_retry_skips_completed_document(self, app: Flask, patch_tenant, patch_scoped_dataset, patch_permission): - api = DocumentRetryApi() - method = unwrap(api.post) - user, tenant_id = patch_tenant - payload = {"document_ids": ["doc-1"]} - req_data = DocumentRetryPayload.model_validate(payload) - document = make_document(indexing_status=IndexingStatus.COMPLETED) - session = self.session - session.add(document) - session.flush() - with ( - app.test_request_context("/"), - patch( - "controllers.console.datasets.datasets_document.DocumentService.retry_document", return_value=None - ) as retry_mock, - ): - response, status = method(api, req_data, session, tenant_id, user, "ds-1") - assert status == 204 - retry_mock.assert_called_once_with("ds-1", [], session) - - def test_retry_foreign_dataset_has_no_side_effects(self, app: Flask, patch_tenant, bypass_knowledge_rate_limit): - api = DocumentRetryApi() - method = unwrap(api.post) - user, tenant_id = patch_tenant - session = self.session - payload = {"document_ids": ["doc-1"]} - req_data = DocumentRetryPayload.model_validate(payload) - - with ( - app.test_request_context("/"), - patch( - "controllers.console.datasets.datasets_document.DatasetService.get_dataset_for_tenant", - return_value=None, - ), - patch("controllers.console.datasets.datasets_document.DocumentService.retry_document") as retry_document, - ): - with pytest.raises(NotFound): - method(api, req_data, session, tenant_id, user, "foreign-dataset") - - bypass_knowledge_rate_limit.assert_not_called() - retry_document.assert_not_called() - - -class TestDocumentPauseRecoverApi(_UsesSQLiteSession): - @pytest.mark.parametrize( - ("api_type", "service_method"), - [(DocumentPauseApi, "pause_document"), (DocumentRecoverApi, "recover_document")], - ) - def test_patch_uses_scoped_document( - self, app: Flask, patch_tenant, bypass_knowledge_rate_limit, api_type, service_method - ): - api = api_type() - method = unwrap(api.patch) - user, tenant_id = patch_tenant - session = self.session - document = make_document() - - with ( - app.test_request_context("/"), - patch.object(api, "get_document", return_value=document) as get_document, - patch("controllers.console.datasets.datasets_document.DocumentService.check_archived", return_value=False), - patch( - f"controllers.console.datasets.datasets_document.DocumentService.{service_method}" - ) as process_document, - ): - response, status = method(api, session, tenant_id, user, "ds-1", "doc-1") - - assert (response, status) == ("", 204) - get_document.assert_called_once_with(session, "ds-1", "doc-1", user, tenant_id) - bypass_knowledge_rate_limit.assert_called_once_with() - process_document.assert_called_once_with(document, session) - - -class TestWebsiteDocumentSyncApi(_UsesSQLiteSession): - def test_get_uses_scoped_dataset_and_document(self, app: Flask, patch_tenant, dataset): - api = WebsiteDocumentSyncApi() - method = unwrap(api.get) - user, tenant_id = patch_tenant - session = self.session - document = make_document(data_source_type=DataSourceType.WEBSITE_CRAWL) - - with ( - app.test_request_context("/"), - patch( - "controllers.console.datasets.datasets_document.DatasetService.get_dataset_for_tenant", - return_value=dataset, - ) as get_dataset, - patch.object(api, "get_document", return_value=document) as get_document, - patch("controllers.console.datasets.datasets_document.DocumentService.check_archived", return_value=False), - patch( - "controllers.console.datasets.datasets_document.DocumentService.sync_website_document" - ) as sync_document, - ): - response, status = method(api, session, tenant_id, user, "ds-1", "doc-1") - - assert status == 200 - assert response["result"] == "success" - get_dataset.assert_called_once_with("ds-1", tenant_id, session=session) - get_document.assert_called_once_with(session, dataset.id, "doc-1", user, tenant_id) - sync_document.assert_called_once_with(dataset, document, session) - - def test_get_rejects_non_editor_before_loading_document(self, app: Flask, dataset): - api = WebsiteDocumentSyncApi() - method = unwrap(api.get) - user = make_account(TenantAccountRole.NORMAL) - session = self.session - - with ( - app.test_request_context("/"), - patch( - "controllers.console.datasets.datasets_document.DatasetService.get_dataset_for_tenant", - return_value=dataset, - ), - patch.object(api, "get_document") as get_document, - patch( - "controllers.console.datasets.datasets_document.DocumentService.sync_website_document" - ) as sync_document, - ): - with pytest.raises(Forbidden): - method(api, session, "tenant-1", user, "ds-1", "doc-1") - - get_document.assert_not_called() - sync_document.assert_not_called() - - -class TestDocumentPipelineExecutionLogApi(_UsesSQLiteSession): - def test_get_log_success(self, app: Flask, patch_tenant): - api = DocumentPipelineExecutionLogApi() - method = unwrap(api.get) - user, tenant_id = patch_tenant - log = DocumentPipelineExecutionLog( - pipeline_id="pipeline-1", - document_id="trusted-doc", - datasource_type="file", - datasource_info="{}", - datasource_node_id="n1", - input_data={}, - created_by="u1", - ) - document = make_document(id="trusted-doc") - session = self.session - session.add(log) - session.flush() - with ( - app.test_request_context("/"), - patch.object(api, "get_document", return_value=document) as get_document, - ): - response, status = method(api, session, tenant_id, user, "ds-1", "doc-1") - assert status == 200 - get_document.assert_called_once_with(session, "ds-1", "doc-1", user, tenant_id) - assert response["datasource_node_id"] == "n1" - - -class TestDocumentGenerateSummaryApi(_UsesSQLiteSession): - def test_generate_summary_missing_documents(self, app: Flask, patch_tenant, patch_permission): - api = DocumentGenerateSummaryApi() - method = unwrap(api.post) - user, _ = patch_tenant - dataset = make_dataset(indexing_technique="high_quality", summary_index_setting={"enable": True}) - payload = {"document_list": ["doc-1", "doc-2"]} - req_data = GenerateSummaryPayload.model_validate(payload) - with ( - app.test_request_context("/"), - patch("controllers.console.datasets.datasets_document.DatasetService.get_dataset", return_value=dataset), - patch( - "controllers.console.datasets.datasets_document.DocumentService.get_documents_by_ids", - return_value=[make_document(id="doc-1")], - ), - ): - with pytest.raises(NotFound): - method(api, req_data, self.session, user, "ds-1") - - def test_generate_not_enabled(self, app: Flask, patch_tenant, patch_permission): - api = DocumentGenerateSummaryApi() - method = unwrap(api.post) - user, _ = patch_tenant - dataset = make_dataset(indexing_technique="high_quality", summary_index_setting={"enable": False}) - payload = {"document_list": ["doc-1"]} - req_data = GenerateSummaryPayload.model_validate(payload) - with ( - app.test_request_context("/"), - patch("controllers.console.datasets.datasets_document.DatasetService.get_dataset", return_value=dataset), - ): - with pytest.raises(ValueError): - method(api, req_data, self.session, user, "ds-1") - - def test_generate_summary_success_with_qa_skip(self, app: Flask, patch_tenant, patch_permission): - api = DocumentGenerateSummaryApi() - method = unwrap(api.post) - user, _ = patch_tenant - dataset = make_dataset(indexing_technique="high_quality", summary_index_setting={"enable": True}) - doc1 = make_document(id="doc-1", doc_form=IndexStructureType.QA_INDEX) - doc2 = make_document(id="doc-2", position=2, doc_form=IndexStructureType.PARAGRAPH_INDEX) - payload = {"document_list": ["doc-1", "doc-2"]} - req_data = GenerateSummaryPayload.model_validate(payload) - with ( - app.test_request_context("/"), - patch("controllers.console.datasets.datasets_document.DatasetService.get_dataset", return_value=dataset), - patch( - "controllers.console.datasets.datasets_document.DocumentService.get_documents_by_ids", - return_value=[doc1, doc2], - ), - patch( - "controllers.console.datasets.datasets_document.generate_summary_index_task.delay", return_value=None - ), - ): - response, status = method(api, req_data, self.session, user, "ds-1") - assert status == 200 - - -class TestDocumentSummaryStatusApi(_UsesSQLiteSession): - def test_get_success(self, app: Flask, patch_tenant, patch_permission): - api = DocumentSummaryStatusApi() - method = unwrap(api.get) - user, _ = patch_tenant - with ( - app.test_request_context("/"), - patch( - "controllers.console.datasets.datasets_document.DatasetService.get_dataset", - return_value=make_dataset(), - ), - patch( - "services.summary_index_service.SummaryIndexService.get_document_summary_status_detail", - return_value={ - "total_segments": 1, - "summary_status": {"timeout": 1}, - "summaries": [ - { - "segment_id": "segment-1", - "segment_position": 1, - "status": "timeout", - } - ], - }, - ), - ): - response, status = method(api, self.session, user, "ds-1", "doc-1") - assert status == 200 - assert response["summary_status"]["timeout"] == 1 - assert response["summaries"][0]["status"] == "timeout" - - -class TestDocumentIndexingEstimateApi(_UsesSQLiteSession): - def test_indexing_estimate_file_not_found(self, app: Flask, patch_tenant): - api = DocumentIndexingEstimateApi() - method = unwrap(api.get) - user, tenant_id = patch_tenant - document = make_document( - indexing_status=IndexingStatus.INDEXING, - data_source_info=json.dumps({"upload_file_id": "file-1"}), - ) - session = self.session - with app.test_request_context("/"), patch.object(api, "get_document", return_value=document): - with pytest.raises(NotFound): - method(api, session, tenant_id, user, "ds-1", "doc-1") - - def test_indexing_estimate_generic_exception(self, app: Flask, patch_tenant): - api = DocumentIndexingEstimateApi() - method = unwrap(api.get) - user, tenant_id = patch_tenant - document = make_document( - indexing_status=IndexingStatus.INDEXING, - data_source_info=json.dumps({"upload_file_id": "file-1"}), - ) - upload_file = make_upload_file() - mock_indexing_runner = MagicMock() - mock_indexing_runner.indexing_estimate.side_effect = RuntimeError("Some indexing error") - session = self.session - session.add(upload_file) - session.flush() - with ( - app.test_request_context("/"), - patch.object(api, "get_document", return_value=document), - patch("controllers.console.datasets.datasets_document.ExtractSetting", return_value=MagicMock()), - patch("controllers.console.datasets.datasets_document.IndexingRunner", return_value=mock_indexing_runner), - ): - with pytest.raises(IndexingEstimateError): - method(api, session, tenant_id, user, "ds-1", "doc-1") - - def test_get_finished(self, app: Flask, patch_tenant): - api = DocumentIndexingEstimateApi() - method = unwrap(api.get) - user, tenant_id = patch_tenant - document = make_document(indexing_status=IndexingStatus.COMPLETED) - with app.test_request_context("/"), patch.object(api, "get_document", return_value=document): - with pytest.raises(DocumentAlreadyFinishedError): - method(api, self.session, tenant_id, user, "ds-1", "doc-1") - - -class TestDocumentBatchDownloadZipApi(_UsesSQLiteSession): - def test_post_no_documents(self, app: Flask, patch_tenant): - api = DocumentBatchDownloadZipApi() - method = unwrap(api.post) - user, tenant_id = patch_tenant - payload: dict[str, list[str]] = {"document_ids": []} - with app.test_request_context("/", json=payload), patch.object(type(console_ns), "payload", payload): - with pytest.raises(ValueError): - method(api, self.session, tenant_id, user, "ds-1") - - -class TestDatasetDocumentListApiDelete(_UsesSQLiteSession): - def test_delete_success(self, app: Flask, patch_tenant, patch_scoped_dataset, patch_permission): - """Test successful deletion of documents""" - api = DatasetDocumentListApi() - method = unwrap(api.delete) - user, tenant_id = patch_tenant - session = self.session - with ( - app.test_request_context("/?document_id=doc-1&document_id=doc-2"), - patch("controllers.console.datasets.datasets_document.DocumentService.delete_documents", return_value=None), - ): - response, status = method(api, session, tenant_id, user, "ds-1") - assert status == 204 - - def test_delete_indexing_error(self, app: Flask, patch_tenant, patch_scoped_dataset, patch_permission): - """Test deletion with indexing error""" - api = DatasetDocumentListApi() - method = unwrap(api.delete) - user, tenant_id = patch_tenant - with ( - app.test_request_context("/?document_id=doc-1"), - patch( - "controllers.console.datasets.datasets_document.DocumentService.delete_documents", - side_effect=services.errors.document.DocumentIndexingError(), - ), - ): - with pytest.raises(DocumentIndexingError): - method(api, self.session, tenant_id, user, "ds-1") - - def test_delete_dataset_not_found(self, app: Flask, patch_tenant, bypass_knowledge_rate_limit): - """Test deletion when dataset not found""" - api = DatasetDocumentListApi() - method = unwrap(api.delete) - user, tenant_id = patch_tenant - with ( - app.test_request_context("/?document_id=doc-1"), - patch( - "controllers.console.datasets.datasets_document.DatasetService.get_dataset_for_tenant", - return_value=None, - ), - patch( - "controllers.console.datasets.datasets_document.DocumentService.delete_documents" - ) as delete_documents, - ): - with pytest.raises(NotFound): - method(api, self.session, tenant_id, user, "foreign-dataset") - - bypass_knowledge_rate_limit.assert_not_called() - delete_documents.assert_not_called() - - -class TestDocumentBatchIndexingEstimateApi(_UsesSQLiteSession): - def test_batch_indexing_estimate_website(self, app: Flask, patch_tenant): - api = DocumentBatchIndexingEstimateApi() - method = unwrap(api.get) - user, tenant_id = patch_tenant - doc = make_document( - indexing_status=IndexingStatus.INDEXING, - data_source_type=DataSourceType.WEBSITE_CRAWL, - data_source_info=json.dumps( - { - "provider": "firecrawl", - "job_id": "j1", - "url": "https://x.com", - "mode": "single", - "only_main_content": True, - } - ), - ) - with ( - app.test_request_context("/"), - patch.object(api, "get_batch_documents", return_value=[doc]), - patch( - "controllers.console.datasets.datasets_document.IndexingRunner.indexing_estimate", - return_value=IndexingEstimate(total_segments=2, preview=[]), - ), - ): - resp, status = method(api, self.session, tenant_id, user, "ds-1", "batch-1") - assert status == 200 - - def test_batch_indexing_estimate_notion(self, app: Flask, patch_tenant): - api = DocumentBatchIndexingEstimateApi() - method = unwrap(api.get) - user, tenant_id = patch_tenant - doc = make_document( - indexing_status=IndexingStatus.INDEXING, - data_source_type=DataSourceType.NOTION_IMPORT, - data_source_info=json.dumps( - { - "credential_id": "c1", - "notion_workspace_id": "w1", - "notion_page_id": "p1", - "type": "page", - } - ), - ) - with ( - app.test_request_context("/"), - patch.object(api, "get_batch_documents", return_value=[doc]), - patch( - "controllers.console.datasets.datasets_document.IndexingRunner.indexing_estimate", - return_value=IndexingEstimate(total_segments=1, preview=[]), - ), - ): - resp, status = method(api, self.session, tenant_id, user, "ds-1", "batch-1") - assert status == 200 - - def test_batch_estimate_unsupported_datasource(self, app: Flask, patch_tenant): - api = DocumentBatchIndexingEstimateApi() - method = unwrap(api.get) - user, tenant_id = patch_tenant - document = make_document( - indexing_status=IndexingStatus.INDEXING, - data_source_type="unknown", - data_source_info="{}", - ) - with app.test_request_context("/"), patch.object(api, "get_batch_documents", return_value=[document]): - with pytest.raises(ValueError): - method(api, self.session, tenant_id, user, "ds-1", "batch-1") - - def test_get_batch_estimate_invalid_batch(self, app: Flask, patch_tenant): - """Test batch estimation with invalid batch""" - api = DocumentBatchIndexingEstimateApi() - method = unwrap(api.get) - user, tenant_id = patch_tenant - with app.test_request_context("/"), patch.object(api, "get_batch_documents", side_effect=NotFound()): - with pytest.raises(NotFound): - method(api, self.session, tenant_id, user, "ds-1", "invalid-batch") - - -class TestDocumentBatchIndexingStatusApi(_UsesSQLiteSession): - def test_get_batch_status_success_serializes_status_shape(self, app: Flask, patch_tenant): - api = DocumentBatchIndexingStatusApi() - method = unwrap(api.get) - user, _ = patch_tenant - error = format_vector_space_admission_error(61, 50) - document = make_document( - indexing_status=IndexingStatus.ERROR, - is_paused=False, - error=error, - ) - session = self.session - session.add_all([make_segment(position=1), make_segment(position=2), make_segment(position=3, completed=False)]) - session.flush() - with app.test_request_context("/"), patch.object(api, "get_batch_documents", return_value=[document]): - response = method(api, session, user, "ds-1", "batch-1") - assert response == { - "data": [ - { - "id": "doc-1", - "indexing_status": "error", - "processing_started_at": None, - "parsing_completed_at": None, - "cleaning_completed_at": None, - "splitting_completed_at": None, - "completed_at": None, - "paused_at": None, - "error": error, - "error_code": VECTOR_SPACE_ADMISSION_ERROR_CODE, - "estimated_vector_space_mb": 61, - "vector_space_limit_mb": 50, - "stopped_at": None, - "completed_segments": 2, - "total_segments": 3, - } - ] - } - - def test_get_batch_status_invalid_batch(self, app: Flask, patch_tenant): - """Test batch status with invalid batch""" - api = DocumentBatchIndexingStatusApi() - method = unwrap(api.get) - user, _ = patch_tenant - with app.test_request_context("/"), patch.object(api, "get_batch_documents", side_effect=NotFound()): - with pytest.raises(NotFound): - method(api, self.session, user, "ds-1", "invalid-batch") - - def test_get_batch_status_uses_one_aggregate_query(self, app: Flask, patch_tenant, sqlite_engine): - api = DocumentBatchIndexingStatusApi() - method = unwrap(api.get) - user, _ = patch_tenant - documents = [make_document(id=f"doc-{index}", position=index) for index in range(1, 6)] - session = self.session - session.add_all( - [ - make_segment( - position=position, - completed=position <= 2, - document_id=document.id, - ) - for document in documents - for position in range(1, 4) - ] - ) - session.flush() - - select_statements: list[str] = [] - - def record_select(_connection, _cursor, statement, _parameters, _context, _executemany): - if statement.lstrip().upper().startswith("SELECT"): - select_statements.append(statement) - - event.listen(sqlite_engine, "before_cursor_execute", record_select) - try: - with app.test_request_context("/"), patch.object(api, "get_batch_documents", return_value=documents): - response = method(api, session, user, "ds-1", "batch-1") - finally: - event.remove(sqlite_engine, "before_cursor_execute", record_select) - - assert len(response["data"]) == 5 - assert all(item["completed_segments"] == 2 for item in response["data"]) - assert all(item["total_segments"] == 3 for item in response["data"]) - assert len(select_statements) == 1 - - -class TestDocumentIndexingStatusApi(_UsesSQLiteSession): - def test_get_status_success_serializes_status_shape(self, app: Flask, patch_tenant): - api = DocumentIndexingStatusApi() - method = unwrap(api.get) - user, tenant_id = patch_tenant - document = make_document( - indexing_status=IndexingStatus.INDEXING, - is_paused=False, - ) - session = self.session - session.add_all([make_segment(position=position, completed=position == 1) for position in range(1, 5)]) - session.flush() - with app.test_request_context("/"), patch.object(api, "get_document", return_value=document): - response = method(api, session, tenant_id, user, "ds-1", "doc-1") - assert response["id"] == "doc-1" - assert response["indexing_status"] == "indexing" - assert response["completed_segments"] == 1 - assert response["total_segments"] == 4 - - def test_get_status_document_not_found(self, app: Flask, patch_tenant): - """Test getting status for non-existent document""" - api = DocumentIndexingStatusApi() - method = unwrap(api.get) - user, tenant_id = patch_tenant - with app.test_request_context("/"), patch.object(api, "get_document", side_effect=NotFound()): - with pytest.raises(NotFound): - method(api, self.session, tenant_id, user, "ds-1", "invalid-doc") - - -class TestDocumentRenameApi(_UsesSQLiteSession): - def test_post_success_serializes_document_shape(self, app: Flask, patch_tenant): - api = DocumentRenameApi() - method = unwrap(api.post) - user, _ = patch_tenant - payload = {"name": "Renamed Document"} - req_data = DocumentRenamePayload.model_validate(payload) - renamed_document = make_document(id="doc-renamed", name="Renamed Document") - session = self.session - with ( - app.test_request_context("/"), - patch( - "controllers.console.datasets.datasets_document.DatasetService.get_dataset", return_value=make_dataset() - ), - patch( - "controllers.console.datasets.datasets_document.DatasetService.check_dataset_operator_permission", - return_value=None, - ), - patch( - "controllers.console.datasets.datasets_document.DocumentService.rename_document", - return_value=renamed_document, - ), - ): - response = method(api, req_data, session, user, "ds-1", "doc-1") - assert response["id"] == "doc-renamed" - assert response["name"] == "Renamed Document" - assert response["data_source_info"] == {} - assert response["doc_metadata"] == [] - assert "data_source_info_dict" not in response - - -class TestDocumentApiMetadata(_UsesSQLiteSession): - def test_get_with_only_option(self, app: Flask, patch_tenant): - """Test get with 'only' metadata option""" - api = DocumentApi() - method = unwrap(api.get) - user, tenant_id = patch_tenant - document = make_document_detail(doc_metadata_details=[]) - with ( - app.test_request_context("/?metadata=only"), - patch.object(api, "get_document", return_value=document), - patch("controllers.console.datasets.datasets_document.DatasetService.get_process_rules", return_value={}), - ): - response, status = method(api, self.session, tenant_id, user, "ds-1", "doc-1") - assert status == 200 - - def test_get_with_without_option(self, app: Flask, patch_tenant): - """Test get with 'without' metadata option""" - api = DocumentApi() - method = unwrap(api.get) - user, tenant_id = patch_tenant - document = make_document_detail() - with ( - app.test_request_context("/?metadata=without"), - patch.object(api, "get_document", return_value=document), - patch("controllers.console.datasets.datasets_document.DatasetService.get_process_rules", return_value={}), - ): - response, status = method(api, self.session, tenant_id, user, "ds-1", "doc-1") - assert status == 200 - - -class TestDocumentGenerateSummaryApiSuccess(_UsesSQLiteSession): - def test_generate_not_enabled_high_quality(self, app: Flask, patch_tenant, patch_permission): - """Test summary generation on non-high-quality dataset""" - api = DocumentGenerateSummaryApi() - method = unwrap(api.post) - user, _ = patch_tenant - dataset = make_dataset(indexing_technique="economy", summary_index_setting={"enable": True}) - payload = {"document_list": ["doc-1"]} - req_data = GenerateSummaryPayload.model_validate(payload) - with ( - app.test_request_context("/"), - patch("controllers.console.datasets.datasets_document.DatasetService.get_dataset", return_value=dataset), - ): - with pytest.raises(ValueError): - method(api, req_data, self.session, user, "ds-1") - - -class TestDocumentProcessingApiResume(_UsesSQLiteSession): - def test_resume_invalid_status(self, app: Flask, patch_tenant): - """Test resume on non-paused document""" - api = DocumentProcessingApi() - method = unwrap(api.patch) - user, tenant_id = patch_tenant - document = make_document(indexing_status=IndexingStatus.COMPLETED, is_paused=False) - with app.test_request_context("/"), patch.object(api, "get_document", return_value=document): - with pytest.raises(InvalidActionError): - method(api, self.session, tenant_id, user, "ds-1", "doc-1", "resume") - - -class TestDocumentPermissionCases(_UsesSQLiteSession): - def test_document_batch_get_permission_denied(self, app: Flask, patch_tenant): - api = DocumentBatchIndexingEstimateApi() - method = unwrap(api.get) - user, tenant_id = patch_tenant - with ( - app.test_request_context("/"), - patch( - "controllers.console.datasets.datasets_document.DatasetService.get_dataset", - return_value=make_dataset(), - ), - patch( - "controllers.console.datasets.datasets_document.DatasetService.check_dataset_permission", - side_effect=services.errors.account.NoPermissionError("No permission"), - ), - ): - with pytest.raises(Forbidden): - method(api, self.session, tenant_id, user, "ds-1", "batch-1") - - def test_document_batch_get_documents_not_found(self, app: Flask, patch_tenant): - api = DocumentBatchIndexingEstimateApi() - method = unwrap(api.get) - user, tenant_id = patch_tenant - with ( - app.test_request_context("/"), - patch( - "controllers.console.datasets.datasets_document.DatasetService.get_dataset", - return_value=make_dataset(), - ), - patch( - "controllers.console.datasets.datasets_document.DatasetService.check_dataset_permission", - return_value=None, - ), - patch.object(api, "get_batch_documents", return_value=None), - ): - response, status = method(api, self.session, tenant_id, user, "ds-1", "batch-1") - assert status == 200 - assert response == {"tokens": 0, "total_price": 0, "currency": "USD", "total_segments": 0, "preview": []} - - def test_process_rule_get_by_document_success(self, app: Flask, patch_tenant): - api = GetProcessRuleApi() - method = unwrap(api.get) - user, _ = patch_tenant - document = make_document() - session = self.session - dataset = make_dataset() - process_rule = DatasetProcessRule(dataset_id="ds-1", mode="custom", rules=json.dumps({"a": 1}), created_by="u1") - session.add(process_rule) - session.flush() - with ( - app.test_request_context("/?document_id=doc-1"), - patch( - "controllers.console.datasets.datasets_document.DocumentService.get_document_by_id", - return_value=document, - ), - patch("controllers.console.datasets.datasets_document.DatasetService.get_dataset", return_value=dataset), - patch( - "controllers.console.datasets.datasets_document.DatasetService.check_dataset_permission", - return_value=None, - ), - ): - result = method(api, session, user) - if isinstance(result, tuple): - response, status = result - else: - response, status = (result, 200) - assert status == 200 - assert response["mode"] == "custom" - - def test_process_rule_permission_denied(self, app: Flask): - api = GetProcessRuleApi() - method = unwrap(api.get) - user = make_account() - document = make_document() - session = self.session - with ( - app.test_request_context("/?document_id=doc-1"), - patch( - "controllers.console.datasets.datasets_document.DocumentService.get_document_by_id", - return_value=document, - ), - patch( - "controllers.console.datasets.datasets_document.DatasetService.get_dataset", - return_value=make_dataset(), - ), - patch( - "controllers.console.datasets.datasets_document.DatasetService.check_dataset_permission", - side_effect=services.errors.account.NoPermissionError("No permission"), - ), - ): - with pytest.raises(Forbidden): - method(api, session, user) - - -class TestDocumentListAdvancedCases(_UsesSQLiteSession): - def test_document_list_with_multiple_sort_options(self, app: Flask, patch_tenant, patch_dataset, patch_permission): - """Test document list with different sort options""" - api = DatasetDocumentListApi() - method = unwrap(api.get) - user, tenant_id = patch_tenant - pagination = MagicMock(items=[make_serializable_document()], total=1) - with ( - app.test_request_context("/?sort=updated_at"), - patch("controllers.console.datasets.datasets_document.paginate_query", return_value=pagination), - patch( - "controllers.console.datasets.datasets_document.DocumentService.enrich_documents_with_summary_index_status", - return_value=None, - ), - ): - response = method(api, self.session, tenant_id, user, "ds-1") - assert response["total"] == 1 - - def test_document_metadata_with_schema_validation(self, app: Flask, patch_tenant): - """Test document metadata update with schema validation""" - api = DocumentMetadataApi() - method = unwrap(api.put) - user, tenant_id = patch_tenant - doc = make_document() - payload = {"doc_type": "contract", "doc_metadata": {"amount": 5000, "currency": "USD", "invalid_field": "x"}} - schema = {"amount": int, "currency": str} - session = self.session - req_data = DocumentMetadataUpdatePayload.model_validate(payload) - with ( - app.test_request_context("/"), - patch.object(api, "get_document", return_value=doc), - patch( - "controllers.console.datasets.datasets_document.DocumentService.DOCUMENT_METADATA_SCHEMA", - {"contract": schema}, - ), - ): - response, status = method(api, req_data, session, tenant_id, user, "ds-1", "doc-1") - assert status == 200 - assert doc.doc_metadata == {"amount": 5000, "currency": "USD"} - - -class TestDocumentIndexingEdgeCases(_UsesSQLiteSession): - def test_document_indexing_with_extraction_setting(self, app: Flask, patch_tenant): - api = DocumentIndexingEstimateApi() - method = unwrap(api.get) - user, tenant_id = patch_tenant - document = make_document( - indexing_status=IndexingStatus.INDEXING, - data_source_info=json.dumps({"upload_file_id": "file-1"}), - ) - upload_file = make_upload_file() - session = self.session - session.add(upload_file) - session.flush() - with ( - app.test_request_context("/"), - patch.object(api, "get_document", return_value=document), - patch("controllers.console.datasets.datasets_document.ExtractSetting", return_value=MagicMock()), - patch( - "controllers.console.datasets.datasets_document.IndexingRunner.indexing_estimate", - return_value=IndexingEstimate(total_segments=5, preview=[]), - ), - ): - response, status = method(api, session, tenant_id, user, "ds-1", "doc-1") - assert status == 200 diff --git a/api/tests/unit_tests/controllers/console/datasets/test_datasets_document_download.py b/api/tests/unit_tests/controllers/console/datasets/test_datasets_document_download.py index 88dab19cd99..92b64d4875f 100644 --- a/api/tests/unit_tests/controllers/console/datasets/test_datasets_document_download.py +++ b/api/tests/unit_tests/controllers/console/datasets/test_datasets_document_download.py @@ -1,461 +1,196 @@ -""" -Unit tests for the dataset document download endpoint. +"""Document downloads through the controller, application service, and SQLite adapter.""" -These tests validate that the controller returns a signed download URL for -upload-file documents, and rejects unsupported or missing file cases. -""" - -from __future__ import annotations - -import importlib import json -import sys -from collections import UserDict -from datetime import UTC, datetime +from datetime import datetime from inspect import unwrap from io import BytesIO +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import create_autospec +from uuid import UUID from zipfile import ZipFile import pytest -from flask import Flask -from sqlalchemy.orm import Session +from flask import send_file +from sqlalchemy import event from werkzeug.exceptions import NotFound -from extensions.storage.storage_type import StorageType -from models.account import Account, TenantAccountRole +from controllers.console.datasets import datasets_document as controller +from machinery.context import RequestContext from models.dataset import Dataset, Document -from models.enums import CreatorUserRole from models.model import UploadFile +from services.knowledge.dataset_access import AccessibleDataset, DatasetAccess +from services.knowledge.documents.adapters import SQLAlchemyDocumentOperations +from services.knowledge.documents.application import DatasetDocumentApplicationService + +CONTEXT = RequestContext("request-1", None, "actor-1", "tenant-1") +DOC_ID = str(UUID(int=1)) +FILE_ID = str(UUID(int=2)) -@pytest.fixture -def app() -> Flask: - """Create a minimal Flask app for request-context based controller tests.""" - app = Flask(__name__) - app.config["TESTING"] = True - return app - - -@pytest.fixture -def datasets_document_module(monkeypatch: pytest.MonkeyPatch): - """ - Reload `controllers.console.datasets.datasets_document` with lightweight decorators. - - We patch auth / setup / rate-limit decorators to no-ops so we can unit test the - controller logic without requiring the full console stack. - """ - - from controllers.console import console_ns, wraps - from libs import login - - def _noop(func): # type: ignore[no-untyped-def] - return func - - # Bypass login/setup/account checks in unit tests. - monkeypatch.setattr(login, "login_required", _noop) - monkeypatch.setattr(wraps, "setup_required", _noop) - monkeypatch.setattr(wraps, "account_initialization_required", _noop) - - # Bypass billing-related decorators used by other endpoints in this module. - monkeypatch.setattr(wraps, "cloud_edition_billing_resource_check", lambda *_args, **_kwargs: lambda f: f) - monkeypatch.setattr(wraps, "cloud_edition_billing_rate_limit_check", lambda *_args, **_kwargs: lambda f: f) - - # Avoid Flask-RESTX route registration side effects during import. - def _noop_route(*_args, **_kwargs): # type: ignore[override] - def _decorator(cls): - return cls - - return _decorator - - monkeypatch.setattr(console_ns, "route", _noop_route) - - module_name = "controllers.console.datasets.datasets_document" - sys.modules.pop(module_name, None) - return importlib.import_module(module_name) - - -def _account(*, is_dataset_editor: bool = True) -> Account: - """Build a real account with the role used by dataset permission checks.""" - account = Account(name="Dataset User", email="dataset-user@example.com") - account.id = "user-123" - account.role = TenantAccountRole.EDITOR if is_dataset_editor else TenantAccountRole.NORMAL - return account - - -def _document( - *, - document_id: str, - tenant_id: str, - data_source_type: str, - upload_file_id: str | None, -) -> Document: - """Build a real document entity used by the controller.""" - document = Document( - id=document_id, - tenant_id=tenant_id, - dataset_id="ds-1", - position=1, - data_source_type=data_source_type, - data_source_info=json.dumps({"upload_file_id": upload_file_id}) if upload_file_id is not None else "{}", - batch="batch-1", - name="document.txt", - created_from="web", - created_by="user-123", +def _document(**values): + return Document( + **{ + "id": DOC_ID, + "tenant_id": "tenant-1", + "dataset_id": "dataset-1", + "position": 1, + "data_source_type": "upload_file", + "data_source_info": json.dumps({"upload_file_id": FILE_ID}), + "batch": "batch-1", + "name": "document.txt", + "created_from": "web", + "created_by": "actor-1", + **values, + } ) - return document -def _dataset() -> Dataset: - return Dataset(id="ds-1", tenant_id="tenant-123", name="Dataset", created_by="user-123") - - -def _upload_file(*, file_id: str, name: str = "document.txt", key: str = "key") -> UploadFile: - upload_file = UploadFile( - tenant_id="tenant-123", - storage_type=StorageType.LOCAL, - key=key, - name=name, - size=1, - extension="txt", - mime_type="text/plain", - created_by_role=CreatorUserRole.ACCOUNT, - created_by="user-123", - created_at=datetime.now(UTC), - used=False, - ) - upload_file.id = file_id - return upload_file - - -def _wire_common_success_mocks( - *, - module, - monkeypatch: pytest.MonkeyPatch, - document_tenant_id: str, - data_source_type: str, - upload_file_id: str | None, - upload_file_exists: bool, - signed_url: str, -) -> None: - """Patch controller dependencies to create a deterministic test environment.""" - import services.dataset_service as dataset_service_module - - # Return a dataset object and allow permission checks to pass. - monkeypatch.setattr( - module.DatasetService, - "get_dataset_for_tenant", - lambda *_args, **_kwargs: _dataset(), - ) - monkeypatch.setattr(module.DatasetService, "check_dataset_permission", lambda *_args, **_kwargs: None) - - # Return a document that will be validated inside DocumentResource.get_document. - document = _document( - document_id="doc-1", - tenant_id=document_tenant_id, - data_source_type=data_source_type, - upload_file_id=upload_file_id, +@pytest.fixture(autouse=True) +def downloads(monkeypatch, sqlite_session_factory): + with sqlite_session_factory.begin() as session: + session.add(Dataset(id="dataset-1", tenant_id="tenant-1", name="Dataset", created_by="actor-1")) + session.add(_document()) + upload = UploadFile( + tenant_id="tenant-1", + storage_type="local", + key="key-1", + name="document.txt", + size=5, + extension="txt", + mime_type="text/plain", + created_by_role="account", + created_by="actor-1", + created_at=datetime(2024, 1, 1), + used=False, + ) + upload.id = FILE_ID + session.add(upload) + access = create_autospec(DatasetAccess, instance=True, spec_set=True) + access.require_accessible.return_value = AccessibleDataset("dataset-1", "tenant-1") + documents = DatasetDocumentApplicationService( + dataset_access=access, + operations=SQLAlchemyDocumentOperations(session_factory=sqlite_session_factory), + metadata_schema={}, ) monkeypatch.setattr( - module.DatasetRefService, - "get_document_by_ref", - lambda *_args, **_kwargs: document if document.tenant_id == "tenant-123" else None, + controller, "application_services", lambda: SimpleNamespace(knowledge=SimpleNamespace(documents=documents)) ) - - # Mock UploadFile lookup via FileService batch helper. - upload_files_by_id: dict[str, object] = {} - if upload_file_exists and upload_file_id is not None: - upload_files_by_id[upload_file_id] = _upload_file(file_id=upload_file_id) - monkeypatch.setattr(module.FileService, "get_upload_files_by_ids", lambda *_args, **_kwargs: upload_files_by_id) - - # Mock signing helper so the returned URL is deterministic. - monkeypatch.setattr(dataset_service_module.file_helpers, "get_signed_file_url", lambda **_kwargs: signed_url) + return documents -def _mock_send_file(obj, **kwargs): # type: ignore[no-untyped-def] - """Return a lightweight representation of `send_file(...)` for unit tests.""" +def test_document_download_returns_signed_url(app, monkeypatch): + from services.knowledge import dataset_service - class _ResponseMock(UserDict): - def __init__(self, sent_file: object, send_file_kwargs: dict[str, object]) -> None: - super().__init__({"_sent_file": sent_file, "_send_file_kwargs": send_file_kwargs}) - self._on_close: object | None = None + signed = [] - def call_on_close(self, func): # type: ignore[no-untyped-def] - self._on_close = func - return func + def sign(**kwargs): + signed.append(kwargs) + return "https://files.example/download" - return _ResponseMock(obj, kwargs) + monkeypatch.setattr(dataset_service.file_helpers, "get_signed_file_url", sign) + with app.test_request_context("/"): + result = unwrap(controller.DocumentDownloadApi.get)( + controller.DocumentDownloadApi(), CONTEXT, "dataset-1", DOC_ID + ) + assert result == {"url": "https://files.example/download"} + assert signed == [{"upload_file_id": FILE_ID, "as_attachment": True}] -def test_batch_download_zip_returns_send_file( - app: Flask, datasets_document_module, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - """Ensure batch ZIP download returns a zip attachment via `send_file`.""" +@pytest.mark.parametrize( + "values", + [ + {"data_source_type": "website_crawl"}, + {"data_source_info": "{}"}, + {"data_source_info": json.dumps({"upload_file_id": "missing"})}, + {"tenant_id": "foreign"}, + {"dataset_id": "foreign"}, + ], +) +def test_invalid_or_foreign_document_is_not_downloadable(app, sqlite_session_factory, values): + with sqlite_session_factory.begin() as session: + document = session.get(Document, DOC_ID) + for key, value in values.items(): + setattr(document, key, value) + with app.test_request_context("/"), pytest.raises(NotFound): + unwrap(controller.DocumentDownloadApi.get)(controller.DocumentDownloadApi(), CONTEXT, "dataset-1", DOC_ID) - monkeypatch.setattr(datasets_document_module.DatasetService, "get_dataset", lambda *_args, **_kwargs: _dataset()) - monkeypatch.setattr( - datasets_document_module.DatasetService, "check_dataset_permission", lambda *_args, **_kwargs: None - ) - # Two upload-file documents, each referencing an UploadFile. - doc1 = _document( - document_id="11111111-1111-1111-1111-111111111111", - tenant_id="tenant-123", - data_source_type="upload_file", - upload_file_id="file-1", - ) - doc2 = _document( - document_id="22222222-2222-2222-2222-222222222222", - tenant_id="tenant-123", - data_source_type="upload_file", - upload_file_id="file-2", - ) - monkeypatch.setattr( - datasets_document_module.DocumentService, - "get_documents_by_ids", - lambda *_args, **_kwargs: [doc1, doc2], - ) - monkeypatch.setattr( - datasets_document_module.FileService, - "get_upload_files_by_ids", - lambda *_args, **_kwargs: { - "file-1": _upload_file(file_id="file-1", name="a.txt", key="k1"), - "file-2": _upload_file(file_id="file-2", name="b.txt", key="k2"), - }, - ) +def test_cross_tenant_upload_file_is_not_signed(app, sqlite_session_factory): + with sqlite_session_factory.begin() as session: + session.get(UploadFile, FILE_ID).tenant_id = "foreign" + with app.test_request_context("/"), pytest.raises(NotFound): + unwrap(controller.DocumentDownloadApi.get)(controller.DocumentDownloadApi(), CONTEXT, "dataset-1", DOC_ID) - # Mock storage streaming content. - import services.file_service as file_service_module - monkeypatch.setattr(file_service_module.storage, "load", lambda _key, stream=True: [b"hello"]) +def test_zip_is_readable_and_cleanup_waits_for_response_close(app, monkeypatch, sqlite_engine): + from services import file_service - # Replace send_file used by the controller to avoid a real Flask response object. - monkeypatch.setattr(datasets_document_module, "send_file", _mock_send_file) + paths = [] + checked_out = set() - # Act - with app.test_request_context( - "/datasets/ds-1/documents/download-zip", - method="POST", - json={"document_ids": ["11111111-1111-1111-1111-111111111111", "22222222-2222-2222-2222-222222222222"]}, + def checkout(_connection, record, _proxy): + checked_out.add(record) + + def checkin(_connection, record): + checked_out.discard(record) + + def load(key, *, stream): + assert not checked_out, "storage I/O must run after the read transaction closes" + assert key == "key-1" + assert stream is True + return iter([b"hello"]) + + def send(path, **kwargs): + paths.append(Path(path)) + return send_file(path, **kwargs) + + monkeypatch.setattr(file_service.storage, "load", load) + monkeypatch.setattr(controller, "send_file", send) + event.listen(sqlite_engine, "checkout", checkout) + event.listen(sqlite_engine, "checkin", checkin) + try: + with app.test_request_context("/", json={"document_ids": [DOC_ID]}): + response = unwrap(controller.DocumentBatchDownloadZipApi.post)( + controller.DocumentBatchDownloadZipApi(), CONTEXT, "dataset-1" + ) + assert paths[0].exists() + response.direct_passthrough = False + with ZipFile(BytesIO(response.get_data())) as archive: + assert archive.read("document.txt") == b"hello" + response.close() + assert not paths[0].exists() + finally: + event.remove(sqlite_engine, "checkout", checkout) + event.remove(sqlite_engine, "checkin", checkin) + + +def test_zip_cleanup_when_response_creation_fails(app, monkeypatch): + from services import file_service + + paths = [] + monkeypatch.setattr(file_service.storage, "load", lambda *_args, **_kwargs: iter([b"hello"])) + + def fail_send(path, **_kwargs): + paths.append(Path(path)) + raise RuntimeError("response failed") + + monkeypatch.setattr(controller, "send_file", fail_send) + with ( + app.test_request_context("/", json={"document_ids": [DOC_ID]}), + pytest.raises(RuntimeError, match="response failed"), ): - api = datasets_document_module.DocumentBatchDownloadZipApi() - method = unwrap(api.post) - result = method(api, sqlite_session, "tenant-123", _account(), dataset_id="ds-1") - - # Assert: we returned via send_file with correct mime type and attachment. - assert result["_send_file_kwargs"]["mimetype"] == "application/zip" - assert result["_send_file_kwargs"]["as_attachment"] is True - assert isinstance(result["_send_file_kwargs"]["download_name"], str) - assert result["_send_file_kwargs"]["download_name"].endswith(".zip") - # Ensure our cleanup hook is registered and execute it to avoid temp file leaks in unit tests. - assert getattr(result, "_on_close", None) is not None - result._on_close() # type: ignore[attr-defined] + unwrap(controller.DocumentBatchDownloadZipApi.post)( + controller.DocumentBatchDownloadZipApi(), CONTEXT, "dataset-1" + ) + assert paths + assert not paths[0].exists() -def test_batch_download_zip_response_is_openable_zip( - app: Flask, datasets_document_module, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - """Ensure the real Flask `send_file` response body is a valid ZIP that can be opened.""" - - # Arrange: same controller mocks as the lightweight send_file test, but we keep the real `send_file`. - monkeypatch.setattr(datasets_document_module.DatasetService, "get_dataset", lambda *_args, **_kwargs: _dataset()) - monkeypatch.setattr( - datasets_document_module.DatasetService, "check_dataset_permission", lambda *_args, **_kwargs: None - ) - - doc1 = _document( - document_id="33333333-3333-3333-3333-333333333333", - tenant_id="tenant-123", - data_source_type="upload_file", - upload_file_id="file-1", - ) - doc2 = _document( - document_id="44444444-4444-4444-4444-444444444444", - tenant_id="tenant-123", - data_source_type="upload_file", - upload_file_id="file-2", - ) - monkeypatch.setattr( - datasets_document_module.DocumentService, - "get_documents_by_ids", - lambda *_args, **_kwargs: [doc1, doc2], - ) - monkeypatch.setattr( - datasets_document_module.FileService, - "get_upload_files_by_ids", - lambda *_args, **_kwargs: { - "file-1": _upload_file(file_id="file-1", name="a.txt", key="k1"), - "file-2": _upload_file(file_id="file-2", name="b.txt", key="k2"), - }, - ) - - # Stream distinct bytes per key so we can verify both ZIP entries. - import services.file_service as file_service_module - - monkeypatch.setattr( - file_service_module.storage, "load", lambda key, stream=True: [b"one"] if key == "k1" else [b"two"] - ) - - # Act - with app.test_request_context( - "/datasets/ds-1/documents/download-zip", - method="POST", - json={"document_ids": ["33333333-3333-3333-3333-333333333333", "44444444-4444-4444-4444-444444444444"]}, - ): - api = datasets_document_module.DocumentBatchDownloadZipApi() - method = unwrap(api.post) - response = method(api, sqlite_session, "tenant-123", _account(), dataset_id="ds-1") - - # Assert: response body is a valid ZIP and contains the expected entries. - response.direct_passthrough = False - data = response.get_data() - response.close() - - with ZipFile(BytesIO(data), mode="r") as zf: - assert zf.namelist() == ["a.txt", "b.txt"] - assert zf.read("a.txt") == b"one" - assert zf.read("b.txt") == b"two" - - -def test_batch_download_zip_rejects_non_upload_file_document( - app: Flask, datasets_document_module, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - """Ensure batch ZIP download rejects non upload-file documents.""" - - monkeypatch.setattr(datasets_document_module.DatasetService, "get_dataset", lambda *_args, **_kwargs: _dataset()) - monkeypatch.setattr( - datasets_document_module.DatasetService, "check_dataset_permission", lambda *_args, **_kwargs: None - ) - - doc = _document( - document_id="55555555-5555-5555-5555-555555555555", - tenant_id="tenant-123", - data_source_type="website_crawl", - upload_file_id="file-1", - ) - monkeypatch.setattr( - datasets_document_module.DocumentService, - "get_documents_by_ids", - lambda *_args, **_kwargs: [doc], - ) - - with app.test_request_context( - "/datasets/ds-1/documents/download-zip", - method="POST", - json={"document_ids": ["55555555-5555-5555-5555-555555555555"]}, - ): - api = datasets_document_module.DocumentBatchDownloadZipApi() - method = unwrap(api.post) - with pytest.raises(NotFound): - method(api, sqlite_session, "tenant-123", _account(), dataset_id="ds-1") - - -def test_document_download_returns_url_for_upload_file_document( - app: Flask, datasets_document_module, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - """Ensure upload-file documents return a `{url}` JSON payload.""" - - _wire_common_success_mocks( - module=datasets_document_module, - monkeypatch=monkeypatch, - document_tenant_id="tenant-123", - data_source_type="upload_file", - upload_file_id="file-123", - upload_file_exists=True, - signed_url="https://example.com/signed", - ) - - # Build a request context then call the resource method directly. - with app.test_request_context("/datasets/ds-1/documents/doc-1/download", method="GET"): - api = datasets_document_module.DocumentDownloadApi() - method = unwrap(api.get) - result = method(api, sqlite_session, "tenant-123", _account(), dataset_id="ds-1", document_id="doc-1") - - assert result == {"url": "https://example.com/signed"} - - -def test_document_download_rejects_non_upload_file_document( - app: Flask, datasets_document_module, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - """Ensure non-upload documents raise 404 (no file to download).""" - - _wire_common_success_mocks( - module=datasets_document_module, - monkeypatch=monkeypatch, - document_tenant_id="tenant-123", - data_source_type="website_crawl", - upload_file_id="file-123", - upload_file_exists=True, - signed_url="https://example.com/signed", - ) - - with app.test_request_context("/datasets/ds-1/documents/doc-1/download", method="GET"): - api = datasets_document_module.DocumentDownloadApi() - method = unwrap(api.get) - with pytest.raises(NotFound): - method(api, sqlite_session, "tenant-123", _account(), dataset_id="ds-1", document_id="doc-1") - - -def test_document_download_rejects_missing_upload_file_id( - app: Flask, datasets_document_module, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - """Ensure missing `upload_file_id` raises 404.""" - - _wire_common_success_mocks( - module=datasets_document_module, - monkeypatch=monkeypatch, - document_tenant_id="tenant-123", - data_source_type="upload_file", - upload_file_id=None, - upload_file_exists=False, - signed_url="https://example.com/signed", - ) - - with app.test_request_context("/datasets/ds-1/documents/doc-1/download", method="GET"): - api = datasets_document_module.DocumentDownloadApi() - method = unwrap(api.get) - with pytest.raises(NotFound): - method(api, sqlite_session, "tenant-123", _account(), dataset_id="ds-1", document_id="doc-1") - - -def test_document_download_rejects_when_upload_file_record_missing( - app: Flask, datasets_document_module, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - """Ensure missing UploadFile row raises 404.""" - - _wire_common_success_mocks( - module=datasets_document_module, - monkeypatch=monkeypatch, - document_tenant_id="tenant-123", - data_source_type="upload_file", - upload_file_id="file-123", - upload_file_exists=False, - signed_url="https://example.com/signed", - ) - - with app.test_request_context("/datasets/ds-1/documents/doc-1/download", method="GET"): - api = datasets_document_module.DocumentDownloadApi() - method = unwrap(api.get) - with pytest.raises(NotFound): - method(api, sqlite_session, "tenant-123", _account(), dataset_id="ds-1", document_id="doc-1") - - -def test_document_download_rejects_document_owner_mismatch( - app: Flask, datasets_document_module, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - """Ensure an owner mismatch is rejected by the shared document resolver.""" - - _wire_common_success_mocks( - module=datasets_document_module, - monkeypatch=monkeypatch, - document_tenant_id="tenant-999", - data_source_type="upload_file", - upload_file_id="file-123", - upload_file_exists=True, - signed_url="https://example.com/signed", - ) - - with app.test_request_context("/datasets/ds-1/documents/doc-1/download", method="GET"): - api = datasets_document_module.DocumentDownloadApi() - method = unwrap(api.get) - with pytest.raises(NotFound): - method(api, sqlite_session, "tenant-123", _account(), dataset_id="ds-1", document_id="doc-1") +def test_zip_rejects_non_upload_document(app, sqlite_session_factory): + with sqlite_session_factory.begin() as session: + session.get(Document, DOC_ID).data_source_type = "website_crawl" + with app.test_request_context("/", json={"document_ids": [DOC_ID]}), pytest.raises(NotFound): + unwrap(controller.DocumentBatchDownloadZipApi.post)( + controller.DocumentBatchDownloadZipApi(), CONTEXT, "dataset-1" + ) 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 00e4737228d..6b5b92b74a5 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 @@ -1,17 +1,14 @@ -from datetime import datetime +from contextlib import AbstractContextManager +from dataclasses import dataclass +from datetime import UTC, datetime from inspect import unwrap -from types import SimpleNamespace -from typing import Any, cast -from unittest.mock import MagicMock, call, patch +from unittest.mock import create_autospec, patch import pytest from flask import Flask -from sqlalchemy.orm import Session 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, @@ -21,1303 +18,360 @@ from controllers.console.datasets.datasets_segments import ( DatasetDocumentSegmentAddApi, DatasetDocumentSegmentApi, DatasetDocumentSegmentBatchImportApi, + DatasetDocumentSegmentBatchImportStatusApi, DatasetDocumentSegmentListApi, DatasetDocumentSegmentUpdateApi, SegmentCreatePayload, SegmentUpdatePayload, + _raise_segment_error, ) from controllers.console.datasets.error import ChildChunkDeleteIndexError, ChildChunkIndexingError, InvalidActionError -from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError -from core.rag.index_processor.constant.index_type import IndexStructureType -from fields.segment_fields import segment_response_with_summary, segment_responses_with_summaries -from libs.datetime_utils import naive_utc_now -from models.account import Account, TenantAccountRole -from models.dataset import ChildChunk, Dataset, Document, DocumentSegment -from models.enums import PermissionEnum, SegmentStatus, SegmentType -from models.model import UploadFile -from services.errors.chunk import ChildChunkDeleteIndexError as ChildChunkDeleteIndexServiceError -from services.errors.chunk import ChildChunkIndexingError as ChildChunkIndexingServiceError -from tests.unit_tests.model_factories import make_account, make_dataset, make_document, make_upload_file +from machinery.context import RequestContext +from services.knowledge.entities.segments import ChildChunkRecord, ChildChunkUpdateArgs, SegmentRecord +from services.knowledge.segments.application import ( + ChildChunkDeleteIndexApplicationError, + ChildChunkIndexingApplicationError, + ChildChunkNotFoundError, + ChildChunkPage, + DatasetSegmentApplicationService, + SegmentBatchImport, + SegmentBatchImportDispatchError, + SegmentBatchImportNotFoundError, + SegmentDatasetModelUnavailableError, + SegmentDatasetNotFoundError, + SegmentDetail, + SegmentDocumentIndexingError, + SegmentDocumentNotFoundError, + SegmentEmbeddingModelUnavailableError, + SegmentInvalidFileTypeError, + SegmentListFilter, + SegmentNotFoundError, + SegmentPage, + SegmentPermissionDeniedError, + SegmentStatusUpdateError, + SegmentUploadFileNotFoundError, +) -def _segment(): - segment = DocumentSegment( - tenant_id="tenant-1", - dataset_id="ds-1", - document_id="doc-1", - position=1, - content="c", - word_count=1, - tokens=1, - created_by="u1", - answer="a", - keywords=["test"], - index_node_id="n1", - index_node_hash="h", - status=SegmentStatus.COMPLETED, - updated_by="u1", - ) - - segment.id = "seg-1" - segment.created_at = naive_utc_now() - segment.updated_at = naive_utc_now() - return segment - - -def _child_chunk(): - child_chunk = ChildChunk( - tenant_id="tenant-1", - dataset_id="ds-1", - document_id="doc-1", - segment_id="seg-1", - position=1, - content="child", - word_count=1, - created_by="u1", - type=SegmentType.CUSTOMIZED, - ) - child_chunk.id = "cc-1" - child_chunk.created_at = naive_utc_now() - child_chunk.updated_at = naive_utc_now() - return child_chunk - - -def _account() -> Account: - return make_account( - account_id="u1", name="Dataset Editor", email="dataset-editor@example.com", role=TenantAccountRole.OWNER +def _context() -> RequestContext: + return RequestContext( + request_id="request-1", + trace_id=None, + account_id="account-1", + active_workspace_id="tenant-1", ) -def _dataset( - *, - dataset_id: str = "ds-1", - tenant_id: str = "tenant-1", - indexing_technique: str = "economy", - embedding_model_provider: str | None = None, - embedding_model: str | None = None, -) -> Dataset: - return make_dataset( - dataset_id=dataset_id, - tenant_id=tenant_id, - description="", - provider="vendor", - permission=PermissionEnum.ONLY_ME, - indexing_technique=indexing_technique, - embedding_model_provider=embedding_model_provider, - embedding_model=embedding_model, - created_by="u1", +def _child_chunk_data() -> ChildChunkRecord: + now = datetime(2026, 1, 1, tzinfo=UTC) + return ChildChunkRecord.model_validate( + { + "id": "child-1", + "segment_id": "segment-1", + "content": "child", + "position": 1, + "word_count": 5, + "type": "customized", + "created_at": now, + "updated_at": now, + } ) -def _document( - *, - document_id: str = "doc-1", - dataset_id: str = "ds-1", - tenant_id: str = "tenant-1", - doc_form: IndexStructureType = IndexStructureType.PARAGRAPH_INDEX, -) -> Document: - return make_document( - document_id=document_id, - tenant_id=tenant_id, - dataset_id=dataset_id, - created_from="api", - created_by="u1", - doc_form=doc_form, +def _segment_data() -> SegmentRecord: + now = datetime(2026, 1, 1, tzinfo=UTC) + return SegmentRecord.model_validate( + { + "id": "segment-1", + "position": 1, + "document_id": "document-1", + "content": "content", + "sign_content": "content", + "answer": None, + "word_count": 7, + "tokens": 2, + "keywords": [], + "index_node_id": "node-1", + "index_node_hash": "hash", + "hit_count": 0, + "enabled": True, + "disabled_at": None, + "disabled_by": None, + "status": "completed", + "created_by": "account-1", + "created_at": now, + "updated_at": now, + "updated_by": None, + "indexing_at": None, + "completed_at": now, + "error": None, + "stopped_at": None, + "child_chunks": [], + "attachments": [], + "summary": None, + } ) -def _upload_file(*, name: str = "test.csv", file_id: str = "test-file-id") -> UploadFile: - return make_upload_file( - file_id=file_id, - storage_type="opendal", - key="test-key", - name=name, - size=0, - extension=name.rsplit(".", maxsplit=1)[-1], - mime_type="text/csv", - created_by="u1", - created_at=datetime.now(), +@dataclass(frozen=True) +class KnowledgeServiceStub: + segments: DatasetSegmentApplicationService + + +@dataclass(frozen=True) +class ApplicationServiceStub: + knowledge: KnowledgeServiceStub + + +def _patch_services(segments: DatasetSegmentApplicationService) -> AbstractContextManager[object]: + return patch( + "controllers.console.datasets.datasets_segments.application_services", + return_value=ApplicationServiceStub(knowledge=KnowledgeServiceStub(segments=segments)), ) -class SQLiteControllerTest: - session: Session +def test_list_segments_delegates_typed_query_and_serializes_page(app: Flask) -> None: + segments = create_autospec(DatasetSegmentApplicationService, instance=True, spec_set=True) + segments.list_segments.return_value = SegmentPage( + items=(_segment_data(),), total=1, total_pages=1, page=2, limit=10 + ) + method = unwrap(DatasetDocumentSegmentListApi.get) - @pytest.fixture(autouse=True) - def _use_sqlite_session(self, sqlite_session: Session) -> None: - self.session = sqlite_session - - -def _segment_response_dict(): - return { - "id": "seg-1", - "position": 1, - "document_id": "doc-1", - "content": "c", - "sign_content": "c", - "answer": "a", - "word_count": 1, - "tokens": 1, - "keywords": ["test"], - "index_node_id": "n1", - "index_node_hash": "h", - "hit_count": 0, - "enabled": True, - "disabled_at": None, - "disabled_by": None, - "status": "completed", - "created_by": "u1", - "created_at": 1779678000, - "updated_at": 1779678000, - "updated_by": "u1", - "indexing_at": None, - "completed_at": None, - "error": None, - "stopped_at": None, - "child_chunks": [], - "attachments": [], - "summary": None, - } - - -def test_segment_response_with_summary(sqlite_session: Session): - segment = _segment() with ( - patch.object(DocumentSegment, "get_child_chunks", autospec=True, return_value=[]) as get_child_chunks, - patch.object(DocumentSegment, "get_attachments", autospec=True, return_value=[]) as get_attachments, + app.test_request_context("/?status=completed&limit=10&page=2&enabled=true&keyword=needle"), + _patch_services(segments), ): - result = segment_response_with_summary(segment, "summary", session=sqlite_session) - assert result.summary == "summary" - assert result.id == segment.id - get_child_chunks.assert_called_once_with(segment, session=sqlite_session, include_full_doc=False) - get_attachments.assert_called_once_with(segment, session=sqlite_session) + response, status = method(DatasetDocumentSegmentListApi(), _context(), "dataset-1", "document-1") + + assert status == 200 + assert response["data"][0]["id"] == "segment-1" + assert response["data"][0]["created_at"] == 1767225600 + assert response["data"][0]["keywords"] == [] + assert response["data"][0]["disabled_at"] is None + query = segments.list_segments.call_args.kwargs["query"] + assert query == SegmentListFilter( + page=2, + limit=10, + statuses=("completed",), + enabled="true", + keyword="needle", + ) -def test_segment_responses_with_summaries_reuses_caller_session(sqlite_session: Session): - segments = [_segment(), _segment()] - segments[1].id = "seg-2" - expected_responses = [MagicMock(), MagicMock()] +def test_delete_segments_delegates_all_selected_ids(app: Flask) -> None: + segments = create_autospec(DatasetSegmentApplicationService, instance=True, spec_set=True) + method = unwrap(DatasetDocumentSegmentListApi.delete) - with patch( - "fields.segment_fields.segment_response_with_summary", side_effect=expected_responses - ) as serialize_segment: - responses = segment_responses_with_summaries( - segments, - {"seg-1": "summary-1", "seg-2": None}, - session=sqlite_session, + with app.test_request_context("/?segment_id=one&segment_id=two"), _patch_services(segments): + response, status = method(DatasetDocumentSegmentListApi(), _context(), "dataset-1", "document-1") + + assert (response, status) == ("", 204) + assert segments.delete_segments.call_args.kwargs["segment_ids"] == ["one", "two"] + + +@pytest.mark.parametrize("action", ["enable", "disable"]) +def test_change_segment_status_delegates_action(app: Flask, action: str) -> None: + segments = create_autospec(DatasetSegmentApplicationService, instance=True, spec_set=True) + method = unwrap(DatasetDocumentSegmentApi.patch) + + with app.test_request_context("/?segment_id=one"), _patch_services(segments): + response, status = method(DatasetDocumentSegmentApi(), _context(), "dataset-1", "document-1", action) + + assert status == 200 + assert response == {"result": "success"} + assert segments.change_segment_status.call_args.kwargs["action"] == action + + +@pytest.mark.parametrize("action", ["typo", "ENABLE", "enable "]) +def test_change_segment_status_rejects_unknown_action_before_service_call(app: Flask, action: str) -> None: + segments = create_autospec(DatasetSegmentApplicationService, instance=True, spec_set=True) + method = unwrap(DatasetDocumentSegmentApi.patch) + + with ( + app.test_request_context("/?segment_id=one"), + _patch_services(segments), + pytest.raises(InvalidActionError) as error, + ): + method(DatasetDocumentSegmentApi(), _context(), "dataset-1", "document-1", action) + + assert error.value.code == 400 + segments.change_segment_status.assert_not_called() + + +def test_create_segment_delegates_payload_and_serializes_detail(app: Flask) -> None: + segments = create_autospec(DatasetSegmentApplicationService, instance=True, spec_set=True) + segments.create_segment.return_value = SegmentDetail(data=_segment_data(), doc_form="text_model") + method = unwrap(DatasetDocumentSegmentAddApi.post) + + with app.test_request_context("/"), _patch_services(segments): + response, status = method( + DatasetDocumentSegmentAddApi(), + SegmentCreatePayload(content="content"), + _context(), + "dataset-1", + "document-1", ) - assert responses == expected_responses - assert serialize_segment.call_args_list == [ - call(segments[0], "summary-1", session=sqlite_session), - call(segments[1], None, session=sqlite_session), - ] + assert status == 200 + assert response["data"]["id"] == "segment-1" + assert segments.create_segment.call_args.kwargs["values"] == {"content": "content"} -class TestDatasetDocumentSegmentListApi(SQLiteControllerTest): - def test_get_success(self, app: Flask): - api = DatasetDocumentSegmentListApi() - method = unwrap(api.get) - dataset = _dataset() - document = _document() - user = _account() - segment = _segment() - pagination = MagicMock() - pagination.items = [segment] - pagination.total = 1 - pagination.pages = 1 - with ( - app.test_request_context("/"), - patch("controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=dataset), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.check_dataset_permission", - return_value=None, - ), - patch("controllers.console.datasets.datasets_segments.DocumentService.get_document", return_value=document), - patch("controllers.console.datasets.datasets_segments.paginate_query", return_value=pagination), - patch("services.summary_index_service.SummaryIndexService.get_segments_summaries", return_value={}), - ): - response, status = method(api, self.session, "tenant-1", user, "ds-1", "doc-1") - assert status == 200 +def test_update_and_delete_segment_delegate_owned_ids(app: Flask) -> None: + segments = create_autospec(DatasetSegmentApplicationService, instance=True, spec_set=True) + segments.update_segment.return_value = SegmentDetail(data=_segment_data(), doc_form="text_model") - def test_get_dataset_not_found(self, app: Flask): - api = DatasetDocumentSegmentListApi() - method = unwrap(api.get) - user = _account() - with ( - app.test_request_context("/"), - patch("controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=None), - ): - with pytest.raises(NotFound): - method(api, self.session, "tenant-1", user, "ds-1", "doc-1") - - def test_get_permission_denied(self, app: Flask): - api = DatasetDocumentSegmentListApi() - method = unwrap(api.get) - dataset = _dataset() - user = _account() - with ( - app.test_request_context("/"), - patch("controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=dataset), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.check_dataset_permission", - side_effect=services.errors.account.NoPermissionError("no access"), - ), - ): - with pytest.raises(Forbidden): - method(api, self.session, "tenant-1", user, "ds-1", "doc-1") - - -class TestDatasetDocumentSegmentApi(SQLiteControllerTest): - def test_patch_success(self, app: Flask): - api = DatasetDocumentSegmentApi() - method = unwrap(api.patch) - user = _account() - dataset = _dataset() - document = _document() - with ( - app.test_request_context("/?segment_id=s1&segment_id=s2"), - patch("controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=dataset), - patch("controllers.console.datasets.datasets_segments.DocumentService.get_document", return_value=document), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.check_dataset_permission", - return_value=None, - ), - patch("controllers.console.datasets.datasets_segments.redis_client.get", return_value=None), - patch( - "controllers.console.datasets.datasets_segments.SegmentService.update_segments_status", - return_value=None, - ), - ): - response, status = method(api, self.session, "tenant-1", user, "ds-1", "doc-1", "enable") - assert status == 200 - assert response["result"] == "success" - - def test_patch_document_indexing_in_progress(self, app: Flask): - api = DatasetDocumentSegmentApi() - method = unwrap(api.patch) - user = _account() - dataset = _dataset() - document = _document() - with ( - app.test_request_context("/"), - patch("controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=dataset), - patch("controllers.console.datasets.datasets_segments.DocumentService.get_document", return_value=document), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.check_dataset_model_setting", - return_value=None, - ), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.check_dataset_permission", - return_value=None, - ), - patch("controllers.console.datasets.datasets_segments.redis_client.get", return_value=b"running"), - ): - with pytest.raises(InvalidActionError): - method(api, self.session, "tenant-1", user, "ds-1", "doc-1", "disable") - - def test_patch_llm_bad_request(self, app: Flask): - api = DatasetDocumentSegmentApi() - method = unwrap(api.patch) - user = _account() - dataset = _dataset( - indexing_technique="high_quality", embedding_model_provider="openai", embedding_model="text-embed" + with app.test_request_context("/"), _patch_services(segments): + update_response, update_status = unwrap(DatasetDocumentSegmentUpdateApi.patch)( + DatasetDocumentSegmentUpdateApi(), + SegmentUpdatePayload(content="updated"), + _context(), + "dataset-1", + "document-1", + "segment-1", ) - document = _document() - with ( - app.test_request_context("/?segment_id=s1"), - patch("controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=dataset), - patch("controllers.console.datasets.datasets_segments.DocumentService.get_document", return_value=document), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.check_dataset_model_setting", - return_value=None, - ), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.check_dataset_permission", - return_value=None, - ), - patch( - "controllers.console.datasets.datasets_segments.ModelManager.get_model_instance", - side_effect=LLMBadRequestError(), - ), - ): - with pytest.raises(ProviderNotInitializeError): - method(api, self.session, "tenant-1", user, "ds-1", "doc-1", "enable") - - def test_patch_provider_token_not_init(self, app: Flask): - api = DatasetDocumentSegmentApi() - method = unwrap(api.patch) - user = _account() - dataset = _dataset( - indexing_technique="high_quality", embedding_model_provider="openai", embedding_model="text-embed" + delete_response, delete_status = unwrap(DatasetDocumentSegmentUpdateApi.delete)( + DatasetDocumentSegmentUpdateApi(), + _context(), + "dataset-1", + "document-1", + "segment-1", ) - document = _document() - with ( - app.test_request_context("/?segment_id=s1"), - patch("controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=dataset), - patch("controllers.console.datasets.datasets_segments.DocumentService.get_document", return_value=document), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.check_dataset_model_setting", - return_value=None, - ), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.check_dataset_permission", - return_value=None, - ), - patch( - "controllers.console.datasets.datasets_segments.ModelManager.get_model_instance", - side_effect=ProviderTokenNotInitError("token missing"), - ), - ): - with pytest.raises(ProviderNotInitializeError): - method(api, self.session, "tenant-1", user, "ds-1", "doc-1", "enable") + + assert update_status == 200 + assert update_response["data"]["content"] == "content" + assert (delete_response, delete_status) == ("", 204) + assert segments.update_segment.call_args.kwargs["segment_id"] == "segment-1" + assert segments.delete_segment.call_args.kwargs["segment_id"] == "segment-1" -class TestDatasetDocumentSegmentAddApi(SQLiteControllerTest): - def test_post_success(self, app: Flask): - api = DatasetDocumentSegmentAddApi() - method = unwrap(api.post) - payload = {"content": "hello"} - user = _account() - dataset = _dataset() - document = _document() - segment = _segment() - with ( - app.test_request_context("/", json=payload), - patch.object(type(console_ns), "payload", payload), - patch("controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=dataset), - patch("controllers.console.datasets.datasets_segments.DocumentService.get_document", return_value=document), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.check_dataset_permission", - return_value=None, - ), - patch( - "controllers.console.datasets.datasets_segments.SegmentService.segment_create_args_validate", - return_value=None, - ), - patch("controllers.console.datasets.datasets_segments.SegmentService.create_segment", return_value=segment), - patch( - "controllers.console.datasets.datasets_segments.SummaryIndexService.get_segment_summary", - return_value=None, - ), - ): - response, status = method( - api, SegmentCreatePayload(content="test content"), self.session, "tenant-1", user, "ds-1", "doc-1" - ) - assert status == 200 - assert response["data"]["id"] == "seg-1" +def test_batch_import_delegates_and_preserves_dispatch_failure_response(app: Flask) -> None: + segments = create_autospec(DatasetSegmentApplicationService, instance=True, spec_set=True) + segments.start_batch_import.return_value = SegmentBatchImport(job_id="job-1", job_status="waiting") + method = unwrap(DatasetDocumentSegmentBatchImportApi.post) - def test_post_llm_bad_request(self, app: Flask): - api = DatasetDocumentSegmentAddApi() - method = unwrap(api.post) - payload = {"content": "x"} - user = _account() - dataset = _dataset( - indexing_technique="high_quality", embedding_model_provider="openai", embedding_model="text-embed" + with app.test_request_context("/"), _patch_services(segments): + response, status = method( + DatasetDocumentSegmentBatchImportApi(), + BatchImportPayload(upload_file_id="file-1"), + _context(), + "dataset-1", + "document-1", ) - document = _document() - with ( - app.test_request_context("/", json=payload), - patch.object(type(console_ns), "payload", payload), - patch("controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=dataset), - patch("controllers.console.datasets.datasets_segments.DocumentService.get_document", return_value=document), - patch( - "controllers.console.datasets.datasets_segments.ModelManager.get_model_instance", - side_effect=LLMBadRequestError(), - ), - ): - with pytest.raises(ProviderNotInitializeError): - method( - api, SegmentCreatePayload(content="test content"), self.session, "tenant-1", user, "ds-1", "doc-1" - ) + assert status == 200 + assert response == {"job_id": "job-1", "job_status": "waiting"} - def test_post_provider_token_not_init(self, app: Flask): - api = DatasetDocumentSegmentAddApi() - method = unwrap(api.post) - payload = {"content": "x"} - user = _account() - dataset = _dataset( - indexing_technique="high_quality", embedding_model_provider="openai", embedding_model="text-embed" + segments.start_batch_import.side_effect = SegmentBatchImportDispatchError("redis down") + with app.test_request_context("/"), _patch_services(segments): + response, status = method( + DatasetDocumentSegmentBatchImportApi(), + BatchImportPayload(upload_file_id="file-1"), + _context(), + "dataset-1", + "document-1", ) - document = _document() - with ( - app.test_request_context("/", json=payload), - patch.object(type(console_ns), "payload", payload), - patch("controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=dataset), - patch("controllers.console.datasets.datasets_segments.DocumentService.get_document", return_value=document), - patch( - "controllers.console.datasets.datasets_segments.ModelManager.get_model_instance", - side_effect=ProviderTokenNotInitError("token missing"), - ), - ): - with pytest.raises(ProviderNotInitializeError): - method( - api, SegmentCreatePayload(content="test content"), self.session, "tenant-1", user, "ds-1", "doc-1" - ) + assert (response, status) == ({"error": "redis down"}, 500) -class TestDatasetDocumentSegmentUpdateApi(SQLiteControllerTest): - def test_patch_success(self, app: Flask): - api = DatasetDocumentSegmentUpdateApi() - method = unwrap(api.patch) - payload = {"content": "updated"} - user = _account() - dataset = _dataset() - document = _document() - segment = _segment() - with ( - app.test_request_context("/", json=payload), - patch.object(type(console_ns), "payload", payload), - patch("controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=dataset), - patch("controllers.console.datasets.datasets_segments.DocumentService.get_document", return_value=document), - patch( - "controllers.console.datasets.datasets_segments.SegmentService.get_segment_by_ref", return_value=segment - ), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.check_dataset_permission", - return_value=None, - ), - patch( - "controllers.console.datasets.datasets_segments.SegmentService.segment_create_args_validate", - return_value=None, - ), - patch("controllers.console.datasets.datasets_segments.SegmentService.update_segment", return_value=segment), - patch( - "controllers.console.datasets.datasets_segments.SummaryIndexService.get_segment_summary", - return_value=None, - ), - ): - response, status = method( - api, - SegmentUpdatePayload(content="test content"), - self.session, - "tenant-1", - user, - "ds-1", - "doc-1", - "seg-1", - ) - assert status == 200 - assert "data" in response +def test_batch_import_status_preserves_missing_job_contract(app: Flask) -> None: + segments = create_autospec(DatasetSegmentApplicationService, instance=True, spec_set=True) + method = unwrap(DatasetDocumentSegmentBatchImportStatusApi.get) - def test_patch_document_outside_dataset_is_not_found(self, app: Flask): - api = DatasetDocumentSegmentUpdateApi() - method = unwrap(api.patch) - payload = {"content": "updated"} - user = _account() - dataset = _dataset() - document = _document(dataset_id="other-dataset") - with ( - app.test_request_context("/", json=payload), - patch.object(type(console_ns), "payload", payload), - patch("controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=dataset), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.check_dataset_model_setting", - return_value=None, - ), - patch("controllers.console.datasets.datasets_segments.DocumentService.get_document", return_value=document), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.check_dataset_permission", - return_value=None, - ), - ): - with pytest.raises(NotFound): - method( - api, - SegmentUpdatePayload(content="test content"), - self.session, - "tenant-1", - user, - "ds-1", - "doc-1", - "seg-1", - ) + segments.get_batch_import_status.side_effect = SegmentBatchImportNotFoundError("The job does not exist.") + with app.test_request_context("/"), _patch_services(segments), pytest.raises(ValueError): + method(DatasetDocumentSegmentBatchImportStatusApi(), _context(), "job-1") - def test_patch_segment_not_found(self, app: Flask): - api = DatasetDocumentSegmentUpdateApi() - method = unwrap(api.patch) - payload = {"content": "updated"} - user = _account() - dataset = _dataset() - document = _document() - with ( - app.test_request_context("/", json=payload), - patch.object(type(console_ns), "payload", payload), - patch("controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=dataset), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.check_dataset_model_setting", - return_value=None, - ), - patch("controllers.console.datasets.datasets_segments.DocumentService.get_document", return_value=document), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.check_dataset_permission", - return_value=None, - ), - patch( - "controllers.console.datasets.datasets_segments.SegmentService.get_segment_by_ref", return_value=None - ), - ): - with pytest.raises(NotFound): - method( - api, - SegmentUpdatePayload(content="test content"), - self.session, - "tenant-1", - user, - "ds-1", - "doc-1", - "seg-1", - ) + segments.get_batch_import_status.side_effect = None + segments.get_batch_import_status.return_value = SegmentBatchImport(job_id="job-1", job_status="completed") + with app.test_request_context("/"), _patch_services(segments): + response, status = method(DatasetDocumentSegmentBatchImportStatusApi(), _context(), "job-1") + assert status == 200 + assert response == {"job_id": "job-1", "job_status": "completed"} - def test_patch_llm_bad_request(self, app: Flask): - api = DatasetDocumentSegmentUpdateApi() - method = unwrap(api.patch) - payload = {"content": "x"} - user = _account() - dataset = _dataset( - indexing_technique="high_quality", embedding_model_provider="openai", embedding_model="text-embed" + +def test_child_chunk_create_and_list_delegate(app: Flask) -> None: + segments = create_autospec(DatasetSegmentApplicationService, instance=True, spec_set=True) + segments.create_child_chunk.return_value = _child_chunk_data() + segments.list_child_chunks.return_value = ChildChunkPage( + items=(_child_chunk_data(),), total=1, total_pages=1, page=1, limit=20 + ) + + with app.test_request_context("/"), _patch_services(segments): + create_response, create_status = unwrap(ChildChunkAddApi.post)( + ChildChunkAddApi(), + ChildChunkCreatePayload(content="child"), + _context(), + "dataset-1", + "document-1", + "segment-1", ) - document = _document() - with ( - app.test_request_context("/", json=payload), - patch.object(type(console_ns), "payload", payload), - patch("controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=dataset), - patch("controllers.console.datasets.datasets_segments.DocumentService.get_document", return_value=document), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.check_dataset_model_setting", - return_value=None, - ), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.check_dataset_permission", - return_value=None, - ), - patch( - "controllers.console.datasets.datasets_segments.ModelManager.get_model_instance", - side_effect=LLMBadRequestError(), - ), - ): - with pytest.raises(ProviderNotInitializeError): - method( - api, - SegmentUpdatePayload(content="test content"), - self.session, - "tenant-1", - user, - "ds-1", - "doc-1", - "seg-1", - ) - - -class TestDatasetDocumentSegmentBatchImportApi(SQLiteControllerTest): - def test_post_success(self, app: Flask): - api = DatasetDocumentSegmentBatchImportApi() - method = unwrap(api.post) - payload = {"upload_file_id": "file-1"} - upload_file = _upload_file() - self.session.add(upload_file) - self.session.commit() - user = _account() - dataset = _dataset() - document = _document() - with ( - app.test_request_context("/", json=payload), - patch.object(type(console_ns), "payload", payload), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.get_dataset_for_tenant", - return_value=dataset, - ) as get_dataset_for_tenant, - patch( - "controllers.console.datasets.datasets_segments.DatasetRefService.get_document_by_ref", - return_value=document, - ) as get_document_by_ref, - patch("controllers.console.datasets.datasets_segments.redis_client.setnx", return_value=True), - patch( - "controllers.console.datasets.datasets_segments.batch_create_segment_to_index_task.delay", - return_value=None, - ), - ): - response, status = method( - api, - BatchImportPayload(upload_file_id="test-file-id"), - self.session, - "tenant-1", - user, - "ds-1", - "doc-1", - ) - assert status == 200 - assert response["job_status"] == "waiting" - get_dataset_for_tenant.assert_called_once_with("ds-1", "tenant-1", session=self.session) - document_ref = get_document_by_ref.call_args.args[0] - assert document_ref.dataset.tenant_id == "tenant-1" - assert document_ref.dataset.dataset_id == "ds-1" - assert document_ref.document_id == "doc-1" - - def test_post_dataset_not_found(self, app: Flask): - api = DatasetDocumentSegmentBatchImportApi() - method = unwrap(api.post) - payload = {"upload_file_id": "file-1"} - user = _account() - with ( - app.test_request_context("/", json=payload), - patch.object(type(console_ns), "payload", payload), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.get_dataset_for_tenant", - return_value=None, - ), - ): - with pytest.raises(NotFound): - method( - api, - BatchImportPayload(upload_file_id="test-file-id"), - self.session, - "tenant-1", - user, - "ds-1", - "doc-1", - ) - - def test_post_document_not_found(self, app: Flask): - api = DatasetDocumentSegmentBatchImportApi() - method = unwrap(api.post) - payload = {"upload_file_id": "file-1"} - user = _account() - with ( - app.test_request_context("/", json=payload), - patch.object(type(console_ns), "payload", payload), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.get_dataset_for_tenant", - return_value=_dataset(), - ), - patch( - "controllers.console.datasets.datasets_segments.DatasetRefService.get_document_by_ref", - return_value=None, - ), - ): - with pytest.raises(NotFound): - method( - api, - BatchImportPayload(upload_file_id="test-file-id"), - self.session, - "tenant-1", - user, - "ds-1", - "doc-1", - ) - - def test_post_upload_file_not_found(self, app: Flask): - api = DatasetDocumentSegmentBatchImportApi() - method = unwrap(api.post) - payload = {"upload_file_id": "file-1"} - user = _account() - with ( - app.test_request_context("/", json=payload), - patch.object(type(console_ns), "payload", payload), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.get_dataset_for_tenant", - return_value=_dataset(), - ), - patch( - "controllers.console.datasets.datasets_segments.DatasetRefService.get_document_by_ref", - return_value=_document(), - ), - ): - with pytest.raises(NotFound): - method( - api, - BatchImportPayload(upload_file_id="test-file-id"), - self.session, - "tenant-1", - user, - "ds-1", - "doc-1", - ) - - def test_post_invalid_file_type(self, app: Flask): - api = DatasetDocumentSegmentBatchImportApi() - method = unwrap(api.post) - payload = {"upload_file_id": "file-1"} - upload_file = _upload_file(name="test.txt") - self.session.add(upload_file) - self.session.commit() - user = _account() - with ( - app.test_request_context("/", json=payload), - patch.object(type(console_ns), "payload", payload), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.get_dataset_for_tenant", - return_value=_dataset(), - ), - patch( - "controllers.console.datasets.datasets_segments.DatasetRefService.get_document_by_ref", - return_value=_document(), - ), - ): - with pytest.raises(ValueError): - method( - api, - BatchImportPayload(upload_file_id="test-file-id"), - self.session, - "tenant-1", - user, - "ds-1", - "doc-1", - ) - - def test_post_async_task_failure(self, app: Flask): - api = DatasetDocumentSegmentBatchImportApi() - method = unwrap(api.post) - payload = {"upload_file_id": "file-1"} - upload_file = _upload_file() - self.session.add(upload_file) - self.session.commit() - user = _account() - with ( - app.test_request_context("/", json=payload), - patch.object(type(console_ns), "payload", payload), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.get_dataset_for_tenant", - return_value=_dataset(), - ), - patch( - "controllers.console.datasets.datasets_segments.DatasetRefService.get_document_by_ref", - return_value=_document(), - ), - patch( - "controllers.console.datasets.datasets_segments.redis_client.setnx", side_effect=Exception("redis down") - ), - ): - response, status = method( - api, - BatchImportPayload(upload_file_id="test-file-id"), - self.session, - "tenant-1", - user, - "ds-1", - "doc-1", - ) - assert status == 500 - assert "error" in response - - def test_get_job_not_found_in_redis(self, app: Flask): - api = DatasetDocumentSegmentBatchImportApi() - method = unwrap(api.get) - with ( - app.test_request_context("/"), - patch("controllers.console.datasets.datasets_segments.redis_client.get", return_value=None), - ): - with pytest.raises(ValueError): - method(api, job_id="job-1") - - -class TestChildChunkAddApi(SQLiteControllerTest): - def test_patch_documents_batch_update_payload(self): - patch_method = cast(Any, ChildChunkAddApi.patch) - api_doc = cast(dict[str, Any], patch_method.__apidoc__) - expected_model = ChildChunkBatchUpdatePayload.__name__ - assert [model.name for model in api_doc["expect"]] == [expected_model] - - def test_get_uses_default_pagination_for_malformed_ints(self, app: Flask): - api = ChildChunkAddApi() - method = unwrap(api.get) - dataset = _dataset() - document = _document() - pagination = MagicMock(items=[], total=0, pages=0) - with ( - app.test_request_context("/?page=bad&limit="), - patch("controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=dataset), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.check_dataset_model_setting", - return_value=None, - ), - patch("controllers.console.datasets.datasets_segments.DocumentService.get_document", return_value=document), - patch( - "controllers.console.datasets.datasets_segments.SegmentService.get_segment_by_ref", - return_value=_segment(), - ), - patch( - "controllers.console.datasets.datasets_segments.SegmentService.get_child_chunks", - return_value=pagination, - ) as get_child_chunks, - ): - response, status = method(api, self.session, "tenant-1", "ds-1", "doc-1", "seg-1") - assert status == 200 - assert response["page"] == 1 - assert response["limit"] == 20 - session = get_child_chunks.call_args.kwargs["session"] - assert session is self.session - assert get_child_chunks.call_args.args == ("seg-1", "doc-1", "ds-1", 1, 20, None) - - def test_post_success(self, app: Flask): - api = ChildChunkAddApi() - method = unwrap(api.post) - payload = {"content": "child"} - user = _account() - dataset = _dataset() - document = _document() - segment = _segment() - child_chunk = _child_chunk() - with ( - app.test_request_context("/", json=payload), - patch.object(type(console_ns), "payload", payload), - patch("controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=dataset), - patch("controllers.console.datasets.datasets_segments.DocumentService.get_document", return_value=document), - patch( - "controllers.console.datasets.datasets_segments.SegmentService.get_segment_by_ref", return_value=segment - ), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.check_dataset_permission", - return_value=None, - ), - patch( - "controllers.console.datasets.datasets_segments.SegmentService.create_child_chunk", - return_value=child_chunk, - ), - ): - response, status = method( - api, - ChildChunkCreatePayload(content="child"), - self.session, - "tenant-1", - user, - "ds-1", - "doc-1", - "seg-1", - ) - assert status == 200 - assert response["data"]["id"] == "cc-1" - - def test_post_child_chunk_indexing_error(self, app: Flask): - api = ChildChunkAddApi() - method = unwrap(api.post) - payload = {"content": "child"} - user = _account() - dataset = _dataset() - document = _document() - segment = _segment() - with ( - app.test_request_context("/", json=payload), - patch.object(type(console_ns), "payload", payload), - patch("controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=dataset), - patch("controllers.console.datasets.datasets_segments.DocumentService.get_document", return_value=document), - patch( - "controllers.console.datasets.datasets_segments.SegmentService.get_segment_by_ref", return_value=segment - ), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.check_dataset_permission", - return_value=None, - ), - patch( - "controllers.console.datasets.datasets_segments.SegmentService.create_child_chunk", - side_effect=ChildChunkIndexingServiceError("fail"), - ), - ): - with pytest.raises(ChildChunkIndexingError): - method( - api, - ChildChunkCreatePayload(content="child"), - self.session, - "tenant-1", - user, - "ds-1", - "doc-1", - "seg-1", - ) - - def test_post_permission_denied(self, app: Flask): - api = ChildChunkAddApi() - method = unwrap(api.post) - payload = {"content": "child"} - user = _account() - dataset = _dataset() - document = _document() - with ( - app.test_request_context("/", json=payload), - patch.object(type(console_ns), "payload", payload), - patch("controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=dataset), - patch("controllers.console.datasets.datasets_segments.DocumentService.get_document", return_value=document), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.check_dataset_permission", - side_effect=services.errors.account.NoPermissionError("no access"), - ), - ): - with pytest.raises(Forbidden): - method( - api, - ChildChunkCreatePayload(content="child"), - self.session, - "tenant-1", - user, - "ds-1", - "doc-1", - "seg-1", - ) - - -class TestChildChunkUpdateApi(SQLiteControllerTest): - def test_delete_success(self, app: Flask): - api = ChildChunkUpdateApi() - method = unwrap(api.delete) - user = _account() - dataset = _dataset() - document = _document() - segment = _segment() - child_chunk = _child_chunk() - with ( - app.test_request_context("/"), - patch("controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=dataset), - patch("controllers.console.datasets.datasets_segments.DocumentService.get_document", return_value=document), - patch( - "controllers.console.datasets.datasets_segments.SegmentService.get_segment_by_ref", return_value=segment - ), - patch( - "controllers.console.datasets.datasets_segments.SegmentService.get_child_chunk_by_segment_ref", - return_value=child_chunk, - ), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.check_dataset_permission", - return_value=None, - ), - patch( - "controllers.console.datasets.datasets_segments.SegmentService.delete_child_chunk", return_value=None - ), - ): - response, status = method(api, self.session, "tenant-1", user, "ds-1", "doc-1", "seg-1", "cc-1") - assert status == 204 - assert response == "" - - def test_delete_child_chunk_index_error(self, app: Flask): - api = ChildChunkUpdateApi() - method = unwrap(api.delete) - user = _account() - dataset = _dataset() - document = _document() - segment = _segment() - child_chunk = _child_chunk() - with ( - app.test_request_context("/"), - patch("controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=dataset), - patch("controllers.console.datasets.datasets_segments.DocumentService.get_document", return_value=document), - patch( - "controllers.console.datasets.datasets_segments.SegmentService.get_segment_by_ref", return_value=segment - ), - patch( - "controllers.console.datasets.datasets_segments.SegmentService.get_child_chunk_by_segment_ref", - return_value=child_chunk, - ), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.check_dataset_permission", - return_value=None, - ), - patch( - "controllers.console.datasets.datasets_segments.SegmentService.delete_child_chunk", - side_effect=ChildChunkDeleteIndexServiceError("fail"), - ), - ): - with pytest.raises(ChildChunkDeleteIndexError): - method(api, self.session, "tenant-1", user, "ds-1", "doc-1", "seg-1", "cc-1") - - def test_delete_child_chunk_not_found(self, app: Flask): - api = ChildChunkUpdateApi() - method = unwrap(api.delete) - user = _account() - dataset = _dataset() - document = _document() - segment = _segment() - with ( - app.test_request_context("/"), - patch("controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=dataset), - patch("controllers.console.datasets.datasets_segments.DocumentService.get_document", return_value=document), - patch( - "controllers.console.datasets.datasets_segments.SegmentService.get_segment_by_ref", return_value=segment - ), - patch( - "controllers.console.datasets.datasets_segments.SegmentService.get_child_chunk_by_segment_ref", - return_value=None, - ), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.check_dataset_permission", - return_value=None, - ), - ): - with pytest.raises(NotFound): - method(api, self.session, "tenant-1", user, "ds-1", "doc-1", "seg-1", "cc-1") - - def test_patch_child_chunk_not_found(self, app: Flask): - api = ChildChunkUpdateApi() - method = unwrap(api.patch) - payload = {"content": "updated child"} - user = _account() - dataset = _dataset() - document = _document() - segment = _segment() - with ( - app.test_request_context("/", json=payload), - patch.object(type(console_ns), "payload", payload), - patch("controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=dataset), - patch("controllers.console.datasets.datasets_segments.DocumentService.get_document", return_value=document), - patch( - "controllers.console.datasets.datasets_segments.SegmentService.get_segment_by_ref", return_value=segment - ), - patch( - "controllers.console.datasets.datasets_segments.SegmentService.get_child_chunk_by_segment_ref", - return_value=None, - ), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.check_dataset_permission", - return_value=None, - ), - ): - with pytest.raises(NotFound): - method( - api, - ChildChunkUpdatePayload(content="updated child"), - self.session, - "tenant-1", - user, - "ds-1", - "doc-1", - "seg-1", - "cc-1", - ) - - -class TestSegmentListAdvancedCases(SQLiteControllerTest): - def test_segment_list_with_keyword_filter(self, app: Flask): - api = DatasetDocumentSegmentListApi() - method = unwrap(api.get) - dataset = _dataset() - document = _document() - user = _account() - segment = _segment() - pagination = MagicMock(items=[segment], total=1, pages=1) - with ( - app.test_request_context("/?keyword=test"), - patch("controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=dataset), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.check_dataset_permission", - return_value=None, - ), - patch("controllers.console.datasets.datasets_segments.DocumentService.get_document", return_value=document), - patch("controllers.console.datasets.datasets_segments.paginate_query", return_value=pagination), - patch("services.summary_index_service.SummaryIndexService.get_segments_summaries", return_value={}), - ): - result = method(api, self.session, "tenant-1", user, "ds-1", "doc-1") - if isinstance(result, tuple): - response, status = result - else: - response, status = (result, 200) - assert status == 200 - assert response["total"] == 1 - - def test_segment_list_postgres_keyword_filter_handles_scalar_keywords(self, app: Flask): - api = DatasetDocumentSegmentListApi() - method = unwrap(api.get) - dataset = _dataset(dataset_id="22222222-2222-2222-2222-222222222222") - document = _document( - document_id="33333333-3333-3333-3333-333333333333", - dataset_id=dataset.id, + list_response, list_status = unwrap(ChildChunkAddApi.get)( + ChildChunkAddApi(), _context(), "dataset-1", "document-1", "segment-1" ) - user = _account() - pagination = MagicMock(items=[], total=0, pages=0) - with ( - app.test_request_context("/?keyword=test"), - patch("controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=dataset), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.check_dataset_permission", - return_value=None, - ), - patch("controllers.console.datasets.datasets_segments.DocumentService.get_document", return_value=document), - patch( - "controllers.console.datasets.datasets_segments.dify_config", - SimpleNamespace(SQLALCHEMY_DATABASE_URI_SCHEME="postgresql"), - ), - patch( - "controllers.console.datasets.datasets_segments.paginate_query", return_value=pagination - ) as paginate_mock, - ): - method( - api, - self.session, - "11111111-1111-1111-1111-111111111111", - user, - "22222222-2222-2222-2222-222222222222", - "33333333-3333-3333-3333-333333333333", - ) - query = paginate_mock.call_args.args[0] - sql = str(query.compile(compile_kwargs={"literal_binds": True})) - assert "jsonb_array_elements_text(CASE" in sql - assert "ELSE CAST('[]' AS JSONB)" in sql - def test_segment_list_permission_denied(self, app: Flask): - """Test segment list with permission denied""" - api = DatasetDocumentSegmentListApi() - method = unwrap(api.get) - user = _account() - with ( - app.test_request_context("/"), - patch("controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=_dataset()), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.check_dataset_permission", - side_effect=services.errors.account.NoPermissionError("No permission"), - ), - ): - with pytest.raises(Forbidden): - method(api, self.session, "tenant-1", user, "ds-1", "doc-1") - - def test_segment_list_dataset_not_found(self, app: Flask): - """Test segment list with dataset not found""" - api = DatasetDocumentSegmentListApi() - method = unwrap(api.get) - user = _account() - with ( - app.test_request_context("/"), - patch("controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=None), - ): - with pytest.raises(NotFound): - method(api, self.session, "tenant-1", user, "ds-1", "doc-1") + assert create_status == list_status == 200 + assert create_response["data"]["id"] == "child-1" + assert create_response["data"]["created_at"] == 1767225600 + assert list_response["data"][0]["id"] == "child-1" -class TestSegmentOperationCases(SQLiteControllerTest): - def test_segment_add_with_provider_token_error(self, app: Flask): - """Test segment add with provider token not initialized""" - api = DatasetDocumentSegmentAddApi() - method = unwrap(api.post) - user = _account() - dataset = _dataset() - document = _document() - payload = {"content": "new content", "answer": None} - with ( - app.test_request_context("/", json=payload), - patch.object(type(console_ns), "payload", payload), - patch("controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=dataset), - patch("controllers.console.datasets.datasets_segments.DocumentService.get_document", return_value=document), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.check_dataset_permission", - return_value=None, - ), - patch( - "controllers.console.datasets.datasets_segments.SegmentService.create_segment", - side_effect=ProviderTokenNotInitError("Token not init"), - ), - ): - with pytest.raises(ProviderTokenNotInitError): - method( - api, - SegmentCreatePayload(content="test content"), - self.session, - "tenant-1", - user, - "ds-1", - "doc-1", - ) +def test_child_chunk_mutations_delegate(app: Flask) -> None: + segments = create_autospec(DatasetSegmentApplicationService, instance=True, spec_set=True) + segments.update_child_chunks.return_value = (_child_chunk_data(),) + segments.update_child_chunk.return_value = _child_chunk_data() + chunks = [ChildChunkUpdateArgs(id="child-1", content="updated")] - def test_batch_import_with_document_not_found(self, app: Flask): - """Test batch import with document not found""" - api = DatasetDocumentSegmentBatchImportApi() - method = unwrap(api.post) - user = _account() - dataset = _dataset() - payload = {"upload_file_id": "file-1"} - with ( - app.test_request_context("/", json=payload), - patch.object(type(console_ns), "payload", payload), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.get_dataset_for_tenant", - return_value=dataset, - ), - patch( - "controllers.console.datasets.datasets_segments.DatasetRefService.get_document_by_ref", - return_value=None, - ), - ): - with pytest.raises(NotFound): - method( - api, - BatchImportPayload(upload_file_id="test-file-id"), - self.session, - "tenant-1", - user, - "ds-1", - "doc-1", - ) + with app.test_request_context("/"), _patch_services(segments): + batch_response, batch_status = unwrap(ChildChunkAddApi.patch)( + ChildChunkAddApi(), + ChildChunkBatchUpdatePayload(chunks=chunks), + _context(), + "dataset-1", + "document-1", + "segment-1", + ) + update_response, update_status = unwrap(ChildChunkUpdateApi.patch)( + ChildChunkUpdateApi(), + ChildChunkUpdatePayload(content="updated"), + _context(), + "dataset-1", + "document-1", + "segment-1", + "child-1", + ) + delete_response, delete_status = unwrap(ChildChunkUpdateApi.delete)( + ChildChunkUpdateApi(), + _context(), + "dataset-1", + "document-1", + "segment-1", + "child-1", + ) - def test_batch_import_with_invalid_file(self, app: Flask): - """Test batch import with invalid file type""" - api = DatasetDocumentSegmentBatchImportApi() - method = unwrap(api.post) - user = _account() - dataset = _dataset() - document = _document() - payload = {"upload_file_id": "file-1"} - with ( - app.test_request_context("/", json=payload), - patch.object(type(console_ns), "payload", payload), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.get_dataset_for_tenant", - return_value=dataset, - ), - patch( - "controllers.console.datasets.datasets_segments.DatasetRefService.get_document_by_ref", - return_value=document, - ), - ): - with pytest.raises(NotFound): - method( - api, - BatchImportPayload(upload_file_id="test-file-id"), - self.session, - "tenant-1", - user, - "ds-1", - "doc-1", - ) + assert batch_status == update_status == 200 + assert batch_response["data"][0]["id"] == "child-1" + assert update_response["data"]["id"] == "child-1" + assert (delete_response, delete_status) == ("", 204) - def test_batch_import_with_async_task_failure(self, app: Flask): - api = DatasetDocumentSegmentBatchImportApi() - method = unwrap(api.post) - user = _account() - dataset = _dataset() - document = _document() - upload_file = _upload_file() - self.session.add(upload_file) - self.session.commit() - payload = {"upload_file_id": "file-1"} - with ( - app.test_request_context("/", json=payload), - patch.object(type(console_ns), "payload", payload), - patch( - "controllers.console.datasets.datasets_segments.DatasetService.get_dataset_for_tenant", - return_value=dataset, - ), - patch( - "controllers.console.datasets.datasets_segments.DatasetRefService.get_document_by_ref", - return_value=document, - ), - patch( - "controllers.console.datasets.datasets_segments.batch_create_segment_to_index_task.delay", - side_effect=Exception("Task failed"), - ), - ): - response, status = method( - api, - BatchImportPayload(upload_file_id="test-file-id"), - self.session, - "tenant-1", - user, - "ds-1", - "doc-1", - ) - assert status == 500 - assert "error" in response - def test_batch_import_get_job_not_found(self, app: Flask): - api = DatasetDocumentSegmentBatchImportApi() - method = unwrap(api.get) - with ( - app.test_request_context("/?job_id=invalid-job"), - patch("controllers.console.datasets.datasets_segments.redis_client.get", return_value=None), - ): - with pytest.raises(ValueError): - method(api, "invalid-job") +@pytest.mark.parametrize( + ("error", "expected"), + [ + (SegmentDatasetNotFoundError(), NotFound), + (SegmentDocumentNotFoundError(), NotFound), + (SegmentNotFoundError(), NotFound), + (ChildChunkNotFoundError(), NotFound), + (SegmentUploadFileNotFoundError(), NotFound), + (SegmentPermissionDeniedError(), Forbidden), + (SegmentEmbeddingModelUnavailableError(), ProviderNotInitializeError), + (SegmentDatasetModelUnavailableError(), ValueError), + (SegmentDocumentIndexingError(), InvalidActionError), + (SegmentStatusUpdateError(), InvalidActionError), + (SegmentInvalidFileTypeError(), ValueError), + (ChildChunkIndexingApplicationError(), ChildChunkIndexingError), + (ChildChunkDeleteIndexApplicationError(), ChildChunkDeleteIndexError), + ], +) +def test_segment_application_errors_map_to_existing_http_contract(error: Exception, expected: type[Exception]) -> None: + with pytest.raises(expected): + _raise_segment_error(error) 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 2ebc0b2e9f8..6d90902ee51 100644 --- a/api/tests/unit_tests/controllers/console/datasets/test_external.py +++ b/api/tests/unit_tests/controllers/console/datasets/test_external.py @@ -1,710 +1,164 @@ -import inspect -import json -from datetime import datetime -from typing import Any -from unittest.mock import ANY, PropertyMock, patch +from inspect import unwrap +from types import SimpleNamespace +from unittest.mock import create_autospec +from uuid import UUID import pytest -from flask import Flask -from sqlalchemy.orm import Session -from werkzeug.exceptions import Forbidden, NotFound +from werkzeug.exceptions import Forbidden, InternalServerError, NotFound -import services -from controllers.console import console_ns +from controllers.console.datasets import external as controller from controllers.console.datasets.error import DatasetNameDuplicateError -from controllers.console.datasets.external import ( - ExternalApiTemplateApi, - ExternalApiTemplateListApi, - ExternalApiTemplateListQuery, - ExternalApiUseCheckApi, - ExternalDatasetCreateApi, - ExternalHitTestingPayload, - ExternalKnowledgeApiPayload, - ExternalKnowledgeHitTestingApi, +from machinery.context import RequestContext +from services.errors.dataset import DatasetNameDuplicateError as DuplicateName +from services.knowledge.dataset_access import DatasetAccessDeniedError, DatasetNotFoundError +from services.knowledge.external.application import ( + ExternalHitTestingError, + ExternalKnowledgeApplicationService, + ExternalTemplateNotFoundError, ) -from fields.dataset_fields import DatasetDetailResponse -from models.account import Account, TenantAccountRole -from models.dataset import Dataset, ExternalKnowledgeApis, ExternalKnowledgeBindings -from services.dataset_service import DatasetService -from services.entities.external_knowledge_entities.external_knowledge_entities import ExternalDatasetCreatePayload -from services.external_knowledge_service import ExternalDatasetService -from services.hit_testing_service import HitTestingService + +CONTEXT = RequestContext("request-1", None, "actor-1", "tenant-1") +ID = UUID(int=1) @pytest.fixture -def app() -> Flask: - app = Flask("test_external_dataset") - app.config["TESTING"] = True - return app +def external(monkeypatch): + service = create_autospec(ExternalKnowledgeApplicationService, instance=True, spec_set=True) + monkeypatch.setattr( + controller, "application_services", lambda: SimpleNamespace(knowledge=SimpleNamespace(external=service)) + ) + return service -@pytest.fixture -def current_user() -> Account: - user = Account(name="Test User", email="user-1@example.com") - user.id = "user-1" - user.role = TenantAccountRole.EDITOR - return user +def test_templates_preserve_pagination(external): + page = {"data": [], "total": 0, "has_more": False, "page": 1, "limit": 100} + external.list_templates.return_value = page + result = unwrap(controller.ExternalApiTemplateListApi.get)( + controller.ExternalApiTemplateListApi(), + controller.ExternalApiTemplateListQuery(page=0, limit=1000, keyword="test"), + CONTEXT, + ) + assert result == (page, 200) + external.list_templates.assert_called_once_with(CONTEXT, page=0, limit=1000, keyword="test") -def _external_api_dict(api_id: str = "api-1") -> dict: - return { - "id": api_id, +@pytest.mark.parametrize( + ("resource", "verb", "operation", "payload"), + [ + ( + controller.ExternalApiTemplateListApi, + controller.ExternalApiTemplateListApi.post, + lambda service: service.create_template, + True, + ), + ( + controller.ExternalApiTemplateApi, + controller.ExternalApiTemplateApi.get, + lambda service: service.get_template, + False, + ), + ( + controller.ExternalApiTemplateApi, + controller.ExternalApiTemplateApi.patch, + lambda service: service.update_template, + True, + ), + ( + controller.ExternalApiTemplateApi, + controller.ExternalApiTemplateApi.delete, + lambda service: service.delete_template, + False, + ), + ], +) +def test_template_missing_is_mapped_and_context_forwarded(external, resource, verb, operation, payload): + method = operation(external) + method.side_effect = ExternalTemplateNotFoundError("API template not found.") + args = ( + [ + controller.ExternalKnowledgeApiPayload( + name="API", settings={"endpoint": "https://example", "api_key": "secret"} + ) + ] + if payload + else [] + ) + args.append(CONTEXT) + if resource is controller.ExternalApiTemplateApi: + args.append(ID) + with pytest.raises(NotFound): + unwrap(verb)(resource(), *args) + assert method.call_args.args == (CONTEXT,) + if resource is controller.ExternalApiTemplateApi: + assert method.call_args.kwargs["template_id"] == str(ID) + + +def test_template_serializes_settings_and_timestamp(external): + external.get_template.return_value = { + "id": str(ID), + "name": "API", "tenant_id": "tenant-1", - "name": f"External API {api_id}", - "description": f"Description for {api_id}", - "settings": { - "endpoint": f"https://external.example.com/{api_id}", - "api_key": "secret", - "headers": {"X-Source": "unit-test"}, - "timeout": 30, - }, - "dataset_bindings": [ - {"id": f"dataset-{api_id}", "name": f"Dataset {api_id}"}, - ], - "created_by": "user-1", + "settings": {"endpoint": "https://example", "api_key": "[__HIDDEN__]"}, + "description": "", + "created_by": "actor-1", "created_at": "2024-01-01T00:00:00", + "dataset_bindings": [], } + result, status = unwrap(controller.ExternalApiTemplateApi.get)(controller.ExternalApiTemplateApi(), CONTEXT, ID) + assert status == 200 + assert result == external.get_template.return_value -def _external_api_object(session: Session, api_id: str = "api-1") -> ExternalKnowledgeApis: - payload = _external_api_dict(api_id) - api = ExternalKnowledgeApis( - name=payload["name"], - description=payload["description"], - tenant_id=payload["tenant_id"], - settings=json.dumps(payload["settings"]), - created_by=payload["created_by"], - updated_by=None, +def test_delete_and_usage_response_contracts(external): + assert unwrap(controller.ExternalApiTemplateApi.delete)(controller.ExternalApiTemplateApi(), CONTEXT, ID) == ( + "", + 204, ) - api.id = api_id - api.created_at = datetime.fromisoformat(payload["created_at"]) - bound_dataset = Dataset( - id=payload["dataset_bindings"][0]["id"], - tenant_id=payload["tenant_id"], - name=payload["dataset_bindings"][0]["name"], - created_by=payload["created_by"], + external.template_usage.return_value = (True, 3) + assert unwrap(controller.ExternalApiUseCheckApi.get)(controller.ExternalApiUseCheckApi(), CONTEXT, ID) == ( + {"is_using": True, "count": 3}, + 200, ) - binding = ExternalKnowledgeBindings( - tenant_id=payload["tenant_id"], - external_knowledge_api_id=api_id, - dataset_id=bound_dataset.id, - external_knowledge_id=f"knowledge-{api_id}", - created_by=payload["created_by"], - ) - session.add_all([api, bound_dataset, binding]) - session.flush() - return api -def _expected_dataset_detail_payload() -> dict[str, Any]: - return { - "id": "dataset-1", - "name": "Support knowledge", - "description": "External support articles", - "provider": "external", - "permission": "only_me", - "data_source_type": "external", - "indexing_technique": "economy", - "app_count": 2, - "document_count": 7, - "word_count": 2048, - "created_by": "user-1", - "author_name": "Test User", - "created_at": 1710000000, - "updated_by": "user-2", - "updated_at": 1710003600, - "embedding_model": None, - "embedding_model_provider": None, - "embedding_available": False, - "retrieval_model_dict": { - "search_method": "semantic_search", - "reranking_enable": False, - "reranking_mode": None, - "reranking_model": {"reranking_provider_name": None, "reranking_model_name": None}, - "weights": None, - "top_k": 4, - "score_threshold_enabled": True, - "score_threshold": 0.5, - }, - "summary_index_setting": { - "enable": True, - "model_name": "summary-model", - "model_provider_name": "provider-a", - "summary_prompt": "Summarize this.", - }, - "tags": [{"id": "tag-1", "name": "Support", "type": "knowledge"}], - "doc_form": "text_model", - "external_knowledge_info": { - "external_knowledge_id": "knowledge-1", - "external_knowledge_api_id": "api-1", - "external_knowledge_api_name": "External API api-1", - "external_knowledge_api_endpoint": "https://external.example.com/api-1", - }, - "external_retrieval_model": { - "top_k": 4, - "score_threshold": 0.5, - "score_threshold_enabled": True, - }, - "doc_metadata": [{"id": "metadata-1", "name": "source", "type": "string"}], - "built_in_field_enabled": True, - "pipeline_id": None, - "runtime_mode": "external", - "chunk_structure": "general", - "icon_info": { - "icon_type": "emoji", - "icon": "book", - "icon_background": "#FFF4ED", - "icon_url": None, - }, - "is_published": True, - "total_documents": 7, - "total_available_documents": 6, - "enable_api": True, - "is_multimodal": False, - "maintainer": None, - "permission_keys": [], +def test_external_dataset_duplicate_name_maps_to_409(external): + payload = controller.ExternalDatasetCreatePayload( + name="Dataset", external_knowledge_api_id="api", external_knowledge_id="knowledge" + ) + external.create_dataset.side_effect = DuplicateName() + with pytest.raises(DatasetNameDuplicateError): + unwrap(controller.ExternalDatasetCreateApi.post)(controller.ExternalDatasetCreateApi(), payload, CONTEXT) + external.create_dataset.assert_called_once_with(CONTEXT, payload=payload) + + +@pytest.mark.parametrize( + ("error", "http_error"), + [ + (DatasetNotFoundError(), NotFound), + (DatasetAccessDeniedError(), Forbidden), + (ExternalHitTestingError("failed"), InternalServerError), + ], +) +def test_retrieval_maps_application_errors(external, error, http_error): + external.hit_testing.side_effect = error + with pytest.raises(http_error): + unwrap(controller.ExternalKnowledgeHitTestingApi.post)( + controller.ExternalKnowledgeHitTestingApi(), + controller.ExternalHitTestingPayload(query="query"), + CONTEXT, + ID, + ) + external.hit_testing.assert_called_once_with( + CONTEXT, dataset_id=str(ID), query="query", retrieval_model=None, metadata_filters=None + ) + + +def test_retrieval_response_is_serialized(external): + external.hit_testing.return_value = { + "query": {"content": "query"}, + "records": [{"content": "answer", "title": "title", "score": 0.8, "metadata": {"source": "doc"}}], } - - -def _dataset() -> Dataset: - payload = _expected_dataset_detail_payload() - return Dataset( - id=payload["id"], - tenant_id="tenant-1", - name=payload["name"], - description=payload["description"], - provider=payload["provider"], - permission=payload["permission"], - data_source_type=payload["data_source_type"], - indexing_technique=payload["indexing_technique"], - created_by=payload["created_by"], + result = unwrap(controller.ExternalKnowledgeHitTestingApi.post)( + controller.ExternalKnowledgeHitTestingApi(), controller.ExternalHitTestingPayload(query="query"), CONTEXT, ID ) - - -def _dataset_detail_response() -> DatasetDetailResponse: - return DatasetDetailResponse.model_validate(_expected_dataset_detail_payload()) - - -class _UsesSQLiteSession: - session: Session - - @pytest.fixture(autouse=True) - def _inject_sqlite_session(self, sqlite_session: Session) -> None: - self.session = sqlite_session - - -class TestExternalApiTemplateListApi(_UsesSQLiteSession): - def test_get_success(self, app: Flask): - api = ExternalApiTemplateListApi() - method = inspect.unwrap(api.get) - - session = self.session - api_item = _external_api_object(session, "api-1") - - with ( - app.test_request_context("/?page=2&limit=1&keyword=vector"), - patch.object( - ExternalDatasetService, - "get_external_knowledge_apis", - return_value=([api_item], 3), - ) as get_external_knowledge_apis, - ): - resp, status = method( - api, ExternalApiTemplateListQuery(page=2, limit=1, keyword="vector"), session, "tenant-1" - ) - - assert status == 200 - assert resp == { - "data": [_external_api_dict("api-1")], - "has_more": True, - "limit": 1, - "total": 3, - "page": 2, - } - get_external_knowledge_apis.assert_called_once_with(2, 1, "tenant-1", "vector", session=ANY) - - def test_post_success_uses_validated_payload_and_returns_template(self, app: Flask, current_user: Account): - api = ExternalApiTemplateListApi() - method = inspect.unwrap(api.post) - - payload = { - "name": "Vendor Search", - "settings": { - "endpoint": "https://external.example.com/search", - "api_key": "secret", - "headers": {"X-Source": "unit-test"}, - "timeout": 30, - }, - } - session = self.session - created = _external_api_object(session, "api-created") - - with ( - app.test_request_context("/", json=payload), - patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload), - patch.object(ExternalDatasetService, "validate_api_list") as validate_api_list, - patch.object( - ExternalDatasetService, - "create_external_knowledge_api", - return_value=created, - ) as create_external_knowledge_api, - ): - resp, status = method( - api, ExternalKnowledgeApiPayload.model_validate(payload), session, "tenant-1", current_user - ) - - assert status == 201 - assert resp == _external_api_dict("api-created") - validate_api_list.assert_called_once_with(payload["settings"]) - create_external_knowledge_api.assert_called_once_with( - tenant_id="tenant-1", - user_id="user-1", - args=payload, - session=session, - ) - - def test_post_forbidden(self, app: Flask, current_user: Account): - current_user.role = TenantAccountRole.NORMAL - api = ExternalApiTemplateListApi() - method = inspect.unwrap(api.post) - - payload = {"name": "x", "settings": {"k": "v"}} - - with ( - app.test_request_context("/"), - patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload), - patch.object(ExternalDatasetService, "validate_api_list"), - ): - with pytest.raises(Forbidden): - method(api, ExternalKnowledgeApiPayload.model_validate(payload), self.session, "tenant-1", current_user) - - def test_post_duplicate_name(self, app: Flask, current_user: Account): - api = ExternalApiTemplateListApi() - method = inspect.unwrap(api.post) - - payload = {"name": "x", "settings": {"k": "v"}} - - with ( - app.test_request_context("/"), - patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload), - patch.object(ExternalDatasetService, "validate_api_list"), - patch.object( - ExternalDatasetService, - "create_external_knowledge_api", - side_effect=services.errors.dataset.DatasetNameDuplicateError(), - ), - ): - with pytest.raises(DatasetNameDuplicateError): - method(api, ExternalKnowledgeApiPayload.model_validate(payload), self.session, "tenant-1", current_user) - - -class TestExternalApiTemplateApi(_UsesSQLiteSession): - def test_get_success_returns_template_contract(self, app: Flask): - api = ExternalApiTemplateApi() - method = inspect.unwrap(api.get) - session = self.session - template = _external_api_object(session, "api-detail") - - with ( - app.test_request_context("/"), - patch.object( - ExternalDatasetService, - "get_external_knowledge_api", - return_value=template, - ) as get_external_knowledge_api, - ): - resp, status = method(api, session, "tenant-1", "api-detail") - - assert status == 200 - assert resp == _external_api_dict("api-detail") - get_external_knowledge_api.assert_called_once_with( - external_knowledge_api_id="api-detail", tenant_id="tenant-1", session=session - ) - - def test_get_not_found(self, app: Flask): - api = ExternalApiTemplateApi() - method = inspect.unwrap(api.get) - - with ( - app.test_request_context("/"), - patch.object( - ExternalDatasetService, - "get_external_knowledge_api", - return_value=None, - ), - ): - with pytest.raises(NotFound): - method(api, self.session, "tenant-1", "api-id") - - def test_patch_success_uses_validated_payload_and_returns_template(self, app: Flask, current_user: Account): - api = ExternalApiTemplateApi() - method = inspect.unwrap(api.patch) - - payload = { - "name": "Updated API", - "settings": { - "endpoint": "https://external.example.com/updated", - "api_key": "new-secret", - "headers": {"X-Version": "2"}, - }, - } - session = self.session - updated = _external_api_object(session, "api-updated") - - with ( - app.test_request_context("/", json=payload), - patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload), - patch.object(ExternalDatasetService, "validate_api_list") as validate_api_list, - patch.object( - ExternalDatasetService, - "update_external_knowledge_api", - return_value=updated, - ) as update_external_knowledge_api, - ): - 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") - validate_api_list.assert_called_once_with(payload["settings"]) - update_external_knowledge_api.assert_called_once_with( - tenant_id="tenant-1", - user_id="user-1", - external_knowledge_api_id="api-updated", - args=payload, - session=session, - ) - - def test_delete_forbidden(self, app: Flask, current_user: Account): - current_user.role = TenantAccountRole.NORMAL - - api = ExternalApiTemplateApi() - method = inspect.unwrap(api.delete) - - with app.test_request_context("/"): - with pytest.raises(Forbidden): - method(api, self.session, "tenant-1", current_user, "api-id") - - -class TestExternalApiUseCheckApi(_UsesSQLiteSession): - def test_get_scopes_usage_check_to_current_tenant(self, app: Flask): - api = ExternalApiUseCheckApi() - method = inspect.unwrap(api.get) - - session = self.session - - with ( - app.test_request_context("/"), - patch.object( - ExternalDatasetService, - "external_knowledge_api_use_check", - return_value=(True, 2), - ) as mock_use_check, - ): - response, status = method(api, session, "tenant-1", "api-id") - - assert status == 200 - assert response == {"is_using": True, "count": 2} - mock_use_check.assert_called_once_with("api-id", "tenant-1", session=ANY) - - -class TestExternalDatasetCreateApi(_UsesSQLiteSession): - def test_create_success(self, app: Flask, current_user: Account): - api = ExternalDatasetCreateApi() - method = inspect.unwrap(api.post) - - payload = { - "external_knowledge_api_id": "api-1", - "external_knowledge_id": "knowledge-1", - "name": "Support knowledge", - "description": "External support articles", - "external_retrieval_model": { - "top_k": 4, - "score_threshold": 0.5, - "score_threshold_enabled": True, - }, - } - - dataset = _dataset() - - with ( - app.test_request_context("/", json=payload), - patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload), - patch.object( - ExternalDatasetService, - "create_external_dataset", - return_value=dataset, - ) as create_external_dataset, - patch( - "controllers.console.datasets.external.dataset_detail_response_source", - return_value=_dataset_detail_response(), - ) as dataset_response_source, - ): - resp, status = method( - api, - ExternalDatasetCreatePayload.model_validate(payload), - self.session, - "tenant-1", - current_user, - ) - - assert status == 201 - assert resp == _expected_dataset_detail_payload() - create_external_dataset.assert_called_once_with( - tenant_id="tenant-1", - user_id="user-1", - args=ExternalDatasetCreatePayload.model_validate(payload), - session=self.session, - ) - dataset_response_source.assert_called_once_with(dataset, session=self.session) - - def test_create_forbidden(self, app: Flask, current_user: Account): - current_user.role = TenantAccountRole.NORMAL - api = ExternalDatasetCreateApi() - method = inspect.unwrap(api.post) - - payload = { - "external_knowledge_api_id": "api", - "external_knowledge_id": "kid", - "name": "dataset", - } - - with ( - app.test_request_context("/"), - patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload), - ): - with pytest.raises(Forbidden): - method( - api, ExternalDatasetCreatePayload.model_validate(payload), self.session, "tenant-1", current_user - ) - - -class TestExternalKnowledgeHitTestingApi(_UsesSQLiteSession): - def test_hit_testing_dataset_not_found(self, app: Flask, current_user: Account): - api = ExternalKnowledgeHitTestingApi() - method = inspect.unwrap(api.post) - - with ( - app.test_request_context("/"), - patch.object( - DatasetService, - "get_dataset", - return_value=None, - ), - ): - with pytest.raises(NotFound): - method(api, ExternalHitTestingPayload(query="test"), self.session, current_user, "dataset-id") - - def test_hit_testing_success(self, app: Flask, current_user: Account): - api = ExternalKnowledgeHitTestingApi() - method = inspect.unwrap(api.post) - - payload = { - "query": "hello", - "external_retrieval_model": { - "top_k": 3, - "score_threshold": 0.25, - "score_threshold_enabled": True, - }, - "metadata_filtering_conditions": { - "logical_operator": "and", - "conditions": [{"name": "source", "comparison_operator": "contains", "value": "external"}], - }, - } - - dataset = _dataset() - retrieve_response = { - "query": {"content": "hello"}, - "records": [ - { - "content": "answer", - "title": "doc", - "score": 0.9, - "metadata": {"source": "external", "page": 2}, - } - ], - } - session = self.session - - with ( - app.test_request_context("/", json=payload), - 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") as check_dataset_permission, - patch.object(HitTestingService, "hit_testing_args_check") as hit_testing_args_check, - patch.object( - HitTestingService, - "external_retrieve", - return_value=retrieve_response, - ) as external_retrieve, - patch("controllers.console.datasets.external.dump_response", side_effect=lambda _model, value: value), - ): - 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) - hit_testing_args_check.assert_called_once_with(payload) - external_retrieve.assert_called_once_with( - session=session, - dataset=dataset, - query="hello", - account=current_user, - external_retrieval_model=payload["external_retrieval_model"], - metadata_filtering_conditions=payload["metadata_filtering_conditions"], - ) - - -class TestExternalApiTemplateListApiAdvanced(_UsesSQLiteSession): - def test_post_duplicate_name_error(self, app: Flask, current_user: Account): - api = ExternalApiTemplateListApi() - method = inspect.unwrap(api.post) - - payload = {"name": "duplicate_api", "settings": {"key": "value"}} - - with ( - app.test_request_context("/", json=payload), - patch.object(type(console_ns), "payload", payload), - patch("controllers.console.datasets.external.ExternalDatasetService.validate_api_list"), - patch( - "controllers.console.datasets.external.ExternalDatasetService.create_external_knowledge_api", - side_effect=services.errors.dataset.DatasetNameDuplicateError("Duplicate"), - ), - ): - with pytest.raises(DatasetNameDuplicateError): - method(api, ExternalKnowledgeApiPayload.model_validate(payload), self.session, "tenant-1", current_user) - - def test_get_with_pagination(self, app: Flask): - api = ExternalApiTemplateListApi() - method = inspect.unwrap(api.get) - - templates = [_external_api_object(self.session, f"api-{i}") for i in range(3)] - - with ( - app.test_request_context("/?page=2&limit=3"), - patch( - "controllers.console.datasets.external.ExternalDatasetService.get_external_knowledge_apis", - return_value=(templates, 25), - ) as get_external_knowledge_apis, - ): - resp, status = method(api, ExternalApiTemplateListQuery(page=2, limit=3), self.session, "tenant-1") - - assert status == 200 - assert resp == { - "data": [_external_api_dict(f"api-{i}") for i in range(3)], - "has_more": True, - "limit": 3, - "total": 25, - "page": 2, - } - get_external_knowledge_apis.assert_called_once_with(2, 3, "tenant-1", None, session=ANY) - - -class TestExternalDatasetCreateApiAdvanced(_UsesSQLiteSession): - def test_create_forbidden(self, app: Flask, current_user: Account): - """Test creating external dataset without permission""" - api = ExternalDatasetCreateApi() - method = inspect.unwrap(api.post) - - current_user.role = TenantAccountRole.NORMAL - - payload = { - "external_knowledge_api_id": "api-1", - "external_knowledge_id": "ek-1", - "name": "new_dataset", - "description": "A dataset", - } - - with app.test_request_context("/", json=payload), patch.object(type(console_ns), "payload", payload): - with pytest.raises(Forbidden): - method( - api, ExternalDatasetCreatePayload.model_validate(payload), self.session, "tenant-1", current_user - ) - - -class TestExternalKnowledgeHitTestingApiAdvanced(_UsesSQLiteSession): - def test_hit_testing_dataset_not_found(self, app: Flask, current_user: Account): - """Test hit testing on non-existent dataset""" - api = ExternalKnowledgeHitTestingApi() - method = inspect.unwrap(api.post) - - payload = { - "query": "test query", - "external_retrieval_model": None, - } - - with ( - app.test_request_context("/", json=payload), - patch.object(type(console_ns), "payload", payload), - patch( - "controllers.console.datasets.external.DatasetService.get_dataset", - return_value=None, - ), - ): - with pytest.raises(NotFound): - method(api, ExternalHitTestingPayload.model_validate(payload), self.session, current_user, "ds-1") - - def test_hit_testing_with_custom_retrieval_model(self, app: Flask, current_user: Account): - api = ExternalKnowledgeHitTestingApi() - method = inspect.unwrap(api.post) - - dataset = _dataset() - payload = { - "query": "test query", - "external_retrieval_model": {"type": "bm25"}, - "metadata_filtering_conditions": {"status": "active"}, - } - session = self.session - - with ( - app.test_request_context("/", json=payload), - patch.object(type(console_ns), "payload", payload), - patch( - "controllers.console.datasets.external.DatasetService.get_dataset", - return_value=dataset, - ), - patch("controllers.console.datasets.external.DatasetService.check_dataset_permission") as check_permission, - patch("controllers.console.datasets.external.HitTestingService.hit_testing_args_check") as args_check, - patch( - "controllers.console.datasets.external.HitTestingService.external_retrieve", - return_value={ - "query": {"content": "test query"}, - "records": [ - { - "content": None, - "title": "metadata-only", - "score": None, - "metadata": {"status": "active"}, - } - ], - }, - ) as external_retrieve, - ): - resp = method(api, ExternalHitTestingPayload.model_validate(payload), session, current_user, "ds-1") - - assert resp == { - "query": {"content": "test query"}, - "records": [ - { - "content": None, - "title": "metadata-only", - "score": None, - "metadata": {"status": "active"}, - } - ], - } - check_permission.assert_called_once_with(dataset, current_user, session) - args_check.assert_called_once_with(payload) - external_retrieve.assert_called_once_with( - session=session, - dataset=dataset, - query="test query", - account=current_user, - external_retrieval_model={"type": "bm25"}, - metadata_filtering_conditions={"status": "active"}, - ) + assert result == external.hit_testing.return_value diff --git a/api/tests/unit_tests/controllers/console/datasets/test_hit_testing_base.py b/api/tests/unit_tests/controllers/console/datasets/test_hit_testing_base.py index 39ddd2ec787..bc7d18b8fc0 100644 --- a/api/tests/unit_tests/controllers/console/datasets/test_hit_testing_base.py +++ b/api/tests/unit_tests/controllers/console/datasets/test_hit_testing_base.py @@ -23,8 +23,8 @@ from core.errors.error import ( from graphon.model_runtime.errors.invoke import InvokeError from models.account import Account, Tenant, TenantAccountRole from models.dataset import Dataset -from services.dataset_service import DatasetService from services.hit_testing_service import HitTestingService +from services.knowledge.dataset_service import DatasetService @pytest.fixture 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 dc528761f74..63c28b60a33 100644 --- a/api/tests/unit_tests/controllers/console/datasets/test_metadata.py +++ b/api/tests/unit_tests/controllers/console/datasets/test_metadata.py @@ -1,340 +1,122 @@ -import uuid -from inspect import unwrap -from unittest.mock import PropertyMock, patch +"""Console metadata routes use admission, explicit context, and response models.""" + +from types import SimpleNamespace +from unittest.mock import create_autospec +from uuid import uuid4 import pytest from flask import Flask -from pytest_mock import MockerFixture -from sqlalchemy.orm import Session -from werkzeug.exceptions import Forbidden, NotFound +from flask_restx import Api -from controllers.common.controller_schemas import MetadataUpdatePayload -from controllers.common.rbac import DatasetId +from controllers.common.rbac import DatasetId, RBACPermission from controllers.console import console_ns -from controllers.console.datasets.metadata import ( - DatasetMetadataApi, - DatasetMetadataBuiltInFieldActionApi, - DatasetMetadataBuiltInFieldApi, - DatasetMetadataCreateApi, - DocumentMetadataEditApi, -) -from controllers.console.wraps import RBACPermission -from models.account import Account -from models.dataset import Dataset -from services.dataset_service import DatasetService -from services.entities.knowledge_entities.knowledge_entities import MetadataArgs, MetadataOperationData -from services.errors.account import NoPermissionError +from controllers.console.datasets.metadata import DatasetMetadataCreateApi +from enums import DeploymentEdition +from libs.login import AccountWithTenant +from models.account import Account, TenantAccountRole from services.errors.metadata import MetadataResourceNotFoundError -from services.metadata_service import MetadataService -from tests.unit_tests.config_override import config_overrides_context +from services.knowledge.dataset_access import DatasetAccessDeniedError, DatasetNotFoundError +from services.knowledge.metadata.application import MetadataService +from services.knowledge.resource_scope import DatasetRef from tests.unit_tests.controllers.rbac_introspection import rbac_checks @pytest.fixture -def app(): - app = Flask("test_dataset_metadata") - - app.config["TESTING"] = True - return app +def routes(monkeypatch: pytest.MonkeyPatch, config_overrides): + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, LOGIN_DISABLED=True, RBAC_ENABLED=False) + app = Flask(__name__) + app.config.update(TESTING=False, PROPAGATE_EXCEPTIONS=False) + Api(app).add_namespace(console_ns, path="/console/api") + account = Account(name="Owner", email="metadata@example.com") + account.id = "account-1" + account.role = TenantAccountRole.OWNER + identity = AccountWithTenant(account=account, tenant_id="tenant-1") + monkeypatch.setattr("controllers.console.flask_admission.current_account_with_tenant", lambda: identity) + monkeypatch.setattr("controllers.console.wraps.current_account_with_tenant", lambda: identity) + service = create_autospec(MetadataService, instance=True, spec_set=True) + service.require_dataset.side_effect = lambda context, dataset_id: DatasetRef( + context.active_workspace_id, dataset_id + ) + app.extensions["application_services"] = SimpleNamespace(knowledge=SimpleNamespace(metadata=service)) + return app.test_client(), service -@pytest.fixture -def current_user() -> Account: - user = Account(name="Test User", email="test@example.com") - user.id = "user-1" - return user +def test_create_serializes_result_and_passes_actor(routes): + client, service = routes + dataset_id = str(uuid4()) + service.create_metadata.return_value = {"id": str(uuid4()), "name": "author", "type": "string"} + response = client.post(f"/console/api/datasets/{dataset_id}/metadata", json={"type": "string", "name": "author"}) + assert response.status_code == 201 + assert response.get_json()["name"] == "author" + assert service.create_metadata.call_args.args[0] == DatasetRef("tenant-1", dataset_id) + assert service.create_metadata.call_args.kwargs == {"actor_id": "account-1"} -@pytest.fixture -def dataset() -> Dataset: - return Dataset(id="dataset-1", tenant_id="tenant-1", name="Test Dataset", created_by="user-1") +@pytest.mark.parametrize("payload", [{"name": None, "type": "string"}, {"name": "author", "type": None}]) +def test_create_rejects_invalid_payload_before_application(routes, payload): + client, service = routes + response = client.post(f"/console/api/datasets/{uuid4()}/metadata", json=payload) + assert response.status_code == 422 + service.create_metadata.assert_not_called() -@pytest.fixture -def dataset_id(): - return uuid.uuid4() +@pytest.mark.parametrize(("error", "status"), [(DatasetNotFoundError(), 404), (DatasetAccessDeniedError(), 403)]) +def test_get_maps_access_errors(routes, error, status): + client, service = routes + service.require_dataset.side_effect = error + assert client.get(f"/console/api/datasets/{uuid4()}/metadata").status_code == status + service.get_dataset_metadatas.assert_not_called() -@pytest.fixture -def metadata_id(): - return uuid.uuid4() +def test_get_serializes_field_counts(routes): + client, service = routes + service.get_dataset_metadatas.return_value = {"doc_metadata": [], "built_in_field_enabled": False} + response = client.get(f"/console/api/datasets/{uuid4()}/metadata") + assert response.status_code == 200 + assert response.get_json() == {"doc_metadata": [], "built_in_field_enabled": False} -@pytest.fixture(autouse=True) -def bypass_decorators(mocker: MockerFixture): - """Bypass setup/login/license decorators.""" - mocker.patch("controllers.console.datasets.metadata.setup_required", lambda f: f) - mocker.patch("controllers.console.datasets.metadata.login_required", lambda f: f) - mocker.patch("controllers.console.datasets.metadata.account_initialization_required", lambda f: f) - mocker.patch("controllers.console.datasets.metadata.enterprise_license_required", lambda f: f) +def test_rename_and_delete_use_explicit_metadata_id(routes): + client, service = routes + dataset_id, metadata_id = str(uuid4()), str(uuid4()) + path = f"/console/api/datasets/{dataset_id}/metadata/{metadata_id}" + service.update_metadata_name.return_value = {"id": metadata_id, "name": "writer", "type": "string"} + assert client.patch(path, json={"name": "writer"}).status_code == 200 + service.update_metadata_name.assert_called_once_with( + DatasetRef("tenant-1", dataset_id), metadata_id, "writer", actor_id="account-1" + ) + response = client.delete(path) + assert response.status_code == 204 + assert response.data == b"" + service.delete_metadata.assert_called_once_with(DatasetRef("tenant-1", dataset_id), metadata_id) -class TestDatasetMetadataCreateApi: - def test_get_requires_dataset_readonly_permission(self): - [check] = rbac_checks(DatasetMetadataCreateApi.get) - - assert check.scene is RBACPermission.DATASET_READONLY - assert isinstance(check.locator, DatasetId) - - def test_create_metadata_success(self, app: Flask, current_user, dataset, dataset_id, sqlite_session: Session): - api = DatasetMetadataCreateApi() - method = unwrap(api.post) - payload = {"name": "author"} - with ( - app.test_request_context("/"), - patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload), - patch.object( - MetadataArgs, - "model_validate", - return_value=MetadataArgs(type="string", name="author"), - ), - patch.object(DatasetService, "get_dataset", return_value=dataset), - patch.object(DatasetService, "check_dataset_permission"), - patch.object( - MetadataService, "create_metadata", return_value={"id": "m1", "type": "string", "name": "author"} - ), - ): - result, status = method( - api, MetadataArgs(type="string", name="author"), sqlite_session, "tenant-1", current_user, dataset_id - ) - assert status == 201 - assert result["type"] == "string" - assert result["name"] == "author" - - def test_create_metadata_dataset_not_found(self, app: Flask, current_user, dataset_id, sqlite_session: Session): - api = DatasetMetadataCreateApi() - method = unwrap(api.post) - valid_payload = {"type": "string", "name": "author"} - with ( - app.test_request_context("/"), - patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=valid_payload), - patch.object( - MetadataArgs, - "model_validate", - return_value=MetadataArgs(type="string", name="author"), - ), - patch.object(DatasetService, "get_dataset", return_value=None), - ): - with pytest.raises(NotFound, match="Dataset not found"): - method( - api, - MetadataArgs(type="string", name="author"), - sqlite_session, - "tenant-1", - current_user, - dataset_id, - ) +@pytest.mark.parametrize("action", ["enable", "disable"]) +def test_builtin_toggle_is_dispatched(routes, action): + client, service = routes + dataset_id = str(uuid4()) + response = client.post(f"/console/api/datasets/{dataset_id}/metadata/built-in/{action}") + assert response.status_code == 204 + operation = service.enable_built_in_field if action == "enable" else service.disable_built_in_field + operation.assert_called_once_with(DatasetRef("tenant-1", dataset_id)) -class TestDatasetMetadataGetApi: - def test_get_metadata_success(self, app: Flask, current_user, dataset, dataset_id, sqlite_session: Session): - api = DatasetMetadataCreateApi() - method = unwrap(api.get) - with ( - app.test_request_context("/"), - patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset) as get_dataset, - patch.object(DatasetService, "check_dataset_permission") as check_permission, - patch.object( - MetadataService, - "get_dataset_metadatas", - return_value={ - "doc_metadata": [{"id": "m1", "name": "author", "type": "string", "count": 0}], - "built_in_field_enabled": False, - }, - ), - ): - result, status = method(api, sqlite_session, "tenant-1", current_user, dataset_id) - assert status == 200 - assert result["doc_metadata"] == [{"id": "m1", "name": "author", "type": "string", "count": 0}] - assert result["built_in_field_enabled"] is False - get_dataset.assert_called_once_with(str(dataset_id), "tenant-1", session=sqlite_session) - check_permission.assert_called_once_with(dataset, current_user, sqlite_session) - - def test_get_metadata_rejects_foreign_tenant_before_read( - self, app: Flask, current_user, dataset_id, sqlite_session: Session - ): - api = DatasetMetadataCreateApi() - method = unwrap(api.get) - with ( - app.test_request_context("/"), - patch.object(DatasetService, "get_dataset_for_tenant", return_value=None) as get_dataset, - patch.object(DatasetService, "check_dataset_permission") as check_permission, - patch.object(MetadataService, "get_dataset_metadatas") as get_metadata, - ): - with pytest.raises(NotFound): - method(api, sqlite_session, "tenant-1", current_user, dataset_id) - - get_dataset.assert_called_once_with(str(dataset_id), "tenant-1", session=sqlite_session) - check_permission.assert_not_called() - get_metadata.assert_not_called() - - def test_get_metadata_delegates_permission_check_in_rbac_mode( - self, app: Flask, current_user, dataset, dataset_id, sqlite_session: Session - ): - api = DatasetMetadataCreateApi() - method = unwrap(api.get) - with ( - app.test_request_context("/"), - config_overrides_context(RBAC_ENABLED=True), - patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset), - patch.object(DatasetService, "check_dataset_permission") as check_permission, - patch.object( - MetadataService, - "get_dataset_metadatas", - return_value={"doc_metadata": [], "built_in_field_enabled": False}, - ), - ): - _, status = method(api, sqlite_session, "tenant-1", current_user, dataset_id) - - assert status == 200 - check_permission.assert_called_once_with(dataset, current_user, sqlite_session) - - def test_get_metadata_rejects_inaccessible_dataset( - self, app: Flask, current_user, dataset, dataset_id, sqlite_session: Session - ): - api = DatasetMetadataCreateApi() - method = unwrap(api.get) - with ( - app.test_request_context("/"), - patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset), - patch.object(DatasetService, "check_dataset_permission", side_effect=NoPermissionError), - patch.object(MetadataService, "get_dataset_metadatas") as get_metadata, - ): - with pytest.raises(Forbidden): - method(api, sqlite_session, "tenant-1", current_user, dataset_id) - - get_metadata.assert_not_called() +def test_builtin_listing_uses_response_model(routes): + client, service = routes + service.get_built_in_fields.return_value = [{"name": "uploader", "type": "string"}] + response = client.get("/console/api/datasets/metadata/built-in") + assert response.status_code == 200 + assert response.get_json() == {"fields": [{"name": "uploader", "type": "string"}]} -class TestDatasetMetadataApi: - def test_update_metadata_success( - self, app: Flask, current_user, dataset, dataset_id, metadata_id, sqlite_session: Session - ): - api = DatasetMetadataApi() - method = unwrap(api.patch) - payload = {"name": "updated-name"} - with ( - app.test_request_context("/"), - patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload), - patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset) as get_dataset, - patch.object(DatasetService, "check_dataset_permission"), - patch.object( - MetadataService, - "update_metadata_name", - return_value={"id": "m1", "type": "string", "name": "updated-name"}, - ) as update_metadata, - ): - result, status = method( - api, - MetadataUpdatePayload(name="updated-name"), - sqlite_session, - "tenant-1", - current_user, - dataset_id, - metadata_id, - ) - assert status == 200 - assert result["type"] == "string" - assert result["name"] == "updated-name" - get_dataset.assert_called_once_with(str(dataset_id), "tenant-1", session=sqlite_session) - update_metadata.assert_called_once_with( - dataset, str(metadata_id), "updated-name", current_user, session=sqlite_session - ) - - def test_delete_metadata_success( - self, app: Flask, current_user, dataset, dataset_id, metadata_id, sqlite_session: Session - ): - api = DatasetMetadataApi() - method = unwrap(api.delete) - with ( - app.test_request_context("/"), - patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset) as get_dataset, - patch.object(DatasetService, "check_dataset_permission"), - patch.object(MetadataService, "delete_metadata") as delete_metadata, - ): - result, status = method(api, sqlite_session, "tenant-1", current_user, dataset_id, metadata_id) - assert status == 204 - assert result == "" - get_dataset.assert_called_once_with(str(dataset_id), "tenant-1", session=sqlite_session) - delete_metadata.assert_called_once_with(dataset, str(metadata_id), sqlite_session) +def test_document_metadata_missing_resource_is_not_found(routes): + client, service = routes + service.update_documents_metadata.side_effect = MetadataResourceNotFoundError("Document not found.") + response = client.post(f"/console/api/datasets/{uuid4()}/documents/metadata", json={"operation_data": []}) + assert response.status_code == 404 -class TestDatasetMetadataBuiltInFieldApi: - def test_get_built_in_fields(self, app: Flask): - api = DatasetMetadataBuiltInFieldApi() - method = unwrap(api.get) - with ( - app.test_request_context("/"), - patch.object( - MetadataService, - "get_built_in_fields", - return_value=[{"name": "document_name", "type": "string"}, {"name": "source", "type": "string"}], - ), - ): - result, status = method(api) - assert status == 200 - assert result["fields"] == [{"name": "document_name", "type": "string"}, {"name": "source", "type": "string"}] - - -class TestDatasetMetadataBuiltInFieldActionApi: - def test_enable_built_in_field(self, app: Flask, current_user, dataset, dataset_id, sqlite_session: Session): - api = DatasetMetadataBuiltInFieldActionApi() - method = unwrap(api.post) - with ( - app.test_request_context("/"), - patch.object(DatasetService, "get_dataset", return_value=dataset), - patch.object(DatasetService, "check_dataset_permission"), - patch.object(MetadataService, "enable_built_in_field"), - ): - result, status = method(api, sqlite_session, current_user, dataset_id, "enable") - assert status == 204 - assert result == "" - - -class TestDocumentMetadataEditApi: - def test_update_document_metadata_success( - self, app: Flask, current_user, dataset, dataset_id, sqlite_session: Session - ): - api = DocumentMetadataEditApi() - method = unwrap(api.post) - payload = {"operation": "add", "metadata": {}} - with ( - app.test_request_context("/"), - patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload), - patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset), - patch.object(DatasetService, "check_dataset_permission"), - patch.object(MetadataService, "update_documents_metadata"), - ): - result, status = method( - api, - MetadataOperationData( - operation_data=[{"document_id": "00000000-0000-0000-0000-000000000001", "metadata_list": []}] - ), - sqlite_session, - dataset.tenant_id, - current_user, - dataset_id, - ) - assert status == 204 - assert result == "" - - def test_update_document_metadata_translates_missing_resource( - self, app: Flask, current_user, dataset, dataset_id, sqlite_session: Session - ): - api = DocumentMetadataEditApi() - method = unwrap(api.post) - request = MetadataOperationData(operation_data=[]) - with ( - app.test_request_context("/"), - patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset), - patch.object(DatasetService, "check_dataset_permission"), - patch.object( - MetadataService, - "update_documents_metadata", - side_effect=MetadataResourceNotFoundError("Metadata not found."), - ), - pytest.raises(NotFound) as exc_info, - ): - method(api, request, sqlite_session, dataset.tenant_id, current_user, dataset_id) - - assert exc_info.value.description == "Metadata not found." +def test_get_requires_dataset_readonly_permission(): + [check] = rbac_checks(DatasetMetadataCreateApi.get) + assert check.scene is RBACPermission.DATASET_READONLY + assert isinstance(check.locator, DatasetId) diff --git a/api/tests/unit_tests/controllers/console/datasets/test_website.py b/api/tests/unit_tests/controllers/console/datasets/test_website.py index b790059f0c7..900b0aa8642 100644 --- a/api/tests/unit_tests/controllers/console/datasets/test_website.py +++ b/api/tests/unit_tests/controllers/console/datasets/test_website.py @@ -11,7 +11,7 @@ from controllers.console.datasets.website import ( WebsiteCrawlStatusApi, WebsiteCrawlStatusQuery, ) -from services.website_service import ( +from services.data_source.website_service import ( WebsiteCrawlApiRequest, WebsiteCrawlStatusApiRequest, WebsiteService, diff --git a/api/tests/unit_tests/controllers/console/test_apikey.py b/api/tests/unit_tests/controllers/console/test_apikey.py index 882bd9bdaeb..a6bea948603 100644 --- a/api/tests/unit_tests/controllers/console/test_apikey.py +++ b/api/tests/unit_tests/controllers/console/test_apikey.py @@ -18,6 +18,7 @@ from controllers.console.agent import roster from controllers.console.datasets import datasets from core.rbac import RBACPermission from enums import DeploymentEdition +from extensions.application_services.datasets import build_dataset_dependencies from extensions.application_services.knowledge import build_dataset_api_key_service from extensions.ext_application_services import ApplicationServices from libs.login import AccountWithTenant @@ -25,6 +26,7 @@ from machinery.context import RequestContext from models.account import Account, AccountStatus, Tenant, TenantAccountJoin, TenantAccountRole from models.dataset import Dataset from models.model import ApiToken +from repositories.workspace_member_query_repository import WorkspaceMemberQueryRepository from services.app.api_key_service import AppApiKeyNotReadyError from services.auth.api_key_contracts import ( ApiKeyLimitExceededError, @@ -282,7 +284,14 @@ def persisted_keys_app( ] ) sqlite_session.commit() - service = build_dataset_api_key_service(database_client=sqlite_session_factory) + dataset_access = build_dataset_dependencies( + database_client=sqlite_session_factory, + workspace_roles=WorkspaceMemberQueryRepository(session_factory=sqlite_session_factory), + ).access + service = build_dataset_api_key_service( + database_client=sqlite_session_factory, + dataset_access=dataset_access, + ) services = Mock(dataset_api_keys=service) monkeypatch.setattr(apikey, "application_services", lambda: services) monkeypatch.setattr(datasets, "application_services", lambda: services) diff --git a/api/tests/unit_tests/controllers/console/test_document_detail_api_data_source_info.py b/api/tests/unit_tests/controllers/console/test_document_detail_api_data_source_info.py index fac5aa79ae7..83a840b0e67 100644 --- a/api/tests/unit_tests/controllers/console/test_document_detail_api_data_source_info.py +++ b/api/tests/unit_tests/controllers/console/test_document_detail_api_data_source_info.py @@ -10,6 +10,7 @@ from typing import Literal, NotRequired, TypedDict from unittest.mock import MagicMock from models.dataset import Document +from services.knowledge.dataset_read_service import get_document_source_detail class LocalFileInfo(TypedDict): @@ -117,7 +118,7 @@ class TestDocumentDetailDataSourceInfo: ) # data_source_detail_dict should return raw data for notion_import - detail_result = document.get_data_source_detail_dict(session=MagicMock()) + detail_result = get_document_source_detail(document, session=MagicMock()) assert detail_result == notion_data # Test website_crawl @@ -128,7 +129,7 @@ class TestDocumentDetailDataSourceInfo: ) # data_source_detail_dict should return raw data for website_crawl - detail_result = document.get_data_source_detail_dict(session=MagicMock()) + detail_result = get_document_source_detail(document, session=MagicMock()) assert detail_result == website_data def test_local_file_data_source_detail_dict_without_db(self): @@ -140,5 +141,5 @@ class TestDocumentDetailDataSourceInfo: ) # Should return empty dict for local_file type (handled in the model) - detail_result = document.get_data_source_detail_dict(session=MagicMock()) + detail_result = get_document_source_detail(document, session=MagicMock()) assert detail_result == {} diff --git a/api/tests/unit_tests/controllers/console/test_workspace_credential_mutation_permissions.py b/api/tests/unit_tests/controllers/console/test_workspace_credential_mutation_permissions.py index f89ccb3e765..7e4888e185b 100644 --- a/api/tests/unit_tests/controllers/console/test_workspace_credential_mutation_permissions.py +++ b/api/tests/unit_tests/controllers/console/test_workspace_credential_mutation_permissions.py @@ -6,7 +6,6 @@ import pytest from controllers.common.rbac import AgentId, RBACPermission, Workspace from controllers.console.agent.composer import AgentComposerApi from controllers.console.agent.roster import AgentAppApi -from controllers.console.datasets.data_source import DataSourceApi from controllers.console.datasets.rag_pipeline.datasource_auth import DatasourceAuth from controllers.console.workspace.model_providers import ModelProviderCredentialApi from controllers.console.workspace.models import ModelProviderModelCredentialApi @@ -22,8 +21,6 @@ from tests.unit_tests.controllers.rbac_introspection import rbac_checks ("method", "permission"), [ (ToolBuiltinProviderAddApi.post, RBACPermission.CREDENTIAL_CREATE), - (DataSourceApi.patch, RBACPermission.CREDENTIAL_MANAGE), - (DataSourceApi.get, RBACPermission.CREDENTIAL_MANAGE), ], ) def test_workspace_credential_mutations_require_management_permission( diff --git a/api/tests/unit_tests/controllers/console/workspace/test_models.py b/api/tests/unit_tests/controllers/console/workspace/test_models.py index d67f7122bef..2afd26c4644 100644 --- a/api/tests/unit_tests/controllers/console/workspace/test_models.py +++ b/api/tests/unit_tests/controllers/console/workspace/test_models.py @@ -137,7 +137,7 @@ class TestDefaultModelApi: with ( app.test_request_context("/"), - patch("services.model_provider_service.create_plugin_provider_manager", return_value=manager), + patch("services.model_provider.service.create_plugin_provider_manager", return_value=manager), patch.object(manager, "get_configurations", return_value=configurations), ): result = unwrap(api.post)(api, payload, "tenant1") @@ -164,7 +164,7 @@ class TestDefaultModelApi: with ( app.test_request_context("/"), - patch("services.model_provider_service.create_plugin_provider_manager", return_value=manager), + patch("services.model_provider.service.create_plugin_provider_manager", return_value=manager), ): result = unwrap(api.post)(api, ParserPostDefault(model_settings=[]), "tenant1") diff --git a/api/tests/unit_tests/controllers/console/workspace/test_tool_provider_apis.py b/api/tests/unit_tests/controllers/console/workspace/test_tool_provider_apis.py index 1f9f4d863b4..1d078671a53 100644 --- a/api/tests/unit_tests/controllers/console/workspace/test_tool_provider_apis.py +++ b/api/tests/unit_tests/controllers/console/workspace/test_tool_provider_apis.py @@ -2,8 +2,10 @@ from __future__ import annotations +from collections.abc import Iterator from inspect import unwrap -from unittest.mock import patch +from types import SimpleNamespace +from unittest.mock import MagicMock, create_autospec, patch import pytest from flask import Flask @@ -54,6 +56,17 @@ from controllers.console.workspace.tool_providers import ( ) from core.tools.entities.api_entities import ToolProviderApiEntity as CoreToolProviderApiEntity from models.account import Account, TenantAccountRole +from services.credentials.query import CredentialQuery + + +@pytest.fixture +def credential_query() -> Iterator[MagicMock]: + query = create_autospec(CredentialQuery, instance=True, spec_set=True) + with patch( + "controllers.console.workspace.tool_providers.application_services", + return_value=SimpleNamespace(credential_queries=query), + ): + yield query def empty_mapping() -> dict[str, object]: @@ -337,7 +350,7 @@ class TestBuiltinProviderApis: req = BuiltinToolUpdatePayload(**payload) assert method(api, req, "t", make_account(), "provider")["result"] == "success" - def test_get_credentials(self, app: Flask) -> None: + def test_get_credentials(self, app: Flask, credential_query: MagicMock) -> None: api = ToolBuiltinProviderGetCredentialsApi() method = unwrap(api.get) @@ -346,9 +359,10 @@ class TestBuiltinProviderApis: patch( "controllers.console.workspace.tool_providers.BuiltinToolManageService.get_builtin_tool_provider_credentials", return_value=[credential_payload()], - ), + ) as get_credentials, ): assert method(api, "t", make_account(id="user-1"), "provider")[0]["id"] == "credential-1" + assert get_credentials.call_args.kwargs["credential_query"] is credential_query def test_icon(self, app: Flask) -> None: api = ToolBuiltinProviderIconApi() @@ -391,7 +405,7 @@ class TestBuiltinProviderApis: req = BuiltinProviderDefaultCredentialPayload(id="c1") assert method(api, req, "t", "provider")["result"] == "success" - def test_get_credential_info(self, app: Flask) -> None: + def test_get_credential_info(self, app: Flask, credential_query: MagicMock) -> None: api = ToolBuiltinProviderGetCredentialInfoApi() method = unwrap(api.get) @@ -400,9 +414,10 @@ class TestBuiltinProviderApis: patch( "controllers.console.workspace.tool_providers.BuiltinToolManageService.get_builtin_tool_provider_credential_info", return_value=credential_info_payload(), - ), + ) as get_info, ): assert method(api, "t", make_account(), "provider")["credentials"][0]["id"] == "credential-1" + assert get_info.call_args.kwargs["credential_query"] is credential_query def test_get_oauth_client_schema(self, app: Flask) -> None: api = ToolBuiltinProviderGetOauthClientSchemaApi() diff --git a/api/tests/unit_tests/controllers/console/workspace/test_tool_providers.py b/api/tests/unit_tests/controllers/console/workspace/test_tool_providers.py index c205f931cc6..cb61284cd35 100644 --- a/api/tests/unit_tests/controllers/console/workspace/test_tool_providers.py +++ b/api/tests/unit_tests/controllers/console/workspace/test_tool_providers.py @@ -7,7 +7,7 @@ import importlib from collections.abc import Callable from contextlib import ExitStack, contextmanager from inspect import unwrap -from types import ModuleType +from types import ModuleType, SimpleNamespace from unittest.mock import MagicMock, patch from uuid import NAMESPACE_URL, uuid5 @@ -24,6 +24,8 @@ from models import Account, BuiltinToolProvider, Tenant, TenantAccountJoin from models.account import TenantAccountRole from models.credential_permission import CredentialPermission from models.enums import PermissionEnum +from repositories.credentials.query_repository import CredentialQueryRepository +from services.credentials.query import CredentialQuery if not hasattr(builtins, "MethodView"): builtins.MethodView = MethodView # type: ignore[attr-defined] @@ -79,6 +81,9 @@ def controller_module(monkeypatch: pytest.MonkeyPatch, config_overrides: Callabl login_module = importlib.import_module("libs.login") monkeypatch.setattr(login_module, "check_csrf_token", lambda *args, **kwargs: None) + monkeypatch.setattr( + module, "application_services", lambda: SimpleNamespace(credential_queries=MagicMock(spec=CredentialQuery)) + ) return module @@ -128,11 +133,18 @@ def _provider_credential( @contextmanager def _bind_database_session(session: Session): - database_session = scoped_session( - sessionmaker(bind=session.get_bind(), expire_on_commit=False), - ) + session_factory = sessionmaker(bind=session.get_bind(), expire_on_commit=False) + database_session = scoped_session(session_factory) try: - with patch("extensions.ext_database.db.session", database_session): + with ( + patch("extensions.ext_database.db.session", database_session), + patch( + "controllers.console.workspace.tool_providers.application_services", + return_value=SimpleNamespace( + credential_queries=CredentialQueryRepository(session_factory=session_factory) + ), + ), + ): yield finally: database_session.remove() diff --git a/api/tests/unit_tests/controllers/console/workspace/test_trigger_provider_apis.py b/api/tests/unit_tests/controllers/console/workspace/test_trigger_provider_apis.py index 074f8d527c2..bf5ab7485af 100644 --- a/api/tests/unit_tests/controllers/console/workspace/test_trigger_provider_apis.py +++ b/api/tests/unit_tests/controllers/console/workspace/test_trigger_provider_apis.py @@ -2,9 +2,11 @@ from __future__ import annotations +from collections.abc import Iterator from datetime import datetime from inspect import unwrap -from unittest.mock import ANY, MagicMock, patch +from types import SimpleNamespace +from unittest.mock import ANY, MagicMock, create_autospec, patch import pytest from flask import Flask @@ -36,6 +38,17 @@ from core.plugin.entities.plugin_daemon import CredentialType from core.trigger.entities.api_entities import SubscriptionBuilderApiEntity, TriggerProviderApiEntity from core.trigger.entities.entities import RequestLog from models.account import Account +from services.credentials.query import CredentialQuery + + +@pytest.fixture +def credential_query() -> Iterator[MagicMock]: + query = create_autospec(CredentialQuery, instance=True, spec_set=True) + with patch( + "controllers.console.workspace.trigger_providers.application_services", + return_value=SimpleNamespace(credential_queries=query), + ): + yield query def mock_user() -> Account: @@ -127,7 +140,7 @@ class TestTriggerProviderApis: class TestTriggerSubscriptionListApi: - def test_list_success(self, app: Flask) -> None: + def test_list_success(self, app: Flask, credential_query: MagicMock) -> None: api = TriggerSubscriptionListApi() method = unwrap(api.get) @@ -136,11 +149,12 @@ class TestTriggerSubscriptionListApi: patch( "controllers.console.workspace.trigger_providers.TriggerProviderService.list_trigger_provider_subscriptions", return_value=[], - ), + ) as list_subscriptions, ): assert method(api, "t1", mock_user(), "github") == [] + assert list_subscriptions.call_args.kwargs["credential_query"] is credential_query - def test_list_invalid_provider(self, app: Flask) -> None: + def test_list_invalid_provider(self, app: Flask, credential_query: MagicMock) -> None: api = TriggerSubscriptionListApi() method = unwrap(api.get) @@ -149,10 +163,11 @@ class TestTriggerSubscriptionListApi: patch( "controllers.console.workspace.trigger_providers.TriggerProviderService.list_trigger_provider_subscriptions", side_effect=ValueError("bad"), - ), + ) as list_subscriptions, ): result, status = method(api, "t1", mock_user(), "bad") assert status == 404 + assert list_subscriptions.call_args.kwargs["credential_query"] is credential_query class TestTriggerSubscriptionBuilderApis: diff --git a/api/tests/unit_tests/controllers/service_api/dataset/rag_pipeline/test_rag_pipeline_workflow.py b/api/tests/unit_tests/controllers/service_api/dataset/rag_pipeline/test_rag_pipeline_workflow.py index 4a87b4a8ff6..4cb3d5ec648 100644 --- a/api/tests/unit_tests/controllers/service_api/dataset/rag_pipeline/test_rag_pipeline_workflow.py +++ b/api/tests/unit_tests/controllers/service_api/dataset/rag_pipeline/test_rag_pipeline_workflow.py @@ -46,12 +46,16 @@ from controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow import ( KnowledgebasePipelineFileUploadApi, PipelineRunApi, ) +from core.app.apps.pipeline.pipeline_generator import PipelineGenerator from core.app.entities.app_invoke_entities import InvokeFrom from extensions.storage.storage_type import StorageType from models.account import Account from models.dataset import Dataset, Pipeline from models.enums import CreatorUserRole from models.model import UploadFile +from repositories.credentials.query_repository import CredentialQueryRepository +from repositories.data_source.credential_repository import SQLAlchemyDatasourceCredentialRepository +from services.data_source.provider_service import DatasourceProviderService from services.errors.file import FileTooLargeError as FileTooLargeServiceError from services.errors.file import UnsupportedFileTypeError from services.rag_pipeline.entity.pipeline_service_api_entities import ( @@ -86,6 +90,15 @@ def _persist_pipeline(session: Session, *, tenant_id: str) -> Pipeline: def _bind_database(sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch): session_proxy = scoped_session(sqlite_session_factory) monkeypatch.setattr(workflow_module.db, "session", session_proxy) + query = CredentialQueryRepository(session_factory=sqlite_session_factory) + providers = DatasourceProviderService( + credentials=SQLAlchemyDatasourceCredentialRepository(session_factory=sqlite_session_factory) + ) + monkeypatch.setattr( + workflow_module, + "application_services", + lambda: SimpleNamespace(credential_queries=query, data_sources=SimpleNamespace(providers=providers)), + ) yield session_proxy.remove() @@ -448,7 +461,11 @@ class TestDatasourcePluginsApiGet: assert status == 200 assert response == datasource_plugins mock_svc_instance.get_datasource_plugins.assert_called_once_with( - tenant_id=tenant_id, dataset_id=dataset_id, is_published=True + tenant_id=tenant_id, + dataset_id=dataset_id, + is_published=True, + credential_query=workflow_module.application_services().credential_queries, + datasource_providers=workflow_module.application_services().data_sources.providers, ) @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.RagPipelineService") @@ -469,7 +486,11 @@ class TestDatasourcePluginsApiGet: assert status == 200 assert response == [] mock_svc_instance.get_datasource_plugins.assert_called_once_with( - tenant_id=tenant_id, dataset_id=dataset_id, is_published=False + tenant_id=tenant_id, + dataset_id=dataset_id, + is_published=False, + credential_query=workflow_module.application_services().credential_queries, + datasource_providers=workflow_module.application_services().data_sources.providers, ) def test_get_plugins_not_found(self, app: Flask, sqlite_session: Session): @@ -607,7 +628,15 @@ class TestPipelineRunApiPost: @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.RagPipelineService") @patch("controllers.service_api.dataset.rag_pipeline.rag_pipeline_workflow.service_api_ns") def test_post_success_streaming( - self, mock_ns, mock_svc_cls, mock_current_user, mock_gen_svc, mock_helper, app, sqlite_session: Session + self, + mock_ns, + mock_svc_cls, + mock_current_user, + mock_gen_svc, + mock_helper, + app: Flask, + sqlite_session: Session, + pipeline_application: PipelineGenerator, ): """Test successful pipeline run with streaming response.""" tenant_id = str(uuid.uuid4()) @@ -639,6 +668,8 @@ class TestPipelineRunApiPost: assert response == {"result": "ok"} mock_svc_cls.assert_called_once_with(sqlite_session) mock_gen_svc.generate.assert_called_once() + assert mock_gen_svc.generate.call_args.kwargs["generator"] is pipeline_application + assert mock_gen_svc.generate.call_args.kwargs["session"] is sqlite_session def test_post_not_found(self, app: Flask, sqlite_session: Session): """Test NotFound when dataset check fails.""" diff --git a/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_segment.py b/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_segment.py index 56aefeda1c7..96eb5c48afa 100644 --- a/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_segment.py +++ b/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_segment.py @@ -17,7 +17,8 @@ Focus on: import inspect import uuid from collections.abc import Callable -from unittest.mock import ANY, Mock, patch +from types import SimpleNamespace +from unittest.mock import ANY, Mock, create_autospec, patch import pytest from flask import Flask @@ -42,10 +43,19 @@ from models.account import Account, Tenant from models.dataset import ChildChunk, Dataset, Document, DocumentSegment, DocumentSegmentSummary from models.enums import IndexingStatus, SegmentType from services.api_token_service import CachedApiToken -from services.dataset_service import DocumentService, SegmentService +from services.knowledge.dataset_service import DocumentService, SegmentService +from services.knowledge.segments.application import SegmentMutationService from tests.unit_tests.model_factories import make_account +@pytest.fixture(autouse=True) +def segment_mutations(): + mutations = create_autospec(SegmentMutationService, instance=True, spec_set=True) + services = SimpleNamespace(knowledge=SimpleNamespace(segments=SimpleNamespace(mutations=mutations))) + with patch("controllers.service_api.dataset.segment.application_services", return_value=services): + yield mutations + + def _segment_response_dict(summary: str | None = None): return { "id": "seg-1", @@ -555,6 +565,7 @@ class TestChildChunkServiceMockedBehavior: document=Document(), dataset=Dataset(), session=unbound_session, + actor_id="account-id", ) assert result == mock_child_chunk @@ -627,6 +638,7 @@ class TestChildChunkServiceMockedBehavior: document=Document(), dataset=Dataset(), session=unbound_session, + actor_id="account-id", ) assert result.content == "Updated content" @@ -722,7 +734,7 @@ class TestSegmentUpdatePayload: def test_payload_with_segment_args(self): """Test payload with SegmentUpdateArgs.""" from controllers.service_api.dataset.segment import SegmentUpdatePayload - from services.entities.knowledge_entities.knowledge_entities import SegmentUpdateArgs + from services.knowledge.entities.segments import SegmentUpdateArgs segment_args = SegmentUpdateArgs(content="Updated content") payload = SegmentUpdatePayload(segment=segment_args) @@ -731,7 +743,7 @@ class TestSegmentUpdatePayload: def test_payload_with_answer_update(self): """Test payload with answer update.""" from controllers.service_api.dataset.segment import SegmentUpdatePayload - from services.entities.knowledge_entities.knowledge_entities import SegmentUpdateArgs + from services.knowledge.entities.segments import SegmentUpdateArgs segment_args = SegmentUpdateArgs(answer="Updated answer") payload = SegmentUpdatePayload(segment=segment_args) @@ -740,7 +752,7 @@ class TestSegmentUpdatePayload: def test_payload_with_keywords_update(self): """Test payload with keywords update.""" from controllers.service_api.dataset.segment import SegmentUpdatePayload - from services.entities.knowledge_entities.knowledge_entities import SegmentUpdateArgs + from services.knowledge.entities.segments import SegmentUpdateArgs segment_args = SegmentUpdateArgs(keywords=["new", "keywords"]) payload = SegmentUpdatePayload(segment=segment_args) @@ -749,7 +761,7 @@ class TestSegmentUpdatePayload: def test_payload_with_enabled_toggle(self): """Test payload with enabled toggle.""" from controllers.service_api.dataset.segment import SegmentUpdatePayload - from services.entities.knowledge_entities.knowledge_entities import SegmentUpdateArgs + from services.knowledge.entities.segments import SegmentUpdateArgs segment_args = SegmentUpdateArgs(enabled=True) payload = SegmentUpdatePayload(segment=segment_args) @@ -758,7 +770,7 @@ class TestSegmentUpdatePayload: def test_payload_with_regenerate_child_chunks(self): """Test payload with regenerate_child_chunks flag.""" from controllers.service_api.dataset.segment import SegmentUpdatePayload - from services.entities.knowledge_entities.knowledge_entities import SegmentUpdateArgs + from services.knowledge.entities.segments import SegmentUpdateArgs segment_args = SegmentUpdateArgs(regenerate_child_chunks=True) payload = SegmentUpdatePayload(segment=segment_args) @@ -770,7 +782,7 @@ class TestSegmentUpdateArgs: def test_args_with_defaults(self): """Test args with default values.""" - from services.entities.knowledge_entities.knowledge_entities import SegmentUpdateArgs + from services.knowledge.entities.segments import SegmentUpdateArgs args = SegmentUpdateArgs() assert args.content is None @@ -781,14 +793,14 @@ class TestSegmentUpdateArgs: def test_args_with_content(self): """Test args with content update.""" - from services.entities.knowledge_entities.knowledge_entities import SegmentUpdateArgs + from services.knowledge.entities.segments import SegmentUpdateArgs args = SegmentUpdateArgs(content="New content here") assert args.content == "New content here" def test_args_with_all_fields(self): """Test args with all fields populated.""" - from services.entities.knowledge_entities.knowledge_entities import SegmentUpdateArgs + from services.knowledge.entities.segments import SegmentUpdateArgs args = SegmentUpdateArgs( content="Full content", @@ -813,7 +825,7 @@ class TestSegmentCreateArgs: def test_args_with_defaults(self): """Test args with default values.""" - from services.entities.knowledge_entities.knowledge_entities import SegmentCreateArgs + from services.knowledge.entities.knowledge_entities import SegmentCreateArgs args = SegmentCreateArgs() assert args.content is None @@ -823,7 +835,7 @@ class TestSegmentCreateArgs: def test_args_with_content_and_answer(self): """Test args with content and answer for Q&A mode.""" - from services.entities.knowledge_entities.knowledge_entities import SegmentCreateArgs + from services.knowledge.entities.knowledge_entities import SegmentCreateArgs args = SegmentCreateArgs(content="Question?", answer="Answer!") assert args.content == "Question?" @@ -831,7 +843,7 @@ class TestSegmentCreateArgs: def test_args_with_keywords(self): """Test args with keywords for search indexing.""" - from services.entities.knowledge_entities.knowledge_entities import SegmentCreateArgs + from services.knowledge.entities.knowledge_entities import SegmentCreateArgs args = SegmentCreateArgs(content="Test content", keywords=["machine learning", "AI", "neural networks"]) assert args.keywords is not None @@ -843,7 +855,7 @@ class TestChildChunkUpdateArgs: def test_args_with_content_only(self): """Test args with content only.""" - from services.entities.knowledge_entities.knowledge_entities import ChildChunkUpdateArgs + from services.knowledge.entities.segments import ChildChunkUpdateArgs args = ChildChunkUpdateArgs(content="Updated chunk content") assert args.content == "Updated chunk content" @@ -851,7 +863,7 @@ class TestChildChunkUpdateArgs: def test_args_with_id_and_content(self): """Test args with both id and content.""" - from services.entities.knowledge_entities.knowledge_entities import ChildChunkUpdateArgs + from services.knowledge.entities.segments import ChildChunkUpdateArgs chunk_id = str(uuid.uuid4()) args = ChildChunkUpdateArgs(id=chunk_id, content="Updated content") @@ -1033,8 +1045,8 @@ class TestSegmentApiGet(SQLiteEndpointTest): ``current_account_with_tenant()`` and response serialization. """ - @patch("controllers.service_api.dataset.segment.segment_responses_with_summaries") - @patch("controllers.service_api.dataset.segment.SummaryIndexService.get_segments_summaries") + @patch("controllers.service_api.dataset.segment.load_segment_details") + @patch("controllers.service_api.dataset.segment.SummaryIndexAdapter.get_segments_summaries") @patch("controllers.service_api.dataset.segment.SegmentService") @patch("controllers.service_api.dataset.segment.DocumentService") @patch("controllers.service_api.dataset.segment.current_account_with_tenant") @@ -1077,8 +1089,8 @@ class TestSegmentApiGet(SQLiteEndpointTest): mock_dump_segments.assert_called_once_with([mock_segment], {}, session=ANY) assert isinstance(mock_dump_segments.call_args.kwargs["session"], Session) - @patch("controllers.service_api.dataset.segment.segment_responses_with_summaries") - @patch("controllers.service_api.dataset.segment.SummaryIndexService.get_segments_summaries") + @patch("controllers.service_api.dataset.segment.load_segment_details") + @patch("controllers.service_api.dataset.segment.SummaryIndexAdapter.get_segments_summaries") @patch("controllers.service_api.dataset.segment.SegmentService") @patch("controllers.service_api.dataset.segment.DocumentService") @patch("controllers.service_api.dataset.segment.current_account_with_tenant") @@ -1119,8 +1131,8 @@ class TestSegmentApiGet(SQLiteEndpointTest): assert response["total"] == page_size assert response["page"] == 1 - @patch("controllers.service_api.dataset.segment.segment_responses_with_summaries") - @patch("controllers.service_api.dataset.segment.SummaryIndexService.get_segments_summaries") + @patch("controllers.service_api.dataset.segment.load_segment_details") + @patch("controllers.service_api.dataset.segment.SummaryIndexAdapter.get_segments_summaries") @patch("controllers.service_api.dataset.segment.SegmentService") @patch("controllers.service_api.dataset.segment.DocumentService") @patch("controllers.service_api.dataset.segment.current_account_with_tenant") @@ -1224,8 +1236,8 @@ class TestSegmentApiPost(SQLiteEndpointTest): mock_rate_limit.enabled = False mock_feature_svc.get_knowledge_rate_limit.return_value = mock_rate_limit - @patch("controllers.service_api.dataset.segment.segment_responses_with_summaries") - @patch("controllers.service_api.dataset.segment.SummaryIndexService.get_segments_summaries") + @patch("controllers.service_api.dataset.segment.load_segment_details") + @patch("controllers.service_api.dataset.segment.SummaryIndexAdapter.get_segments_summaries") @patch("controllers.service_api.dataset.segment.SegmentService") @patch("controllers.service_api.dataset.segment.DocumentService") @patch("controllers.service_api.dataset.segment.current_account_with_tenant") @@ -1426,7 +1438,9 @@ class TestDatasetSegmentApiDelete(SQLiteEndpointTest): # Assert assert response == ("", 204) - mock_seg_svc.delete_segment.assert_called_once_with(mock_segment, mock_doc, mock_dataset, self.session) + mock_seg_svc.delete_segment.assert_called_once_with( + mock_segment, mock_doc, mock_dataset, self.session, mutations=ANY + ) @patch("controllers.service_api.dataset.segment.SegmentService") @patch("controllers.service_api.dataset.segment.DocumentService") @@ -1564,8 +1578,8 @@ class TestDatasetSegmentApiUpdate(SQLiteEndpointTest): mock_rate_limit.enabled = False mock_feature_svc.get_knowledge_rate_limit.return_value = mock_rate_limit - @patch("controllers.service_api.dataset.segment.segment_response_with_summary") - @patch("controllers.service_api.dataset.segment.SummaryIndexService.get_segment_summary") + @patch("controllers.service_api.dataset.segment.load_segment_detail") + @patch("controllers.service_api.dataset.segment.SummaryIndexAdapter.get_segment_summary") @patch("controllers.service_api.dataset.segment.SegmentService") @patch("controllers.service_api.dataset.segment.DocumentService") @patch("controllers.service_api.dataset.segment.DatasetService") @@ -1717,8 +1731,8 @@ class TestDatasetSegmentApiGetSingle(SQLiteEndpointTest): ``current_account_with_tenant()`` and response serialization. """ - @patch("controllers.service_api.dataset.segment.segment_response_with_summary") - @patch("controllers.service_api.dataset.segment.SummaryIndexService.get_segment_summary") + @patch("controllers.service_api.dataset.segment.load_segment_detail") + @patch("controllers.service_api.dataset.segment.SummaryIndexAdapter.get_segment_summary") @patch("controllers.service_api.dataset.segment.SegmentService") @patch("controllers.service_api.dataset.segment.DocumentService") @patch("controllers.service_api.dataset.segment.DatasetService") @@ -1764,8 +1778,8 @@ class TestDatasetSegmentApiGetSingle(SQLiteEndpointTest): mock_dump_segment.assert_called_once_with(mock_segment, None, session=ANY) assert isinstance(mock_dump_segment.call_args.kwargs["session"], Session) - @patch("controllers.service_api.dataset.segment.segment_response_with_summary") - @patch("controllers.service_api.dataset.segment.SummaryIndexService.get_segment_summary") + @patch("controllers.service_api.dataset.segment.load_segment_detail") + @patch("controllers.service_api.dataset.segment.SummaryIndexAdapter.get_segment_summary") @patch("controllers.service_api.dataset.segment.SegmentService") @patch("controllers.service_api.dataset.segment.DocumentService") @patch("controllers.service_api.dataset.segment.DatasetService") @@ -1783,7 +1797,7 @@ class TestDatasetSegmentApiGetSingle(SQLiteEndpointTest): mock_dataset, mock_segment, ): - """Test that single segment response includes summary content from SummaryIndexService.""" + """Test that single segment response includes summary content from SummaryIndexAdapter.""" mock_account_fn.return_value = (_account(), mock_tenant.id) self._persist_dataset(mock_dataset, mock_tenant.id) mock_dataset_svc.check_dataset_model_setting.return_value = None diff --git a/api/tests/unit_tests/controllers/service_api/dataset/test_document.py b/api/tests/unit_tests/controllers/service_api/dataset/test_document.py index 665d4cb0607..159dddfbbc2 100644 --- a/api/tests/unit_tests/controllers/service_api/dataset/test_document.py +++ b/api/tests/unit_tests/controllers/service_api/dataset/test_document.py @@ -60,10 +60,10 @@ from models.enums import ( SegmentStatus, ) from models.model import ApiToken, UploadFile -from services.dataset_ref_service import DatasetRef -from services.dataset_service import DocumentService -from services.entities.knowledge_entities.knowledge_entities import ProcessRule, RetrievalModel from services.errors.file import FileTooLargeError as FileTooLargeServiceError +from services.knowledge.dataset_service import DocumentService +from services.knowledge.entities.knowledge_entities import ProcessRule, RetrievalModel +from services.knowledge.resource_scope import DatasetRef from tests.unit_tests.config_override import config_overrides_context from tests.unit_tests.model_factories import make_account, make_upload_file @@ -639,8 +639,8 @@ class TestDocumentServiceBatchMethods: class TestDocumentServiceFileOperations: """Test DocumentService file related operations.""" - @patch("services.dataset_service.file_helpers.get_signed_file_url") - @patch("services.dataset_service.DocumentService._get_upload_file_for_upload_file_document") + @patch("services.knowledge.dataset_service.file_helpers.get_signed_file_url") + @patch("services.knowledge.dataset_service.DocumentService._get_upload_file_for_upload_file_document") def test_get_document_download_url(self, mock_get_file, mock_signed_url, sqlite_session: Session): """Test generation of download URL.""" mock_doc = make_serializable_document() @@ -661,8 +661,8 @@ class TestDocumentServiceSaveValidation: """Test validations during document saving.""" @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - @patch("services.dataset_service.DatasetService.check_doc_form") - @patch("services.dataset_service.FeatureService.get_features") + @patch("services.knowledge.dataset_service.DatasetService.check_doc_form") + @patch("services.knowledge.dataset_service.FeatureService.get_features") def test_save_document_validates_doc_form(self, mock_features, mock_check_form, sqlite_session: Session): """Test that doc_form is validated during save.""" dataset = make_dataset(tenant_id="tenant_id") diff --git a/api/tests/unit_tests/controllers/service_api/dataset/test_hit_testing.py b/api/tests/unit_tests/controllers/service_api/dataset/test_hit_testing.py index f47de0c8d50..b1c5a2a6250 100644 --- a/api/tests/unit_tests/controllers/service_api/dataset/test_hit_testing.py +++ b/api/tests/unit_tests/controllers/service_api/dataset/test_hit_testing.py @@ -25,7 +25,7 @@ import services from controllers.service_api.dataset.hit_testing import HitTestingApi, HitTestingPayload from models.account import Account, Tenant, TenantAccountRole from models.dataset import Dataset -from services.entities.knowledge_entities.knowledge_entities import RetrievalModel +from services.knowledge.entities.knowledge_entities import RetrievalModel # --------------------------------------------------------------------------- # HitTestingPayload Model Tests diff --git a/api/tests/unit_tests/controllers/service_api/dataset/test_metadata.py b/api/tests/unit_tests/controllers/service_api/dataset/test_metadata.py index e138a5b3c77..55b42a4e355 100644 --- a/api/tests/unit_tests/controllers/service_api/dataset/test_metadata.py +++ b/api/tests/unit_tests/controllers/service_api/dataset/test_metadata.py @@ -38,8 +38,9 @@ from controllers.service_api.dataset.metadata import ( from models.account import Account, Tenant from models.dataset import Dataset from models.enums import PermissionEnum -from services.entities.knowledge_entities.knowledge_entities import MetadataArgs, MetadataOperationData from services.errors.metadata import MetadataResourceNotFoundError +from services.knowledge.entities.knowledge_entities import MetadataArgs, MetadataOperationData +from services.knowledge.resource_scope import DatasetRef @pytest.fixture @@ -93,6 +94,30 @@ class _UsesSQLiteSession: self.session = sqlite_session +@pytest.mark.parametrize("method", ["patch", "delete"]) +def test_metadata_mutation_translates_missing_resource(method, sqlite_session, mock_tenant, mock_dataset): + with ( + patch.object(metadata_module, "DatasetService") as datasets, + patch.object(metadata_module, "application_services") as services, + ): + datasets.get_dataset_for_tenant.return_value = mock_dataset + metadata = services.return_value.knowledge.metadata + operation = metadata.update_metadata_name if method == "patch" else metadata.delete_metadata + operation.side_effect = MetadataResourceNotFoundError("Metadata not found.") + resource = DatasetMetadataServiceApi() + arguments = { + "session": sqlite_session, + "tenant_id": mock_tenant.id, + "dataset_id": mock_dataset.id, + "metadata_id": uuid.uuid4(), + } + if method == "patch": + arguments["payload"] = MetadataUpdatePayload(name="renamed") + endpoint = DatasetMetadataServiceApi.patch if method == "patch" else DatasetMetadataServiceApi.delete + with pytest.raises(NotFound, match="Metadata not found"): + unwrap(endpoint)(resource, **arguments) + + class TestDatasetMetadataCreatePost(_UsesSQLiteSession): """Tests for DatasetMetadataCreateServiceApi.post(). @@ -107,7 +132,7 @@ class TestDatasetMetadataCreatePost(_UsesSQLiteSession): metadata_args = MetadataArgs.model_validate(request.get_json() or {}) return unwrap(api.post)(api, metadata_args, session, **kwargs) - @patch("controllers.service_api.dataset.metadata.MetadataService") + @patch("controllers.service_api.dataset.metadata.application_services") @patch("controllers.service_api.dataset.metadata.DatasetService") def test_create_metadata_success( self, @@ -118,7 +143,8 @@ class TestDatasetMetadataCreatePost(_UsesSQLiteSession): mock_dataset, ): """Test successful metadata creation.""" - mock_dataset_svc.get_dataset.return_value = mock_dataset + mock_meta_svc = mock_meta_svc.return_value.knowledge.metadata + mock_dataset_svc.get_dataset_for_tenant.return_value = mock_dataset mock_dataset_svc.check_dataset_permission.return_value = None mock_metadata = {"id": "meta-1", "type": "string", "name": "Author"} mock_meta_svc.create_metadata.return_value = mock_metadata @@ -150,7 +176,7 @@ class TestDatasetMetadataCreatePost(_UsesSQLiteSession): mock_dataset, ): """Test 404 when dataset not found.""" - mock_dataset_svc.get_dataset.return_value = None + mock_dataset_svc.get_dataset_for_tenant.return_value = None with app.test_request_context( f"/datasets/{mock_dataset.id}/metadata", @@ -171,7 +197,7 @@ class TestDatasetMetadataCreatePost(_UsesSQLiteSession): class TestDatasetMetadataCreateGet(_UsesSQLiteSession): """Tests for DatasetMetadataCreateServiceApi.get().""" - @patch("controllers.service_api.dataset.metadata.MetadataService") + @patch("controllers.service_api.dataset.metadata.application_services") @patch("controllers.service_api.dataset.metadata.DatasetService") def test_get_metadata_success( self, @@ -182,7 +208,8 @@ class TestDatasetMetadataCreateGet(_UsesSQLiteSession): mock_dataset, ): """Test successful metadata list retrieval.""" - mock_dataset_svc.get_dataset.return_value = mock_dataset + mock_meta_svc = mock_meta_svc.return_value.knowledge.metadata + mock_dataset_svc.get_dataset_for_tenant.return_value = mock_dataset mock_meta_svc.get_dataset_metadatas.return_value = { "doc_metadata": [{"id": "m1", "name": "Author", "type": "string", "count": 0}], "built_in_field_enabled": False, @@ -213,7 +240,7 @@ class TestDatasetMetadataCreateGet(_UsesSQLiteSession): mock_dataset, ): """Test 404 when dataset not found.""" - mock_dataset_svc.get_dataset.return_value = None + mock_dataset_svc.get_dataset_for_tenant.return_value = None with app.test_request_context( f"/datasets/{mock_dataset.id}/metadata", @@ -240,7 +267,7 @@ class TestDatasetMetadataServiceApiPatch(_UsesSQLiteSession): payload = MetadataUpdatePayload.model_validate(request.get_json() or {}) return unwrap(api.patch)(api, payload, session, **kwargs) - @patch("controllers.service_api.dataset.metadata.MetadataService") + @patch("controllers.service_api.dataset.metadata.application_services") @patch("controllers.service_api.dataset.metadata.DatasetService") def test_update_metadata_name_success( self, @@ -252,6 +279,7 @@ class TestDatasetMetadataServiceApiPatch(_UsesSQLiteSession): account: Account, ): """Test successful metadata name update.""" + mock_meta_svc = mock_meta_svc.return_value.knowledge.metadata metadata_id = str(uuid.uuid4()) mock_dataset_svc.get_dataset_for_tenant.return_value = mock_dataset mock_dataset_svc.check_dataset_permission.return_value = None @@ -278,7 +306,7 @@ class TestDatasetMetadataServiceApiPatch(_UsesSQLiteSession): str(mock_dataset.id), mock_tenant.id, session=session ) mock_meta_svc.update_metadata_name.assert_called_once_with( - mock_dataset, metadata_id, "New Name", account, session=session + DatasetRef(mock_tenant.id, mock_dataset.id), metadata_id, "New Name", actor_id=account.id ) @patch("controllers.service_api.dataset.metadata.DatasetService") @@ -320,7 +348,7 @@ class TestDatasetMetadataServiceApiDelete(_UsesSQLiteSession): def _call_delete(api, session: Session, **kwargs): return unwrap(api.delete)(api, session, **kwargs) - @patch("controllers.service_api.dataset.metadata.MetadataService") + @patch("controllers.service_api.dataset.metadata.application_services") @patch("controllers.service_api.dataset.metadata.DatasetService") def test_delete_metadata_success( self, @@ -331,6 +359,7 @@ class TestDatasetMetadataServiceApiDelete(_UsesSQLiteSession): mock_dataset, ): """Test successful metadata deletion.""" + mock_meta_svc = mock_meta_svc.return_value.knowledge.metadata metadata_id = str(uuid.uuid4()) mock_dataset_svc.get_dataset_for_tenant.return_value = mock_dataset mock_dataset_svc.check_dataset_permission.return_value = None @@ -354,7 +383,7 @@ class TestDatasetMetadataServiceApiDelete(_UsesSQLiteSession): mock_dataset_svc.get_dataset_for_tenant.assert_called_once_with( str(mock_dataset.id), mock_tenant.id, session=session ) - mock_meta_svc.delete_metadata.assert_called_once_with(mock_dataset, metadata_id, session) + mock_meta_svc.delete_metadata.assert_called_once_with(DatasetRef(mock_tenant.id, mock_dataset.id), metadata_id) @patch("controllers.service_api.dataset.metadata.DatasetService") def test_delete_metadata_dataset_not_found( @@ -392,7 +421,7 @@ class TestDatasetMetadataServiceApiDelete(_UsesSQLiteSession): class TestDatasetMetadataBuiltInFieldGet(_UsesSQLiteSession): """Tests for DatasetMetadataBuiltInFieldServiceApi.get().""" - @patch("controllers.service_api.dataset.metadata.MetadataService") + @patch("controllers.service_api.dataset.metadata.application_services") def test_get_built_in_fields_success( self, mock_meta_svc, @@ -401,6 +430,7 @@ class TestDatasetMetadataBuiltInFieldGet(_UsesSQLiteSession): mock_dataset, ): """Test successful built-in fields retrieval.""" + mock_meta_svc = mock_meta_svc.return_value.knowledge.metadata mock_meta_svc.get_built_in_fields.return_value = [ {"name": "source", "type": "string"}, ] @@ -434,7 +464,7 @@ class TestDatasetMetadataBuiltInFieldAction(_UsesSQLiteSession): def _call_post(api, session: Session, **kwargs): return unwrap(api.post)(api, session, **kwargs) - @patch("controllers.service_api.dataset.metadata.MetadataService") + @patch("controllers.service_api.dataset.metadata.application_services") @patch("controllers.service_api.dataset.metadata.DatasetService") def test_enable_built_in_field( self, @@ -445,7 +475,8 @@ class TestDatasetMetadataBuiltInFieldAction(_UsesSQLiteSession): mock_dataset, ): """Test enabling built-in metadata field.""" - mock_dataset_svc.get_dataset.return_value = mock_dataset + mock_meta_svc = mock_meta_svc.return_value.knowledge.metadata + mock_dataset_svc.get_dataset_for_tenant.return_value = mock_dataset mock_dataset_svc.check_dataset_permission.return_value = None with app.test_request_context( @@ -464,9 +495,9 @@ class TestDatasetMetadataBuiltInFieldAction(_UsesSQLiteSession): assert status == 200 assert response["result"] == "success" - mock_meta_svc.enable_built_in_field.assert_called_once_with(mock_dataset, session) + mock_meta_svc.enable_built_in_field.assert_called_once_with(DatasetRef(mock_tenant.id, mock_dataset.id)) - @patch("controllers.service_api.dataset.metadata.MetadataService") + @patch("controllers.service_api.dataset.metadata.application_services") @patch("controllers.service_api.dataset.metadata.DatasetService") def test_disable_built_in_field( self, @@ -477,7 +508,8 @@ class TestDatasetMetadataBuiltInFieldAction(_UsesSQLiteSession): mock_dataset, ): """Test disabling built-in metadata field.""" - mock_dataset_svc.get_dataset.return_value = mock_dataset + mock_meta_svc = mock_meta_svc.return_value.knowledge.metadata + mock_dataset_svc.get_dataset_for_tenant.return_value = mock_dataset mock_dataset_svc.check_dataset_permission.return_value = None with app.test_request_context( @@ -495,7 +527,7 @@ class TestDatasetMetadataBuiltInFieldAction(_UsesSQLiteSession): ) assert status == 200 - mock_meta_svc.disable_built_in_field.assert_called_once_with(mock_dataset, session) + mock_meta_svc.disable_built_in_field.assert_called_once_with(DatasetRef(mock_tenant.id, mock_dataset.id)) @patch("controllers.service_api.dataset.metadata.DatasetService") def test_action_dataset_not_found( @@ -506,7 +538,7 @@ class TestDatasetMetadataBuiltInFieldAction(_UsesSQLiteSession): mock_dataset, ): """Test 404 when dataset not found.""" - mock_dataset_svc.get_dataset.return_value = None + mock_dataset_svc.get_dataset_for_tenant.return_value = None with app.test_request_context( f"/datasets/{mock_dataset.id}/metadata/built-in/enable", @@ -540,7 +572,7 @@ class TestDocumentMetadataEditPost(_UsesSQLiteSession): metadata_args = MetadataOperationData.model_validate(request.get_json() or {}) return unwrap(api.post)(api, metadata_args, session, **kwargs) - @patch("controllers.service_api.dataset.metadata.MetadataService") + @patch("controllers.service_api.dataset.metadata.application_services") @patch("controllers.service_api.dataset.metadata.DatasetService") def test_update_documents_metadata_success( self, @@ -552,6 +584,7 @@ class TestDocumentMetadataEditPost(_UsesSQLiteSession): account: Account, ): """Test successful documents metadata update.""" + mock_meta_svc = mock_meta_svc.return_value.knowledge.metadata mock_dataset_svc.get_dataset_for_tenant.return_value = mock_dataset mock_dataset_svc.check_dataset_permission.return_value = None mock_meta_svc.update_documents_metadata.return_value = None @@ -573,10 +606,9 @@ class TestDocumentMetadataEditPost(_UsesSQLiteSession): assert status == 200 assert response["result"] == "success" mock_meta_svc.update_documents_metadata.assert_called_once_with( - mock_dataset, + DatasetRef(mock_tenant.id, mock_dataset.id), ANY, - account, - session=session, + actor_id=account.id, ) @patch("controllers.service_api.dataset.metadata.DatasetService") @@ -605,7 +637,7 @@ class TestDocumentMetadataEditPost(_UsesSQLiteSession): dataset_id=mock_dataset.id, ) - @patch("controllers.service_api.dataset.metadata.MetadataService") + @patch("controllers.service_api.dataset.metadata.application_services") @patch("controllers.service_api.dataset.metadata.DatasetService") def test_update_documents_metadata_translates_missing_resource( self, @@ -615,6 +647,7 @@ class TestDocumentMetadataEditPost(_UsesSQLiteSession): mock_tenant, mock_dataset, ): + mock_meta_svc = mock_meta_svc.return_value.knowledge.metadata mock_dataset_svc.get_dataset_for_tenant.return_value = mock_dataset mock_meta_svc.update_documents_metadata.side_effect = MetadataResourceNotFoundError("Document not found.") diff --git a/api/tests/unit_tests/controllers/test_rbac_route_resource_contract.py b/api/tests/unit_tests/controllers/test_rbac_route_resource_contract.py index 0fd2a205b2f..60a21436242 100644 --- a/api/tests/unit_tests/controllers/test_rbac_route_resource_contract.py +++ b/api/tests/unit_tests/controllers/test_rbac_route_resource_contract.py @@ -27,15 +27,7 @@ NAMESPACES = (console_ns, openapi_ns) HTTP_METHODS = ("delete", "get", "head", "options", "patch", "post", "put") -# Known violations tracked separately: DatasetDocumentSegmentBatchImportApi binds one class -# to both the dataset-scoped import route and the job-scoped status route, so every method -# on it is reachable at a path carrying only a job id. Its permission points are genuinely -# per-dataset, so it needs the route split rather than a Workspace() locator. Remove these -# entries with that fix. -KNOWN_VIOLATIONS = { - ("DatasetDocumentSegmentBatchImportApi", "post"), - ("DatasetDocumentSegmentBatchImportApi", "get"), -} +KNOWN_VIOLATIONS: set[tuple[str, str]] = set() # The generated access-permission endpoints, spelled out here so this guard fails loudly if # the factory stops registering them (or renames a URL) rather than silently checking less. diff --git a/api/tests/unit_tests/controllers/test_swagger.py b/api/tests/unit_tests/controllers/test_swagger.py index 63933381138..33857795895 100644 --- a/api/tests/unit_tests/controllers/test_swagger.py +++ b/api/tests/unit_tests/controllers/test_swagger.py @@ -262,6 +262,38 @@ def test_agent_tts_routes_document_voice_queries_and_binary_audio(openapi_json: assert _response_content_types(preview, status) == {"application/json"} +def test_console_segment_mutation_contracts_survive_admission(openapi_json: OpenAPIJSONLoader) -> None: + from controllers.console import bp as console_bp + + payload = json.loads(openapi_json(console_bp, "/console/api/openapi.json", False)) + paths = payload["paths"] + document_path = "/datasets/{dataset_id}/documents/{document_id}" + mutations = ( + (f"{document_path}/segment", "post", "SegmentCreatePayload", "SegmentDetailResponse"), + (f"{document_path}/segments/{{segment_id}}", "patch", "SegmentUpdatePayload", "SegmentDetailResponse"), + ( + f"{document_path}/segments/{{segment_id}}/child_chunks", + "patch", + "ChildChunkBatchUpdatePayload", + "ChildChunkBatchUpdateResponse", + ), + ) + for path, method, request_schema, response_schema in mutations: + operation = paths[path][method] + assert operation["requestBody"]["content"]["application/json"]["schema"] == { + "$ref": f"#/components/schemas/{request_schema}" + } + assert operation["responses"]["200"]["content"]["application/json"]["schema"] == { + "$ref": f"#/components/schemas/{response_schema}" + } + + delete_operation = paths[f"{document_path}/segments"]["delete"] + assert set(delete_operation["responses"]) == {"204"} + parameters = _parameters_by_name(delete_operation) + assert parameters["segment_id"]["in"] == "query" + assert parameters["segment_id"]["schema"]["type"] == "array" + + def test_service_document_file_routes_document_multipart_form_data(openapi_json: OpenAPIJSONLoader): from controllers.service_api import bp as service_api_bp diff --git a/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_generator.py b/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_generator.py index 841a003b84a..de16f2dc540 100644 --- a/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_generator.py +++ b/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_generator.py @@ -1,13 +1,13 @@ import contextlib import json from types import SimpleNamespace -from unittest.mock import MagicMock, PropertyMock +from unittest.mock import MagicMock, PropertyMock, create_autospec import pytest from pytest_mock import MockerFixture from sqlalchemy import select from sqlalchemy.engine import Engine -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, sessionmaker import core.app.apps.pipeline.pipeline_generator as module from core.app.apps.exc import GenerateTaskStoppedError @@ -17,6 +17,9 @@ from models.dataset import Dataset, Document, DocumentPipelineExecutionLog, Pipe from models.enums import DataSourceType, EndUserType from models.model import EndUser from models.workflow import Workflow, WorkflowType +from repositories.knowledge.document_repository import SQLAlchemyDocumentRepository +from services.data_source.credential_gateway import DatasourceProviderCredentialStore +from services.data_source.provider_service import DatasourceProviderService TENANT_ID = "00000000-0000-0000-0000-000000000001" PIPELINE_ID = "00000000-0000-0000-0000-000000000002" @@ -38,7 +41,12 @@ class FakeRagPipelineGenerateEntity(SimpleNamespace): @pytest.fixture def generator(mocker: MockerFixture, sqlite_engine: Engine): - gen = module.PipelineGenerator() + gen = module.PipelineGenerator( + documents=SQLAlchemyDocumentRepository(session_factory=sessionmaker(bind=sqlite_engine)), + datasource_providers=DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), + ) _patch_sqlite_engine(mocker, sqlite_engine) mocker.patch.object(module, "RagPipelineGenerateEntity", FakeRagPipelineGenerateEntity) @@ -228,10 +236,10 @@ def test_generate_published_pipeline_creates_documents_and_delay( ) mocker.patch.object(generator, "_prepare_user_inputs", return_value={"k": "v"}) - mocker.patch("services.dataset_service.DocumentService.get_documents_position", return_value=1) + mocker.patch("services.knowledge.dataset_service.DocumentService.get_documents_position", return_value=1) features = SimpleNamespace() get_features = mocker.patch("services.feature_service.FeatureService.get_features", return_value=features) - check_limits = mocker.patch("services.dataset_service.DocumentService.check_document_creation_limits") + check_limits = mocker.patch("services.knowledge.dataset_service.DocumentService.check_document_creation_limits") mocker.patch.object( module.DifyCoreRepositoryFactory, @@ -292,7 +300,7 @@ def test_generate_published_pipeline_rejects_when_document_creation_limits_excee features = SimpleNamespace() mocker.patch("services.feature_service.FeatureService.get_features", return_value=features) check_limits = mocker.patch( - "services.dataset_service.DocumentService.check_document_creation_limits", + "services.knowledge.dataset_service.DocumentService.check_document_creation_limits", side_effect=ValueError("document limit exceeded"), ) diff --git a/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_runner.py b/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_runner.py index 37e3912ba68..e2117219e00 100644 --- a/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_runner.py +++ b/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_runner.py @@ -19,7 +19,7 @@ import pytest from pytest_mock import MockerFixture from sqlalchemy import event from sqlalchemy.engine import Engine -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, sessionmaker import core.app.apps.pipeline.pipeline_runner as module from core.app.apps.pipeline.pipeline_runner import PipelineRunner @@ -29,6 +29,7 @@ from models.dataset import Dataset, Document, Pipeline from models.enums import DocumentCreatedFrom from models.model import EndUser from models.workflow import Workflow, WorkflowType +from repositories.knowledge.document_repository import SQLAlchemyDocumentRepository from tests.unit_tests.model_factories import make_dataset, make_document, make_end_user, make_workflow @@ -117,7 +118,7 @@ def _build_app_generate_entity() -> SimpleNamespace: @pytest.fixture -def runner(): +def runner(sqlite_engine: Engine): app_generate_entity = _build_app_generate_entity() queue_manager = MagicMock() variable_loader = MagicMock() @@ -133,6 +134,7 @@ def runner(): system_user_id="sys", workflow_execution_repository=workflow_execution_repository, workflow_node_execution_repository=workflow_node_execution_repository, + documents=SQLAlchemyDocumentRepository(session_factory=sessionmaker(bind=sqlite_engine)), ) @@ -174,12 +176,15 @@ def test_init_rag_pipeline_graph_not_found(mocker, runner): def test_update_document_status_on_failure(runner, sqlite_session: Session): document = _document() _, dataset, _ = _persist_scope(sqlite_session, documents=(document,)) - dataset_ref = module.DatasetRefService.create_dataset_ref(dataset) - document_ref = module.DatasetRefService.create_document_ref_from_id(dataset_ref, document.id) event = GraphRunFailedEvent(error="boom") - runner._update_document_status(event, document_ref) + runner._update_document_status( + event, + workspace_id=dataset.tenant_id, + dataset_id=dataset.id, + document_id=document.id, + ) sqlite_session.expire_all() updated = sqlite_session.get(Document, document.id) @@ -190,10 +195,13 @@ def test_update_document_status_on_failure(runner, sqlite_session: Session): def test_update_document_status_skips_when_document_not_found(runner, sqlite_session: Session): _, dataset, _ = _persist_scope(sqlite_session) - dataset_ref = module.DatasetRefService.create_dataset_ref(dataset) - document_ref = module.DatasetRefService.create_document_ref_from_id(dataset_ref, "missing") - runner._update_document_status(GraphRunFailedEvent(error="boom"), document_ref) + runner._update_document_status( + GraphRunFailedEvent(error="boom"), + workspace_id=dataset.tenant_id, + dataset_id=dataset.id, + document_id="missing", + ) assert sqlite_session.get(Document, "missing") is None @@ -207,14 +215,19 @@ def test_update_document_status_skips_without_document_ref(runner, sqlite_engine event.listen(sqlite_engine, "checkout", record_checkout) try: - runner._update_document_status(GraphRunFailedEvent(error="boom"), None) + runner._update_document_status( + GraphRunFailedEvent(error="boom"), + workspace_id="workspace-1", + dataset_id="dataset-1", + document_id=None, + ) finally: event.remove(sqlite_engine, "checkout", record_checkout) assert checkouts == 0 -def test_run_pipeline_not_found(): +def test_run_pipeline_not_found(sqlite_engine: Engine): app_generate_entity = _build_app_generate_entity() app_generate_entity.invoke_from = InvokeFrom.WEB_APP app_generate_entity.single_iteration_run = None @@ -228,6 +241,7 @@ def test_run_pipeline_not_found(): system_user_id="sys", workflow_execution_repository=MagicMock(), workflow_node_execution_repository=MagicMock(), + documents=SQLAlchemyDocumentRepository(session_factory=sessionmaker(bind=sqlite_engine)), ) with pytest.raises(ValueError): @@ -299,7 +313,7 @@ def test_run_rejects_original_document_outside_pipeline_dataset_after_async_boun runner.get_workflow.assert_not_called() -def test_run_workflow_not_initialized(sqlite_session: Session): +def test_run_workflow_not_initialized(sqlite_session: Session, sqlite_engine: Engine): app_generate_entity = _build_app_generate_entity() pipeline = _pipeline() @@ -316,18 +330,17 @@ def test_run_workflow_not_initialized(sqlite_session: Session): system_user_id="sys", workflow_execution_repository=MagicMock(), workflow_node_execution_repository=MagicMock(), + documents=SQLAlchemyDocumentRepository(session_factory=sessionmaker(bind=sqlite_engine)), ) with pytest.raises(ValueError): runner.run() -def test_run_single_iteration_path(mocker: MockerFixture, sqlite_session: Session): +def test_run_single_iteration_path(mocker: MockerFixture, sqlite_session: Session, sqlite_engine: Engine): app_generate_entity = _build_app_generate_entity() app_generate_entity.single_iteration_run = MagicMock() _, dataset, _ = _persist_scope(sqlite_session, documents=(_document(),)) - dataset_ref = module.DatasetRefService.create_dataset_ref(dataset) - document_ref = module.DatasetRefService.create_document_ref_from_id(dataset_ref, "doc") runner = PipelineRunner( application_generate_entity=app_generate_entity, @@ -337,6 +350,7 @@ def test_run_single_iteration_path(mocker: MockerFixture, sqlite_session: Sessio system_user_id="sys", workflow_execution_repository=MagicMock(), workflow_node_execution_repository=MagicMock(), + documents=SQLAlchemyDocumentRepository(session_factory=sessionmaker(bind=sqlite_engine)), ) runner._resolve_user_from = MagicMock(return_value=UserFrom.ACCOUNT) @@ -355,7 +369,12 @@ def test_run_single_iteration_path(mocker: MockerFixture, sqlite_session: Sessio runner.run() runner._prepare_single_node_execution.assert_called_once() - runner._update_document_status.assert_called_once_with(event, document_ref) + runner._update_document_status.assert_called_once_with( + event, + workspace_id=dataset.tenant_id, + dataset_id=dataset.id, + document_id="doc", + ) runner._handle_event.assert_called() @@ -388,6 +407,7 @@ def test_run_normal_path_builds_graph(mocker: MockerFixture, sqlite_session: Ses system_user_id="sys", workflow_execution_repository=MagicMock(), workflow_node_execution_repository=MagicMock(), + documents=SQLAlchemyDocumentRepository(session_factory=sessionmaker(bind=sqlite_engine)), ) runner._resolve_user_from = MagicMock(return_value=UserFrom.ACCOUNT) diff --git a/api/tests/unit_tests/core/datasource/test_datasource_manager.py b/api/tests/unit_tests/core/datasource/test_datasource_manager.py index 5e38166a5cd..da2c1591b65 100644 --- a/api/tests/unit_tests/core/datasource/test_datasource_manager.py +++ b/api/tests/unit_tests/core/datasource/test_datasource_manager.py @@ -236,10 +236,6 @@ def test_stream_online_results_yields_messages_online_document(mocker: MockerFix fake_runtime = mocker.Mock() fake_runtime.get_online_document_page_content.side_effect = _doc_messages mocker.patch.object(DatasourceManager, "get_datasource_runtime", return_value=fake_runtime) - mocker.patch( - "core.datasource.datasource_manager.DatasourceProviderService.get_datasource_credentials", - return_value=None, - ) gen = DatasourceManager.stream_online_results( user_id="u1", @@ -247,11 +243,9 @@ def test_stream_online_results_yields_messages_online_document(mocker: MockerFix datasource_type="online_document", provider_id="p/x", tenant_id="t1", - provider="prov", - plugin_id="plug", - credential_id="", datasource_param=types.SimpleNamespace(workspace_id="w", page_id="pg", type="t"), online_drive_request=None, + credentials={}, ) msgs = list(gen) assert len(msgs) == 1 @@ -268,10 +262,6 @@ def test_stream_online_results_sets_credentials_and_returns_empty_dict_online_do runtime = _Runtime() mocker.patch.object(DatasourceManager, "get_datasource_runtime", return_value=runtime) - mocker.patch( - "core.datasource.datasource_manager.DatasourceProviderService.get_datasource_credentials", - return_value={"token": "t"}, - ) gen = DatasourceManager.stream_online_results( user_id="u1", @@ -279,11 +269,9 @@ def test_stream_online_results_sets_credentials_and_returns_empty_dict_online_do datasource_type="online_document", provider_id="p/x", tenant_id="t1", - provider="prov", - plugin_id="plug", - credential_id="cred", datasource_param=types.SimpleNamespace(workspace_id="w", page_id="pg", type="t"), online_drive_request=None, + credentials={"token": "t"}, ) messages, final_value = _drain_generator(gen) @@ -304,10 +292,6 @@ def test_stream_online_results_raises_when_missing_params(mocker: MockerFixture) yield from _gen_messages_text_only("never") mocker.patch.object(DatasourceManager, "get_datasource_runtime", return_value=_Runtime()) - mocker.patch( - "core.datasource.datasource_manager.DatasourceProviderService.get_datasource_credentials", - return_value={}, - ) with pytest.raises(ValueError, match="datasource_param is required for ONLINE_DOCUMENT streaming"): list( @@ -317,11 +301,9 @@ def test_stream_online_results_raises_when_missing_params(mocker: MockerFixture) datasource_type="online_document", provider_id="p/x", tenant_id="t1", - provider="prov", - plugin_id="plug", - credential_id="", datasource_param=None, online_drive_request=None, + credentials={}, ) ) @@ -333,11 +315,9 @@ def test_stream_online_results_raises_when_missing_params(mocker: MockerFixture) datasource_type="online_drive", provider_id="p/x", tenant_id="t1", - provider="prov", - plugin_id="plug", - credential_id="", datasource_param=None, online_drive_request=None, + credentials={}, ) ) @@ -352,10 +332,6 @@ def test_stream_online_results_yields_messages_and_returns_empty_dict_online_dri runtime = _Runtime() mocker.patch.object(DatasourceManager, "get_datasource_runtime", return_value=runtime) - mocker.patch( - "core.datasource.datasource_manager.DatasourceProviderService.get_datasource_credentials", - return_value={"token": "t"}, - ) gen = DatasourceManager.stream_online_results( user_id="u1", @@ -363,11 +339,9 @@ def test_stream_online_results_yields_messages_and_returns_empty_dict_online_dri datasource_type="online_drive", provider_id="p/x", tenant_id="t1", - provider="prov", - plugin_id="plug", - credential_id="cred", datasource_param=None, online_drive_request=types.SimpleNamespace(id="fid", bucket="b"), + credentials={"token": "t"}, ) messages, final_value = _drain_generator(gen) @@ -378,10 +352,6 @@ def test_stream_online_results_yields_messages_and_returns_empty_dict_online_dri def test_stream_online_results_raises_for_unsupported_stream_type(mocker: MockerFixture): mocker.patch.object(DatasourceManager, "get_datasource_runtime", return_value=mocker.Mock()) - mocker.patch( - "core.datasource.datasource_manager.DatasourceProviderService.get_datasource_credentials", - return_value={}, - ) with pytest.raises(ValueError, match="Unsupported datasource type for streaming"): list( @@ -391,11 +361,9 @@ def test_stream_online_results_raises_for_unsupported_stream_type(mocker: Mocker datasource_type="website_crawl", provider_id="p/x", tenant_id="t1", - provider="prov", - plugin_id="plug", - credential_id="", datasource_param=None, online_drive_request=None, + credentials={}, ) ) @@ -416,14 +384,12 @@ def test_stream_node_events_emits_events_online_document(mocker: MockerFixture): datasource_type="online_document", provider_id="p/x", tenant_id="t1", - provider="prov", - plugin_id="plug", - credential_id="", parameters_for_log={"k": "v"}, datasource_info={"user_id": "u1"}, variable_pool=mocker.Mock(), datasource_param=types.SimpleNamespace(workspace_id="w", page_id="pg", type="t"), online_drive_request=None, + credentials={}, ) ) # should contain one StreamChunkEvent then a final chunk (empty) and a completed event @@ -503,14 +469,12 @@ def test_stream_node_events_builds_file_and_variables_from_messages(mocker: Mock datasource_type="online_document", provider_id="p/x", tenant_id="t1", - provider="prov", - plugin_id="plug", - credential_id="", parameters_for_log={"k": "v"}, datasource_info={"info": "x"}, variable_pool=variable_pool, datasource_param=types.SimpleNamespace(workspace_id="w", page_id="pg", type="t"), online_drive_request=None, + credentials={}, ) ) @@ -554,14 +518,12 @@ def test_stream_node_events_raises_when_toolfile_missing(mocker: MockerFixture, datasource_type="online_document", provider_id="p/x", tenant_id="t1", - provider="prov", - plugin_id="plug", - credential_id="", parameters_for_log={}, datasource_info={}, variable_pool=mocker.Mock(), datasource_param=types.SimpleNamespace(workspace_id="w", page_id="pg", type="t"), online_drive_request=None, + credentials={}, ) ) @@ -599,14 +561,12 @@ def test_stream_node_events_online_drive_sets_variable_pool_file_and_outputs(moc datasource_type="online_drive", provider_id="p/x", tenant_id="t1", - provider="prov", - plugin_id="plug", - credential_id="", parameters_for_log={}, datasource_info={"k": "v"}, variable_pool=variable_pool, datasource_param=None, online_drive_request=types.SimpleNamespace(id="id", bucket="b"), + credentials={}, ) ) @@ -644,14 +604,12 @@ def test_stream_node_events_skips_file_build_for_non_online_types(mocker: Mocker datasource_type="website_crawl", provider_id="p/x", tenant_id="t1", - provider="prov", - plugin_id="plug", - credential_id="", parameters_for_log={}, datasource_info={}, variable_pool=mocker.Mock(), datasource_param=None, online_drive_request=None, + credentials={}, ) ) diff --git a/api/tests/unit_tests/core/datasource/test_notion_provider.py b/api/tests/unit_tests/core/datasource/test_notion_provider.py index 1e4c5bbb355..d1e232f8f57 100644 --- a/api/tests/unit_tests/core/datasource/test_notion_provider.py +++ b/api/tests/unit_tests/core/datasource/test_notion_provider.py @@ -35,6 +35,8 @@ from core.rag.models.document import Document from models.base import TypeBase from models.dataset import Document as DocumentModel from models.enums import DataSourceType, DocumentCreatedFrom +from services.data_source.credential_gateway import DatasourceProviderCredentialStore +from services.data_source.provider_service import DatasourceProviderService @dataclass(frozen=True) @@ -102,6 +104,7 @@ class TestNotionExtractorAuthentication: tenant_id="tenant-789", notion_access_token="explicit-token-abc", document_model=mock_document_model, + notion_token_loader=lambda: "token", ) # Assert @@ -110,13 +113,11 @@ class TestNotionExtractorAuthentication: assert extractor._notion_obj_id == "page-456" assert extractor._notion_page_type == "page" - @patch("core.rag.extractor.notion_extractor.DatasourceProviderService") - def test_init_with_credential_id(self, mock_service_class, mock_document_model): + def test_init_with_credential_id(self, mock_document_model): """Test NotionExtractor initialization with credential ID retrieval.""" # Arrange mock_service = Mock() - mock_service.get_datasource_credentials.return_value = {"integration_secret": "credential-token-xyz"} - mock_service_class.return_value = mock_service + mock_service.get_stored_notion_access_token.return_value = "credential-token-xyz" # Act extractor = NotionExtractor( @@ -126,23 +127,25 @@ class TestNotionExtractorAuthentication: tenant_id="tenant-789", credential_id="cred-123", document_model=mock_document_model, + notion_token_loader=lambda: mock_service.get_stored_notion_access_token( + tenant_id="tenant-789", credential_id="cred-123" + ), ) # Assert assert extractor._notion_access_token == "credential-token-xyz" - mock_service.get_datasource_credentials.assert_called_once_with( + mock_service.get_stored_notion_access_token.assert_called_once_with( tenant_id="tenant-789", credential_id="cred-123", - provider="notion_datasource", - plugin_id="langgenius/notion_datasource", ) - @patch("core.rag.extractor.notion_extractor.NotionExtractor._get_access_token") + @patch("services.data_source.provider_service.DatasourceProviderService.get_datasource_credentials") def test_init_with_integration_token_fallback(self, mock_get_token, mock_document_model, config_overrides): """Test NotionExtractor falls back to integration token when credential not found.""" # Arrange mock_get_token.side_effect = Exception("No credential id found") config_overrides(NOTION_INTEGRATION_TOKEN="integration-token-fallback") + providers = DatasourceProviderService(credentials=Mock(spec=DatasourceProviderCredentialStore)) # Act extractor = NotionExtractor( @@ -152,17 +155,21 @@ class TestNotionExtractorAuthentication: tenant_id="tenant-789", credential_id=None, document_model=mock_document_model, + notion_token_loader=lambda: providers.get_stored_notion_access_token( + tenant_id="tenant-789", credential_id=None + ), ) # Assert assert extractor._notion_access_token == "integration-token-fallback" - @patch("core.rag.extractor.notion_extractor.NotionExtractor._get_access_token") + @patch("services.data_source.provider_service.DatasourceProviderService.get_datasource_credentials") def test_init_missing_credentials_raises_error(self, mock_get_token, mock_document_model, config_overrides): """Test NotionExtractor raises error when no credentials available.""" # Arrange mock_get_token.side_effect = Exception("No credential id found") config_overrides(NOTION_INTEGRATION_TOKEN=None) + providers = DatasourceProviderService(credentials=Mock(spec=DatasourceProviderCredentialStore)) # Act & Assert with pytest.raises(ValueError) as exc_info: @@ -173,6 +180,9 @@ class TestNotionExtractorAuthentication: tenant_id="tenant-789", credential_id=None, document_model=mock_document_model, + notion_token_loader=lambda: providers.get_stored_notion_access_token( + tenant_id="tenant-789", credential_id=None + ), ) assert "Must specify `integration_token`" in str(exc_info.value) @@ -196,6 +206,7 @@ class TestNotionExtractorPageRetrieval: notion_page_type="page", tenant_id="tenant-789", notion_access_token="test-token", + notion_token_loader=lambda: "token", ) def _create_mock_response(self, data: dict[str, Any], status_code: int = 200) -> Mock: @@ -396,6 +407,7 @@ class TestNotionExtractorDatabaseRetrieval: notion_page_type="database", tenant_id="tenant-789", notion_access_token="test-token", + notion_token_loader=lambda: "token", ) def _create_database_page(self, page_id: str, properties: dict[str, Any]) -> dict[str, Any]: @@ -600,6 +612,7 @@ class TestNotionExtractorTableParsing: notion_page_type="page", tenant_id="tenant-789", notion_access_token="test-token", + notion_token_loader=lambda: "token", ) @patch("httpx.request") @@ -751,6 +764,7 @@ class TestNotionExtractorLastEditedTime: tenant_id="tenant-789", notion_access_token="test-token", document_model=mock_document_model, + notion_token_loader=lambda: "token", ) @pytest.fixture @@ -763,6 +777,7 @@ class TestNotionExtractorLastEditedTime: tenant_id="tenant-789", notion_access_token="test-token", document_model=mock_document_model, + notion_token_loader=lambda: "token", ) @patch("httpx.request") @@ -869,6 +884,7 @@ class TestNotionExtractorIntegration: tenant_id="tenant-789", notion_access_token="test-token", document_model=persisted_document, + notion_token_loader=lambda: "token", ) # Mock last edited time request @@ -940,6 +956,7 @@ class TestNotionExtractorIntegration: tenant_id="tenant-789", notion_access_token="test-token", document_model=persisted_document, + notion_token_loader=lambda: "token", ) # Mock last edited time request @@ -990,6 +1007,7 @@ class TestNotionExtractorIntegration: notion_page_type="invalid_type", tenant_id="tenant-789", notion_access_token="test-token", + notion_token_loader=lambda: "token", ) # Act & Assert @@ -1016,6 +1034,7 @@ class TestNotionExtractorReadBlock: notion_page_type="page", tenant_id="tenant-789", notion_access_token="test-token", + notion_token_loader=lambda: "token", ) @patch("httpx.request") @@ -1172,6 +1191,7 @@ class TestNotionExtractorAdvancedBlockTypes: notion_page_type="page", tenant_id="tenant-789", notion_access_token="test-token", + notion_token_loader=lambda: "token", ) def _create_block_with_rich_text( @@ -1343,6 +1363,7 @@ class TestNotionExtractorDatabaseAdvanced: notion_page_type="database", tenant_id="tenant-789", notion_access_token="test-token", + notion_token_loader=lambda: "token", ) def _create_database_page_with_properties(self, page_id: str, properties: dict[str, Any]) -> dict[str, Any]: @@ -1535,6 +1556,7 @@ class TestNotionExtractorErrorScenarios: notion_page_type="page", tenant_id="tenant-789", notion_access_token="test-token", + notion_token_loader=lambda: "token", ) @pytest.mark.parametrize( @@ -1680,6 +1702,7 @@ class TestNotionExtractorTableAdvanced: notion_page_type="page", tenant_id="tenant-789", notion_access_token="test-token", + notion_token_loader=lambda: "token", ) @patch("httpx.request") diff --git a/api/tests/unit_tests/core/datasource/test_stream_node_events.py b/api/tests/unit_tests/core/datasource/test_stream_node_events.py index 888c1a8d9ee..c05f9092a25 100644 --- a/api/tests/unit_tests/core/datasource/test_stream_node_events.py +++ b/api/tests/unit_tests/core/datasource/test_stream_node_events.py @@ -33,14 +33,12 @@ def test_stream_node_events_accumulates_variables(mocker: MockerFixture) -> None datasource_type="online_document", provider_id="p/x", tenant_id="t", - provider="prov", - plugin_id="plug", - credential_id="", parameters_for_log={}, datasource_info={"user_id": "u"}, variable_pool=mocker.Mock(), datasource_param=type("P", (), {"workspace_id": "w", "page_id": "pg", "type": "t"})(), online_drive_request=None, + credentials={}, ) ) assert isinstance(events[-1], StreamCompletedEvent) diff --git a/api/tests/unit_tests/core/datasource/test_website_crawl.py b/api/tests/unit_tests/core/datasource/test_website_crawl.py index 53000881ddc..d589acb747b 100644 --- a/api/tests/unit_tests/core/datasource/test_website_crawl.py +++ b/api/tests/unit_tests/core/datasource/test_website_crawl.py @@ -12,11 +12,12 @@ The tests cover multiple crawl providers (Firecrawl, WaterCrawl, JinaReader) and ensure proper handling of crawl options, status checking, and data retrieval. """ -from unittest.mock import Mock, patch +from unittest.mock import Mock, create_autospec, patch import pytest from pytest_mock import MockerFixture +import services.data_source.website_service as website_service_module from core.datasource.entities.datasource_entities import ( DatasourceEntity, DatasourceIdentity, @@ -27,7 +28,8 @@ from core.datasource.entities.datasource_entities import ( from core.datasource.website_crawl.website_crawl_plugin import WebsiteCrawlDatasourcePlugin from core.datasource.website_crawl.website_crawl_provider import WebsiteCrawlDatasourcePluginProviderController from core.rag.extractor.watercrawl.provider import WaterCrawlProvider -from services.website_service import CrawlOptions, CrawlRequest, WebsiteService +from services.data_source.credential_gateway import DatasourceProviderCredentialStore +from services.data_source.website_service import CrawlOptions, CrawlRequest, WebsiteService # ============================================================================ # Fixtures @@ -489,8 +491,8 @@ class TestCrawlStatus: class TestWebsiteService: """Test suite for WebsiteService with multiple providers.""" - @patch("services.website_service.current_user") - @patch("services.website_service.DatasourceProviderService") + @patch("services.data_source.website_service.current_user") + @patch("services.data_source.website_service.DatasourceProviderService") def test_crawl_url_firecrawl(self, mock_provider_service: Mock, mock_current_user: Mock, mocker: MockerFixture): """Test crawling with Firecrawl provider.""" # Setup mocks @@ -500,14 +502,14 @@ class TestWebsiteService: "base_url": "https://api.firecrawl.dev", } - mock_firecrawl = mocker.patch("services.website_service.FirecrawlApp") + mock_firecrawl = mocker.patch("services.data_source.website_service.FirecrawlApp") mock_firecrawl_instance = mock_firecrawl.return_value mock_firecrawl_instance.crawl_url.return_value = "job-123" # Mock redis - mocker.patch("services.website_service.redis_client") + mocker.patch("services.data_source.website_service.redis_client") - from services.website_service import WebsiteCrawlApiRequest + from services.data_source.website_service import WebsiteCrawlApiRequest api_request = WebsiteCrawlApiRequest( provider="firecrawl", @@ -515,13 +517,17 @@ class TestWebsiteService: options={"limit": 10, "crawl_sub_pages": True, "only_main_content": True}, ) - result = WebsiteService.crawl_url(api_request) + result = WebsiteService( + providers=website_service_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ) + ).crawl_url(api_request) assert result["status"] == "active" assert result["job_id"] == "job-123" - @patch("services.website_service.current_user") - @patch("services.website_service.DatasourceProviderService") + @patch("services.data_source.website_service.current_user") + @patch("services.data_source.website_service.DatasourceProviderService") def test_crawl_url_watercrawl(self, mock_provider_service: Mock, mock_current_user: Mock, mocker: MockerFixture): """Test crawling with WaterCrawl provider.""" # Setup mocks @@ -531,11 +537,11 @@ class TestWebsiteService: "base_url": "https://app.watercrawl.dev", } - mock_watercrawl = mocker.patch("services.website_service.WaterCrawlProvider") + mock_watercrawl = mocker.patch("services.data_source.website_service.WaterCrawlProvider") mock_watercrawl_instance = mock_watercrawl.return_value mock_watercrawl_instance.crawl_url.return_value = {"status": "active", "job_id": "job-456"} - from services.website_service import WebsiteCrawlApiRequest + from services.data_source.website_service import WebsiteCrawlApiRequest api_request = WebsiteCrawlApiRequest( provider="watercrawl", @@ -543,13 +549,17 @@ class TestWebsiteService: options={"limit": 20, "crawl_sub_pages": True, "max_depth": 2}, ) - result = WebsiteService.crawl_url(api_request) + result = WebsiteService( + providers=website_service_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ) + ).crawl_url(api_request) assert result["status"] == "active" assert result["job_id"] == "job-456" - @patch("services.website_service.current_user") - @patch("services.website_service.DatasourceProviderService") + @patch("services.data_source.website_service.current_user") + @patch("services.data_source.website_service.DatasourceProviderService") def test_crawl_url_jinareader(self, mock_provider_service: Mock, mock_current_user: Mock, mocker: MockerFixture): """Test crawling with JinaReader provider.""" # Setup mocks @@ -561,11 +571,11 @@ class TestWebsiteService: mock_response = Mock() mock_response.json.return_value = {"code": 200, "data": {"taskId": "task-789"}} mock_httpx_post = mocker.patch( - "services.website_service._adaptive_http_client.post", + "services.data_source.website_service._adaptive_http_client.post", return_value=mock_response, ) - from services.website_service import WebsiteCrawlApiRequest + from services.data_source.website_service import WebsiteCrawlApiRequest api_request = WebsiteCrawlApiRequest( provider="jinareader", @@ -573,7 +583,11 @@ class TestWebsiteService: options={"limit": 15, "crawl_sub_pages": True, "use_sitemap": True}, ) - result = WebsiteService.crawl_url(api_request) + result = WebsiteService( + providers=website_service_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ) + ).crawl_url(api_request) assert result["status"] == "active" assert result["job_id"] == "task-789" @@ -855,11 +869,11 @@ class TestContentExtraction: class TestErrorHandling: """Test suite for error handling in crawl operations.""" - @patch("services.website_service.current_user") - @patch("services.website_service.DatasourceProviderService") + @patch("services.data_source.website_service.current_user") + @patch("services.data_source.website_service.DatasourceProviderService") def test_invalid_provider_error(self, mock_provider_service: Mock, mock_current_user: Mock): """Test that invalid provider raises ValueError.""" - from services.website_service import WebsiteCrawlApiRequest + from services.data_source.website_service import WebsiteCrawlApiRequest # Setup mocks mock_current_user.current_tenant_id = "test_tenant" @@ -873,7 +887,11 @@ class TestErrorHandling: # The error should be raised when trying to crawl with invalid provider with pytest.raises(ValueError, match="Invalid provider"): - WebsiteService.crawl_url(api_request) + WebsiteService( + providers=website_service_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ) + ).crawl_url(api_request) def test_missing_api_key_error(self, mocker: MockerFixture): """Test that missing API key is handled properly at the httpx client level.""" @@ -1267,8 +1285,8 @@ class TestProviderSpecificFeatures: unique features and API behaviors that require specific testing. """ - @patch("services.website_service.current_user") - @patch("services.website_service.DatasourceProviderService") + @patch("services.data_source.website_service.current_user") + @patch("services.data_source.website_service.DatasourceProviderService") def test_firecrawl_with_prompt_parameter( self, mock_provider_service: Mock, mock_current_user: Mock, mocker: MockerFixture ): @@ -1285,14 +1303,14 @@ class TestProviderSpecificFeatures: "base_url": "https://api.firecrawl.dev", } - mock_firecrawl = mocker.patch("services.website_service.FirecrawlApp") + mock_firecrawl = mocker.patch("services.data_source.website_service.FirecrawlApp") mock_firecrawl_instance = mock_firecrawl.return_value mock_firecrawl_instance.crawl_url.return_value = "prompt-job-123" # Mock redis - mocker.patch("services.website_service.redis_client") + mocker.patch("services.data_source.website_service.redis_client") - from services.website_service import WebsiteCrawlApiRequest + from services.data_source.website_service import WebsiteCrawlApiRequest # Include a prompt for AI-guided extraction api_request = WebsiteCrawlApiRequest( @@ -1306,7 +1324,11 @@ class TestProviderSpecificFeatures: }, ) - result = WebsiteService.crawl_url(api_request) + result = WebsiteService( + providers=website_service_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ) + ).crawl_url(api_request) assert result["status"] == "active" assert result["job_id"] == "prompt-job-123" @@ -1317,8 +1339,8 @@ class TestProviderSpecificFeatures: assert "prompt" in params assert params["prompt"] == "Extract only technical documentation and API references" - @patch("services.website_service.current_user") - @patch("services.website_service.DatasourceProviderService") + @patch("services.data_source.website_service.current_user") + @patch("services.data_source.website_service.DatasourceProviderService") def test_jinareader_single_page_mode( self, mock_provider_service: Mock, mock_current_user: Mock, mocker: MockerFixture ): @@ -1343,23 +1365,27 @@ class TestProviderSpecificFeatures: "url": "https://example.com/page", }, } - mocker.patch("services.website_service._jina_http_client.get", return_value=mock_response) + mocker.patch("services.data_source.website_service._jina_http_client.get", return_value=mock_response) - from services.website_service import WebsiteCrawlApiRequest + from services.data_source.website_service import WebsiteCrawlApiRequest # Single page mode (crawl_sub_pages = False) api_request = WebsiteCrawlApiRequest( provider="jinareader", url="https://example.com/page", options={"crawl_sub_pages": False, "limit": 1} ) - result = WebsiteService.crawl_url(api_request) + result = WebsiteService( + providers=website_service_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ) + ).crawl_url(api_request) # In single-page mode, JinaReader returns data immediately assert result["status"] == "active" assert "data" in result - @patch("services.website_service.current_user") - @patch("services.website_service.DatasourceProviderService") + @patch("services.data_source.website_service.current_user") + @patch("services.data_source.website_service.DatasourceProviderService") def test_watercrawl_with_tag_filtering( self, mock_provider_service: Mock, mock_current_user: Mock, mocker: MockerFixture ): @@ -1376,11 +1402,11 @@ class TestProviderSpecificFeatures: "base_url": "https://app.watercrawl.dev", } - mock_watercrawl = mocker.patch("services.website_service.WaterCrawlProvider") + mock_watercrawl = mocker.patch("services.data_source.website_service.WaterCrawlProvider") mock_watercrawl_instance = mock_watercrawl.return_value mock_watercrawl_instance.crawl_url.return_value = {"status": "active", "job_id": "tag-filter-job"} - from services.website_service import WebsiteCrawlApiRequest + from services.data_source.website_service import WebsiteCrawlApiRequest # Configure with tag filtering api_request = WebsiteCrawlApiRequest( @@ -1394,7 +1420,11 @@ class TestProviderSpecificFeatures: }, ) - result = WebsiteService.crawl_url(api_request) + result = WebsiteService( + providers=website_service_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ) + ).crawl_url(api_request) assert result["status"] == "active" assert result["job_id"] == "tag-filter-job" @@ -1451,7 +1481,7 @@ class TestDataStructureValidation: This conversion ensures that external API parameters are properly mapped to internal data structures. """ - from services.website_service import WebsiteCrawlApiRequest + from services.data_source.website_service import WebsiteCrawlApiRequest # Create API request with all options api_request = WebsiteCrawlApiRequest( @@ -1634,7 +1664,7 @@ class TestEdgeCasesAndBoundaries: Empty or invalid URLs should fail validation before attempting to crawl. """ - from services.website_service import WebsiteCrawlApiRequest + from services.data_source.website_service import WebsiteCrawlApiRequest # Empty URL should raise ValueError during validation with pytest.raises(ValueError, match="URL is required"): diff --git a/api/tests/unit_tests/core/rag/datasource/keyword/jieba/test_jieba.py b/api/tests/unit_tests/core/rag/datasource/keyword/jieba/test_jieba.py index 2cfbfd5db68..41edd48db92 100644 --- a/api/tests/unit_tests/core/rag/datasource/keyword/jieba/test_jieba.py +++ b/api/tests/unit_tests/core/rag/datasource/keyword/jieba/test_jieba.py @@ -1,407 +1,275 @@ -import json -from types import SimpleNamespace -from typing import Any -from unittest.mock import MagicMock +from collections.abc import Generator +from contextlib import contextmanager +from dataclasses import dataclass, field import pytest -from sqlalchemy import event, select +from sqlalchemy import select from sqlalchemy.orm import Session -import core.rag.datasource.keyword.jieba.jieba as jieba_module -from core.rag.datasource.keyword.jieba.jieba import Jieba, dumps_with_sets, set_orjson_default +from core.rag.datasource.keyword.jieba import jieba as jieba_module +from core.rag.datasource.keyword.jieba.jieba import Jieba from core.rag.models.document import Document from models.dataset import Dataset, DatasetKeywordTable, DocumentSegment +from repositories.knowledge.dataset_read_repository import get_dataset_keyword_table +from repositories.knowledge.keyword_table_repository import load_keyword_table +from tests.unit_tests.config_override import apply_config_overrides from tests.unit_tests.model_factories import make_dataset -class _DummyLock: - def __enter__(self): - return self +@dataclass +class KeywordStorage: + files: dict[str, bytes] = field(default_factory=dict) + reads: list[str] = field(default_factory=list) + writes: list[str] = field(default_factory=list) - def __exit__(self, exc_type, exc, tb): - return False + def load_once(self, key: str) -> bytes: + self.reads.append(key) + if key not in self.files: + raise FileNotFoundError(key) + return self.files[key] + + def exists(self, key: str) -> bool: + return key in self.files + + def save(self, key: str, data: bytes) -> None: + self.files[key] = data + self.writes.append(key) + + def delete(self, key: str) -> None: + del self.files[key] -def _dataset_keyword_table( - data_source_type: str = "database", keyword_table_dict: dict[str, Any] | None = None -) -> DatasetKeywordTable: - keyword_table = DatasetKeywordTable( - dataset_id="dataset-1", - data_source_type=data_source_type, - keyword_table="", - ) - keyword_table.get_keyword_table_dict = MagicMock(return_value=keyword_table_dict) - return keyword_table +@dataclass +class KeywordLocks: + acquired: list[str] = field(default_factory=list) + held: bool = False + + @contextmanager + def lock(self, name: str, *, timeout: int) -> Generator[None]: + assert timeout == 600 + assert not self.held + self.acquired.append(name) + self.held = True + try: + yield + finally: + self.held = False -def _dataset(dataset_keyword_table: DatasetKeywordTable | None = None, keyword_number: int | None = None) -> Dataset: - dataset = make_dataset(keyword_number=keyword_number) - dataset.get_dataset_keyword_table = MagicMock(return_value=dataset_keyword_table) - return dataset +@dataclass +class KeywordRuntime: + session: Session + dataset: Dataset + storage: KeywordStorage + locks: KeywordLocks + extracted: list[tuple[str, int]] + + def table(self) -> dict[str, set[str]]: + row = get_dataset_keyword_table(self.dataset, session=self.session) + assert row is not None + payload = load_keyword_table( + storage=self.storage, + tenant_id=self.dataset.tenant_id, + dataset_id=self.dataset.id, + storage_type=row.data_source_type, + data=row.keyword_table, + ) + assert payload is not None + return payload["__data__"]["table"] -@pytest.fixture -def patched_runtime(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session): - storage = MagicMock() - lock = MagicMock(return_value=_DummyLock()) - redis_client = SimpleNamespace(lock=lock) +@pytest.fixture(params=["database", "file"]) +def runtime(request: pytest.FixtureRequest, sqlite_session: Session, monkeypatch: pytest.MonkeyPatch) -> KeywordRuntime: + storage = KeywordStorage() + locks = KeywordLocks() + extracted: list[tuple[str, int]] = [] + class KeywordExtractor: + def extract_keywords(self, text: str, keyword_number: int = 10) -> set[str]: + extracted.append((text, keyword_number)) + return set(text.split()[:keyword_number]) + + monkeypatch.setattr(jieba_module, "JiebaKeywordTableHandler", KeywordExtractor) + monkeypatch.setattr(jieba_module, "redis_client", locks) monkeypatch.setattr(jieba_module, "storage", storage) - monkeypatch.setattr(jieba_module, "redis_client", redis_client) - - return SimpleNamespace(session=sqlite_session, storage=storage, lock=lock) + apply_config_overrides(monkeypatch, KEYWORD_DATA_SOURCE_TYPE=request.param) + dataset = make_dataset(name="Test", created_by="author", keyword_number=2) + sqlite_session.add(dataset) + sqlite_session.commit() + return KeywordRuntime(sqlite_session, dataset, storage, locks, extracted) -def _segment(*, index_node_id: str = "node-2") -> DocumentSegment: - segment = DocumentSegment( - tenant_id="tenant-1", +def segment(node_id: str, *, tenant_id: str = "tenant-1") -> DocumentSegment: + return DocumentSegment( + tenant_id=tenant_id, dataset_id="dataset-1", - document_id="doc-2", + document_id="document-1", position=1, - content="segment-content", - word_count=1, - tokens=1, - created_by="user-1", - enabled=True, - keywords=[], - answer=None, - index_node_id=index_node_id, - index_node_hash="hash-2", - status="completed", + content=node_id, + word_count=len(node_id), + tokens=0, + created_by="author", + index_node_id=node_id, ) - segment.id = "segment-1" - return segment -def test_create_indexes_documents_and_returns_self(monkeypatch: pytest.MonkeyPatch, patched_runtime): - dataset = _dataset(_dataset_keyword_table(), keyword_number=2) - keyword = Jieba(dataset) - handler = MagicMock() - handler.extract_keywords.return_value = {"kw1", "kw2"} - - monkeypatch.setattr(jieba_module, "JiebaKeywordTableHandler", lambda: handler) - monkeypatch.setattr(keyword, "_get_dataset_keyword_table", MagicMock(return_value={})) - monkeypatch.setattr(keyword, "_update_segment_keywords", MagicMock()) - monkeypatch.setattr(keyword, "_save_dataset_keyword_table", MagicMock()) - +def test_create_extracts_keywords_and_persists_the_shared_format(runtime: KeywordRuntime) -> None: + chunk = segment("node-1") + foreign = segment("node-1", tenant_id="tenant-2") + runtime.session.add_all([chunk, foreign]) + keyword = Jieba(runtime.dataset) result = keyword.create( - [ - Document(page_content="alpha", metadata={"doc_id": "node-1"}), - SimpleNamespace(page_content="ignored", metadata=None), - ], - patched_runtime.session, + [Document(page_content="中文 alpha ignored", metadata={"doc_id": "node-1"})], runtime.session ) + runtime.session.commit() assert result is keyword - keyword._update_segment_keywords.assert_called_once() - call_args = keyword._update_segment_keywords.call_args.args - assert call_args[0] == "dataset-1" - assert call_args[1] == "node-1" - assert set(call_args[2]) == {"kw1", "kw2"} - assert call_args[3] is patched_runtime.session - saved_table = keyword._save_dataset_keyword_table.call_args.args[0] - assert saved_table["kw1"] == {"node-1"} - assert saved_table["kw2"] == {"node-1"} - patched_runtime.lock.assert_called_once_with("keyword_indexing_lock_dataset-1", timeout=600) - - -def test_add_texts_supports_keywords_list_and_extract_fallback(monkeypatch: pytest.MonkeyPatch, patched_runtime): - keyword = Jieba(_dataset(_dataset_keyword_table(), keyword_number=3)) - handler = MagicMock() - handler.extract_keywords.return_value = {"auto"} - - monkeypatch.setattr(jieba_module, "JiebaKeywordTableHandler", lambda: handler) - monkeypatch.setattr(keyword, "_get_dataset_keyword_table", MagicMock(return_value={})) - monkeypatch.setattr(keyword, "_update_segment_keywords", MagicMock()) - monkeypatch.setattr(keyword, "_save_dataset_keyword_table", MagicMock()) - - texts = [ - Document(page_content="extract-this", metadata={"doc_id": "node-1"}), - Document(page_content="use-manual", metadata={"doc_id": "node-2"}), - ] - keyword.add_texts(texts, patched_runtime.session, keywords_list=[[], ["manual"]]) - - assert keyword._update_segment_keywords.call_count == 2 - first_call = keyword._update_segment_keywords.call_args_list[0].args - second_call = keyword._update_segment_keywords.call_args_list[1].args - assert set(first_call[2]) == {"auto"} - assert second_call[2] == ["manual"] - assert first_call[3] is patched_runtime.session - assert second_call[3] is patched_runtime.session - keyword._save_dataset_keyword_table.assert_called_once_with( - {"auto": {"node-1"}, "manual": {"node-2"}}, patched_runtime.session + assert set(chunk.keywords) == {"中文", "alpha"} + assert foreign.keywords is None + assert runtime.table() == {"中文": {"node-1"}, "alpha": {"node-1"}} + assert runtime.extracted == [("中文 alpha ignored", 2)] + assert runtime.locks.acquired == ["keyword_indexing_lock_dataset-1"] + assert runtime.locks.held is False + row = get_dataset_keyword_table(runtime.dataset, session=runtime.session) + assert row is not None + payload = load_keyword_table( + storage=runtime.storage, + tenant_id=runtime.dataset.tenant_id, + dataset_id=runtime.dataset.id, + storage_type=row.data_source_type, + data=row.keyword_table, ) + assert payload is not None + assert payload["__type__"] == "keyword_table" + assert payload["__data__"]["index_id"] == "dataset-1" + if row.data_source_type == "file": + assert row.keyword_table == "" + assert runtime.storage.writes == ["keyword_files/tenant-1/dataset-1.txt"] + else: + assert runtime.storage.reads == [] + assert runtime.storage.writes == [] -def test_add_texts_without_keywords_list_always_uses_extractor(monkeypatch: pytest.MonkeyPatch, patched_runtime): - keyword = Jieba(_dataset(_dataset_keyword_table(), keyword_number=1)) - handler = MagicMock() - handler.extract_keywords.return_value = {"from-extractor"} - - monkeypatch.setattr(jieba_module, "JiebaKeywordTableHandler", lambda: handler) - monkeypatch.setattr(keyword, "_get_dataset_keyword_table", MagicMock(return_value={})) - monkeypatch.setattr(keyword, "_update_segment_keywords", MagicMock()) - monkeypatch.setattr(keyword, "_save_dataset_keyword_table", MagicMock()) - - keyword.add_texts([Document(page_content="content", metadata={"doc_id": "node-1"})], patched_runtime.session) - - handler.extract_keywords.assert_called_once_with("content", 1) - assert set(keyword._update_segment_keywords.call_args.args[2]) == {"from-extractor"} - assert keyword._update_segment_keywords.call_args.args[3] is patched_runtime.session - - -def test_text_exists_handles_missing_and_existing_keyword_table( - monkeypatch: pytest.MonkeyPatch, unbound_session: Session -): - keyword = Jieba(_dataset(_dataset_keyword_table(keyword_table_dict=None))) - session = unbound_session - assert keyword.text_exists("node-1", session=session) is False - - keyword = Jieba( - _dataset( - _dataset_keyword_table( - keyword_table_dict={"__type__": "keyword_table", "__data__": {"table": {"k": {"node-1", "node-2"}}}} - ) - ) - ) - assert keyword.text_exists("node-2", session=session) is True - assert keyword.text_exists("node-x", session=session) is False - - -def test_delete_by_ids_updates_table_when_present(monkeypatch: pytest.MonkeyPatch, patched_runtime): - keyword = Jieba( - _dataset( - _dataset_keyword_table( - keyword_table_dict={"__type__": "keyword_table", "__data__": {"table": {"k": {"node-1", "node-2"}}}} - ) - ) - ) - monkeypatch.setattr(keyword, "_get_dataset_keyword_table", MagicMock(return_value={"k": {"node-1", "node-2"}})) - monkeypatch.setattr(keyword, "_delete_ids_from_keyword_table", MagicMock(return_value={"k": {"node-2"}})) - monkeypatch.setattr(keyword, "_save_dataset_keyword_table", MagicMock()) - - keyword.delete_by_ids(["node-1"], patched_runtime.session) - - keyword._get_dataset_keyword_table.assert_called_once_with(patched_runtime.session) - keyword._delete_ids_from_keyword_table.assert_called_once_with({"k": {"node-1", "node-2"}}, ["node-1"]) - keyword._save_dataset_keyword_table.assert_called_once_with({"k": {"node-2"}}, patched_runtime.session) - - -def test_delete_by_ids_saves_none_when_keyword_table_is_missing(monkeypatch: pytest.MonkeyPatch, patched_runtime): - keyword = Jieba(_dataset(_dataset_keyword_table())) - monkeypatch.setattr(keyword, "_get_dataset_keyword_table", MagicMock(return_value=None)) - monkeypatch.setattr(keyword, "_delete_ids_from_keyword_table", MagicMock()) - monkeypatch.setattr(keyword, "_save_dataset_keyword_table", MagicMock()) - - keyword.delete_by_ids(["node-1"], patched_runtime.session) - - keyword._get_dataset_keyword_table.assert_called_once_with(patched_runtime.session) - keyword._delete_ids_from_keyword_table.assert_not_called() - keyword._save_dataset_keyword_table.assert_called_once_with(None, patched_runtime.session) - - -def test_search_returns_documents_in_rank_order_and_applies_filter(monkeypatch: pytest.MonkeyPatch, patched_runtime): - keyword = Jieba(_dataset(_dataset_keyword_table())) - patched_runtime.session.add(_segment()) - patched_runtime.session.flush() - monkeypatch.setattr(keyword, "_retrieve_ids_by_query", MagicMock(return_value=["node-1", "node-2"])) - - documents = keyword.search("query", session=patched_runtime.session, top_k=2, document_ids_filter=["doc-2"]) - - assert len(documents) == 1 - assert documents[0].page_content == "segment-content" - assert documents[0].metadata["doc_id"] == "node-2" - assert documents[0].metadata["doc_hash"] == "hash-2" - - -def test_delete_removes_keyword_table_and_optional_file(patched_runtime): - db_keyword = DatasetKeywordTable(dataset_id="dataset-1", keyword_table="", data_source_type="database") - patched_runtime.session.add(db_keyword) - patched_runtime.session.commit() - commits: list[str] = [] - event.listen(patched_runtime.session, "after_commit", lambda _session: commits.append("commit")) - - keyword_db = Jieba(_dataset(db_keyword)) - keyword_db.delete(session=patched_runtime.session) - patched_runtime.storage.delete.assert_not_called() - assert patched_runtime.session.get(DatasetKeywordTable, db_keyword.id) is None - - file_keyword = DatasetKeywordTable(dataset_id="dataset-1", keyword_table="", data_source_type="object_storage") - patched_runtime.session.add(file_keyword) - patched_runtime.session.commit() - keyword_file = Jieba(_dataset(file_keyword)) - keyword_file.delete(session=patched_runtime.session) - - patched_runtime.storage.delete.assert_called_once_with("keyword_files/tenant-1/dataset-1.txt") - assert patched_runtime.session.get(DatasetKeywordTable, file_keyword.id) is None - assert commits == ["commit", "commit", "commit"] - - -def test_save_dataset_keyword_table_to_database(patched_runtime): - dataset_keyword_table = DatasetKeywordTable(dataset_id="dataset-1", keyword_table="", data_source_type="database") - patched_runtime.session.add(dataset_keyword_table) - patched_runtime.session.flush() - keyword = Jieba(_dataset(dataset_keyword_table)) - - keyword._save_dataset_keyword_table({"kw": {"node-1"}}, patched_runtime.session) - - assert '"__type__":"keyword_table"' in dataset_keyword_table.keyword_table - assert '"index_id":"dataset-1"' in dataset_keyword_table.keyword_table - - -def test_save_dataset_keyword_table_to_file_storage(patched_runtime): - dataset_keyword_table = DatasetKeywordTable(dataset_id="dataset-1", keyword_table="", data_source_type="file") - patched_runtime.session.add(dataset_keyword_table) - patched_runtime.session.flush() - keyword = Jieba(_dataset(dataset_keyword_table)) - patched_runtime.storage.exists.return_value = True - - keyword._save_dataset_keyword_table({"kw": {"node-1"}}, patched_runtime.session) - - patched_runtime.storage.delete.assert_called_once_with("keyword_files/tenant-1/dataset-1.txt") - patched_runtime.storage.save.assert_called_once() - save_args = patched_runtime.storage.save.call_args.args - assert save_args[0] == "keyword_files/tenant-1/dataset-1.txt" - assert isinstance(save_args[1], bytes) - - -def test_get_dataset_keyword_table_returns_existing_table_data(patched_runtime): - existing = DatasetKeywordTable( - dataset_id="dataset-1", - keyword_table="", - data_source_type="database", - ) - existing.get_keyword_table_dict = MagicMock( - return_value={"__type__": "keyword_table", "__data__": {"table": {"kw": ["node-1"]}}} - ) - patched_runtime.session.add(existing) - patched_runtime.session.flush() - keyword = Jieba(_dataset(existing)) - assert keyword._get_dataset_keyword_table(patched_runtime.session) == {"kw": ["node-1"]} - - existing.get_keyword_table_dict = MagicMock(return_value=None) - keyword_with_missing_payload = Jieba(_dataset(existing)) - assert keyword_with_missing_payload._get_dataset_keyword_table(patched_runtime.session) == {} - - -def test_get_dataset_keyword_table_creates_table_when_missing(monkeypatch: pytest.MonkeyPatch, patched_runtime): - keyword = Jieba(_dataset(dataset_keyword_table=None)) - from tests.unit_tests.config_override import apply_config_overrides - - apply_config_overrides(monkeypatch, KEYWORD_DATA_SOURCE_TYPE="database") - result = keyword._get_dataset_keyword_table(patched_runtime.session) - - assert result == {} - created_table = patched_runtime.session.scalar( - select(DatasetKeywordTable).where(DatasetKeywordTable.dataset_id == "dataset-1") - ) - assert created_table is not None - assert created_table.dataset_id == "dataset-1" - assert created_table.data_source_type == "database" - assert '"index_id":"dataset-1"' in created_table.keyword_table - - -def test_add_and_delete_ids_from_keyword_table_helpers(): - keyword = Jieba(_dataset(_dataset_keyword_table())) - keyword_table = {"kw1": {"node-1"}, "kw2": {"node-1", "node-2"}} - - updated = keyword._add_text_to_keyword_table(keyword_table, "node-3", ["kw1", "kw3"]) - assert updated["kw1"] == {"node-1", "node-3"} - assert updated["kw3"] == {"node-3"} - - deleted = keyword._delete_ids_from_keyword_table(updated, ["node-1", "node-3"]) - assert "kw3" not in deleted - assert "kw1" not in deleted - assert deleted["kw2"] == {"node-2"} - - -def test_retrieve_ids_by_query_ranks_by_keyword_frequency(monkeypatch: pytest.MonkeyPatch): - keyword = Jieba(_dataset(_dataset_keyword_table())) - handler = MagicMock() - handler.extract_keywords.return_value = ["kw-a", "kw-b"] - monkeypatch.setattr(jieba_module, "JiebaKeywordTableHandler", lambda: handler) - - ranked_ids = keyword._retrieve_ids_by_query( - {"kw-a": {"node-1", "node-2"}, "kw-b": {"node-2"}, "kw-c": {"node-3"}}, - "query", - k=1, - ) - - assert ranked_ids == ["node-2"] - - -def test_update_segment_keywords_updates_when_segment_exists(patched_runtime): - keyword = Jieba(_dataset(_dataset_keyword_table())) - segment = _segment(index_node_id="node-1") - patched_runtime.session.add(segment) - patched_runtime.session.flush() - - keyword._update_segment_keywords("dataset-1", "node-1", ["kw1", "kw2"], patched_runtime.session) - - assert segment.keywords == ["kw1", "kw2"] - - keyword._update_segment_keywords("dataset-1", "node-missing", ["kw3"], patched_runtime.session) - assert segment.keywords == ["kw1", "kw2"] - - -def test_create_segment_keywords_and_update_segment_keywords_index(monkeypatch: pytest.MonkeyPatch, patched_runtime): - keyword = Jieba(_dataset(_dataset_keyword_table())) - monkeypatch.setattr(keyword, "_get_dataset_keyword_table", MagicMock(return_value={})) - monkeypatch.setattr(keyword, "_update_segment_keywords", MagicMock()) - monkeypatch.setattr(keyword, "_save_dataset_keyword_table", MagicMock()) - - keyword.create_segment_keywords("node-1", ["kw"], patched_runtime.session) - keyword._get_dataset_keyword_table.assert_called_once_with(patched_runtime.session) - keyword._update_segment_keywords.assert_called_once_with("dataset-1", "node-1", ["kw"], patched_runtime.session) - keyword._save_dataset_keyword_table.assert_called_once_with({"kw": {"node-1"}}, patched_runtime.session) - - keyword._get_dataset_keyword_table.reset_mock() - keyword._save_dataset_keyword_table.reset_mock() - keyword.update_segment_keywords_index("node-2", ["kw2"], patched_runtime.session) - keyword._get_dataset_keyword_table.assert_called_once_with(patched_runtime.session) - keyword._save_dataset_keyword_table.assert_called_once_with({"kw2": {"node-2"}}, patched_runtime.session) - - -def test_multi_create_segment_keywords_uses_provided_and_extracted_keywords( - monkeypatch: pytest.MonkeyPatch, patched_runtime -): - keyword = Jieba(_dataset(_dataset_keyword_table(), keyword_number=2)) - handler = MagicMock() - handler.extract_keywords.return_value = {"auto"} - monkeypatch.setattr(jieba_module, "JiebaKeywordTableHandler", lambda: handler) - monkeypatch.setattr(keyword, "_get_dataset_keyword_table", MagicMock(return_value={})) - monkeypatch.setattr(keyword, "_save_dataset_keyword_table", MagicMock()) - - first_segment = _segment(index_node_id="node-1") - first_segment.content = "first content" - first_segment.keywords = None - second_segment = _segment(index_node_id="node-2") - second_segment.content = "second content" - second_segment.keywords = None - - keyword.multi_create_segment_keywords( +def test_add_uses_manual_keywords_and_extracts_empty_selections(runtime: KeywordRuntime) -> None: + chunks = [segment("node-1"), segment("node-2")] + runtime.session.add_all(chunks) + keyword = Jieba(runtime.dataset) + keyword.add_texts( [ - {"segment": first_segment, "keywords": ["manual"]}, - {"segment": second_segment, "keywords": []}, + Document(page_content="automatic", metadata={"doc_id": "node-1"}), + Document(page_content="ignored", metadata={"doc_id": "node-2"}), ], - patched_runtime.session, + runtime.session, + keywords_list=[[], ["manual"]], ) - - assert first_segment.keywords == ["manual"] - assert second_segment.keywords == ["auto"] - saved_table = keyword._save_dataset_keyword_table.call_args.args[0] - assert saved_table["manual"] == {"node-1"} - assert saved_table["auto"] == {"node-2"} - assert keyword._save_dataset_keyword_table.call_args.args[1] is patched_runtime.session + runtime.session.commit() + assert runtime.table() == {"automatic": {"node-1"}, "manual": {"node-2"}} + assert runtime.extracted == [("automatic", 2)] + assert chunks[0].keywords == ["automatic"] + assert chunks[1].keywords == ["manual"] -def test_set_orjson_default_and_dumps_with_sets(): - assert set(set_orjson_default({"a", "b"})) == {"a", "b"} +def test_delete_ids_preserves_unrelated_entries(runtime: KeywordRuntime) -> None: + keyword = Jieba(runtime.dataset) + keyword.add_texts( + [ + Document(page_content="shared first", metadata={"doc_id": "node-1"}), + Document(page_content="shared second", metadata={"doc_id": "node-2"}), + ], + runtime.session, + ) + runtime.session.commit() + keyword.delete_by_ids(["node-1"], runtime.session) + runtime.session.commit() + assert runtime.table() == {"shared": {"node-2"}, "second": {"node-2"}} + assert keyword.text_exists("node-1", session=runtime.session) is False + assert keyword.text_exists("node-2", session=runtime.session) is True + keyword.delete_by_ids(["node-2"], runtime.session) + runtime.session.commit() + assert runtime.table() == {} + assert keyword.text_exists("node-2", session=runtime.session) is False - with pytest.raises(TypeError, match="is not JSON serializable"): - set_orjson_default(("not", "a", "set")) - payload = {"items": {"a", "b"}} - json_payload = dumps_with_sets(payload) - decoded = json.loads(json_payload) - assert set(decoded["items"]) == {"a", "b"} +def test_delete_without_an_existing_table_is_idempotent(runtime: KeywordRuntime) -> None: + keyword = Jieba(runtime.dataset) + assert keyword.text_exists("missing", session=runtime.session) is False + keyword.delete_by_ids(["missing"], runtime.session) + runtime.session.commit() + keyword.delete_by_ids(["missing"], runtime.session) + runtime.session.commit() + assert runtime.table() == {} + + +def test_search_preserves_ranking_and_document_filters(runtime: KeywordRuntime) -> None: + first, second = segment("node-1"), segment("node-2") + second.document_id = "document-2" + runtime.session.add_all([first, second]) + keyword = Jieba(runtime.dataset) + keyword.add_texts( + [ + Document(page_content="alpha", metadata={"doc_id": "node-1"}), + Document(page_content="alpha beta", metadata={"doc_id": "node-2"}), + ], + runtime.session, + ) + runtime.session.commit() + documents = keyword.search("alpha beta", session=runtime.session, top_k=2) + assert [doc.metadata["doc_id"] for doc in documents] == ["node-2", "node-1"] + documents = keyword.search("alpha beta", session=runtime.session, top_k=2, document_ids_filter=["document-1"]) + assert [doc.metadata["doc_id"] for doc in documents] == ["node-1"] + + +def test_delete_removes_the_database_record_and_file(runtime: KeywordRuntime) -> None: + keyword = Jieba(runtime.dataset) + keyword.create([Document(page_content="alpha", metadata={"doc_id": "node-1"})], runtime.session) + runtime.session.commit() + keyword.delete(session=runtime.session) + assert runtime.session.scalar(select(DatasetKeywordTable)) is None + assert runtime.storage.files == {} + + +def test_mutation_does_not_replace_an_unreadable_table( + runtime: KeywordRuntime, monkeypatch: pytest.MonkeyPatch +) -> None: + keyword = Jieba(runtime.dataset) + keyword.create([Document(page_content="original", metadata={"doc_id": "node-1"})], runtime.session) + runtime.session.commit() + row = get_dataset_keyword_table(runtime.dataset, session=runtime.session) + assert row is not None + stored_files = dict(runtime.storage.files) + writes = list(runtime.storage.writes) + if row.data_source_type == "database": + row.keyword_table = "invalid json" + runtime.session.commit() + expected_error = ValueError + else: + + def fail_read(_key: str) -> bytes: + raise OSError("storage unavailable") + + monkeypatch.setattr(runtime.storage, "load_once", fail_read) + expected_error = OSError + with pytest.raises(expected_error): + keyword.add_texts([Document(page_content="replacement", metadata={"doc_id": "node-1"})], runtime.session) + runtime.session.rollback() + assert runtime.storage.files == stored_files + assert runtime.storage.writes == writes + assert runtime.locks.held is False + if row.data_source_type == "database": + assert row.keyword_table == "invalid json" + + +def test_deletion_does_not_initialize_the_keyword_extractor( + runtime: KeywordRuntime, monkeypatch: pytest.MonkeyPatch +) -> None: + keyword = Jieba(runtime.dataset) + keyword.create([Document(page_content="original", metadata={"doc_id": "node-1"})], runtime.session) + runtime.session.commit() + + def unexpected_extractor() -> None: + pytest.fail("Index cleanup must not initialize the text extractor") + + monkeypatch.setattr(jieba_module, "JiebaKeywordTableHandler", unexpected_extractor) + keyword.delete_by_ids(["node-1"], runtime.session) + runtime.session.commit() + assert runtime.table() == {} diff --git a/api/tests/unit_tests/core/rag/datasource/test_datasource_retrieval.py b/api/tests/unit_tests/core/rag/datasource/test_datasource_retrieval.py index 4c78edf6615..4d81e4424ed 100644 --- a/api/tests/unit_tests/core/rag/datasource/test_datasource_retrieval.py +++ b/api/tests/unit_tests/core/rag/datasource/test_datasource_retrieval.py @@ -899,7 +899,7 @@ class TestRetrievalServiceInternals: monkeypatch.setattr( RetrievalService, "get_segment_attachment_infos", - lambda attachment_ids, session: [ + lambda attachment_ids, *, dataset_refs, session: [ { "attachment_id": "attach-node-1", "attachment_info": { @@ -1172,7 +1172,9 @@ class TestRetrievalServiceInternals: assert result is None def test_get_segment_attachment_infos_returns_empty_when_upload_files_missing(self, sqlite_session: Session): - result = RetrievalService.get_segment_attachment_infos(["upload-1"], sqlite_session) + result = RetrievalService.get_segment_attachment_infos( + ["upload-1"], dataset_refs=(("tenant-id", "dataset-id"),), session=sqlite_session + ) assert result == [] @@ -1180,7 +1182,9 @@ class TestRetrievalServiceInternals: sqlite_session.add(_upload_file(file_id="upload-1", name="file-name")) sqlite_session.commit() - result = RetrievalService.get_segment_attachment_infos(["upload-1"], sqlite_session) + result = RetrievalService.get_segment_attachment_infos( + ["upload-1"], dataset_refs=(("tenant-id", "dataset-id"),), session=sqlite_session + ) assert result == [] @@ -1200,7 +1204,9 @@ class TestRetrievalServiceInternals: sqlite_session.add_all([upload_file_1, upload_file_2, binding]) sqlite_session.commit() - result = RetrievalService.get_segment_attachment_infos(["upload-1", "upload-2"], sqlite_session) + result = RetrievalService.get_segment_attachment_infos( + ["upload-1", "upload-2"], dataset_refs=(("tenant-id", "dataset-id"),), session=sqlite_session + ) assert result == [ { diff --git a/api/tests/unit_tests/core/rag/datasource/test_retrieval_attachment_access.py b/api/tests/unit_tests/core/rag/datasource/test_retrieval_attachment_access.py index 6ebfbf4c053..8f89aca59d0 100644 --- a/api/tests/unit_tests/core/rag/datasource/test_retrieval_attachment_access.py +++ b/api/tests/unit_tests/core/rag/datasource/test_retrieval_attachment_access.py @@ -55,6 +55,7 @@ def test_file_access_grants_ignore_empty_inputs_and_missing_scope() -> None: @pytest.mark.parametrize("sqlite_session", [(UploadFile, SegmentAttachmentBinding)], indirect=True) def test_segment_attachment_lookup_grants_returned_upload_files_to_current_scope(sqlite_session: Session) -> None: tenant_id = str(uuid4()) + dataset_id = str(uuid4()) upload_file_id = str(uuid4()) segment_id = str(uuid4()) user_id = str(uuid4()) @@ -74,7 +75,7 @@ def test_segment_attachment_lookup_grants_returned_upload_files_to_current_scope upload_file.id = upload_file_id binding = SegmentAttachmentBinding( tenant_id=tenant_id, - dataset_id=str(uuid4()), + dataset_id=dataset_id, document_id=str(uuid4()), segment_id=segment_id, attachment_id=upload_file_id, @@ -89,7 +90,11 @@ def test_segment_attachment_lookup_grants_returned_upload_files_to_current_scope ) with bind_file_access_scope(scope): - result = RetrievalService.get_segment_attachment_infos([upload_file_id], sqlite_session) + result = RetrievalService.get_segment_attachment_infos( + [upload_file_id], + dataset_refs=((tenant_id, dataset_id),), + session=sqlite_session, + ) scoped_stmt = DatabaseFileAccessController().apply_upload_file_filters( select(UploadFile).where(UploadFile.id == upload_file_id) ) @@ -100,6 +105,59 @@ def test_segment_attachment_lookup_grants_returned_upload_files_to_current_scope assert "upload_files.id IN" in whereclause +@pytest.mark.parametrize("sqlite_session", [(UploadFile, SegmentAttachmentBinding)], indirect=True) +def test_segment_attachment_lookup_rejects_another_tenants_upload(sqlite_session: Session) -> None: + tenant_id = str(uuid4()) + foreign_tenant_id = str(uuid4()) + upload_file_id = str(uuid4()) + foreign_dataset_id = str(uuid4()) + upload_file = UploadFile( + tenant_id=foreign_tenant_id, + storage_type=StorageType.LOCAL, + key="uploads/foreign.png", + name="foreign.png", + size=1024, + extension="png", + mime_type="image/png", + created_by_role=CreatorUserRole.END_USER, + created_by=str(uuid4()), + created_at=datetime.now(UTC), + used=False, + ) + upload_file.id = upload_file_id + sqlite_session.add_all( + [ + upload_file, + SegmentAttachmentBinding( + tenant_id=foreign_tenant_id, + dataset_id=foreign_dataset_id, + document_id=str(uuid4()), + segment_id=str(uuid4()), + attachment_id=upload_file_id, + ), + ] + ) + sqlite_session.commit() + scope = FileAccessScope( + tenant_id=tenant_id, + user_id=str(uuid4()), + user_from=UserFrom.END_USER, + invoke_from=InvokeFrom.WEB_APP, + ) + + with bind_file_access_scope(scope): + result = RetrievalService.get_segment_attachment_infos( + [upload_file_id], + dataset_refs=((tenant_id, str(uuid4())),), + session=sqlite_session, + ) + current_scope = get_current_file_access_scope() + + assert result == [] + assert current_scope is not None + assert current_scope.granted_upload_file_ids == frozenset() + + @pytest.mark.parametrize("sqlite_session", [(Dataset, DatasetDocument, DocumentSegment)], indirect=True) def test_knowledge_retrieval_grants_returned_segments_to_current_scope( monkeypatch: pytest.MonkeyPatch, diff --git a/api/tests/unit_tests/core/rag/datasource/vdb/test_vector_factory.py b/api/tests/unit_tests/core/rag/datasource/vdb/test_vector_factory.py index 5ed354d15ac..9c4dc266ad1 100644 --- a/api/tests/unit_tests/core/rag/datasource/vdb/test_vector_factory.py +++ b/api/tests/unit_tests/core/rag/datasource/vdb/test_vector_factory.py @@ -431,6 +431,7 @@ def test_create_multimodal_filters_missing_uploads( vector._embeddings.embed_multimodal_documents.return_value = [[0.1, 0.2]] vector._vector_processor = MagicMock() vector._session = sqlite_session + vector._dataset = Dataset(tenant_id=upload_file.tenant_id, name="dataset", created_by=upload_file.created_by) monkeypatch.setattr(vector_factory_module.storage, "load_once", MagicMock(return_value=b"abc")) docs = [ diff --git a/api/tests/unit_tests/core/rag/extractor/firecrawl/test_firecrawl.py b/api/tests/unit_tests/core/rag/extractor/firecrawl/test_firecrawl.py index 3bf1441c0af..cce929d3780 100644 --- a/api/tests/unit_tests/core/rag/extractor/firecrawl/test_firecrawl.py +++ b/api/tests/unit_tests/core/rag/extractor/firecrawl/test_firecrawl.py @@ -3,7 +3,7 @@ import json from collections.abc import Mapping from typing import Any -from unittest.mock import MagicMock +from unittest.mock import MagicMock, create_autospec import pytest from pytest_mock import MockerFixture @@ -11,6 +11,8 @@ from pytest_mock import MockerFixture import core.rag.extractor.firecrawl.firecrawl_app as firecrawl_module from core.rag.extractor.firecrawl.firecrawl_app import FirecrawlApp from core.rag.extractor.firecrawl.firecrawl_web_extractor import FirecrawlWebExtractor +from services.data_source.provider_service import DatasourceProviderService +from services.data_source.website_service import WebsiteService def _response(status_code: int, json_data: Mapping[str, Any] | None = None, text: str = "") -> MagicMock: @@ -389,7 +391,13 @@ class TestFirecrawlWebExtractor: }, ) - extractor = FirecrawlWebExtractor("https://example.com", "job-1", "tenant-1", mode="crawl") + extractor = FirecrawlWebExtractor( + "https://example.com", + "job-1", + "tenant-1", + mode="crawl", + website_service=WebsiteService(providers=create_autospec(DatasourceProviderService, instance=True)), + ) docs = extractor.extract() assert len(docs) == 1 @@ -402,7 +410,13 @@ class TestFirecrawlWebExtractor: return_value=None, ) - extractor = FirecrawlWebExtractor("https://example.com", "job-1", "tenant-1", mode="crawl") + extractor = FirecrawlWebExtractor( + "https://example.com", + "job-1", + "tenant-1", + mode="crawl", + website_service=WebsiteService(providers=create_autospec(DatasourceProviderService, instance=True)), + ) assert extractor.extract() == [] def test_extract_scrape_mode_returns_document(self, mocker: MockerFixture): @@ -417,7 +431,12 @@ class TestFirecrawlWebExtractor: ) extractor = FirecrawlWebExtractor( - "https://example.com", "job-1", "tenant-1", mode="scrape", only_main_content=False + "https://example.com", + "job-1", + "tenant-1", + mode="scrape", + only_main_content=False, + website_service=WebsiteService(providers=create_autospec(DatasourceProviderService, instance=True)), ) docs = extractor.extract() @@ -426,5 +445,11 @@ class TestFirecrawlWebExtractor: mock_scrape.assert_called_once_with("firecrawl", "https://example.com", "tenant-1", False) def test_extract_unknown_mode_returns_empty(self): - extractor = FirecrawlWebExtractor("https://example.com", "job-1", "tenant-1", mode="unknown") + extractor = FirecrawlWebExtractor( + "https://example.com", + "job-1", + "tenant-1", + mode="unknown", + website_service=WebsiteService(providers=create_autospec(DatasourceProviderService, instance=True)), + ) assert extractor.extract() == [] diff --git a/api/tests/unit_tests/core/rag/extractor/test_extract_processor.py b/api/tests/unit_tests/core/rag/extractor/test_extract_processor.py index abfc01d671b..344ee8d9d55 100644 --- a/api/tests/unit_tests/core/rag/extractor/test_extract_processor.py +++ b/api/tests/unit_tests/core/rag/extractor/test_extract_processor.py @@ -2,11 +2,13 @@ from datetime import datetime from pathlib import Path from types import SimpleNamespace +import httpx import pytest from sqlalchemy.orm import Session import core.rag.extractor.extract_processor as processor_module from core.rag.extractor.entity.datasource_type import DatasourceType +from core.rag.extractor.entity.extract_setting import ExtractSetting, NotionInfo, WebsiteInfo from core.rag.extractor.extract_processor import ExtractProcessor from core.rag.models.document import Document from models.model import UploadFile @@ -76,8 +78,6 @@ def _patch_all_extractors(monkeypatch) -> _ExtractorFactory: class TestExtractProcessorLoaders: def test_load_from_upload_file_return_docs_and_text(self, monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr(processor_module, "ExtractSetting", lambda **kwargs: SimpleNamespace(**kwargs)) - monkeypatch.setattr( ExtractProcessor, "extract", @@ -126,9 +126,8 @@ class TestExtractProcessorLoaders: def test_load_from_url_builds_temp_file_with_correct_suffix( self, monkeypatch: pytest.MonkeyPatch, url, headers, expected_suffix ): - response = SimpleNamespace(headers=headers, content=b"body") + response = httpx.Response(200, headers=headers, content=b"body") monkeypatch.setattr(processor_module.remote_fetcher, "make_request", lambda *args, **kwargs: response) - monkeypatch.setattr(processor_module, "ExtractSetting", lambda **kwargs: SimpleNamespace(**kwargs)) captured = {} @@ -150,7 +149,7 @@ class TestExtractProcessorLoaders: def test_load_from_url_extracts_long_text_without_upload_file(self, monkeypatch: pytest.MonkeyPatch): content = "a" * 100_000 - response = SimpleNamespace(headers={"Content-Type": "text/plain"}, content=content.encode()) + response = httpx.Response(200, headers={"Content-Type": "text/plain"}, content=content.encode()) monkeypatch.setattr(processor_module.remote_fetcher, "make_request", lambda *args, **kwargs: response) apply_config_overrides(monkeypatch, ETL_TYPE="SelfHosted") @@ -215,7 +214,7 @@ class TestExtractProcessorFileRouting: monkeypatch.setattr(processor_module.storage, "download", fake_download) monkeypatch.setattr(processor_module.tempfile, "_get_candidate_names", lambda: iter(["candidate-name"])) - setting = SimpleNamespace( + setting = ExtractSetting( datasource_type=DatasourceType.FILE, upload_file=_upload_file(key=f"uploaded{extension}"), ) @@ -296,7 +295,7 @@ class TestExtractProcessorFileRouting: assert kwargs["session"] is unbound_session def test_extract_requires_upload_file_when_file_path_not_provided(self): - setting = SimpleNamespace(datasource_type=DatasourceType.FILE, upload_file=None) + setting = ExtractSetting(datasource_type=DatasourceType.FILE, upload_file=None) with pytest.raises(AssertionError, match="upload_file is required"): ExtractProcessor.extract(setting) @@ -306,20 +305,22 @@ class TestExtractProcessorDatasourceRouting: def test_extract_routes_notion_datasource(self, monkeypatch: pytest.MonkeyPatch): factory = _patch_all_extractors(monkeypatch) - notion_info = SimpleNamespace( + notion_info = NotionInfo( notion_workspace_id="ws", notion_obj_id="obj", notion_page_type="page", - document="doc", + document=None, tenant_id="tenant", credential_id="cred", + notion_access_token="secret", ) - setting = SimpleNamespace(datasource_type=DatasourceType.NOTION, notion_info=notion_info) + setting = ExtractSetting(datasource_type=DatasourceType.NOTION, notion_info=notion_info) docs = ExtractProcessor.extract(setting) assert docs[0].page_content == "extracted-by-NotionExtractor" assert factory.calls[-1][0] == "NotionExtractor" + assert factory.calls[-1][2]["notion_access_token"] == "secret" @pytest.mark.parametrize( ("provider", "expected"), @@ -334,7 +335,7 @@ class TestExtractProcessorDatasourceRouting: ): factory = _patch_all_extractors(monkeypatch) - website_info = SimpleNamespace( + website_info = WebsiteInfo( provider=provider, url="https://example.com", job_id="job", @@ -342,14 +343,14 @@ class TestExtractProcessorDatasourceRouting: mode="crawl", only_main_content=True, ) - setting = SimpleNamespace(datasource_type=DatasourceType.WEBSITE, website_info=website_info) + setting = ExtractSetting(datasource_type=DatasourceType.WEBSITE, website_info=website_info) docs = ExtractProcessor.extract(setting) assert docs[0].page_content == f"extracted-by-{expected}" assert factory.calls[-1][0] == expected def test_extract_unsupported_website_provider(self): - bad_provider = SimpleNamespace( + bad_provider = WebsiteInfo( provider="unknown", url="https://example.com", job_id="job", @@ -357,19 +358,19 @@ class TestExtractProcessorDatasourceRouting: mode="crawl", only_main_content=True, ) - setting = SimpleNamespace(datasource_type=DatasourceType.WEBSITE, website_info=bad_provider) + setting = ExtractSetting(datasource_type=DatasourceType.WEBSITE, website_info=bad_provider) with pytest.raises(ValueError, match="Unsupported website provider"): ExtractProcessor.extract(setting) def test_extract_unsupported_datasource_type(self): with pytest.raises(ValueError, match="Unsupported datasource type"): - ExtractProcessor.extract(SimpleNamespace(datasource_type="unknown")) + ExtractProcessor.extract(ExtractSetting(datasource_type="unknown")) def test_extract_requires_notion_info(self): with pytest.raises(AssertionError, match="notion_info is required"): - ExtractProcessor.extract(SimpleNamespace(datasource_type=DatasourceType.NOTION, notion_info=None)) + ExtractProcessor.extract(ExtractSetting(datasource_type=DatasourceType.NOTION, notion_info=None)) def test_extract_requires_website_info(self): with pytest.raises(AssertionError, match="website_info is required"): - ExtractProcessor.extract(SimpleNamespace(datasource_type=DatasourceType.WEBSITE, website_info=None)) + ExtractProcessor.extract(ExtractSetting(datasource_type=DatasourceType.WEBSITE, website_info=None)) diff --git a/api/tests/unit_tests/core/rag/extractor/test_jina_reader_extractor.py b/api/tests/unit_tests/core/rag/extractor/test_jina_reader_extractor.py index 0b4c9bd809e..79fcac2f88b 100644 --- a/api/tests/unit_tests/core/rag/extractor/test_jina_reader_extractor.py +++ b/api/tests/unit_tests/core/rag/extractor/test_jina_reader_extractor.py @@ -1,6 +1,10 @@ +from unittest.mock import create_autospec + from pytest_mock import MockerFixture from core.rag.extractor.jina_reader_extractor import JinaReaderWebExtractor +from services.data_source.provider_service import DatasourceProviderService +from services.data_source.website_service import WebsiteService class TestJinaReaderWebExtractor: @@ -15,7 +19,13 @@ class TestJinaReaderWebExtractor: }, ) - extractor = JinaReaderWebExtractor("https://example.com", "job-1", "tenant-1", mode="crawl") + extractor = JinaReaderWebExtractor( + "https://example.com", + "job-1", + "tenant-1", + mode="crawl", + website_service=WebsiteService(providers=create_autospec(DatasourceProviderService, instance=True)), + ) docs = extractor.extract() assert len(docs) == 1 @@ -32,7 +42,13 @@ class TestJinaReaderWebExtractor: return_value=None, ) - extractor = JinaReaderWebExtractor("https://example.com", "job-1", "tenant-1", mode="crawl") + extractor = JinaReaderWebExtractor( + "https://example.com", + "job-1", + "tenant-1", + mode="crawl", + website_service=WebsiteService(providers=create_autospec(DatasourceProviderService, instance=True)), + ) assert extractor.extract() == [] @@ -41,7 +57,13 @@ class TestJinaReaderWebExtractor: "core.rag.extractor.jina_reader_extractor.WebsiteService.get_crawl_url_data", return_value={"content": "unused"}, ) - extractor = JinaReaderWebExtractor("https://example.com", "job-1", "tenant-1", mode="scrape") + extractor = JinaReaderWebExtractor( + "https://example.com", + "job-1", + "tenant-1", + mode="scrape", + website_service=WebsiteService(providers=create_autospec(DatasourceProviderService, instance=True)), + ) assert extractor.extract() == [] mock_get_crawl.assert_not_called() diff --git a/api/tests/unit_tests/core/rag/extractor/test_notion_extractor.py b/api/tests/unit_tests/core/rag/extractor/test_notion_extractor.py index 56c9e254e38..2eac243de8c 100644 --- a/api/tests/unit_tests/core/rag/extractor/test_notion_extractor.py +++ b/api/tests/unit_tests/core/rag/extractor/test_notion_extractor.py @@ -13,11 +13,11 @@ from sqlalchemy import Engine from sqlalchemy.orm import Session, sessionmaker from core.rag.extractor import notion_extractor +from core.rag.extractor.entity.extract_setting import StoredDocumentExtractionInput from core.rag.index_processor.constant.index_type import IndexStructureType from models.base import TypeBase from models.dataset import Document as DocumentModel from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus -from tests.unit_tests.config_override import apply_config_overrides @pytest.fixture @@ -65,36 +65,27 @@ class TestNotionExtractorInitAndPublicMethods: notion_page_type="page", tenant_id="tenant", notion_access_token="token", + notion_token_loader=lambda: "token", ) assert extractor._notion_access_token == "token" - def test_init_falls_back_to_env_token_when_credential_lookup_fails(self, monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr( - notion_extractor.NotionExtractor, - "_get_access_token", - classmethod(lambda cls, tenant_id, credential_id: (_ for _ in ()).throw(Exception("credential error"))), - ) - apply_config_overrides(monkeypatch, NOTION_INTEGRATION_TOKEN="env-token") - + def test_init_uses_injected_token_loader(self): + loader = mock.Mock(return_value="resolved-token") extractor = notion_extractor.NotionExtractor( notion_workspace_id="ws", notion_obj_id="obj", notion_page_type="page", tenant_id="tenant", credential_id="cred", + notion_token_loader=loader, ) - assert extractor._notion_access_token == "env-token" - - def test_init_raises_if_no_credential_and_no_env_token(self, monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr( - notion_extractor.NotionExtractor, - "_get_access_token", - classmethod(lambda cls, tenant_id, credential_id: (_ for _ in ()).throw(Exception("credential error"))), - ) - apply_config_overrides(monkeypatch, NOTION_INTEGRATION_TOKEN=None) + assert extractor._notion_access_token == "resolved-token" + loader.assert_called_once_with() + def test_init_propagates_missing_token_from_loader(self): + loader = mock.Mock(side_effect=ValueError("Must specify `integration_token`")) with pytest.raises(ValueError, match="Must specify `integration_token`"): notion_extractor.NotionExtractor( notion_workspace_id="ws", @@ -102,6 +93,7 @@ class TestNotionExtractorInitAndPublicMethods: notion_page_type="page", tenant_id="tenant", credential_id="cred", + notion_token_loader=loader, ) def test_extract_updates_last_edited_and_loads_documents(self, monkeypatch: pytest.MonkeyPatch): @@ -111,6 +103,7 @@ class TestNotionExtractorInitAndPublicMethods: notion_page_type="page", tenant_id="tenant", notion_access_token="token", + notion_token_loader=lambda: "token", ) update_mock = mock.Mock() @@ -131,6 +124,7 @@ class TestNotionExtractorInitAndPublicMethods: notion_page_type="page", tenant_id="tenant", notion_access_token="token", + notion_token_loader=lambda: "token", ) monkeypatch.setattr(extractor, "_get_notion_block_data", lambda _: ["line1", "line2"]) @@ -153,6 +147,7 @@ class TestNotionDatabase: notion_page_type="database", tenant_id="tenant", notion_access_token="token", + notion_token_loader=lambda: "token", ) first_page = { @@ -204,6 +199,7 @@ class TestNotionDatabase: notion_page_type="database", tenant_id="tenant", notion_access_token="token", + notion_token_loader=lambda: "token", ) page = { @@ -246,6 +242,7 @@ class TestNotionDatabase: notion_page_type="database", tenant_id="tenant", notion_access_token="token", + notion_token_loader=lambda: "token", ) mocker.patch("httpx.post", return_value=_mock_response({"results": None})) @@ -259,6 +256,7 @@ class TestNotionDatabase: notion_page_type="database", tenant_id="tenant", notion_access_token="token", + notion_token_loader=lambda: "token", ) mock_post = mocker.patch("httpx.post", return_value=_mock_response({"results": None})) @@ -278,6 +276,7 @@ class TestNotionDatabase: notion_page_type="database", tenant_id="tenant", notion_access_token="token", + notion_token_loader=lambda: "token", ) extractor._notion_access_token = None @@ -293,6 +292,7 @@ class TestNotionBlocks: notion_page_type="page", tenant_id="tenant", notion_access_token="token", + notion_token_loader=lambda: "token", ) first_response = { @@ -349,6 +349,7 @@ class TestNotionBlocks: notion_page_type="page", tenant_id="tenant", notion_access_token="token", + notion_token_loader=lambda: "token", ) mocker.patch("httpx.request", side_effect=httpx.HTTPError("network")) @@ -370,6 +371,7 @@ class TestNotionBlocks: notion_page_type="page", tenant_id="tenant", notion_access_token="token", + notion_token_loader=lambda: "token", ) root_payload = { @@ -419,6 +421,7 @@ class TestNotionBlocks: notion_page_type="page", tenant_id="tenant", notion_access_token="token", + notion_token_loader=lambda: "token", ) mocker.patch("httpx.request", return_value=_mock_response({"results": None, "next_cursor": None})) @@ -431,6 +434,7 @@ class TestNotionBlocks: notion_page_type="page", tenant_id="tenant", notion_access_token="token", + notion_token_loader=lambda: "token", ) page_one = { @@ -492,6 +496,7 @@ class TestNotionBlocks: notion_page_type="page", tenant_id="tenant", notion_access_token="token", + notion_token_loader=lambda: "token", ) # A cell with mixed formatting arrives as multiple rich text segments. @@ -534,6 +539,7 @@ class TestNotionBlocks: notion_page_type="page", tenant_id="tenant", notion_access_token="token", + notion_token_loader=lambda: "token", ) page = { @@ -577,6 +583,7 @@ class TestNotionMetadataAndCredentialMethods: notion_page_type="page", tenant_id="tenant", notion_access_token="token", + notion_token_loader=lambda: "token", ) assert extractor.update_last_edited_time(None) is None @@ -592,11 +599,13 @@ class TestNotionMetadataAndCredentialMethods: notion_page_type="page", tenant_id="tenant", notion_access_token="token", + notion_token_loader=lambda: "token", ) monkeypatch.setattr(extractor, "get_notion_last_edited_time", lambda: "2026-01-01T00:00:00.000Z") session_maker, document = persisted_document - extractor.update_last_edited_time(document) + snapshot = StoredDocumentExtractionInput.model_validate(document) + extractor.update_last_edited_time(snapshot) # Closing the writer session rolls back an uncommitted update before the independent read below. notion_extractor.db.session.close() @@ -608,6 +617,28 @@ class TestNotionMetadataAndCredentialMethods: "last_edited_time": "2026-01-01T00:00:00.000Z", } + def test_last_edited_time_update_rejects_a_foreign_owner_snapshot( + self, + monkeypatch: pytest.MonkeyPatch, + persisted_document: tuple[sessionmaker[Session], DocumentModel], + ): + session_maker, document = persisted_document + extractor = notion_extractor.NotionExtractor( + notion_workspace_id="ws", + notion_obj_id="obj", + notion_page_type="page", + tenant_id=document.tenant_id, + notion_access_token="token", + notion_token_loader=lambda: "token", + ) + monkeypatch.setattr(extractor, "get_notion_last_edited_time", lambda: "2026-01-01T00:00:00.000Z") + snapshot = StoredDocumentExtractionInput.model_validate(document).model_copy(update={"tenant_id": "foreign"}) + extractor.update_last_edited_time(snapshot) + with session_maker() as session: + stored = session.get(DocumentModel, document.id) + assert stored is not None + assert stored.data_source_info_dict["last_edited_time"] == "2025-01-01T00:00:00.000Z" + def test_get_notion_last_edited_time_uses_page_and_database_urls(self, mocker: MockerFixture): extractor_page = notion_extractor.NotionExtractor( notion_workspace_id="ws", @@ -615,6 +646,7 @@ class TestNotionMetadataAndCredentialMethods: notion_page_type="page", tenant_id="tenant", notion_access_token="token", + notion_token_loader=lambda: "token", ) request_mock = mocker.patch( "httpx.request", return_value=_mock_response({"last_edited_time": "2025-05-01T00:00:00.000Z"}) @@ -629,6 +661,7 @@ class TestNotionMetadataAndCredentialMethods: notion_page_type="database", tenant_id="tenant", notion_access_token="token", + notion_token_loader=lambda: "token", ) request_mock = mocker.patch( "httpx.request", return_value=_mock_response({"last_edited_time": "2025-06-01T00:00:00.000Z"}) @@ -644,32 +677,13 @@ class TestNotionMetadataAndCredentialMethods: notion_page_type="page", tenant_id="tenant", notion_access_token="token", + notion_token_loader=lambda: "token", ) extractor._notion_access_token = None with pytest.raises(AssertionError, match="Notion access token is required"): extractor.get_notion_last_edited_time() - def test_get_access_token_success_and_errors(self, monkeypatch: pytest.MonkeyPatch): - with pytest.raises(Exception, match="No credential id found"): - notion_extractor.NotionExtractor._get_access_token("tenant", None) - - class FakeProviderServiceMissing: - def get_datasource_credentials(self, **kwargs): - return None - - monkeypatch.setattr(notion_extractor, "DatasourceProviderService", FakeProviderServiceMissing) - with pytest.raises(Exception, match="No notion credential found"): - notion_extractor.NotionExtractor._get_access_token("tenant", "cred") - - class FakeProviderServiceFound: - def get_datasource_credentials(self, **kwargs): - return {"integration_secret": "token-from-credential"} - - monkeypatch.setattr(notion_extractor, "DatasourceProviderService", FakeProviderServiceFound) - - assert notion_extractor.NotionExtractor._get_access_token("tenant", "cred") == "token-from-credential" - def test_get_cell_text_uses_plain_text_for_mention_and_equation_segments(): # Notion rich text segments of type mention or equation carry plain_text diff --git a/api/tests/unit_tests/core/rag/extractor/test_notion_page_rich_text.py b/api/tests/unit_tests/core/rag/extractor/test_notion_page_rich_text.py index 4e403e311a7..15cbdc124f2 100644 --- a/api/tests/unit_tests/core/rag/extractor/test_notion_page_rich_text.py +++ b/api/tests/unit_tests/core/rag/extractor/test_notion_page_rich_text.py @@ -18,6 +18,7 @@ def _extractor() -> notion_extractor.NotionExtractor: notion_page_type="page", tenant_id="tenant", notion_access_token="token", + notion_token_loader=lambda: "token", ) diff --git a/api/tests/unit_tests/core/rag/extractor/watercrawl/test_watercrawl.py b/api/tests/unit_tests/core/rag/extractor/watercrawl/test_watercrawl.py index 701198e968e..0b13ebc84c5 100644 --- a/api/tests/unit_tests/core/rag/extractor/watercrawl/test_watercrawl.py +++ b/api/tests/unit_tests/core/rag/extractor/watercrawl/test_watercrawl.py @@ -2,7 +2,7 @@ import json from typing import Any -from unittest.mock import MagicMock +from unittest.mock import MagicMock, create_autospec import httpx import pytest @@ -16,6 +16,8 @@ from core.rag.extractor.watercrawl.exceptions import ( ) from core.rag.extractor.watercrawl.extractor import WaterCrawlWebExtractor from core.rag.extractor.watercrawl.provider import WaterCrawlProvider +from services.data_source.provider_service import DatasourceProviderService +from services.data_source.website_service import WebsiteService def _response( @@ -454,7 +456,7 @@ class TestWaterCrawlWebExtractor: def test_extract_crawl_and_scrape_modes(self, monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr( "core.rag.extractor.watercrawl.extractor.WebsiteService.get_crawl_url_data", - lambda job_id, provider, url, tenant_id: { + lambda _self, job_id, provider, url, tenant_id: { "markdown": "crawl", "source_url": url, "description": "d", @@ -463,7 +465,7 @@ class TestWaterCrawlWebExtractor: ) monkeypatch.setattr( "core.rag.extractor.watercrawl.extractor.WebsiteService.get_scrape_url_data", - lambda provider, url, tenant_id, only_main_content: { + lambda _self, provider, url, tenant_id, only_main_content: { "markdown": "scrape", "source_url": url, "description": "d", @@ -471,8 +473,20 @@ class TestWaterCrawlWebExtractor: }, ) - crawl_extractor = WaterCrawlWebExtractor("https://example.com", "job-1", "tenant-1", mode="crawl") - scrape_extractor = WaterCrawlWebExtractor("https://example.com", "job-1", "tenant-1", mode="scrape") + crawl_extractor = WaterCrawlWebExtractor( + "https://example.com", + "job-1", + "tenant-1", + mode="crawl", + website_service=WebsiteService(providers=create_autospec(DatasourceProviderService, instance=True)), + ) + scrape_extractor = WaterCrawlWebExtractor( + "https://example.com", + "job-1", + "tenant-1", + mode="scrape", + website_service=WebsiteService(providers=create_autospec(DatasourceProviderService, instance=True)), + ) assert crawl_extractor.extract()[0].page_content == "crawl" assert scrape_extractor.extract()[0].page_content == "scrape" @@ -480,14 +494,26 @@ class TestWaterCrawlWebExtractor: def test_extract_crawl_returns_empty_when_service_returns_none(self, monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr( "core.rag.extractor.watercrawl.extractor.WebsiteService.get_crawl_url_data", - lambda job_id, provider, url, tenant_id: None, + lambda _self, job_id, provider, url, tenant_id: None, ) - extractor = WaterCrawlWebExtractor("https://example.com", "job-1", "tenant-1", mode="crawl") + extractor = WaterCrawlWebExtractor( + "https://example.com", + "job-1", + "tenant-1", + mode="crawl", + website_service=WebsiteService(providers=create_autospec(DatasourceProviderService, instance=True)), + ) assert extractor.extract() == [] def test_extract_unknown_mode_returns_empty(self): - extractor = WaterCrawlWebExtractor("https://example.com", "job-1", "tenant-1", mode="other") + extractor = WaterCrawlWebExtractor( + "https://example.com", + "job-1", + "tenant-1", + mode="other", + website_service=WebsiteService(providers=create_autospec(DatasourceProviderService, instance=True)), + ) assert extractor.extract() == [] diff --git a/api/tests/unit_tests/core/rag/indexing/processor/test_paragraph_index_processor.py b/api/tests/unit_tests/core/rag/indexing/processor/test_paragraph_index_processor.py index f9f8740cbf1..2f330e00159 100644 --- a/api/tests/unit_tests/core/rag/indexing/processor/test_paragraph_index_processor.py +++ b/api/tests/unit_tests/core/rag/indexing/processor/test_paragraph_index_processor.py @@ -184,7 +184,12 @@ class TestParagraphIndexProcessor: processor, "_get_content_files", return_value=[AttachmentDocument(page_content="image", metadata={})] ), ): - documents = processor.transform([source_document], process_rule=process_rule, session=session) + documents = processor.transform( + [source_document], + process_rule=process_rule, + tenant_id="tenant-1", + session=session, + ) assert len(documents) == 1 assert documents[0].page_content == "first" @@ -219,6 +224,7 @@ class TestParagraphIndexProcessor: processor.transform( [Document(page_content="text", metadata={})], process_rule={"mode": "automatic"}, + tenant_id="tenant-1", session=session, ) @@ -285,7 +291,7 @@ class TestParagraphIndexProcessor: with ( patch( - "core.rag.index_processor.processor.paragraph_index_processor.SummaryIndexService.delete_summaries_for_segments" + "core.rag.index_processor.processor.paragraph_index_processor.SummaryIndexAdapter.delete_summaries_for_segments" ) as mock_summary, patch("core.rag.index_processor.processor.paragraph_index_processor.Vector") as mock_vector_cls, ): @@ -303,7 +309,7 @@ class TestParagraphIndexProcessor: with ( patch( - "core.rag.index_processor.processor.paragraph_index_processor.SummaryIndexService.delete_summaries_for_segments" + "core.rag.index_processor.processor.paragraph_index_processor.SummaryIndexAdapter.delete_summaries_for_segments" ) as mock_summary, patch("core.rag.index_processor.processor.paragraph_index_processor.Keyword") as mock_keyword_cls, ): diff --git a/api/tests/unit_tests/core/rag/indexing/processor/test_parent_child_index_processor.py b/api/tests/unit_tests/core/rag/indexing/processor/test_parent_child_index_processor.py index 2919c71e0a0..32b09e69877 100644 --- a/api/tests/unit_tests/core/rag/indexing/processor/test_parent_child_index_processor.py +++ b/api/tests/unit_tests/core/rag/indexing/processor/test_parent_child_index_processor.py @@ -139,12 +139,13 @@ class TestParentChildIndexProcessor: patch.object( processor, "_get_content_files", return_value=[AttachmentDocument(page_content="image", metadata={})] ), - patch.object(processor, "_split_child_nodes", return_value=child_docs), + patch.object(processor, "split_child_nodes", return_value=child_docs), ): result = processor.transform( [parent_document], process_rule={"mode": "custom", "rules": {"enabled": True}}, preview=False, + tenant_id="tenant-1", session=self.session, ) @@ -176,12 +177,13 @@ class TestParentChildIndexProcessor: return_value="hash", ), patch.object(processor, "_get_content_files", return_value=[]), - patch.object(processor, "_split_child_nodes", return_value=[]), + patch.object(processor, "split_child_nodes", return_value=[]), ): result = processor.transform( documents, process_rule={"mode": "custom", "rules": {"enabled": True}}, preview=True, + tenant_id="tenant-1", session=self.session, ) @@ -202,7 +204,7 @@ class TestParentChildIndexProcessor: patch.object( processor, "_get_content_files", return_value=[AttachmentDocument(page_content="image", metadata={})] ), - patch.object(processor, "_split_child_nodes", return_value=child_docs), + patch.object(processor, "split_child_nodes", return_value=child_docs), patch( "core.rag.index_processor.processor.parent_child_index_processor.helper.generate_text_hash", return_value="hash", @@ -213,6 +215,7 @@ class TestParentChildIndexProcessor: docs, process_rule={"mode": "hierarchical", "rules": {"enabled": True}}, preview=True, + tenant_id="tenant-1", session=self.session, ) @@ -336,7 +339,7 @@ class TestParentChildIndexProcessor: with ( patch( - "core.rag.index_processor.processor.parent_child_index_processor.SummaryIndexService.delete_summaries_for_segments" + "core.rag.index_processor.processor.parent_child_index_processor.SummaryIndexAdapter.delete_summaries_for_segments" ) as mock_summary, patch("core.rag.index_processor.processor.parent_child_index_processor.Vector"), ): @@ -349,7 +352,7 @@ class TestParentChildIndexProcessor: ) -> None: with ( patch( - "core.rag.index_processor.processor.parent_child_index_processor.SummaryIndexService.delete_summaries_for_segments" + "core.rag.index_processor.processor.parent_child_index_processor.SummaryIndexAdapter.delete_summaries_for_segments" ) as mock_summary, patch("core.rag.index_processor.processor.parent_child_index_processor.Vector"), ): @@ -362,7 +365,7 @@ class TestParentChildIndexProcessor: rules = Rule(subchunk_segmentation=None) with pytest.raises(ValueError, match="No subchunk segmentation found"): - processor._split_child_nodes(Document(page_content="parent", metadata={}), rules, "custom", None) + processor.split_child_nodes(Document(page_content="parent", metadata={}), rules, "custom", None) def test_split_child_nodes_generates_child_documents(self, processor: ParentChildIndexProcessor) -> None: rules = Rule(subchunk_segmentation=Segmentation(max_tokens=200, chunk_overlap=10, separator="\n")) @@ -379,7 +382,7 @@ class TestParentChildIndexProcessor: return_value="hash", ), ): - child_docs = processor._split_child_nodes( + child_docs = processor.split_child_nodes( Document(page_content="parent", metadata={}), rules, "custom", None ) diff --git a/api/tests/unit_tests/core/rag/indexing/processor/test_qa_index_processor.py b/api/tests/unit_tests/core/rag/indexing/processor/test_qa_index_processor.py index d95a201c1fc..57cbc443b59 100644 --- a/api/tests/unit_tests/core/rag/indexing/processor/test_qa_index_processor.py +++ b/api/tests/unit_tests/core/rag/indexing/processor/test_qa_index_processor.py @@ -289,7 +289,7 @@ class TestQAIndexProcessor: with ( patch( - "core.rag.index_processor.processor.qa_index_processor.SummaryIndexService.delete_summaries_for_segments" + "core.rag.index_processor.processor.qa_index_processor.SummaryIndexAdapter.delete_summaries_for_segments" ) as mock_summary, patch("core.rag.index_processor.processor.qa_index_processor.Vector") as mock_vector_cls, ): @@ -305,7 +305,7 @@ class TestQAIndexProcessor: session = sqlite_session with ( patch( - "core.rag.index_processor.processor.qa_index_processor.SummaryIndexService.delete_summaries_for_segments" + "core.rag.index_processor.processor.qa_index_processor.SummaryIndexAdapter.delete_summaries_for_segments" ) as mock_summary, patch("core.rag.index_processor.processor.qa_index_processor.Vector") as mock_vector_cls, ): diff --git a/api/tests/unit_tests/core/rag/indexing/test_index_processor_base.py b/api/tests/unit_tests/core/rag/indexing/test_index_processor_base.py index faa6417900c..98cd0479635 100644 --- a/api/tests/unit_tests/core/rag/indexing/test_index_processor_base.py +++ b/api/tests/unit_tests/core/rag/indexing/test_index_processor_base.py @@ -24,9 +24,14 @@ from models.tools import ToolFile from tests.unit_tests.config_override import config_overrides_context -def _upload(*, upload_id: str, name: str) -> UploadFile: +def _upload( + *, + upload_id: str, + name: str, + tenant_id: str = "tenant-1", +) -> UploadFile: upload = UploadFile( - tenant_id=str(uuid4()), + tenant_id=tenant_id, storage_type=StorageType.LOCAL, key=f"uploads/{name}", name=name, @@ -42,8 +47,8 @@ def _upload(*, upload_id: str, name: str) -> UploadFile: return upload -def _persist_upload(session: Session, *, upload_id: str, name: str) -> UploadFile: - upload = _upload(upload_id=upload_id, name=name) +def _persist_upload(session: Session, *, upload_id: str, name: str, tenant_id: str = "tenant-1") -> UploadFile: + upload = _upload(upload_id=upload_id, name=name, tenant_id=tenant_id) session.add(upload) return upload @@ -170,7 +175,7 @@ class TestBaseIndexProcessor: self, processor: _ForwardingBaseIndexProcessor, unbound_session: Session ) -> None: document = Document(page_content="no image markdown", metadata={"document_id": "doc-1", "dataset_id": "ds-1"}) - assert processor._get_content_files(document, session=unbound_session) == [] + assert processor._get_content_files(document, tenant_id="tenant-1", session=unbound_session) == [] def test_get_content_files_handles_all_sources_and_duplicates( self, processor: _ForwardingBaseIndexProcessor, sqlite_session: Session @@ -195,7 +200,12 @@ class TestBaseIndexProcessor: patch.object(processor, "_download_tool_file", return_value=tool_upload.id) as mock_tool_download, patch.object(processor, "_download_image", return_value=remote_upload.id) as mock_image_download, ): - files = processor._get_content_files(document, current_user=current_user, session=sqlite_session) + files = processor._get_content_files( + document, + current_user=current_user, + tenant_id="tenant-1", + session=sqlite_session, + ) assert len(files) == 5 assert all(isinstance(file, AttachmentDocument) for file in files) @@ -207,6 +217,7 @@ class TestBaseIndexProcessor: mock_tool_download.assert_called_once_with( "cccccccc-cccc-cccc-cccc-cccccccccccc", current_user, + tenant_id="tenant-1", session=sqlite_session, ) mock_image_download.assert_called_once() @@ -218,7 +229,12 @@ class TestBaseIndexProcessor: images = ["/files/tools/cccccccc-cccc-cccc-cccc-cccccccccccc.png", "https://example.com/remote.png"] with patch.object(processor, "_extract_markdown_images", return_value=images): - files = processor._get_content_files(document, current_user=None, session=unbound_session) + files = processor._get_content_files( + document, + current_user=None, + tenant_id="tenant-1", + session=unbound_session, + ) assert files == [] @@ -238,7 +254,9 @@ class TestBaseIndexProcessor: patch.object(processor, "_extract_markdown_images", return_value=images), patch.object(processor, "_download_image") as mock_image_download, ): - files = processor._get_content_files(document, current_user=Mock(), session=unbound_session) + files = processor._get_content_files( + document, current_user=_account(), tenant_id="tenant-1", session=unbound_session + ) assert files == [] mock_image_download.assert_not_called() @@ -251,7 +269,22 @@ class TestBaseIndexProcessor: with ( patch.object(processor, "_extract_markdown_images", return_value=images), ): - files = processor._get_content_files(document, session=sqlite_session) + files = processor._get_content_files(document, tenant_id="tenant-1", session=sqlite_session) + + assert files == [] + + def test_get_content_files_ignores_upload_from_another_tenant( + self, processor: _ForwardingBaseIndexProcessor, sqlite_session: Session + ) -> None: + upload_id = "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa" + _persist_upload(sqlite_session, upload_id=upload_id, name="foreign.png", tenant_id="tenant-2") + sqlite_session.commit() + document = Document( + page_content=f"![foreign](/files/{upload_id}/file-preview)", + metadata={"document_id": "doc-1", "dataset_id": "ds-1"}, + ) + + files = processor._get_content_files(document, tenant_id="tenant-1", session=sqlite_session) assert files == [] @@ -333,7 +366,15 @@ class TestBaseIndexProcessor: def test_download_tool_file_returns_none_when_not_found( self, processor: _ForwardingBaseIndexProcessor, sqlite_session: Session ) -> None: - assert processor._download_tool_file(str(uuid4()), current_user=_account(), session=sqlite_session) is None + assert ( + processor._download_tool_file( + str(uuid4()), + current_user=_account(), + tenant_id="tenant-1", + session=sqlite_session, + ) + is None + ) def test_download_tool_file_uploads_file_when_found( self, processor: _ForwardingBaseIndexProcessor, sqlite_session: Session, sqlite_engine: Engine @@ -360,8 +401,39 @@ class TestBaseIndexProcessor: patch("services.file_service.FileService") as mock_file_service, ): mock_file_service.return_value.upload_file.return_value = upload_result - result = processor._download_tool_file(tool_file.id, current_user=_account(), session=sqlite_session) + result = processor._download_tool_file( + tool_file.id, + current_user=_account(), + tenant_id=tool_file.tenant_id, + session=sqlite_session, + ) assert result == upload_result.id mock_load.assert_called_once_with("k1") mock_file_service.return_value.upload_file.assert_called_once() + + def test_download_tool_file_rejects_file_from_another_tenant( + self, processor: _ForwardingBaseIndexProcessor, sqlite_session: Session + ) -> None: + tool_file = ToolFile( + user_id=str(uuid4()), + tenant_id="tenant-2", + conversation_id=None, + file_key="foreign-key", + mimetype="image/png", + name="foreign.png", + size=4, + ) + sqlite_session.add(tool_file) + sqlite_session.commit() + + with patch("core.rag.index_processor.index_processor_base.storage.load_once") as mock_load: + result = processor._download_tool_file( + tool_file.id, + current_user=_account(), + tenant_id="tenant-1", + session=sqlite_session, + ) + + assert result is None + mock_load.assert_not_called() diff --git a/api/tests/unit_tests/core/rag/indexing/test_indexing_runner.py b/api/tests/unit_tests/core/rag/indexing/test_indexing_runner.py deleted file mode 100644 index 7bea0a9fe34..00000000000 --- a/api/tests/unit_tests/core/rag/indexing/test_indexing_runner.py +++ /dev/null @@ -1,2018 +0,0 @@ -"""Comprehensive unit tests for IndexingRunner. - -This test module provides complete coverage of the IndexingRunner class, which is responsible -for orchestrating the document indexing pipeline in the Dify RAG system. - -Test Coverage Areas: -================== -1. **Document Parsing Pipeline (Extract Phase)** - - Tests extraction from various data sources (upload files, Notion, websites) - - Validates metadata preservation and document status updates - - Ensures proper error handling for missing or invalid sources - -2. **Chunk Creation Logic (Transform Phase)** - - Tests document splitting with different segmentation strategies - - Validates embedding model integration for high-quality indexing - - Tests text cleaning and preprocessing rules - -3. **Embedding Generation Orchestration** - - Tests parallel processing of document chunks - - Validates token counting and embedding generation - - Tests integration with various embedding model providers - -4. **Vector Storage Integration (Load Phase)** - - Tests vector index creation and updates - - Validates keyword index generation for economy mode - - Tests parent-child index structures - -5. **Retry Logic & Error Handling** - - Tests pause/resume functionality - - Validates error recovery and status updates - - Tests handling of provider token errors and deleted documents - -6. **Document Status Management** - - Tests status transitions (parsing → splitting → indexing → completed) - - Validates timestamp updates and error state persistence - - Tests concurrent document processing - -Testing Approach: -================ -- All tests use mocking to avoid external dependencies (database, storage, Redis) -- Tests follow the Arrange-Act-Assert (AAA) pattern for clarity -- Each test is isolated and can run independently -- Fixtures provide reusable test data and mock objects -- Comprehensive docstrings explain the purpose and assertions of each test - -Note: These tests focus on unit testing the IndexingRunner logic. Integration tests -for the full indexing pipeline are handled separately in the integration test suite. -""" - -import json -import uuid -from datetime import UTC, datetime -from typing import Any -from unittest.mock import MagicMock, Mock, patch - -import pytest -from sqlalchemy import event -from sqlalchemy.orm import Session, sessionmaker -from sqlalchemy.orm.exc import ObjectDeletedError - -from core.entities.knowledge_entities import PreviewDetail -from core.errors.error import ProviderTokenNotInitError -from core.indexing_runner import ( - DocumentIsDeletedPausedError, - DocumentIsPausedError, - IndexingRunner, -) -from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType -from core.rag.models.document import ChildDocument, Document -from enums import DeploymentEdition -from extensions.storage.storage_type import StorageType -from graphon.model_runtime.entities.model_entities import ModelType -from libs.datetime_utils import naive_utc_now -from models.dataset import ChildChunk, Dataset, DatasetProcessRule, DocumentSegment -from models.dataset import Document as DatasetDocument -from models.enums import CreatorUserRole, DataSourceType, DocumentCreatedFrom, SegmentStatus -from models.model import Account, UploadFile -from services.vector_space_admission_service import VectorSpaceAdmissionError - -# ============================================================================ -# Helper Functions -# ============================================================================ - - -def create_mock_dataset( - dataset_id: str | None = None, - tenant_id: str | None = None, - indexing_technique: str = IndexTechniqueType.HIGH_QUALITY, - embedding_provider: str = "openai", - embedding_model: str = "text-embedding-ada-002", -) -> Dataset: - """Create a Dataset object with configurable parameters. - - This helper function creates a properly configured Dataset object that can be - used across multiple tests, ensuring consistency in test data. - - Args: - dataset_id: Optional dataset ID. If None, generates a new UUID. - tenant_id: Optional tenant ID. If None, generates a new UUID. - indexing_technique: The indexing technique ("high_quality" or "economy"). - embedding_provider: The embedding model provider name. - embedding_model: The embedding model name. - - Returns: - Dataset: A configured Dataset object with all required attributes. - - Example: - >>> dataset = create_mock_dataset(indexing_technique="economy") - >>> assert dataset.indexing_technique == "economy" - """ - dataset = Dataset( - id=dataset_id or str(uuid.uuid4()), - tenant_id=tenant_id or str(uuid.uuid4()), - indexing_technique=indexing_technique, - embedding_model_provider=embedding_provider, - embedding_model=embedding_model, - ) - return dataset - - -def create_mock_dataset_document( - document_id: str | None = None, - dataset_id: str | None = None, - tenant_id: str | None = None, - doc_form: str = IndexStructureType.PARAGRAPH_INDEX, - data_source_type: str = "upload_file", - doc_language: str = "English", -) -> DatasetDocument: - """Create a DatasetDocument object with configurable parameters. - - This helper function creates a properly configured DatasetDocument object, - reducing boilerplate code in individual tests. - - Args: - document_id: Optional document ID. If None, generates a new UUID. - dataset_id: Optional dataset ID. If None, generates a new UUID. - tenant_id: Optional tenant ID. If None, generates a new UUID. - doc_form: The document form/index type (e.g., PARAGRAPH_INDEX, QA_INDEX). - data_source_type: The data source type ("upload_file", "notion_import", etc.). - doc_language: The document language. - - Returns: - DatasetDocument: A configured DatasetDocument object with all required attributes. - - Example: - >>> doc = create_mock_dataset_document(doc_form=IndexStructureType.QA_INDEX) - >>> assert doc.doc_form == IndexStructureType.QA_INDEX - """ - return DatasetDocument( - id=document_id or str(uuid.uuid4()), - dataset_id=dataset_id or str(uuid.uuid4()), - tenant_id=tenant_id or str(uuid.uuid4()), - doc_form=doc_form, - doc_language=doc_language, - data_source_type=data_source_type, - data_source_info=json.dumps({"upload_file_id": str(uuid.uuid4())}), - dataset_process_rule_id=str(uuid.uuid4()), - created_by=str(uuid.uuid4()), - ) - - -def persist_indexing_scope(session: Session) -> tuple[Dataset, DatasetDocument]: - tenant_id = str(uuid.uuid4()) - created_by = str(uuid.uuid4()) - dataset = Dataset( - id=str(uuid.uuid4()), - tenant_id=tenant_id, - name="Dataset", - description="", - created_by=created_by, - indexing_technique=IndexTechniqueType.HIGH_QUALITY, - ) - document = DatasetDocument( - id=str(uuid.uuid4()), - tenant_id=tenant_id, - dataset_id=dataset.id, - position=1, - data_source_type=DataSourceType.UPLOAD_FILE, - batch="batch", - name="Document", - created_from=DocumentCreatedFrom.API, - created_by=created_by, - ) - session.add_all([dataset, document]) - session.commit() - return dataset, document - - -def create_sample_documents( - count: int = 3, - include_children: bool = False, - base_content: str = "Sample chunk content", -) -> list[Document]: - """Create a list of sample Document objects for testing. - - This helper function generates test documents with proper metadata, - optionally including child documents for hierarchical indexing tests. - - Args: - count: Number of documents to create. - include_children: Whether to add child documents to each parent. - base_content: Base content string for documents. - - Returns: - list[Document]: A list of Document objects with metadata. - - Example: - >>> docs = create_sample_documents(count=2, include_children=True) - >>> assert len(docs) == 2 - >>> assert docs[0].children is not None - """ - documents = [] - for i in range(count): - doc = Document( - page_content=f"{base_content} {i + 1}", - metadata={ - "doc_id": f"chunk{i + 1}", - "doc_hash": f"hash{i + 1}", - "document_id": "doc1", - "dataset_id": "dataset1", - }, - ) - - # Add child documents if requested (for parent-child indexing) - if include_children: - doc.children = [ - ChildDocument( - page_content=f"Child of {base_content} {i + 1}", - metadata={ - "doc_id": f"child_chunk{i + 1}", - "doc_hash": f"child_hash{i + 1}", - }, - ) - ] - - documents.append(doc) - - return documents - - -def create_mock_process_rule( - mode: str = "automatic", - max_tokens: int = 500, - chunk_overlap: int = 50, - separator: str = "\\n\\n", -) -> dict[str, Any]: - """Create a mock processing rule dictionary. - - This helper function creates a processing rule configuration that matches - the structure expected by the IndexingRunner. - - Args: - mode: Processing mode ("automatic", "custom", or "hierarchical"). - max_tokens: Maximum tokens per chunk. - chunk_overlap: Number of overlapping tokens between chunks. - separator: Separator string for splitting. - - Returns: - dict: A processing rule configuration dictionary. - - Example: - >>> rule = create_mock_process_rule(mode="custom", max_tokens=1000) - >>> assert rule["mode"] == "custom" - >>> assert rule["rules"]["segmentation"]["max_tokens"] == 1000 - """ - return { - "mode": mode, - "rules": { - "segmentation": { - "max_tokens": max_tokens, - "chunk_overlap": chunk_overlap, - "separator": separator, - }, - "pre_processing_rules": [{"id": "remove_extra_spaces", "enabled": True}], - }, - } - - -# ============================================================================ -# Test Classes -# ============================================================================ - - -class TestIndexingRunnerExtract: - """Unit tests for IndexingRunner._extract method. - - Tests cover: - - Upload file extraction - - Notion import extraction - - Website crawl extraction - - Document status updates during extraction - - Error handling for missing data sources - """ - - @pytest.fixture - def mock_dependencies(self, sqlite_session: Session): - """Mock all external dependencies for extract tests.""" - with ( - patch("core.indexing_runner.IndexProcessorFactory") as mock_factory, - patch("core.indexing_runner.storage") as mock_storage, - ): - yield { - "session": sqlite_session, - "factory": mock_factory, - "storage": mock_storage, - } - - @pytest.fixture - def sample_dataset_document(self): - """Create a sample dataset document for testing.""" - return DatasetDocument( - id=str(uuid.uuid4()), - dataset_id=str(uuid.uuid4()), - tenant_id=str(uuid.uuid4()), - doc_form=IndexStructureType.PARAGRAPH_INDEX, - data_source_type="upload_file", - data_source_info=json.dumps({"upload_file_id": str(uuid.uuid4())}), - ) - - @pytest.fixture - def sample_process_rule(self): - """Create a sample processing rule.""" - return { - "mode": "automatic", - "rules": { - "segmentation": {"max_tokens": 500, "chunk_overlap": 50, "separator": "\\n\\n"}, - "pre_processing_rules": [{"id": "remove_extra_spaces", "enabled": True}], - }, - } - - def test_extract_upload_file_success(self, mock_dependencies, sample_dataset_document, sample_process_rule): - """Test successful extraction from uploaded file. - - This test verifies that the IndexingRunner can successfully extract content - from an uploaded file and properly update document metadata. It ensures: - - The processor's extract method is called with correct parameters - - Document and dataset IDs are properly added to metadata - - The document status is updated during extraction - - Expected behavior: - - Extract should return documents with updated metadata - - Each document should have document_id and dataset_id in metadata - - The processor's extract method should be called exactly once - """ - # Arrange: Set up the test environment with mocked dependencies - runner = IndexingRunner() - mock_processor = MagicMock() - mock_dependencies["factory"].return_value.init_index_processor.return_value = mock_processor - - # Create mock extracted documents that simulate PDF page extraction - extracted_docs = [ - Document( - page_content="Test content 1", - metadata={"doc_id": "doc1", "source": "test.pdf", "page": 1}, - ), - Document( - page_content="Test content 2", - metadata={"doc_id": "doc2", "source": "test.pdf", "page": 2}, - ), - ] - mock_processor.extract.return_value = extracted_docs - file_id = json.loads(sample_dataset_document.data_source_info)["upload_file_id"] - upload_file = UploadFile( - tenant_id=sample_dataset_document.tenant_id, - storage_type=StorageType.LOCAL, - key="uploads/test.pdf", - name="test.pdf", - size=10, - extension="pdf", - mime_type="application/pdf", - created_by_role=CreatorUserRole.ACCOUNT, - created_by="user-id", - created_at=datetime.now(UTC), - used=True, - ) - upload_file.id = file_id - mock_dependencies["session"].add(upload_file) - mock_dependencies["session"].commit() - - with patch.object(runner, "_update_document_index_status"), patch("core.indexing_runner.ExtractSetting"): - result = runner._extract( - mock_processor, sample_dataset_document, sample_process_rule, mock_dependencies["session"] - ) - - # Assert: Verify the extraction results - assert len(result) == 2, "Should extract 2 documents from the PDF" - assert result[0].page_content == "Test content 1", "First document content should match" - # Verify metadata was properly updated with document and dataset IDs - assert result[0].metadata["document_id"] == sample_dataset_document.id - assert result[0].metadata["dataset_id"] == sample_dataset_document.dataset_id - assert result[1].page_content == "Test content 2", "Second document content should match" - # Verify the processor was called exactly once (not multiple times) - mock_processor.extract.assert_called_once() - assert mock_processor.extract.call_args.kwargs["session"] is mock_dependencies["session"] - - def test_extract_notion_import_success(self, mock_dependencies, sample_dataset_document, sample_process_rule): - """Test successful extraction from Notion import.""" - # Arrange - runner = IndexingRunner() - sample_dataset_document.data_source_type = "notion_import" - sample_dataset_document.data_source_info = json.dumps( - { - "credential_id": str(uuid.uuid4()), - "notion_workspace_id": "workspace123", - "notion_page_id": "page123", - "type": "page", - } - ) - - mock_processor = MagicMock() - mock_dependencies["factory"].return_value.init_index_processor.return_value = mock_processor - - extracted_docs = [Document(page_content="Notion content", metadata={"doc_id": "notion1", "source": "notion"})] - mock_processor.extract.return_value = extracted_docs - - # Mock update_document_index_status to avoid database calls - with patch.object(runner, "_update_document_index_status"): - # Act - result = runner._extract( - mock_processor, sample_dataset_document, sample_process_rule, mock_dependencies["session"] - ) - - # Assert - assert len(result) == 1 - assert result[0].page_content == "Notion content" - assert result[0].metadata["document_id"] == sample_dataset_document.id - - def test_extract_website_crawl_success(self, mock_dependencies, sample_dataset_document, sample_process_rule): - """Test successful extraction from website crawl.""" - # Arrange - runner = IndexingRunner() - sample_dataset_document.data_source_type = "website_crawl" - sample_dataset_document.data_source_info = json.dumps( - { - "provider": "firecrawl", - "url": "https://example.com", - "job_id": "job123", - "mode": "crawl", - "only_main_content": True, - } - ) - - mock_processor = MagicMock() - mock_dependencies["factory"].return_value.init_index_processor.return_value = mock_processor - - extracted_docs = [ - Document(page_content="Website content", metadata={"doc_id": "web1", "url": "https://example.com"}) - ] - mock_processor.extract.return_value = extracted_docs - - # Mock update_document_index_status to avoid database calls - with patch.object(runner, "_update_document_index_status"): - # Act - result = runner._extract( - mock_processor, sample_dataset_document, sample_process_rule, mock_dependencies["session"] - ) - - # Assert - assert len(result) == 1 - assert result[0].page_content == "Website content" - assert result[0].metadata["document_id"] == sample_dataset_document.id - - def test_extract_missing_upload_file(self, mock_dependencies, sample_dataset_document, sample_process_rule): - """Test extraction fails when upload file is missing.""" - # Arrange - runner = IndexingRunner() - sample_dataset_document.data_source_info = "{}" - - mock_processor = MagicMock() - mock_dependencies["factory"].return_value.init_index_processor.return_value = mock_processor - - # Act & Assert - with pytest.raises(ValueError, match="no upload file found"): - runner._extract(mock_processor, sample_dataset_document, sample_process_rule, mock_dependencies["session"]) - - def test_extract_unsupported_data_source(self, mock_dependencies, sample_dataset_document, sample_process_rule): - """Test extraction returns empty list for unsupported data sources.""" - # Arrange - runner = IndexingRunner() - sample_dataset_document.data_source_type = "unsupported_type" - - mock_processor = MagicMock() - - # Act - result = runner._extract( - mock_processor, sample_dataset_document, sample_process_rule, mock_dependencies["session"] - ) - - # Assert - assert result == [] - - -class TestIndexingRunnerTransform: - """Unit tests for IndexingRunner._transform method. - - Tests cover: - - Document chunking with different splitters - - Embedding model instance retrieval - - Text cleaning and preprocessing - - Metadata preservation - - Child chunk generation for hierarchical indexing - """ - - @pytest.fixture - def mock_dependencies(self, unbound_session: Session): - """Mock all external dependencies for transform tests.""" - with ( - patch("core.indexing_runner.ModelManager.for_tenant") as mock_model_manager, - ): - yield { - "session": unbound_session, - "model_manager": mock_model_manager, - } - - @pytest.fixture - def sample_dataset(self): - """Create a sample dataset for testing.""" - dataset = Dataset( - id=str(uuid.uuid4()), - tenant_id=str(uuid.uuid4()), - indexing_technique=IndexTechniqueType.HIGH_QUALITY, - embedding_model_provider="openai", - embedding_model="text-embedding-ada-002", - ) - return dataset - - @pytest.fixture - def sample_text_docs(self): - """Create sample text documents for transformation.""" - return [ - Document( - page_content="This is a long document that needs to be split into multiple chunks. " * 10, - metadata={"doc_id": "doc1", "source": "test.pdf"}, - ), - Document( - page_content="Another document with different content. " * 5, - metadata={"doc_id": "doc2", "source": "test.pdf"}, - ), - ] - - def test_transform_with_high_quality_indexing(self, mock_dependencies, sample_dataset, sample_text_docs): - """Test transformation with high quality indexing (embeddings).""" - # Arrange - runner = IndexingRunner() - mock_embedding_instance = MagicMock() - model_manager = mock_dependencies["model_manager"].return_value - model_manager.get_model_instance.return_value = mock_embedding_instance - - mock_processor = MagicMock() - transformed_docs = [ - Document( - page_content="Chunk 1", - metadata={"doc_id": "chunk1", "doc_hash": "hash1", "document_id": "doc1"}, - ), - Document( - page_content="Chunk 2", - metadata={"doc_id": "chunk2", "doc_hash": "hash2", "document_id": "doc1"}, - ), - ] - mock_processor.transform.return_value = transformed_docs - - process_rule = { - "mode": "automatic", - "rules": {"segmentation": {"max_tokens": 500, "chunk_overlap": 50}}, - } - - # Act - result = runner._transform( - mock_processor, - sample_dataset, - sample_text_docs, - "English", - process_rule, - session=mock_dependencies["session"], - ) - - # Assert - assert len(result) == 2 - assert result[0].page_content == "Chunk 1" - assert result[1].page_content == "Chunk 2" - model_manager.get_model_instance.assert_called_once_with( - tenant_id=sample_dataset.tenant_id, - provider=sample_dataset.embedding_model_provider, - model_type=ModelType.TEXT_EMBEDDING, - model=sample_dataset.embedding_model, - ) - mock_processor.transform.assert_called_once() - - def test_transform_with_economy_indexing(self, mock_dependencies, sample_dataset, sample_text_docs): - """Test transformation with economy indexing (no embeddings).""" - # Arrange - runner = IndexingRunner() - model_manager = mock_dependencies["model_manager"].return_value - sample_dataset.indexing_technique = IndexTechniqueType.ECONOMY - - mock_processor = MagicMock() - transformed_docs = [ - Document( - page_content="Chunk 1", - metadata={"doc_id": "chunk1", "doc_hash": "hash1"}, - ) - ] - mock_processor.transform.return_value = transformed_docs - - process_rule = {"mode": "automatic", "rules": {}} - - # Act - result = runner._transform( - mock_processor, - sample_dataset, - sample_text_docs, - "English", - process_rule, - session=mock_dependencies["session"], - ) - - # Assert - assert len(result) == 1 - model_manager.get_model_instance.assert_not_called() - - def test_transform_with_custom_segmentation(self, mock_dependencies, sample_dataset, sample_text_docs): - """Test transformation with custom segmentation rules.""" - # Arrange - runner = IndexingRunner() - mock_embedding_instance = MagicMock() - model_manager = mock_dependencies["model_manager"].return_value - model_manager.get_model_instance.return_value = mock_embedding_instance - - mock_processor = MagicMock() - transformed_docs = [Document(page_content="Custom chunk", metadata={"doc_id": "custom1", "doc_hash": "hash1"})] - mock_processor.transform.return_value = transformed_docs - - process_rule = { - "mode": "custom", - "rules": {"segmentation": {"max_tokens": 1000, "chunk_overlap": 100, "separator": "\\n"}}, - } - - # Act - result = runner._transform( - mock_processor, - sample_dataset, - sample_text_docs, - "Chinese", - process_rule, - session=mock_dependencies["session"], - ) - - # Assert - assert len(result) == 1 - assert result[0].page_content == "Custom chunk" - # Verify transform was called with correct parameters - call_args = mock_processor.transform.call_args - assert call_args[1]["doc_language"] == "Chinese" - assert call_args[1]["process_rule"] == process_rule - - -class TestIndexingRunnerLoad: - """Unit tests for IndexingRunner._load method. - - Tests cover: - - Vector index creation - - Keyword index creation - - Multi-threaded processing - - Document segment status updates - - Precomputed token totals - - Error handling during loading - """ - - @pytest.fixture - def mock_dependencies(self, unbound_session: Session): - """Mock all external dependencies for load tests.""" - with ( - patch("core.indexing_runner.ModelManager.for_tenant") as mock_model_manager, - patch("core.indexing_runner.current_app") as mock_app, - patch("core.indexing_runner.threading.Thread") as mock_thread, - patch("core.indexing_runner.concurrent.futures.ThreadPoolExecutor") as mock_executor, - ): - mock_app._get_current_object = Mock(return_value=Mock()) - yield { - "session": unbound_session, - "model_manager": mock_model_manager, - "app": mock_app, - "thread": mock_thread, - "executor": mock_executor, - } - - @pytest.fixture - def sample_dataset(self): - """Create a sample dataset for testing.""" - dataset = Dataset( - id=str(uuid.uuid4()), - tenant_id=str(uuid.uuid4()), - indexing_technique=IndexTechniqueType.HIGH_QUALITY, - embedding_model_provider="openai", - embedding_model="text-embedding-ada-002", - ) - return dataset - - @pytest.fixture - def sample_dataset_document(self): - """Create a sample dataset document for testing.""" - return DatasetDocument( - id=str(uuid.uuid4()), - dataset_id=str(uuid.uuid4()), - doc_form=IndexStructureType.PARAGRAPH_INDEX, - ) - - @pytest.fixture - def sample_documents(self): - """Create sample documents for loading.""" - return [ - Document( - page_content="Chunk 1 content", - metadata={"doc_id": "chunk1", "doc_hash": "hash1", "document_id": "doc1"}, - ), - Document( - page_content="Chunk 2 content", - metadata={"doc_id": "chunk2", "doc_hash": "hash2", "document_id": "doc1"}, - ), - Document( - page_content="Chunk 3 content", - metadata={"doc_id": "chunk3", "doc_hash": "hash3", "document_id": "doc1"}, - ), - ] - - def test_load_with_high_quality_indexing( - self, mock_dependencies, sample_dataset, sample_dataset_document, sample_documents - ): - """Test loading with high quality indexing (vector embeddings).""" - # Arrange - runner = IndexingRunner() - - # Mock ThreadPoolExecutor - mock_future = MagicMock() - mock_future.result.return_value = None - mock_executor_instance = MagicMock() - mock_executor_instance.__enter__.return_value = mock_executor_instance - mock_executor_instance.__exit__.return_value = None - mock_executor_instance.submit.return_value = mock_future - mock_dependencies["executor"].return_value = mock_executor_instance - - # Mock update_document_index_status to avoid database calls - with patch.object(runner, "_update_document_index_status") as mock_update_status: - # Act - runner._load( - session=mock_dependencies["session"], - dataset=sample_dataset, - dataset_document=sample_dataset_document, - documents=sample_documents, - total_tokens=300, - ) - - # Assert - mock_dependencies["model_manager"].assert_not_called() - # Verify executor was used for parallel processing - assert mock_executor_instance.submit.called - for submit_call in mock_executor_instance.submit.call_args_list: - assert submit_call.args[0].__name__ == runner._process_chunk.__name__ - assert len(submit_call.args) == 6 - mock_future.result.assert_called() - assert mock_update_status.call_args.kwargs["extra_update_params"][DatasetDocument.tokens] == 300 - - def test_load_propagates_worker_errors( - self, mock_dependencies, sample_dataset, sample_dataset_document, sample_documents - ): - runner = IndexingRunner() - mock_future = MagicMock() - mock_future.result.side_effect = RuntimeError("index failed") - mock_executor_instance = MagicMock() - mock_executor_instance.__enter__.return_value = mock_executor_instance - mock_executor_instance.__exit__.return_value = None - mock_executor_instance.submit.return_value = mock_future - mock_dependencies["executor"].return_value = mock_executor_instance - - with ( - patch.object(runner, "_update_document_index_status") as mock_update_status, - pytest.raises(RuntimeError, match="index failed"), - ): - runner._load( - session=mock_dependencies["session"], - dataset=sample_dataset, - dataset_document=sample_dataset_document, - documents=sample_documents, - total_tokens=300, - ) - - mock_update_status.assert_not_called() - - def test_load_with_economy_indexing( - self, mock_dependencies, sample_dataset, sample_dataset_document, sample_documents - ): - """Test loading with economy indexing (keyword only).""" - # Arrange - runner = IndexingRunner() - sample_dataset.indexing_technique = IndexTechniqueType.ECONOMY - - # Mock thread for keyword indexing - mock_thread_instance = MagicMock() - mock_thread_instance.join = MagicMock() - mock_dependencies["thread"].return_value = mock_thread_instance - - # Mock update_document_index_status to avoid database calls - with patch.object(runner, "_update_document_index_status"): - # Act - runner._load( - session=mock_dependencies["session"], - dataset=sample_dataset, - dataset_document=sample_dataset_document, - documents=sample_documents, - total_tokens=0, - ) - - # Assert - # Verify keyword thread was created and joined - mock_dependencies["thread"].assert_called_once() - mock_thread_instance.start.assert_called_once() - mock_thread_instance.join.assert_called_once() - - def test_load_with_parent_child_index( - self, mock_dependencies, sample_dataset, sample_dataset_document, sample_documents - ): - """Test loading with parent-child index structure.""" - # Arrange - runner = IndexingRunner() - sample_dataset_document.doc_form = IndexStructureType.PARENT_CHILD_INDEX - sample_dataset.indexing_technique = IndexTechniqueType.HIGH_QUALITY - - # Add child documents - for doc in sample_documents: - doc.children = [ - ChildDocument( - page_content=f"Child of {doc.page_content}", - metadata={"doc_id": f"child_{doc.metadata['doc_id']}", "doc_hash": "child_hash"}, - ) - ] - - # Mock ThreadPoolExecutor - mock_future = MagicMock() - mock_future.result.return_value = None - mock_executor_instance = MagicMock() - mock_executor_instance.__enter__.return_value = mock_executor_instance - mock_executor_instance.__exit__.return_value = None - mock_executor_instance.submit.return_value = mock_future - mock_dependencies["executor"].return_value = mock_executor_instance - - # Mock update_document_index_status to avoid database calls - with patch.object(runner, "_update_document_index_status"): - # Act - runner._load( - session=mock_dependencies["session"], - dataset=sample_dataset, - dataset_document=sample_dataset_document, - documents=sample_documents, - total_tokens=150, - ) - - # Assert - mock_dependencies["model_manager"].assert_not_called() - # Verify no keyword thread for parent-child index - mock_dependencies["thread"].assert_not_called() - - -class TestIndexingRunnerRun: - """Unit tests for IndexingRunner.run method. - - Tests cover: - - Complete end-to-end indexing flow - - Error handling and recovery - - Document status transitions - - Pause detection - - Multiple document processing - """ - - @pytest.fixture - def mock_dependencies(self, sqlite_session: Session): - """Mock all external dependencies for run tests.""" - with ( - patch("core.indexing_runner.IndexProcessorFactory") as mock_factory, - patch("core.indexing_runner.ModelManager.for_tenant") as mock_model_manager, - patch("core.indexing_runner.storage") as mock_storage, - patch("core.indexing_runner.threading.Thread") as mock_thread, - ): - yield { - "session": sqlite_session, - "factory": mock_factory, - "model_manager": mock_model_manager, - "storage": mock_storage, - "thread": mock_thread, - } - - @pytest.fixture - def sample_dataset_documents(self, sqlite_session: Session): - """Create sample dataset documents for testing.""" - docs = [] - for i in range(2): - tenant_id = str(uuid.uuid4()) - dataset_id = str(uuid.uuid4()) - account = Account(name=f"Account {i}", email=f"account-{i}@example.com") - account.id = str(uuid.uuid4()) - dataset = Dataset( - id=dataset_id, - tenant_id=tenant_id, - name=f"Dataset {i}", - description="", - created_by=account.id, - indexing_technique=IndexTechniqueType.HIGH_QUALITY, - ) - process_rule = DatasetProcessRule( - dataset_id=dataset_id, - mode="automatic", - rules="{}", - created_by=account.id, - ) - document = DatasetDocument( - id=str(uuid.uuid4()), - dataset_id=dataset_id, - tenant_id=tenant_id, - position=1, - doc_form=IndexStructureType.PARAGRAPH_INDEX, - doc_language="English", - data_source_type=DataSourceType.UPLOAD_FILE, - data_source_info=json.dumps({"upload_file_id": str(uuid.uuid4())}), - dataset_process_rule_id=process_rule.id, - batch="batch", - name=f"Document {i}", - created_from=DocumentCreatedFrom.API, - created_by=account.id, - ) - sqlite_session.add_all([account, dataset, process_rule, document]) - docs.append(document) - sqlite_session.commit() - return docs - - def test_run_in_indexing_status_loads_child_chunks_with_caller_session( - self, mock_dependencies, sample_dataset_documents - ): - runner = IndexingRunner() - dataset_document = sample_dataset_documents[0] - dataset_document.doc_form = IndexStructureType.PARENT_CHILD_INDEX - session = mock_dependencies["session"] - dataset = session.get(Dataset, dataset_document.dataset_id) - assert dataset is not None - process_rule = session.get(DatasetProcessRule, dataset_document.dataset_process_rule_id) - assert process_rule is not None - process_rule.mode = "hierarchical" - process_rule.rules = json.dumps({"parent_mode": "paragraph"}) - segment = DocumentSegment( - tenant_id=dataset_document.tenant_id, - dataset_id=dataset_document.dataset_id, - document_id=dataset_document.id, - position=1, - content="parent", - word_count=0, - tokens=12, - created_by="account-id", - status="waiting", - index_node_id="parent-node", - index_node_hash="parent-hash", - ) - child_chunk = ChildChunk( - tenant_id=dataset_document.tenant_id, - dataset_id=dataset_document.dataset_id, - document_id=dataset_document.id, - segment_id=segment.id, - position=1, - content="child", - word_count=1, - created_by=dataset_document.created_by, - index_node_id="child-node", - index_node_hash="child-hash", - ) - session.add_all([segment, child_chunk]) - session.commit() - - with ( - patch.object(runner, "_load") as load, - ): - runner.run_in_indexing_status(dataset_document, session) - - assert load.call_args.kwargs["documents"][0].children[0].page_content == "child" - assert load.call_args.kwargs["total_tokens"] == 12 - - def test_run_in_indexing_status_uses_tokens_from_all_segments(self, mock_dependencies, sample_dataset_documents): - runner = IndexingRunner() - dataset_document = sample_dataset_documents[0] - session = mock_dependencies["session"] - dataset = session.get(Dataset, dataset_document.dataset_id) - assert dataset is not None - completed_segment = DocumentSegment( - tenant_id=dataset_document.tenant_id, - dataset_id=dataset_document.dataset_id, - document_id=dataset_document.id, - position=1, - content="", - word_count=0, - tokens=10, - created_by="account-id", - status=SegmentStatus.COMPLETED, - ) - incomplete_segment = DocumentSegment( - tenant_id="tenant-id", - dataset_id=dataset_document.dataset_id, - document_id=dataset_document.id, - position=1, - content="pending", - word_count=0, - tokens=20, - created_by="account-id", - status=SegmentStatus.WAITING, - index_node_id="pending-node", - index_node_hash="pending-hash", - ) - session.add_all([completed_segment, incomplete_segment]) - session.commit() - - with patch.object(runner, "_load") as load: - runner.run_in_indexing_status(dataset_document, session) - - assert load.call_args.kwargs["documents"][0].page_content == "pending" - assert load.call_args.kwargs["total_tokens"] == 30 - - @patch.object(Account, "set_tenant_id_with_session", autospec=True) - def test_run_success_single_document(self, set_tenant_id, mock_dependencies, sample_dataset_documents): - """Test successful run with single document.""" - # Arrange - runner = IndexingRunner() - doc = sample_dataset_documents[0] - session = mock_dependencies["session"] - mock_dataset = session.get(Dataset, doc.dataset_id) - mock_current_user = session.get(Account, doc.created_by) - assert mock_dataset is not None - assert mock_current_user is not None - mock_dataset.indexing_technique = IndexTechniqueType.ECONOMY - session.commit() - - # Mock processor - mock_processor = MagicMock() - mock_dependencies["factory"].return_value.init_index_processor.return_value = mock_processor - - # Mock extract, transform, load - mock_processor.extract.return_value = [Document(page_content="Test content", metadata={"doc_id": "doc1"})] - mock_processor.transform.return_value = [ - Document( - page_content="Chunk 1", - metadata={"doc_id": "chunk1", "doc_hash": "hash1"}, - ) - ] - - # Mock thread for keyword indexing - mock_thread_instance = MagicMock() - mock_dependencies["thread"].return_value = mock_thread_instance - - commits = 0 - - def record_commit(_session) -> None: - nonlocal commits - commits += 1 - - event.listen(session, "after_commit", record_commit) - with ( - patch.object(runner, "_extract", return_value=[Document(page_content="Test", metadata={})]) as mock_extract, - patch.object( - runner, - "_transform", - return_value=[Document(page_content="Chunk", metadata={"doc_id": "c1", "doc_hash": "h1"})], - ) as mock_transform, - patch.object(runner, "_load_segments") as mock_load_segments, - patch.object(runner, "_load") as mock_load, - ): - # Act - runner.run([doc], session) - event.remove(session, "after_commit", record_commit) - - # Assert - verify the methods were called - mock_extract.assert_called_once() - mock_transform.assert_called_once() - mock_load_segments.assert_called_once() - mock_load.assert_called_once() - assert commits == 2 - set_tenant_id.assert_called_once_with(mock_current_user, mock_dataset.tenant_id, session=session) - - mock_processor = MagicMock() - mock_dependencies["factory"].return_value.init_index_processor.return_value = mock_processor - - # Mock _extract to raise DocumentIsPausedError - with patch.object(runner, "_extract", side_effect=DocumentIsPausedError("Document paused")): - # Act & Assert - with pytest.raises(DocumentIsPausedError): - runner.run([doc], mock_dependencies["session"]) - - @patch.object(Account, "set_tenant_id_with_session", autospec=True) - def test_run_counts_each_transformed_document_once( - self, set_tenant_id, mock_dependencies, sample_dataset_documents - ): - runner = IndexingRunner() - dataset_document = sample_dataset_documents[0] - session = mock_dependencies["session"] - dataset = session.get(Dataset, dataset_document.dataset_id) - current_user = session.get(Account, dataset_document.created_by) - assert dataset is not None - assert current_user is not None - transformed_documents = [ - Document(page_content="first", metadata={"doc_id": "first", "doc_hash": "hash-first"}), - Document(page_content="second", metadata={"doc_id": "second", "doc_hash": "hash-second"}), - ] - - with ( - patch.object(runner, "_extract", return_value=[Document(page_content="source", metadata={})]), - patch.object(runner, "_transform", return_value=transformed_documents), - patch.object(runner, "_load_segments") as load_segments, - patch.object(runner, "_load") as load, - patch( - "core.indexing_runner.calculate_segment_token_counts", - return_value=[11, 22], - ) as calculate_token_counts, - ): - runner.run([dataset_document], mock_dependencies["session"]) - - calculate_token_counts.assert_called_once_with(dataset=dataset, documents=transformed_documents) - load_segments.assert_called_once_with( - session=mock_dependencies["session"], - dataset=dataset, - dataset_document=dataset_document, - documents=transformed_documents, - token_counts=[11, 22], - ) - assert load.call_args.kwargs["total_tokens"] == 33 - set_tenant_id.assert_called_once_with( - current_user, - dataset.tenant_id, - session=mock_dependencies["session"], - ) - - @patch.object(Account, "set_tenant_id_with_session", autospec=True) - def test_run_rejects_before_segment_or_vector_writes( - self, set_tenant_id, mock_dependencies, sample_dataset_documents - ): - runner = IndexingRunner(enforce_vector_space_admission=True) - dataset_document = sample_dataset_documents[0] - dataset_document.need_summary = False - session = mock_dependencies["session"] - dataset = session.get(Dataset, dataset_document.dataset_id) - current_user = session.get(Account, dataset_document.created_by) - assert dataset is not None - assert current_user is not None - transformed_documents = [Document(page_content="Chunk", metadata={"doc_id": "c1", "doc_hash": "h1"})] - admission_error = VectorSpaceAdmissionError("estimated storage exceeds capacity") - admission_service = Mock() - admission_service.ensure_document_can_be_indexed.side_effect = admission_error - - with ( - patch("core.indexing_runner.VectorSpaceAdmissionService", return_value=admission_service), - patch.object(runner, "_extract", return_value=[Document(page_content="source", metadata={})]), - patch.object( - runner, - "_transform", - return_value=transformed_documents, - ), - patch.object(runner, "_load_segments") as load_segments, - patch.object(runner, "_load") as load, - patch.object(runner, "_handle_indexing_error") as handle_error, - ): - runner.run([dataset_document], mock_dependencies["session"]) - - load_segments.assert_not_called() - load.assert_not_called() - admission_service.ensure_document_can_be_indexed.assert_called_once_with( - dataset=dataset, - document_id=dataset_document.id, - doc_form=dataset_document.doc_form, - documents=transformed_documents, - include_summaries=False, - session=mock_dependencies["session"], - ) - handle_error.assert_called_once_with(dataset_document.id, admission_error, mock_dependencies["session"]) - set_tenant_id.assert_called_once_with( - current_user, - dataset.tenant_id, - session=mock_dependencies["session"], - ) - - @patch.object(Account, "set_tenant_id_with_session", autospec=True) - def test_run_in_splitting_status_counts_each_transformed_document_once( - self, set_tenant_id, mock_dependencies, sample_dataset_documents - ): - runner = IndexingRunner() - dataset_document = sample_dataset_documents[0] - session = mock_dependencies["session"] - dataset = session.get(Dataset, dataset_document.dataset_id) - current_user = session.get(Account, dataset_document.created_by) - assert dataset is not None - assert current_user is not None - transformed_documents = [ - Document(page_content="first", metadata={"doc_id": "first", "doc_hash": "hash-first"}), - Document(page_content="second", metadata={"doc_id": "second", "doc_hash": "hash-second"}), - ] - - with ( - patch.object(runner, "_extract", return_value=[Document(page_content="source", metadata={})]), - patch.object(runner, "_transform", return_value=transformed_documents), - patch.object(runner, "_load_segments") as load_segments, - patch.object(runner, "_load") as load, - patch( - "core.indexing_runner.calculate_segment_token_counts", - return_value=[11, 22], - ) as calculate_token_counts, - ): - runner.run_in_splitting_status(dataset_document, mock_dependencies["session"]) - - calculate_token_counts.assert_called_once_with(dataset=dataset, documents=transformed_documents) - load_segments.assert_called_once_with( - session=mock_dependencies["session"], - dataset=dataset, - dataset_document=dataset_document, - documents=transformed_documents, - token_counts=[11, 22], - ) - assert load.call_args.kwargs["total_tokens"] == 33 - set_tenant_id.assert_called_once_with( - current_user, - dataset.tenant_id, - session=mock_dependencies["session"], - ) - - def test_run_handles_provider_token_error(self, mock_dependencies, sample_dataset_documents): - """Test run handles ProviderTokenNotInitError and updates document status.""" - # Arrange - runner = IndexingRunner() - doc = sample_dataset_documents[0] - - mock_processor = MagicMock() - mock_dependencies["factory"].return_value.init_index_processor.return_value = mock_processor - mock_processor.extract.side_effect = ProviderTokenNotInitError("Token not initialized") - - # Act - with patch.object(runner, "_extract", side_effect=ProviderTokenNotInitError("Token not initialized")): - runner.run([doc], mock_dependencies["session"]) - - mock_dependencies["session"].refresh(doc) - assert doc.indexing_status == "error" - - def test_run_handles_object_deleted_error(self, mock_dependencies, sample_dataset_documents): - """Test run handles ObjectDeletedError gracefully.""" - # Arrange - runner = IndexingRunner() - doc = sample_dataset_documents[0] - - mock_processor = MagicMock() - mock_dependencies["factory"].return_value.init_index_processor.return_value = mock_processor - - # Mock _extract to raise ObjectDeletedError - with patch.object(runner, "_extract", side_effect=ObjectDeletedError(state=None, msg="Object deleted")): - # Act - runner.run([doc], mock_dependencies["session"]) - - # Assert - should not raise, just log warning - # No exception should be raised - - @patch.object(Account, "set_tenant_id_with_session", autospec=True) - def test_run_processes_multiple_documents(self, set_tenant_id, mock_dependencies, sample_dataset_documents): - """Test run processes multiple documents sequentially.""" - # Arrange - runner = IndexingRunner() - docs = sample_dataset_documents - for doc in docs: - dataset = mock_dependencies["session"].get(Dataset, doc.dataset_id) - assert dataset is not None - dataset.indexing_technique = IndexTechniqueType.ECONOMY - mock_dependencies["session"].commit() - - mock_processor = MagicMock() - mock_dependencies["factory"].return_value.init_index_processor.return_value = mock_processor - - # Mock thread - mock_thread_instance = MagicMock() - mock_dependencies["thread"].return_value = mock_thread_instance - - # Mock all internal methods - with ( - patch.object(runner, "_extract", return_value=[Document(page_content="Test", metadata={})]) as mock_extract, - patch.object( - runner, - "_transform", - return_value=[Document(page_content="Chunk", metadata={"doc_id": "c1", "doc_hash": "h1"})], - ), - patch.object(runner, "_load_segments"), - patch.object(runner, "_load"), - ): - # Act - runner.run(docs, mock_dependencies["session"]) - - # Assert - # Verify extract was called for each document - assert mock_extract.call_count == len(docs) - assert set_tenant_id.call_count == len(docs) - - -class TestIndexingRunnerRetryLogic: - """Unit tests for retry logic and error handling. - - Tests cover: - - Document pause status checking - - Document status updates - - Error state persistence - - Deleted document handling - """ - - @pytest.fixture - def mock_dependencies(self, sqlite_session: Session): - """Mock all external dependencies.""" - with ( - patch("core.indexing_runner.redis_client") as mock_redis, - ): - yield { - "session": sqlite_session, - "redis": mock_redis, - } - - def test_check_document_paused_status_not_paused(self, mock_dependencies): - """Test document pause check when document is not paused.""" - # Arrange - mock_dependencies["redis"].get.return_value = None - document_id = str(uuid.uuid4()) - - # Act & Assert - should not raise - IndexingRunner._check_document_paused_status(document_id) - - def test_check_document_paused_status_is_paused(self, mock_dependencies): - """Test document pause check when document is paused.""" - # Arrange - mock_dependencies["redis"].get.return_value = "1" - document_id = str(uuid.uuid4()) - - # Act & Assert - with pytest.raises(DocumentIsPausedError): - IndexingRunner._check_document_paused_status(document_id) - - def test_update_document_index_status_success(self, mock_dependencies): - """Test successful document status update.""" - document_id = str(uuid.uuid4()) - session = mock_dependencies["session"] - document = DatasetDocument( - id=document_id, - tenant_id=str(uuid.uuid4()), - dataset_id=str(uuid.uuid4()), - position=1, - data_source_type=DataSourceType.UPLOAD_FILE, - batch="batch", - name="Document", - created_from=DocumentCreatedFrom.API, - created_by=str(uuid.uuid4()), - is_paused=False, - ) - session.add(document) - session.commit() - - # Act - IndexingRunner._update_document_index_status( - document_id, - "completed", - {"tokens": 100, "completed_at": naive_utc_now()}, - session=session, - ) - - session.refresh(document) - assert document.indexing_status == "completed" - assert document.tokens == 100 - - def test_update_document_index_status_paused(self, mock_dependencies): - """Test document status update when document is paused.""" - document_id = str(uuid.uuid4()) - session = mock_dependencies["session"] - document = DatasetDocument( - id=document_id, - tenant_id=str(uuid.uuid4()), - dataset_id=str(uuid.uuid4()), - position=1, - data_source_type=DataSourceType.UPLOAD_FILE, - batch="batch", - name="Paused document", - created_from=DocumentCreatedFrom.API, - created_by=str(uuid.uuid4()), - is_paused=True, - ) - session.add(document) - session.commit() - - # Act & Assert - with pytest.raises(DocumentIsPausedError): - IndexingRunner._update_document_index_status(document_id, "completed", session=session) - - def test_update_document_index_status_deleted(self, mock_dependencies): - """Test document status update when document is deleted.""" - # Arrange - document_id = str(uuid.uuid4()) - # Act & Assert - with pytest.raises(DocumentIsDeletedPausedError): - IndexingRunner._update_document_index_status(document_id, "completed", session=mock_dependencies["session"]) - - -class TestIndexingRunnerDocumentCleaning: - """Unit tests for document cleaning and preprocessing. - - Tests cover: - - Text cleaning rules - - Whitespace normalization - - Special character handling - - Custom preprocessing rules - """ - - @pytest.fixture - def sample_process_rule_automatic(self): - """Create automatic processing rule.""" - rule = DatasetProcessRule( - dataset_id="dataset-id", - mode="automatic", - rules=None, - created_by="account-id", - ) - return rule - - @pytest.fixture - def sample_process_rule_custom(self): - """Create custom processing rule.""" - rule = DatasetProcessRule( - dataset_id="dataset-id", - mode="custom", - rules=json.dumps( - { - "pre_processing_rules": [ - {"id": "remove_extra_spaces", "enabled": True}, - {"id": "remove_urls_emails", "enabled": True}, - ] - } - ), - created_by="account-id", - ) - return rule - - def test_document_clean_automatic_mode(self, sample_process_rule_automatic): - """Test document cleaning with automatic mode.""" - # Arrange - text = "This is a test document with extra spaces." - - # Act - with patch("core.indexing_runner.CleanProcessor.clean") as mock_clean: - mock_clean.return_value = "This is a test document with extra spaces." - result = IndexingRunner._document_clean(text, sample_process_rule_automatic) - - # Assert - assert "extra spaces" in result - mock_clean.assert_called_once() - - def test_document_clean_custom_mode(self, sample_process_rule_custom): - """Test document cleaning with custom rules.""" - # Arrange - text = "Visit https://example.com or email test@example.com for more info." - - # Act - with patch("core.indexing_runner.CleanProcessor.clean") as mock_clean: - mock_clean.return_value = "Visit or email for more info." - result = IndexingRunner._document_clean(text, sample_process_rule_custom) - - # Assert - assert "https://" not in result - assert "@" not in result - mock_clean.assert_called_once() - - def test_filter_string_removes_special_characters(self): - """Test filter_string removes special control characters.""" - # Arrange - text = "Normal text\x00with\x08control\x1fcharacters\x7f" - - # Act - result = IndexingRunner.filter_string(text) - - # Assert - assert "\x00" not in result - assert "\x08" not in result - assert "\x1f" not in result - assert "\x7f" not in result - assert "Normal text" in result - - def test_filter_string_handles_unicode_fffe(self): - """Test filter_string removes Unicode U+FFFE.""" - # Arrange - text = "Text with \ufffe unicode issue" - - # Act - result = IndexingRunner.filter_string(text) - - # Assert - assert "\ufffe" not in result - assert "Text with" in result - - def test_filter_string_preserves_valid_extended_characters(self): - """filter_string must keep valid printable characters like 'ï', '¿', '¾'.""" - # Arrange - text = "naïve ¿Cómo? ¾ done" - - # Act - result = IndexingRunner.filter_string(text) - - # Assert - assert result == text - # The U+FFFE noncharacter is still stripped. - assert IndexingRunner.filter_string("keep\ufffedrop") == "keepdrop" - - -class TestIndexingRunnerSplitter: - """Unit tests for text splitter configuration. - - Tests cover: - - Custom segmentation rules - - Automatic segmentation - - Chunk size validation - - Separator handling - """ - - @pytest.fixture - def mock_embedding_instance(self): - """Create mock embedding model instance.""" - instance = MagicMock() - instance.get_text_embedding_num_tokens.return_value = 100 - return instance - - def test_get_splitter_custom_mode(self, mock_embedding_instance): - """Test splitter creation with custom mode.""" - # Arrange - with patch("core.indexing_runner.FixedRecursiveCharacterTextSplitter") as mock_splitter_class: - mock_splitter = MagicMock() - mock_splitter_class.from_encoder.return_value = mock_splitter - - # Act - result = IndexingRunner._get_splitter( - processing_rule_mode="custom", - max_tokens=500, - chunk_overlap=50, - separator="\\n\\n", - embedding_model_instance=mock_embedding_instance, - ) - - # Assert - assert result == mock_splitter - mock_splitter_class.from_encoder.assert_called_once() - call_kwargs = mock_splitter_class.from_encoder.call_args[1] - assert call_kwargs["chunk_size"] == 500 - assert call_kwargs["chunk_overlap"] == 50 - assert call_kwargs["fixed_separator"] == "\n\n" - - def test_get_splitter_automatic_mode(self, mock_embedding_instance): - """Test splitter creation with automatic mode.""" - # Arrange - with patch("core.indexing_runner.EnhanceRecursiveCharacterTextSplitter") as mock_splitter_class: - mock_splitter = MagicMock() - mock_splitter_class.from_encoder.return_value = mock_splitter - - # Act - result = IndexingRunner._get_splitter( - processing_rule_mode="automatic", - max_tokens=500, - chunk_overlap=50, - separator="", - embedding_model_instance=mock_embedding_instance, - ) - - # Assert - assert result == mock_splitter - mock_splitter_class.from_encoder.assert_called_once() - - def test_get_splitter_validates_max_tokens_too_small(self, mock_embedding_instance): - """Test splitter validation rejects max_tokens below minimum.""" - # Act & Assert - with pytest.raises(ValueError, match="Custom segment length should be between"): - IndexingRunner._get_splitter( - processing_rule_mode="custom", - max_tokens=30, # Below minimum of 50 - chunk_overlap=10, - separator="\\n", - embedding_model_instance=mock_embedding_instance, - ) - - def test_get_splitter_validates_max_tokens_too_large(self, mock_embedding_instance, config_overrides): - """Test splitter validation rejects max_tokens above maximum.""" - config_overrides(INDEXING_MAX_SEGMENTATION_TOKENS_LENGTH=5000) - with pytest.raises(ValueError, match="Custom segment length should be between"): - IndexingRunner._get_splitter( - processing_rule_mode="custom", - max_tokens=10000, # Above maximum - chunk_overlap=100, - separator="\\n", - embedding_model_instance=mock_embedding_instance, - ) - - -class TestIndexingRunnerLoadSegments: - """Unit tests for segment loading and storage. - - Tests cover: - - Segment creation in database - - Child chunk handling - - Document status updates - - Word count calculation - """ - - @pytest.fixture - def mock_dependencies(self, unbound_session: Session): - """Mock all external dependencies.""" - with ( - patch("core.indexing_runner.DatasetDocumentStore") as mock_docstore, - ): - yield { - "session": unbound_session, - "docstore": mock_docstore, - } - - @pytest.fixture - def sample_dataset(self): - """Create sample dataset.""" - dataset = Dataset( - id=str(uuid.uuid4()), - tenant_id=str(uuid.uuid4()), - ) - return dataset - - @pytest.fixture - def sample_dataset_document(self): - """Create sample dataset document.""" - return DatasetDocument( - id=str(uuid.uuid4()), - dataset_id=str(uuid.uuid4()), - created_by=str(uuid.uuid4()), - doc_form=IndexStructureType.PARAGRAPH_INDEX, - ) - - @pytest.fixture - def sample_documents(self): - """Create sample documents.""" - return [ - Document( - page_content="This is chunk 1 with some content.", - metadata={"doc_id": "chunk1", "doc_hash": "hash1"}, - ), - Document( - page_content="This is chunk 2 with different content.", - metadata={"doc_id": "chunk2", "doc_hash": "hash2"}, - ), - ] - - def test_load_segments_paragraph_index( - self, mock_dependencies, sample_dataset, sample_dataset_document, sample_documents - ): - """Test loading segments for paragraph index.""" - # Arrange - runner = IndexingRunner() - mock_docstore_instance = MagicMock() - mock_dependencies["docstore"].return_value = mock_docstore_instance - - # Mock update methods to avoid database calls - with ( - patch.object(runner, "_update_document_index_status"), - patch.object(runner, "_update_segments_by_document"), - ): - # Act - runner._load_segments( - session=mock_dependencies["session"], - dataset=sample_dataset, - dataset_document=sample_dataset_document, - documents=sample_documents, - token_counts=[10, 20], - ) - - # Assert - mock_dependencies["docstore"].assert_called_once_with( - dataset=sample_dataset, - user_id=sample_dataset_document.created_by, - document_id=sample_dataset_document.id, - ) - mock_docstore_instance.add_documents.assert_called_once_with( - session=mock_dependencies["session"], - docs=sample_documents, - save_child=False, - token_counts=[10, 20], - ) - - def test_load_segments_parent_child_index( - self, mock_dependencies, sample_dataset, sample_dataset_document, sample_documents - ): - """Test loading segments for parent-child index.""" - # Arrange - runner = IndexingRunner() - sample_dataset_document.doc_form = IndexStructureType.PARENT_CHILD_INDEX - - # Add child documents - for doc in sample_documents: - doc.children = [ - ChildDocument( - page_content=f"Child of {doc.page_content}", - metadata={"doc_id": f"child_{doc.metadata['doc_id']}", "doc_hash": "child_hash"}, - ) - ] - - mock_docstore_instance = MagicMock() - mock_dependencies["docstore"].return_value = mock_docstore_instance - - # Mock update methods to avoid database calls - with ( - patch.object(runner, "_update_document_index_status"), - patch.object(runner, "_update_segments_by_document"), - ): - # Act - runner._load_segments( - session=mock_dependencies["session"], - dataset=sample_dataset, - dataset_document=sample_dataset_document, - documents=sample_documents, - token_counts=[10, 20], - ) - - # Assert - mock_docstore_instance.add_documents.assert_called_once_with( - session=mock_dependencies["session"], - docs=sample_documents, - save_child=True, - token_counts=[10, 20], - ) - - def test_load_segments_updates_word_count( - self, mock_dependencies, sample_dataset, sample_dataset_document, sample_documents - ): - """Test load segments calculates and updates word count.""" - # Arrange - runner = IndexingRunner() - mock_docstore_instance = MagicMock() - mock_dependencies["docstore"].return_value = mock_docstore_instance - - # Calculate expected word count - expected_word_count = sum(len(doc.page_content.split()) for doc in sample_documents) - - # Mock update methods to avoid database calls - with ( - patch.object(runner, "_update_document_index_status") as mock_update_status, - patch.object(runner, "_update_segments_by_document"), - ): - # Act - runner._load_segments( - session=mock_dependencies["session"], - dataset=sample_dataset, - dataset_document=sample_dataset_document, - documents=sample_documents, - token_counts=[10, 20], - ) - - # Assert - # Verify word count was calculated correctly and passed to status update - mock_update_status.assert_called_once() - call_kwargs = mock_update_status.call_args.kwargs - assert "extra_update_params" in call_kwargs - - -class TestIndexingRunnerEstimate: - """Unit tests for indexing estimation. - - Tests cover: - - Token estimation - - Segment count estimation - - Batch upload limit enforcement - """ - - @pytest.fixture - def mock_dependencies(self, sqlite_session: Session): - """Mock all external dependencies.""" - with ( - patch("core.indexing_runner.IndexProcessorFactory") as mock_factory, - ): - yield { - "session": sqlite_session, - "factory": mock_factory, - } - - def test_indexing_estimate_respects_batch_limit(self, mock_dependencies, config_overrides): - """Test indexing estimate enforces batch upload limit.""" - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, BATCH_UPLOAD_LIMIT=10) - # Arrange - runner = IndexingRunner() - tenant_id = str(uuid.uuid4()) - - extract_settings = [MagicMock() for _ in range(15)] - - with pytest.raises(ValueError, match="batch upload limit"): - runner.indexing_estimate( - tenant_id=tenant_id, - extract_settings=extract_settings, - tmp_processing_rule={"mode": "automatic", "rules": {}}, - doc_form=IndexStructureType.PARAGRAPH_INDEX, - session=mock_dependencies["session"], - ) - - def test_indexing_estimate_limits_qa_preview(self, mock_dependencies, config_overrides): - """Test indexing estimate returns no more than ten QA preview items.""" - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - runner = IndexingRunner() - mock_processor = MagicMock() - mock_dependencies["factory"].return_value.init_index_processor.return_value = mock_processor - - qa_documents = [ - Document( - page_content=f"Question {index}", - metadata={"answer": f"Answer {index}"}, - ) - for index in range(11) - ] - mock_processor.extract.return_value = [Document(page_content="Source content")] - mock_processor.transform.return_value = qa_documents - - result = runner.indexing_estimate( - tenant_id=str(uuid.uuid4()), - extract_settings=[MagicMock()], - tmp_processing_rule=create_mock_process_rule(), - doc_form=IndexStructureType.QA_INDEX, - session=mock_dependencies["session"], - ) - - assert result.total_segments == 220 - assert result.qa_preview is not None - assert len(result.qa_preview) == 10 - assert result.qa_preview[-1].question == "Question 9" - - def test_indexing_estimate_commits_preview_cleanup_before_summary_workers( - self, mock_dependencies, config_overrides - ): - """Test preview cleanup is visible before summary workers use independent sessions.""" - config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - runner = IndexingRunner() - tenant_id = str(uuid.uuid4()) - mock_processor = MagicMock() - mock_dependencies["factory"].return_value.init_index_processor.return_value = mock_processor - phase_events: list[str] = [] - session = mock_dependencies["session"] - event.listen(session, "after_commit", lambda _session: phase_events.append("commit")) - - preview_doc = Document( - page_content="![image](http://files.local/files/image-1/file-preview)", - metadata={}, - ) - mock_processor.extract.return_value = [preview_doc] - mock_processor.transform.return_value = [preview_doc] - mock_processor.generate_summary_preview.side_effect = lambda *_args, **_kwargs: ( - phase_events.append("summary") or [PreviewDetail(content=preview_doc.page_content)] - ) - - image_file = UploadFile( - tenant_id=tenant_id, - storage_type=StorageType.LOCAL, - key="image_files/tenant-1/source-file-1/image.png", - name="image.png", - size=10, - extension="png", - mime_type="image/png", - created_by_role=CreatorUserRole.ACCOUNT, - created_by="user-id", - created_at=datetime.now(UTC), - used=True, - ) - image_file.id = "image-1" - session.add(image_file) - session.commit() - phase_events.clear() - - with ( - patch("core.indexing_runner.get_image_upload_file_ids", return_value=["image-1"]), - patch("core.indexing_runner.storage") as mock_storage, - ): - result = runner.indexing_estimate( - tenant_id=tenant_id, - extract_settings=[MagicMock()], - tmp_processing_rule={ - "mode": "automatic", - "rules": {}, - "summary_index_setting": {"enable": True}, - }, - doc_form=IndexStructureType.PARAGRAPH_INDEX, - session=session, - ) - - assert result.total_segments == 1 - mock_storage.delete.assert_called_once_with(image_file.key) - assert session.get(UploadFile, image_file.id) is None - assert phase_events == ["commit", "summary"] - - -class TestIndexingRunnerProcessChunk: - """Unit tests for chunk processing in parallel. - - Tests cover: - - Vector index creation - - Segment status updates - - Pause detection during processing - """ - - @pytest.fixture - def mock_dependencies(self, sqlite_session: Session): - """Mock all external dependencies.""" - with ( - patch("core.indexing_runner.redis_client") as mock_redis, - ): - yield { - "session": sqlite_session, - "redis": mock_redis, - } - - @pytest.fixture - def mock_flask_app(self): - """Create mock Flask app context.""" - app = MagicMock() - app.app_context.return_value.__enter__ = MagicMock() - app.app_context.return_value.__exit__ = MagicMock() - return app - - def test_process_chunk_loads_index_and_completes_segments( - self, - mock_dependencies, - mock_flask_app, - sqlite_session_factory: sessionmaker[Session], - ): - """Test process chunk loads the index and completes segments without counting tokens.""" - # Arrange - from core.indexing_runner import IndexingRunner - - runner = IndexingRunner() - mock_processor = MagicMock() - chunk_documents = [ - Document(page_content="Chunk 1", metadata={"doc_id": "c1"}), - Document(page_content="Chunk 2", metadata={"doc_id": "c2"}), - ] - - session = mock_dependencies["session"] - mock_dataset, mock_dataset_document = persist_indexing_scope(session) - segments = [ - DocumentSegment( - tenant_id=mock_dataset.tenant_id, - dataset_id=mock_dataset.id, - document_id=mock_dataset_document.id, - position=position, - content=f"Chunk {position}", - word_count=1, - tokens=1, - created_by=mock_dataset.created_by, - index_node_id=doc.metadata["doc_id"], - status=SegmentStatus.INDEXING, - ) - for position, doc in enumerate(chunk_documents, start=1) - ] - session.add_all(segments) - session.commit() - - mock_dependencies["redis"].get.return_value = None - - # Create a proper context manager mock - mock_context = MagicMock() - mock_context.__enter__ = MagicMock(return_value=None) - mock_context.__exit__ = MagicMock(return_value=None) - mock_flask_app.app_context.return_value = mock_context - - with ( - patch("core.indexing_runner.session_factory.create_session", sqlite_session_factory), - patch("core.indexing_runner.IndexProcessorFactory") as mock_factory, - ): - mock_factory.return_value.init_index_processor.return_value = mock_processor - - # Act - the method creates its own app_context and session - result = runner._process_chunk( - mock_flask_app, - IndexStructureType.PARAGRAPH_INDEX, - chunk_documents, - mock_dataset.id, - mock_dataset_document.id, - ) - - # Assert - assert result is None - mock_processor.load.assert_called_once() - session.expire_all() - assert all(session.get(DocumentSegment, segment.id).status == SegmentStatus.COMPLETED for segment in segments) - - def test_process_chunk_detects_pause( - self, - mock_dependencies, - mock_flask_app, - sqlite_session_factory: sessionmaker[Session], - ): - """Test process chunk detects document pause.""" - # Arrange - from core.indexing_runner import IndexingRunner - - runner = IndexingRunner() - chunk_documents = [Document(page_content="Chunk", metadata={"doc_id": "c1"})] - - mock_dataset, mock_dataset_document = persist_indexing_scope(mock_dependencies["session"]) - - # Mock Redis to return paused status - mock_dependencies["redis"].get.return_value = "1" - # Create a proper context manager mock - mock_context = MagicMock() - mock_context.__enter__ = MagicMock(return_value=None) - mock_context.__exit__ = MagicMock(return_value=None) - mock_flask_app.app_context.return_value = mock_context - - with patch("core.indexing_runner.session_factory.create_session", sqlite_session_factory): - # Act & Assert - the method creates its own app_context and session - with pytest.raises(DocumentIsPausedError): - runner._process_chunk( - mock_flask_app, - IndexStructureType.PARAGRAPH_INDEX, - chunk_documents, - mock_dataset.id, - mock_dataset_document.id, - ) diff --git a/api/tests/unit_tests/core/rag/retrieval/test_dataset_retrieval_methods.py b/api/tests/unit_tests/core/rag/retrieval/test_dataset_retrieval_methods.py index 4bf883cba6d..01bc89c2831 100644 --- a/api/tests/unit_tests/core/rag/retrieval/test_dataset_retrieval_methods.py +++ b/api/tests/unit_tests/core/rag/retrieval/test_dataset_retrieval_methods.py @@ -8,6 +8,8 @@ import pytest from sqlalchemy import select from sqlalchemy.orm import Session, sessionmaker +from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom +from core.app.file_access import FileAccessScope, bind_file_access_scope, get_current_file_access_scope from core.rag.index_processor.constant.index_type import IndexTechniqueType from core.rag.models.document import Document from core.rag.retrieval import dataset_retrieval as retrieval_module @@ -322,17 +324,22 @@ class TestDatasetRetrievalKnowledgeRetrieval: "format_retrieval_documents", MagicMock(return_value=[record]), ) - grant_access = MagicMock() - monkeypatch.setattr(retrieval_module, "grant_retriever_segment_access", grant_access) - - with retrieval_database.session_maker() as caller_session: + scope = FileAccessScope( + tenant_id=dataset.tenant_id, + user_id="visitor", + user_from=UserFrom.END_USER, + invoke_from=InvokeFrom.WEB_APP, + ) + with bind_file_access_scope(scope), retrieval_database.session_maker() as caller_session: result = retrieval.knowledge_retrieval(caller_session, _request(dataset_ids=[dataset.id])) + granted_scope = get_current_file_access_scope() assert len(result) == 1 assert result[0].metadata.dataset_id == dataset.id assert result[0].metadata.document_id == document.id assert result[0].title == document.name - grant_access.assert_called_once_with([segment.id]) + assert granted_scope is not None + assert granted_scope.granted_retriever_segment_ids == frozenset([segment.id]) def test_metadata_filtering_disabled_skips_filter_builder( self, diff --git a/api/tests/unit_tests/core/tools/utils/test_misc_utils_extra.py b/api/tests/unit_tests/core/tools/utils/test_misc_utils_extra.py index 1c04b710806..f19950b5c12 100644 --- a/api/tests/unit_tests/core/tools/utils/test_misc_utils_extra.py +++ b/api/tests/unit_tests/core/tools/utils/test_misc_utils_extra.py @@ -363,7 +363,11 @@ def test_single_dataset_retriever_non_economy_run_sorts_context_and_resources(sq "format_retrieval_documents", return_value=records, ), - patch.object(DocumentSegment, "get_sign_content", lambda segment: segment.content.replace("raw", "signed")), + patch.object( + single_retriever_module, + "sign_segment_content", + lambda segment, **_kwargs: segment.content.replace("raw", "signed"), + ), ): result = tool.run(session=sqlite_session, query="hello") @@ -560,7 +564,11 @@ def test_multi_dataset_retriever_run_orders_segments_and_returns_resources(sqlit fake_current_app = _FakeCurrentApp() with ( - patch.object(DocumentSegment, "get_sign_content", lambda segment: segment.content.replace("raw", "signed")), + patch.object( + multi_retriever_module, + "sign_segment_content", + lambda segment, **_kwargs: segment.content.replace("raw", "signed"), + ), patch.object(tool, "_retriever", side_effect=fake_retriever) as retriever_mock, patch.object(multi_retriever_module, "current_app", fake_current_app), patch.object(multi_retriever_module.threading, "Thread", _ImmediateThread), diff --git a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_validators.py b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_validators.py index 3024db1440c..328832171e8 100644 --- a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_validators.py +++ b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_validators.py @@ -687,7 +687,7 @@ def test_publish_validation_rejects_missing_or_out_of_scope_knowledge_datasets( captured["tenant_id"] = tenant_id return [], 0 - import services.dataset_service as dataset_service_module + import services.knowledge.dataset_service as dataset_service_module monkeypatch.setattr(dataset_service_module.DatasetService, "get_datasets_by_ids", fake_get_datasets_by_ids) diff --git a/api/tests/unit_tests/core/workflow/nodes/datasource/test_datasource_node.py b/api/tests/unit_tests/core/workflow/nodes/datasource/test_datasource_node.py index 235d56e9897..d6987f0d49a 100644 --- a/api/tests/unit_tests/core/workflow/nodes/datasource/test_datasource_node.py +++ b/api/tests/unit_tests/core/workflow/nodes/datasource/test_datasource_node.py @@ -92,6 +92,7 @@ def test_datasource_node_delegates_to_manager_stream(mocker: MockerFixture): ), graph_init_params=gp, graph_runtime_state=gs, + datasource_credentials=lambda _tenant, _provider, _plugin, _credential: {}, ) evts = list(node._run()) diff --git a/api/tests/unit_tests/core/workflow/nodes/datasource/test_datasource_node_stream.py b/api/tests/unit_tests/core/workflow/nodes/datasource/test_datasource_node_stream.py index a5ca12e851a..4a26d3d29e1 100644 --- a/api/tests/unit_tests/core/workflow/nodes/datasource/test_datasource_node_stream.py +++ b/api/tests/unit_tests/core/workflow/nodes/datasource/test_datasource_node_stream.py @@ -2,11 +2,16 @@ import time from collections.abc import Generator +from types import SimpleNamespace from typing import Never +import pytest from pytest_mock import MockerFixture +from sqlalchemy import Engine +from sqlalchemy.orm import sessionmaker from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom +from core.workflow.node_factory import DifyNodeFactory from core.workflow.nodes.datasource.datasource_node import DatasourceNode from core.workflow.nodes.datasource.entities import DatasourceNodeData from core.workflow.system_variables import build_system_variables @@ -16,14 +21,15 @@ from graphon.runtime import GraphRuntimeState, VariablePool from tests.workflow_test_utils import build_test_graph_init_params -def test_node_integration_minimal_stream(mocker: MockerFixture) -> None: +@pytest.mark.parametrize("use_factory", [False, True]) +def test_node_integration_minimal_stream(mocker: MockerFixture, sqlite_engine: Engine, use_factory: bool) -> None: variable_pool = VariablePool.from_bootstrap( system_variables=build_system_variables( datasource_type="online_document", datasource_info={ "workspace_id": "w", "page": {"page_id": "pg", "type": "t"}, - "credential_id": "", + "credential_id": "credential-1", }, ) ) @@ -43,7 +49,9 @@ def test_node_integration_minimal_stream(mocker: MockerFixture) -> None: return "icon" @classmethod - def stream_node_events(cls, **_: object) -> Generator[StreamCompletedEvent, None, None]: + def stream_node_events(cls, **kwargs: object) -> Generator[StreamCompletedEvent, None, None]: + assert kwargs["tenant_id"] == "t1" + assert kwargs["credentials"] == {"integration_secret": "token"} yield from () yield StreamCompletedEvent(node_run_result=NodeRunResult(status=WorkflowNodeExecutionStatus.SUCCEEDED)) @@ -53,20 +61,36 @@ def test_node_integration_minimal_stream(mocker: MockerFixture) -> None: mocker.patch("core.workflow.nodes.datasource.datasource_node.DatasourceManager", new=_Mgr) - node = DatasourceNode( - node_id="n", - data=DatasourceNodeData( - type="datasource", - version="1", - title="Datasource", - provider_type="plugin", - provider_name="p", - plugin_id="plug", - datasource_name="ds", - ), - graph_init_params=graph_init_params, - graph_runtime_state=graph_runtime_state, + load_credentials = mocker.Mock(return_value={"integration_secret": "token"}) + node_data = DatasourceNodeData( + type="datasource", + version="1", + title="Datasource", + provider_type="plugin", + provider_name="p", + plugin_id="plug", + datasource_name="ds", ) + if use_factory: + database_client = sessionmaker(sqlite_engine) + mocker.patch("core.workflow.node_factory.get_session_maker", return_value=database_client) + build_credentials = mocker.patch( + "core.workflow.node_factory.build_data_source_credentials", + return_value=SimpleNamespace(providers=SimpleNamespace(get_datasource_credentials=load_credentials)), + ) + factory = DifyNodeFactory(graph_init_params=graph_init_params, graph_runtime_state=graph_runtime_state) + node = factory.create_node({"id": "n", "data": node_data.model_dump()}) + build_credentials.assert_called_once_with(database_client=database_client) + assert isinstance(node, DatasourceNode) + else: + node = DatasourceNode( + node_id="n", + data=node_data, + graph_init_params=graph_init_params, + graph_runtime_state=graph_runtime_state, + datasource_credentials=load_credentials, + ) out = list(node._run()) assert isinstance(out[-1], StreamCompletedEvent) + load_credentials.assert_called_once_with("t1", "p", "plug", "credential-1") diff --git a/api/tests/unit_tests/events/event_handlers/test_create_document_index.py b/api/tests/unit_tests/events/event_handlers/test_create_document_index.py index fad8932f9e2..c381796ff6b 100644 --- a/api/tests/unit_tests/events/event_handlers/test_create_document_index.py +++ b/api/tests/unit_tests/events/event_handlers/test_create_document_index.py @@ -1,13 +1,15 @@ import logging +from collections.abc import Sequence from unittest.mock import MagicMock, patch import pytest from sqlalchemy.orm import Session -from core.indexing_runner import DocumentIsPausedError from events.event_handlers import create_document_index as handler_module from models.dataset import Document from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus +from services.knowledge.indexing.errors import DocumentIsPausedError +from services.knowledge.resource_scope import DocumentRef @pytest.fixture @@ -45,7 +47,7 @@ def test_handle_logs_document_pause( ) -> None: mock_indexing_runner.run.side_effect = DocumentIsPausedError("Document is paused") - with patch.object(handler_module, "IndexingRunner", return_value=mock_indexing_runner): + with patch.object(handler_module, "build_document_indexing_service", return_value=mock_indexing_runner): with caplog.at_level(logging.INFO, logger=handler_module.logger.name): handler_module.handle("dataset-1", document_ids=["doc-1"]) @@ -64,7 +66,7 @@ def test_handle_logs_unexpected_indexing_errors( ) -> None: mock_indexing_runner.run.side_effect = RuntimeError("Indexing failed") - with patch.object(handler_module, "IndexingRunner", return_value=mock_indexing_runner): + with patch.object(handler_module, "build_document_indexing_service", return_value=mock_indexing_runner): with caplog.at_level(logging.ERROR, logger=handler_module.logger.name): handler_module.handle("dataset-1", document_ids=["doc-1"]) @@ -83,12 +85,16 @@ def test_handle_runs_indexing_on_success( mock_indexing_runner: MagicMock, caplog: pytest.LogCaptureFixture, ) -> None: - def assert_status_committed(_documents: list[Document], session: Session) -> None: - assert not session.in_transaction() + def assert_status_committed(refs: Sequence[DocumentRef]) -> None: + with Session(sqlite_session.get_bind()) as observer: + row = observer.get(Document, refs[0].document_id) + assert row is not None + assert row.indexing_status == IndexingStatus.PARSING + assert row.processing_started_at is not None mock_indexing_runner.run.side_effect = assert_status_committed - with patch.object(handler_module, "IndexingRunner", return_value=mock_indexing_runner): + with patch.object(handler_module, "build_document_indexing_service", return_value=mock_indexing_runner): with caplog.at_level(logging.INFO, logger=handler_module.logger.name): handler_module.handle("dataset-1", document_ids=["doc-1"]) diff --git a/api/tests/unit_tests/extensions/test_ext_application_services.py b/api/tests/unit_tests/extensions/test_ext_application_services.py index 4ad63dcc7f5..1362d844e47 100644 --- a/api/tests/unit_tests/extensions/test_ext_application_services.py +++ b/api/tests/unit_tests/extensions/test_ext_application_services.py @@ -41,6 +41,7 @@ from models.model import ( Site, TrialApp, ) +from models.provider import ProviderCredential from repositories.account_activation_repository import SQLAlchemyAccountActivationRepository from repositories.account_integration_repository import SQLAlchemyAccountIntegrationRepository from repositories.account_oauth_repository import ( @@ -96,10 +97,14 @@ from services.app_site_service import AppSiteService from services.app_tracing_config_gateway import OpsTraceManagerGateway from services.app_tracing_config_service import AppTracingConfigService from services.audio_types import AudioAppRef, AudioOutput, AudioUpload -from services.auth.data_source_api_key_auth_service import DataSourceApiKeyAuthService from services.billing_portal_service import BillingPortalService from services.billing_service import BillingService from services.compliance_download_service import ComplianceDownloadService +from services.data_source.auth.api_key_service import DataSourceApiKeyAuthService +from services.data_source.binding_application_service import DataSourceBindingApplicationService +from services.data_source.credential_gateway import ActorAwareDatasourceCredentialGateway +from services.data_source.notion_import_adapters import PluginNotionSourceGateway +from services.data_source.notion_import_application_service import NotionImportApplicationService from services.errors.enterprise import EnterpriseAPIError, EnterpriseAPINotFoundError, EnterpriseServiceError from services.file_service import FileService from services.human_input_file_upload_service import HumanInputFileUploadService @@ -108,6 +113,15 @@ from services.installed_app_access_service import InstalledAppAccessDeniedError, from services.installed_app_generation_adapters import AppGenerateServiceRuntime as InstalledAppGenerateServiceRuntime from services.installed_app_generation_service import InstalledAppGenerationService from services.knowledge.api_key_service import DatasetApiKeyService +from services.knowledge.dataset_access import DatasetAccessService +from services.knowledge.datasets.application import DatasetApplicationService +from services.knowledge.document_sync import DocumentSyncApplicationService +from services.knowledge.documents.application import DatasetDocumentApplicationService +from services.knowledge.external.application import ExternalKnowledgeApplicationService +from services.knowledge.indexing.adapters.estimate import IndexingEstimateAdapter, SQLAlchemyProcessRuleReader +from services.knowledge.indexing.adapters.sources import NotionSourceResolver +from services.knowledge.indexing.estimate import IndexingEstimateApplicationService +from services.knowledge.segments.application import DatasetSegmentApplicationService from services.message_file_preview_service import MessageFilePreviewService from services.oauth_device_application_service import OAuthDeviceApplicationService from services.partner_tenant_binding_service import PartnerTenantBindingService @@ -239,6 +253,7 @@ def test_build_application_services_preserves_composed_boundaries( assert isinstance(services.app_api_keys, AppApiKeyService) assert isinstance(services.dataset_api_keys, DatasetApiKeyService) + assert services.dataset_api_keys._access is services.knowledge.datasets._dataset_access assert isinstance(services.oauth_device, OAuthDeviceApplicationService) assert register_script.call_count == 3 @@ -768,7 +783,73 @@ def test_build_application_services_wires_data_source_api_key_auth( redis=_redis(), ) - assert isinstance(services.data_source_api_key_auth, DataSourceApiKeyAuthService) + assert isinstance(services.data_sources.api_key_auth, DataSourceApiKeyAuthService) + + +def test_build_application_services_groups_dataset_services_and_reuses_repositories( + sqlite_session_factory: sessionmaker[Session], +) -> None: + services = ext_application_services.build_application_services( + database_client=sqlite_session_factory, + deployment_edition=DeploymentEdition.COMMUNITY, + initialization_password="", + redis=MagicMock(spec=RedisClientWrapper), + ) + + assert isinstance(services.data_sources.bindings, DataSourceBindingApplicationService) + assert isinstance(services.data_sources.notion_imports, NotionImportApplicationService) + assert isinstance(services.knowledge.document_sync, DocumentSyncApplicationService) + assert isinstance(services.knowledge.documents, DatasetDocumentApplicationService) + assert isinstance(services.knowledge.datasets, DatasetApplicationService) + assert isinstance(services.knowledge.external, ExternalKnowledgeApplicationService) + assert services.knowledge.datasets._dataset_access is services.knowledge.document_sync._dataset_access + assert services.knowledge.external._dataset_access is services.knowledge.document_sync._dataset_access + assert services.knowledge.documents._dataset_access is services.knowledge.document_sync._dataset_access + assert isinstance(services.knowledge.indexing_estimates, IndexingEstimateApplicationService) + assert isinstance(services.knowledge.segments, DatasetSegmentApplicationService) + assert services.data_sources.bindings._bindings is services.data_sources.oauth["notion"]._bindings + assert services.data_sources.notion_imports._dataset_access is services.knowledge.document_sync._dataset_access + assert services.data_sources.notion_imports._dataset_access is services.knowledge.indexing_estimates._dataset_access + assert services.data_sources.notion_imports._documents is services.knowledge.document_sync._documents + assert services.data_sources.notion_imports._documents is services.knowledge.indexing_estimates._documents + notion_source = services.data_sources.notion_imports._source + estimate_service = services.knowledge.indexing_estimates + dataset_access = services.data_sources.notion_imports._dataset_access + assert isinstance(notion_source, PluginNotionSourceGateway) + assert isinstance(estimate_service._runner, IndexingEstimateAdapter) + assert isinstance(estimate_service._process_rules, SQLAlchemyProcessRuleReader) + assert isinstance(dataset_access, DatasetAccessService) + notion_resolver = estimate_service._notion + assert isinstance(notion_resolver, NotionSourceResolver) + actor_credentials = notion_source._credentials + assert isinstance(actor_credentials, ActorAwareDatasourceCredentialGateway) + assert actor_credentials is notion_resolver._actor_credentials + assert services.data_sources.providers._credentials is actor_credentials._credentials + assert services.knowledge.pipeline_generator._datasource_providers is services.data_sources.providers + assert notion_resolver._stored_credentials is not notion_resolver._actor_credentials + assert services.workspace_member_queries._members is dataset_access._workspace_roles + + +def test_build_application_services_wires_credential_query( + sqlite_session_factory: sessionmaker[Session], +) -> None: + services = ext_application_services.build_application_services( + database_client=sqlite_session_factory, + deployment_edition=DeploymentEdition.COMMUNITY, + initialization_password="", + redis=MagicMock(spec=RedisClientWrapper), + ) + tenant_id, actor_id = str(uuid4()), str(uuid4()) + with sqlite_session_factory.begin() as session: + credential = ProviderCredential( + tenant_id=tenant_id, provider_name="openai", credential_name="Team", encrypted_config="encrypted" + ) + session.add(credential) + credential_id = credential.id + + records = services.credential_queries.list_models(workspace_id=tenant_id, provider="openai", actor_id=actor_id) + + assert [(record.id, record.name) for record in records] == [(credential_id, "Team")] def test_build_application_services_uses_supplied_redis_for_both_workflow_stop_signals( @@ -851,6 +932,19 @@ def installed_app_ref(sqlite_session_factory: sessionmaker[Session]) -> Installe return result +def test_build_application_services_reuses_installed_app_generation_dependencies( + sqlite_session_factory: sessionmaker[Session], +) -> None: + services = ext_application_services.build_application_services( + database_client=sqlite_session_factory, + deployment_edition=DeploymentEdition.COMMUNITY, + initialization_password="", + redis=MagicMock(spec=RedisClientWrapper), + ) + + assert services.installed_apps.access._installed_apps is services.installed_apps.generation._usage + + @pytest.mark.parametrize( ("deployment_edition", "permission_result"), [ diff --git a/api/tests/unit_tests/fields/test_dataset_fields.py b/api/tests/unit_tests/fields/test_dataset_fields.py index ecde29d57fe..02f0692079b 100644 --- a/api/tests/unit_tests/fields/test_dataset_fields.py +++ b/api/tests/unit_tests/fields/test_dataset_fields.py @@ -2,16 +2,12 @@ import json from collections.abc import Iterator from contextlib import contextmanager -import pytest from sqlalchemy import event from sqlalchemy.orm import Session from core.rag.index_processor.constant.index_type import IndexStructureType from fields.dataset_fields import ( - DatasetDetailPrefetch, DatasetDetailResponse, - build_dataset_detail_prefetch, - dataset_detail_response_source, ) from models.account import Account from models.dataset import ( @@ -25,6 +21,7 @@ from models.dataset import ( ) from models.enums import DatasetMetadataType, DataSourceType, DocumentCreatedFrom, TagType from models.model import App, AppMode, IconType, Tag, TagBinding +from services.knowledge.dataset_read_service import load_dataset_detail, load_dataset_details def _dataset_detail_payload(**overrides): @@ -207,8 +204,7 @@ def test_dataset_detail_expands_missing_weighted_score_nested_fields(): } -@pytest.mark.parametrize("sqlite_session", [(Dataset, Account, App, AppDatasetJoin)], indirect=True) -def test_dataset_detail_response_source_uses_caller_session_for_database_fields(sqlite_session: Session): +def test_dataset_detail_serializes_after_session_closes(sqlite_session: Session) -> None: account = Account(name="Ada", email="ada@example.com") account.id = "account-1" dataset = Dataset( @@ -247,10 +243,10 @@ def test_dataset_detail_response_source_uses_caller_session_for_database_fields( sqlite_session.add_all([account, dataset, decoy_app, decoy_join]) sqlite_session.flush() - response = DatasetDetailResponse.model_validate( - dataset_detail_response_source(dataset, session=sqlite_session), - from_attributes=True, - ) + detail = load_dataset_detail(dataset, session=sqlite_session) + sqlite_session.expire_all() + sqlite_session.close() + response = DatasetDetailResponse.model_validate(detail) assert response.app_count == 0 assert response.document_count == 0 @@ -390,11 +386,8 @@ def _seed_dataset_page(session: Session) -> list[Dataset]: return datasets -def _dump(dataset: Dataset, *, session: Session, prefetch: DatasetDetailPrefetch | None = None) -> dict[str, object]: - return DatasetDetailResponse.model_validate( - dataset_detail_response_source(dataset, session=session, prefetch=prefetch), - from_attributes=True, - ).model_dump(mode="json") +def _dump(dataset: Dataset, *, session: Session) -> dict[str, object]: + return DatasetDetailResponse.model_validate(load_dataset_detail(dataset, session=session)).model_dump(mode="json") @contextmanager @@ -419,8 +412,10 @@ def test_dataset_detail_prefetch_matches_per_dataset_lookups(sqlite_session: Ses expected = [_dump(dataset, session=sqlite_session) for dataset in datasets] - prefetch = build_dataset_detail_prefetch(datasets, session=sqlite_session) - prefetched = [_dump(dataset, session=sqlite_session, prefetch=prefetch) for dataset in datasets] + details = load_dataset_details(datasets, session=sqlite_session) + sqlite_session.expire_all() + sqlite_session.close() + prefetched = [DatasetDetailResponse.model_validate(detail).model_dump(mode="json") for detail in details] assert prefetched == expected @@ -452,9 +447,8 @@ def test_dataset_detail_prefetch_keeps_query_count_independent_of_page_size(sqli def select_count(page: list[Dataset]) -> int: with _count_selects(sqlite_session) as statements: - prefetch = build_dataset_detail_prefetch(page, session=sqlite_session) - for dataset in page: - _dump(dataset, session=sqlite_session, prefetch=prefetch) + for detail in load_dataset_details(page, session=sqlite_session): + DatasetDetailResponse.model_validate(detail).model_dump(mode="json") return len(statements) small_page = select_count(datasets[:2]) diff --git a/api/tests/unit_tests/fields/test_document_fields.py b/api/tests/unit_tests/fields/test_document_fields.py index c097de67b65..7f4e82a7583 100644 --- a/api/tests/unit_tests/fields/test_document_fields.py +++ b/api/tests/unit_tests/fields/test_document_fields.py @@ -1,18 +1,17 @@ import json from datetime import datetime -import pytest from sqlalchemy.orm import Session from extensions.storage.storage_type import StorageType -from fields.document_fields import DocumentWithSession +from fields.document_fields import DocumentResponse from models.dataset import Document, DocumentSegment from models.enums import CreatorUserRole, DataSourceType, DocumentCreatedFrom from models.model import UploadFile +from services.knowledge.dataset_read_service import load_document_details -@pytest.mark.parametrize("sqlite_session", [(Document, DocumentSegment, UploadFile)], indirect=True) -def test_document_with_session_uses_explicit_getters(sqlite_session: Session) -> None: +def test_document_detail_serializes_after_session_closes(sqlite_session: Session) -> None: upload = UploadFile( tenant_id="tenant-1", storage_type=StorageType.LOCAL, @@ -52,8 +51,11 @@ def test_document_with_session_uses_explicit_getters(sqlite_session: Session) -> ) sqlite_session.add_all([upload, document, segment]) sqlite_session.flush() - source = DocumentWithSession(document=document, session=sqlite_session) + source = load_document_details([document], session=sqlite_session)[0] + sqlite_session.expire_all() + sqlite_session.close() + response = DocumentResponse.model_validate(source).model_dump(mode="json") - assert source.data_source_detail_dict["upload_file"]["name"] == "source.txt" - assert source.hit_count == 3 - assert source.doc_metadata_details is None + assert response["data_source_detail_dict"]["upload_file"]["name"] == "source.txt" + assert response["hit_count"] == 3 + assert response["doc_metadata"] == [] diff --git a/api/tests/unit_tests/fields/test_segment_fields.py b/api/tests/unit_tests/fields/test_segment_fields.py new file mode 100644 index 00000000000..71beaa42fc0 --- /dev/null +++ b/api/tests/unit_tests/fields/test_segment_fields.py @@ -0,0 +1,87 @@ +import json +from datetime import datetime + +from sqlalchemy.orm import Session + +from extensions.storage.storage_type import StorageType +from fields.segment_fields import SegmentResponse +from models.dataset import ChildChunk, DatasetProcessRule, Document, DocumentSegment, SegmentAttachmentBinding +from models.enums import CreatorUserRole, DataSourceType, DocumentCreatedFrom, ProcessRuleMode +from models.model import UploadFile +from services.knowledge.dataset_read_service import load_segment_details + + +def test_segment_detail_serializes_nested_relations_after_session_closes(sqlite_session: Session) -> None: + rule = DatasetProcessRule( + dataset_id="dataset-1", + mode=ProcessRuleMode.HIERARCHICAL, + rules=json.dumps({"parent_mode": "paragraph"}), + created_by="account-1", + ) + document = Document( + id="document-1", + tenant_id="tenant-1", + dataset_id="dataset-1", + dataset_process_rule_id=rule.id, + position=1, + data_source_type=DataSourceType.UPLOAD_FILE, + batch="batch-1", + name="source.txt", + created_from=DocumentCreatedFrom.WEB, + created_by="account-1", + ) + segment = DocumentSegment( + tenant_id=document.tenant_id, + dataset_id=document.dataset_id, + document_id=document.id, + position=1, + content="hello world", + word_count=2, + tokens=2, + created_by="account-1", + ) + chunk = ChildChunk( + tenant_id=document.tenant_id, + dataset_id=document.dataset_id, + document_id=document.id, + segment_id=segment.id, + position=1, + content="hello", + word_count=1, + created_by="account-1", + ) + attachment = UploadFile( + tenant_id=document.tenant_id, + storage_type=StorageType.LOCAL, + key="documents/image.png", + name="image.png", + size=12, + extension="png", + mime_type="image/png", + created_by_role=CreatorUserRole.ACCOUNT, + created_by="account-1", + created_at=datetime(2024, 1, 1), + used=True, + ) + binding = SegmentAttachmentBinding( + tenant_id=document.tenant_id, + dataset_id=document.dataset_id, + document_id=document.id, + segment_id=segment.id, + attachment_id=attachment.id, + ) + sqlite_session.add_all([rule, document, segment, chunk, attachment, binding]) + sqlite_session.flush() + details = load_segment_details([segment], {segment.id: "A summary"}, session=sqlite_session) + + # Expire even nested ORM objects so a retained proxy cannot pass by using cached attributes. + sqlite_session.expire_all() + sqlite_session.close() + response = SegmentResponse.model_validate(details[0]).model_dump(mode="json") + + assert response["summary"] == "A summary" + assert response["sign_content"] == "hello world" + assert response["child_chunks"][0]["content"] == "hello" + assert isinstance(response["child_chunks"][0]["created_at"], int) + assert response["attachments"][0]["name"] == "image.png" + assert "&sign=" in response["attachments"][0]["source_url"] diff --git a/api/tests/unit_tests/models/test_dataset_models.py b/api/tests/unit_tests/models/test_dataset_models.py index 77e36f10f63..e2fa9ead2d4 100644 --- a/api/tests/unit_tests/models/test_dataset_models.py +++ b/api/tests/unit_tests/models/test_dataset_models.py @@ -1,3 +1,5 @@ +from repositories.knowledge.keyword_table_repository import load_keyword_table + """ Comprehensive unit tests for Dataset models. @@ -12,7 +14,7 @@ This test suite covers: import json import pickle from datetime import UTC, datetime -from unittest.mock import patch +from unittest.mock import Mock, patch from urllib.parse import parse_qs, urlparse from uuid import uuid4 @@ -46,6 +48,37 @@ from models.enums import ( SegmentStatus, ) from models.model import App, AppMode, IconType, UploadFile +from repositories.knowledge.dataset_read_repository import ( + get_child_dataset, + get_child_document, + get_child_segment, + get_dataset_app_count, + get_dataset_author_name, + get_dataset_available_document_count, + get_dataset_creator, + get_dataset_doc_form, + get_dataset_doc_metadata, + get_dataset_document_count, + get_dataset_is_published, + get_dataset_keyword_table, + get_dataset_tags, + get_dataset_word_count, + get_document_dataset, + get_document_hit_count, + get_document_process_rule, + get_document_segment_count, + get_external_api_dataset_bindings, + get_external_knowledge_info, + get_joined_app, + get_latest_dataset_process_rule, + get_next_segment, + get_previous_segment, + get_segment_child_chunks, + get_segment_dataset, + get_segment_document, +) +from repositories.knowledge.segment_read_adapter import get_segment_attachments +from services.knowledge.dataset_read_service import get_dataset_queries from tests.unit_tests.config_override import apply_config_overrides @@ -186,9 +219,9 @@ class TestDatasetModelValidation: sqlite_session.add_all([dataset, account, process_rule, document]) sqlite_session.flush() - assert dataset.get_created_by_account(session=sqlite_session) is account - assert dataset.get_latest_process_rule(session=sqlite_session) is process_rule - assert dataset.get_doc_form(session=sqlite_session) == IndexStructureType.PARAGRAPH_INDEX + assert get_dataset_creator(dataset, session=sqlite_session) is account + assert get_latest_dataset_process_rule(dataset, session=sqlite_session) is process_rule + assert get_dataset_doc_form(dataset, session=sqlite_session) == IndexStructureType.PARAGRAPH_INDEX @pytest.mark.parametrize("sqlite_session", [(Dataset, Document)], indirect=True) def test_get_doc_form_ignores_foreign_tenant_document(self, sqlite_session: Session) -> None: @@ -201,7 +234,7 @@ class TestDatasetModelValidation: sqlite_session.add_all([dataset, foreign_document]) sqlite_session.flush() - assert dataset.get_doc_form(session=sqlite_session) is None + assert get_dataset_doc_form(dataset, session=sqlite_session) is None @pytest.mark.parametrize("sqlite_session", [(Dataset, DatasetKeywordTable)], indirect=True) def test_get_dataset_keyword_table_uses_caller_session(self, sqlite_session: Session): @@ -213,7 +246,7 @@ class TestDatasetModelValidation: sqlite_session.add_all([dataset, keyword_table]) sqlite_session.flush() - result = dataset.get_dataset_keyword_table(session=sqlite_session) + result = get_dataset_keyword_table(dataset, session=sqlite_session) assert result is keyword_table @@ -250,17 +283,17 @@ class TestDatasetModelValidation: ) sqlite_session.flush() - assert dataset.get_total_documents(session=sqlite_session) == 2 - assert dataset.get_total_available_documents(session=sqlite_session) == 1 - assert dataset.get_app_count(session=sqlite_session) == 1 - assert dataset.get_document_count(session=sqlite_session) == 2 - assert dataset.get_word_count(session=sqlite_session) == 500 - assert dataset.get_author_name(session=sqlite_session) == "Ada" - assert dataset.get_tags(session=sqlite_session) == [] - assert dataset.get_doc_form(session=sqlite_session) == IndexStructureType.PARAGRAPH_INDEX - assert dataset.get_external_knowledge_info(session=sqlite_session) is None - assert dataset.get_doc_metadata(session=sqlite_session) == [] - assert dataset.get_is_published(session=sqlite_session) is False + assert get_dataset_document_count(dataset, session=sqlite_session) == 2 + assert get_dataset_available_document_count(dataset, session=sqlite_session) == 1 + assert get_dataset_app_count(dataset, session=sqlite_session) == 1 + assert get_dataset_document_count(dataset, session=sqlite_session) == 2 + assert get_dataset_word_count(dataset, session=sqlite_session) == 500 + assert get_dataset_author_name(dataset, session=sqlite_session) == "Ada" + assert get_dataset_tags(dataset, session=sqlite_session) == [] + assert get_dataset_doc_form(dataset, session=sqlite_session) == IndexStructureType.PARAGRAPH_INDEX + assert get_external_knowledge_info(dataset, session=sqlite_session) is None + assert get_dataset_doc_metadata(dataset, session=sqlite_session) == [] + assert get_dataset_is_published(dataset, session=sqlite_session) is False def test_dataset_indexing_technique_validation(self): """Test dataset indexing technique values.""" @@ -387,7 +420,7 @@ class TestDatasetModelValidation: sqlite_session.add_all([dataset, external_api, binding]) sqlite_session.flush() - assert dataset.get_external_knowledge_info(session=sqlite_session) is None + assert get_external_knowledge_info(dataset, session=sqlite_session) is None @pytest.mark.parametrize( "sqlite_session", [(ExternalKnowledgeApis, ExternalKnowledgeBindings, Dataset)], indirect=True @@ -433,7 +466,7 @@ class TestDatasetModelValidation: ) sqlite_session.flush() - result = external_api.get_dataset_bindings(session=sqlite_session) + result = get_external_api_dataset_bindings(external_api, session=sqlite_session) assert result == [{"id": dataset.id, "name": dataset.name}] @@ -461,11 +494,11 @@ class TestDatasetModelValidation: used=False, ) upload_file.id = "file-1" - sqlite_session.add_all([dataset_query, upload_file]) + sqlite_session.add_all([_make_dataset(dataset_id=dataset_query.dataset_id), dataset_query, upload_file]) sqlite_session.flush() - with patch("models.dataset.sign_upload_file_preview_url", return_value="signed-url"): - queries = dataset_query.get_queries(session=sqlite_session) + with patch("services.knowledge.dataset_read_service.sign_upload_file_preview_url", return_value="signed-url"): + queries = get_dataset_queries(dataset_query, session=sqlite_session) assert queries == [ { @@ -604,9 +637,9 @@ class TestDocumentModelRelationships: sqlite_session.add_all([process_rule, document, *segments]) sqlite_session.flush() - assert document.get_dataset_process_rule(session=sqlite_session) is process_rule - assert document.get_segment_count(session=sqlite_session) == 3 - assert document.get_hit_count(session=sqlite_session) == 7 + assert get_document_process_rule(document, session=sqlite_session) is process_rule + assert get_document_segment_count(document, session=sqlite_session) == 3 + assert get_document_hit_count(document, session=sqlite_session) == 7 def test_document_display_status_queuing(self): """Test document display_status property for queuing state.""" @@ -818,7 +851,59 @@ class TestDocumentModelRelationships: sqlite_session.add_all([dataset, document]) sqlite_session.flush() - assert document.get_dataset(session=sqlite_session) is dataset + assert get_document_dataset(document, session=sqlite_session) is dataset + + @pytest.mark.parametrize("sqlite_session", [(Document, DocumentSegment)], indirect=True) + def test_document_average_segment_length( + self, + sqlite_session: Session, + ): + """Test average_segment_length property calculation.""" + # Arrange + document = Document( + tenant_id=str(uuid4()), + dataset_id=str(uuid4()), + position=1, + data_source_type=DataSourceType.UPLOAD_FILE, + batch="batch_001", + name="test.pdf", + created_from=DocumentCreatedFrom.WEB, + created_by=str(uuid4()), + word_count=1000, + ) + sqlite_session.add(document) + sqlite_session.flush() + sqlite_session.add_all(_make_segments(document, [0] * 10)) + sqlite_session.commit() + + # Act + result = document.to_dict( + segment_count=get_document_segment_count(document, session=sqlite_session), hit_count=0, process_rule=None + )["average_segment_length"] + + # Assert + assert result == 100 + + def test_document_average_segment_length_zero(self): + """Test average_segment_length property when word_count is zero.""" + # Arrange + document = Document( + tenant_id=str(uuid4()), + dataset_id=str(uuid4()), + position=1, + data_source_type=DataSourceType.UPLOAD_FILE, + batch="batch_001", + name="test.pdf", + created_from=DocumentCreatedFrom.WEB, + created_by=str(uuid4()), + word_count=0, + ) + + # Act + result = document.to_dict(segment_count=0, hit_count=0, process_rule=None)["average_segment_length"] + + # Assert + assert result == 0 class TestDocumentSegmentIndexing: @@ -858,7 +943,7 @@ class TestDocumentSegmentIndexing: sqlite_session.add_all([process_rule, document, segment, child_chunk]) sqlite_session.flush() - result = segment.get_child_chunks(session=sqlite_session) + result = get_segment_child_chunks(segment, session=sqlite_session) assert result == [child_chunk] @pytest.mark.parametrize( @@ -895,8 +980,8 @@ class TestDocumentSegmentIndexing: sqlite_session.add_all([process_rule, document, segment, child_chunk]) sqlite_session.flush() - result = segment.get_child_chunks(session=sqlite_session) - response_result = segment.get_child_chunks(session=sqlite_session, include_full_doc=False) + result = get_segment_child_chunks(segment, session=sqlite_session) + response_result = get_segment_child_chunks(segment, session=sqlite_session, include_full_doc=False) assert result == [child_chunk] assert response_result == [] @@ -917,8 +1002,8 @@ class TestDocumentSegmentIndexing: sqlite_session.add_all([dataset, document, segment]) sqlite_session.flush() - assert segment.get_dataset(session=sqlite_session) is dataset - assert segment.get_document(session=sqlite_session) is document + assert get_segment_dataset(segment, session=sqlite_session) is dataset + assert get_segment_document(segment, session=sqlite_session) is document def test_document_segment_creation_with_required_fields(self): """Test creating a document segment with all required fields.""" @@ -1116,8 +1201,8 @@ class TestDocumentSegmentIndexing: sqlite_session.add_all([segment, attachment, binding]) sqlite_session.flush() - monkeypatch.setattr("models.dataset.time.time", lambda: 1700000000) - monkeypatch.setattr("models.dataset.os.urandom", lambda _: b"\x01" * 16) + monkeypatch.setattr("repositories.knowledge.segment_read_adapter.time.time", lambda: 1700000000) + monkeypatch.setattr("repositories.knowledge.segment_read_adapter.os.urandom", lambda _: b"\x01" * 16) apply_config_overrides( monkeypatch, SECRET_KEY="unit-secret", @@ -1126,7 +1211,7 @@ class TestDocumentSegmentIndexing: ) # Act - attachments = segment.get_attachments(session=sqlite_session) + attachments = get_segment_attachments(segment, session=sqlite_session) # Assert assert len(attachments) == 1 @@ -1376,7 +1461,13 @@ class TestDatasetKeywordTable: sqlite_session.add_all([dataset, keyword_table]) sqlite_session.flush() - result = keyword_table.get_keyword_table_dict(session=sqlite_session) + result = load_keyword_table( + storage=Mock(), + tenant_id=dataset.tenant_id, + dataset_id=dataset.id, + storage_type=keyword_table.data_source_type, + data=keyword_table.keyword_table, + ) assert result == {"__data__": {"table": {"keyword": {"node-1"}}}} @@ -1519,6 +1610,50 @@ class TestModelIntegration: assert document.word_count == 100 assert segment.status == SegmentStatus.COMPLETED + @pytest.mark.parametrize("sqlite_session", [(Document, DocumentSegment)], indirect=True) + def test_document_to_dict_serialization( + self, + sqlite_session: Session, + ): + """Test document to_dict method for serialization.""" + # Arrange + tenant_id = str(uuid4()) + dataset_id = str(uuid4()) + created_by = str(uuid4()) + + document = Document( + tenant_id=tenant_id, + dataset_id=dataset_id, + position=1, + data_source_type=DataSourceType.UPLOAD_FILE, + batch="batch_001", + name="test.pdf", + created_from=DocumentCreatedFrom.WEB, + created_by=created_by, + word_count=100, + indexing_status=IndexingStatus.COMPLETED, + ) + sqlite_session.add(document) + sqlite_session.flush() + sqlite_session.add_all(_make_segments(document, [2, 2, 2, 2, 2])) + sqlite_session.commit() + + # Act + result = document.to_dict( + segment_count=get_document_segment_count(document, session=sqlite_session), + hit_count=get_document_hit_count(document, session=sqlite_session), + process_rule=get_document_process_rule(document, session=sqlite_session), + ) + + # Assert + assert result["tenant_id"] == tenant_id + assert result["dataset_id"] == dataset_id + assert result["name"] == "test.pdf" + assert result["word_count"] == 100 + assert result["indexing_status"] == IndexingStatus.COMPLETED + assert result["segment_count"] == 5 + assert result["hit_count"] == 10 + class TestChildChunkSessionAccessors: """Regression coverage for the ``ChildChunk`` accessors refactored to take a caller-provided session. @@ -1544,9 +1679,9 @@ class TestChildChunkSessionAccessors: sqlite_session.add_all([dataset, document, segment, child_chunk]) sqlite_session.flush() - assert child_chunk.dataset(session=sqlite_session) is dataset - assert child_chunk.document(session=sqlite_session) is document - assert child_chunk.segment(session=sqlite_session) is segment + assert get_child_dataset(child_chunk, session=sqlite_session) is dataset + assert get_child_document(child_chunk, session=sqlite_session) is document + assert get_child_segment(child_chunk, session=sqlite_session) is segment def test_accessors_return_none_when_related_rows_are_absent(self, sqlite_session: Session): child_chunk = ChildChunk( @@ -1560,9 +1695,9 @@ class TestChildChunkSessionAccessors: created_by="account-1", ) - assert child_chunk.dataset(session=sqlite_session) is None - assert child_chunk.document(session=sqlite_session) is None - assert child_chunk.segment(session=sqlite_session) is None + assert get_child_dataset(child_chunk, session=sqlite_session) is None + assert get_child_document(child_chunk, session=sqlite_session) is None + assert get_child_segment(child_chunk, session=sqlite_session) is None class TestDatasetAvailableDocumentCount: @@ -1590,7 +1725,7 @@ class TestDatasetAvailableDocumentCount: sqlite_session.add_all([dataset, completed, waiting]) sqlite_session.flush() - assert dataset.get_available_document_count(session=sqlite_session) == 1 + assert get_dataset_available_document_count(dataset, session=sqlite_session) == 1 class TestDocumentSegmentNeighborAccessors: @@ -1609,8 +1744,8 @@ class TestDocumentSegmentNeighborAccessors: sqlite_session.flush() middle_segment = segments[1] - assert middle_segment.previous_segment(session=sqlite_session) is segments[0] - assert middle_segment.next_segment(session=sqlite_session) is segments[2] + assert get_previous_segment(middle_segment, session=sqlite_session) is segments[0] + assert get_next_segment(middle_segment, session=sqlite_session) is segments[2] def test_accessors_return_none_when_neighbors_are_absent(self, sqlite_session: Session): dataset = _make_dataset(dataset_id=str(uuid4()), tenant_id=str(uuid4())) @@ -1619,8 +1754,8 @@ class TestDocumentSegmentNeighborAccessors: sqlite_session.add_all([dataset, document, segment]) sqlite_session.flush() - assert segment.previous_segment(session=sqlite_session) is None - assert segment.next_segment(session=sqlite_session) is None + assert get_previous_segment(segment, session=sqlite_session) is None + assert get_next_segment(segment, session=sqlite_session) is None class TestAppDatasetJoinSessionAccessors: @@ -1639,7 +1774,7 @@ class TestAppDatasetJoinSessionAccessors: sqlite_session.add_all([app, join]) sqlite_session.flush() - assert join.app(session=sqlite_session) is app + assert get_joined_app(join, session=sqlite_session) is app def test_app_accessor_returns_none_when_app_absent(self, sqlite_session: Session): join = AppDatasetJoin( @@ -1647,4 +1782,4 @@ class TestAppDatasetJoinSessionAccessors: dataset_id=str(uuid4()), ) - assert join.app(session=sqlite_session) is None + assert get_joined_app(join, session=sqlite_session) is None diff --git a/api/tests/unit_tests/pyrefly.toml b/api/tests/unit_tests/pyrefly.toml index 01f5847d4db..21274d824f5 100644 --- a/api/tests/unit_tests/pyrefly.toml +++ b/api/tests/unit_tests/pyrefly.toml @@ -477,7 +477,6 @@ project-excludes = [ "core/rag/indexing/processor/test_qa_index_processor.py", "core/rag/indexing/test_index_processor.py", "core/rag/indexing/test_index_processor_base.py", - "core/rag/indexing/test_indexing_runner.py", "core/rag/pipeline/test_queue.py", "core/rag/rerank/test_reranker.py", "core/rag/retrieval/test_dataset_retrieval.py", @@ -843,7 +842,6 @@ project-excludes = [ "services/test_telemetry_service.py", "services/test_trigger_provider_service.py", "services/test_variable_truncator.py", - "services/test_vector_service.py", "services/test_webhook_service.py", "services/test_webhook_service_additional.py", "services/test_website_service.py", diff --git a/api/tests/unit_tests/repositories/credentials/__init__.py b/api/tests/unit_tests/repositories/credentials/__init__.py new file mode 100644 index 00000000000..6c42b60d5a7 --- /dev/null +++ b/api/tests/unit_tests/repositories/credentials/__init__.py @@ -0,0 +1 @@ +"""Tests for credentials persistence adapters.""" diff --git a/api/tests/unit_tests/repositories/credentials/test_query_repository.py b/api/tests/unit_tests/repositories/credentials/test_query_repository.py new file mode 100644 index 00000000000..9bb0ffab543 --- /dev/null +++ b/api/tests/unit_tests/repositories/credentials/test_query_repository.py @@ -0,0 +1,451 @@ +"""Unit tests for persisted credential permission policy. + +Tests visibility filtering and complete domain reads against real SQLite rows. +""" + +from datetime import datetime +from uuid import uuid4 + +import pytest +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker + +from core.plugin.entities.plugin_daemon import CredentialType as TriggerCredentialType +from models.account import Account, TenantAccountRole +from models.credential_permission import CredentialPermission, CredentialType +from models.enums import PermissionEnum +from models.oauth import DatasourceProvider +from models.provider import ProviderCredential +from models.tools import BuiltinToolProvider +from models.trigger import TriggerSubscription, WorkflowPluginTrigger +from repositories.credentials.query_repository import ( + CredentialQueryRepository, + apply_credential_visibility_filter_for_actor, +) + + +@pytest.fixture +def tenant_id() -> str: + return str(uuid4()) + + +@pytest.fixture +def user_id() -> str: + return str(uuid4()) + + +@pytest.fixture +def other_user_id() -> str: + return str(uuid4()) + + +@pytest.fixture +def repository(sqlite_session_factory: sessionmaker[Session]) -> CredentialQueryRepository: + return CredentialQueryRepository(session_factory=sqlite_session_factory) + + +@pytest.mark.parametrize("provider", ["openai", "langgenius/openai/openai"]) +def test_model_aliases_preserve_visibility_and_tenant_scope( + sqlite_session: Session, repository: CredentialQueryRepository, provider: str +) -> None: + tenant, actor, teammate = (str(uuid4()) for _ in range(3)) + rows = [ + ProviderCredential( + tenant_id=tenant, + provider_name="openai" if index % 2 else "langgenius/openai/openai", + credential_name=name, + encrypted_config="encrypted", + user_id=owner, + visibility=visibility, + ) + for index, (name, owner, visibility) in enumerate( + [ + ("team", teammate, PermissionEnum.ALL_TEAM), + ("mine", actor, PermissionEnum.ONLY_ME), + ("legacy", None, PermissionEnum.ONLY_ME), + ("shared", teammate, PermissionEnum.PARTIAL_TEAM), + ("disabled", teammate, PermissionEnum.PARTIAL_TEAM), + ("wrong-tenant-grant", teammate, PermissionEnum.PARTIAL_TEAM), + ("wrong-type-grant", teammate, PermissionEnum.PARTIAL_TEAM), + ("private", teammate, PermissionEnum.ONLY_ME), + ] + ) + ] + for index, row in enumerate(rows): + row.created_at = datetime(2026, 1, index + 1) + sqlite_session.add_all(rows) + sqlite_session.add_all( + [ + CredentialPermission( + credential_id=rows[index].id, + credential_type=CredentialType.BUILTIN_TOOL_PROVIDER + if index == 6 + else CredentialType.PROVIDER_CREDENTIAL, + account_id=actor, + tenant_id=str(uuid4()) if index == 5 else tenant, + has_permission=index != 4, + ) + for index in (3, 4, 5, 6) + ] + ) + sqlite_session.add_all( + [ + ProviderCredential( + tenant_id=str(uuid4()), provider_name="openai", credential_name="foreign", encrypted_config="secret" + ), + ProviderCredential( + tenant_id=tenant, + provider_name="anthropic", + credential_name="different-provider", + encrypted_config="secret", + ), + ] + ) + sqlite_session.commit() + + records = repository.list_models(workspace_id=tenant, provider=provider, actor_id=actor) + + assert [record.name for record in records] == ["shared", "legacy", "mine", "team"] + internal = repository.list_models(workspace_id=tenant, provider=provider, actor_id=None) + assert {record.id for record in internal} == {row.id for row in rows} + + +def test_datasource_hidden_default_and_complete_owner_scope( + sqlite_session: Session, repository: CredentialQueryRepository +) -> None: + tenant, actor, teammate = (str(uuid4()) for _ in range(3)) + visible = DatasourceProvider( + tenant_id=tenant, + name="visible", + provider="notion", + plugin_id="plugin", + auth_type="oauth2", + encrypted_credentials={"access_token": "encrypted"}, + user_id=actor, + visibility=PermissionEnum.ONLY_ME, + ) + hidden_default = DatasourceProvider( + tenant_id=tenant, + name="hidden", + provider="notion", + plugin_id="plugin", + auth_type="oauth2", + encrypted_credentials={}, + user_id=teammate, + visibility=PermissionEnum.ONLY_ME, + is_default=True, + ) + sqlite_session.add_all([visible, hidden_default]) + for owner, provider, plugin in [ + (str(uuid4()), "notion", "plugin"), + (tenant, "other", "plugin"), + (tenant, "notion", "other"), + ]: + sqlite_session.add( + DatasourceProvider( + tenant_id=owner, + name="unrelated", + provider=provider, + plugin_id=plugin, + auth_type="oauth2", + encrypted_credentials={}, + is_default=True, + ) + ) + sqlite_session.commit() + + records = repository.list_datasources(workspace_id=tenant, provider="notion", plugin_id="plugin", actor_id=actor) + + assert len(records) == 1 + assert records[0].id == visible.id + assert records[0].encrypted_credentials == {"access_token": "encrypted"} + assert records[0].is_default is False + internal = repository.list_datasources(workspace_id=tenant, provider="notion", plugin_id="plugin", actor_id=None) + assert {record.id for record in internal if record.is_default} == {hidden_default.id} + + +def test_tool_borrowing_members_and_default_do_not_escape_scope_or_write_state( + sqlite_session: Session, repository: CredentialQueryRepository +) -> None: + tenant, actor, teammate = (str(uuid4()) for _ in range(3)) + visible = BuiltinToolProvider( + name="visible", + tenant_id=tenant, + user_id=actor, + provider="tool", + visibility=PermissionEnum.PARTIAL_TEAM, + encrypted_credentials='{"api_key": "encrypted"}', + ) + borrowed = BuiltinToolProvider( + name="borrowed", + tenant_id=tenant, + user_id=teammate, + provider="tool", + visibility=PermissionEnum.ONLY_ME, + is_default=True, + ) + foreign = BuiltinToolProvider(name="foreign", tenant_id=str(uuid4()), user_id=teammate, provider="tool") + other_provider = BuiltinToolProvider(name="other", tenant_id=tenant, user_id=teammate, provider="other") + sqlite_session.add_all([visible, borrowed, foreign, other_provider]) + sqlite_session.add_all( + [ + CredentialPermission( + credential_id=visible.id, + credential_type=kind, + account_id=member, + tenant_id=owner, + has_permission=enabled, + ) + for member, owner, kind, enabled in [ + (teammate, tenant, CredentialType.BUILTIN_TOOL_PROVIDER, True), + (str(uuid4()), tenant, CredentialType.BUILTIN_TOOL_PROVIDER, False), + (str(uuid4()), str(uuid4()), CredentialType.BUILTIN_TOOL_PROVIDER, True), + (str(uuid4()), tenant, CredentialType.PROVIDER_CREDENTIAL, True), + ] + ] + ) + sqlite_session.commit() + + records = repository.list_tools( + workspace_id=tenant, + provider="tool", + actor_id=actor, + include_credential_ids=[visible.id, borrowed.id, borrowed.id, foreign.id, other_provider.id, ""], + ) + + assert [record.id for record in records] == [visible.id, borrowed.id] + assert records[0].partial_member_ids == (teammate,) + assert records[0].credentials == {"api_key": "encrypted"} + assert records[0].is_default is True + assert records[0].from_other_member is False + assert records[1].from_other_member is True + sqlite_session.refresh(visible) + assert visible.is_default is False + + +def test_trigger_visibility_and_usage_counts_are_tenant_scoped( + sqlite_session: Session, repository: CredentialQueryRepository +) -> None: + tenant, actor, teammate, app_id = (str(uuid4()) for _ in range(4)) + rows = [ + TriggerSubscription( + name=name, + tenant_id=owner, + user_id=teammate, + provider_id=provider, + endpoint_id=f"{name}-endpoint", + parameters={"event": "created"}, + properties={"label": name}, + credentials={"api_key": "encrypted"}, + credential_type=TriggerCredentialType.API_KEY, + visibility=visibility, + ) + for name, owner, provider, visibility in [ + ("shared", tenant, "trigger", PermissionEnum.PARTIAL_TEAM), + ("private", tenant, "trigger", PermissionEnum.ONLY_ME), + ("foreign", str(uuid4()), "trigger", PermissionEnum.ALL_TEAM), + ("other-provider", tenant, "other", PermissionEnum.ALL_TEAM), + ] + ] + sqlite_session.add_all(rows) + sqlite_session.add( + CredentialPermission( + credential_id=rows[0].id, + credential_type=CredentialType.TRIGGER_SUBSCRIPTION, + tenant_id=tenant, + account_id=actor, + ) + ) + sqlite_session.add_all( + [ + WorkflowPluginTrigger( + app_id=app, + node_id=node, + tenant_id=owner, + provider_id="trigger", + event_name="created", + subscription_id=rows[0].id, + ) + for owner, app, node in [ + (tenant, app_id, "node-1"), + (tenant, app_id, "node-2"), + (tenant, str(uuid4()), "node-3"), + (str(uuid4()), str(uuid4()), "foreign-node"), + ] + ] + ) + sqlite_session.commit() + + records = repository.list_trigger_subscriptions(workspace_id=tenant, provider="trigger", actor_id=actor) + + assert len(records) == 1 + assert records[0].id == rows[0].id + assert records[0].workflows_in_use == 2 + assert records[0].parameters == {"event": "created"} + assert records[0].endpoint_id == "shared-endpoint" + assert records[0].credentials == {"api_key": "encrypted"} + + +def test_empty_lists_do_not_require_related_rows(repository: CredentialQueryRepository) -> None: + assert repository.list_models(workspace_id="empty", provider="openai", actor_id="actor") == [] + assert ( + repository.list_datasources(workspace_id="empty", provider="notion", plugin_id="plugin", actor_id="actor") == [] + ) + assert repository.list_tools(workspace_id="empty", provider="tool", actor_id="actor") == [] + assert repository.list_trigger_subscriptions(workspace_id="empty", provider="trigger", actor_id="actor") == [] + + +def _subscription( + *, + tenant_id: str, + owner_id: str, + name: str, + visibility: PermissionEnum, +) -> TriggerSubscription: + return TriggerSubscription( + name=name, + tenant_id=tenant_id, + user_id=owner_id, + provider_id="test/provider", + endpoint_id=f"{name}-endpoint", + parameters={}, + properties={}, + credentials={}, + credential_type=TriggerCredentialType.API_KEY, + visibility=visibility, + ) + + +def _user(user_id: str, *, is_admin: bool) -> Account: + user = Account(name="Credential User", email=f"{user_id}@example.com") + user.id = user_id + user.role = TenantAccountRole.ADMIN if is_admin else TenantAccountRole.NORMAL + return user + + +class TestApplyVisibilityFilter: + def test_admin_does_not_bypass_personal_visibility( + self, + sqlite_session: Session, + tenant_id: str, + user_id: str, + other_user_id: str, + ) -> None: + private_subscription = _subscription( + tenant_id=tenant_id, + owner_id=other_user_id, + name="private", + visibility=PermissionEnum.ONLY_ME, + ) + sqlite_session.add(private_subscription) + sqlite_session.commit() + + query = apply_credential_visibility_filter_for_actor( + select(TriggerSubscription).where(TriggerSubscription.tenant_id == tenant_id), + tenant_id=tenant_id, + model_id_column=TriggerSubscription.id, + model_user_id_column=TriggerSubscription.user_id, + model_visibility_column=TriggerSubscription.visibility, + credential_type=CredentialType.TRIGGER_SUBSCRIPTION, + actor_id=_user(user_id, is_admin=True).id, + ) + + assert sqlite_session.scalars(query).all() == [] + + def test_non_admin_sees_team_owned_and_partial_member_subscriptions( + self, + sqlite_session: Session, + tenant_id: str, + user_id: str, + other_user_id: str, + ) -> None: + team_subscription = _subscription( + tenant_id=tenant_id, + owner_id=other_user_id, + name="team", + visibility=PermissionEnum.ALL_TEAM, + ) + owned_subscription = _subscription( + tenant_id=tenant_id, + owner_id=user_id, + name="owned", + visibility=PermissionEnum.ONLY_ME, + ) + shared_subscription = _subscription( + tenant_id=tenant_id, + owner_id=other_user_id, + name="shared", + visibility=PermissionEnum.PARTIAL_TEAM, + ) + private_subscription = _subscription( + tenant_id=tenant_id, + owner_id=other_user_id, + name="private", + visibility=PermissionEnum.ONLY_ME, + ) + sqlite_session.add_all( + [ + team_subscription, + owned_subscription, + shared_subscription, + private_subscription, + CredentialPermission( + credential_id=shared_subscription.id, + credential_type=CredentialType.TRIGGER_SUBSCRIPTION, + account_id=user_id, + tenant_id=tenant_id, + ), + ] + ) + sqlite_session.commit() + + query = apply_credential_visibility_filter_for_actor( + select(TriggerSubscription).where(TriggerSubscription.tenant_id == tenant_id), + tenant_id=tenant_id, + model_id_column=TriggerSubscription.id, + model_user_id_column=TriggerSubscription.user_id, + model_visibility_column=TriggerSubscription.visibility, + credential_type=CredentialType.TRIGGER_SUBSCRIPTION, + actor_id=_user(user_id, is_admin=False).id, + ) + + visible_ids = {subscription.id for subscription in sqlite_session.scalars(query)} + assert visible_ids == {team_subscription.id, owned_subscription.id, shared_subscription.id} + + def test_partial_permissions_are_scoped_by_tenant( + self, + sqlite_session: Session, + tenant_id: str, + user_id: str, + other_user_id: str, + ) -> None: + other_tenant_id = str(uuid4()) + shared_subscription = _subscription( + tenant_id=tenant_id, + owner_id=other_user_id, + name="shared", + visibility=PermissionEnum.PARTIAL_TEAM, + ) + sqlite_session.add(shared_subscription) + sqlite_session.flush() + sqlite_session.add( + CredentialPermission( + credential_id=shared_subscription.id, + credential_type=CredentialType.TRIGGER_SUBSCRIPTION, + account_id=user_id, + tenant_id=other_tenant_id, + ) + ) + sqlite_session.commit() + + query = apply_credential_visibility_filter_for_actor( + select(TriggerSubscription).where(TriggerSubscription.tenant_id == tenant_id), + tenant_id=tenant_id, + model_id_column=TriggerSubscription.id, + model_user_id_column=TriggerSubscription.user_id, + model_visibility_column=TriggerSubscription.visibility, + credential_type=CredentialType.TRIGGER_SUBSCRIPTION, + actor_id=_user(user_id, is_admin=False).id, + ) + + assert sqlite_session.scalars(query).all() == [] diff --git a/api/tests/unit_tests/repositories/data_source/__init__.py b/api/tests/unit_tests/repositories/data_source/__init__.py new file mode 100644 index 00000000000..a32a4aee0cd --- /dev/null +++ b/api/tests/unit_tests/repositories/data_source/__init__.py @@ -0,0 +1 @@ +"""Tests for data-source persistence adapters.""" diff --git a/api/tests/unit_tests/repositories/test_data_source_api_key_auth_repository.py b/api/tests/unit_tests/repositories/data_source/test_api_key_auth_repository.py similarity index 94% rename from api/tests/unit_tests/repositories/test_data_source_api_key_auth_repository.py rename to api/tests/unit_tests/repositories/data_source/test_api_key_auth_repository.py index d2a4742aa3e..802fd19f400 100644 --- a/api/tests/unit_tests/repositories/test_data_source_api_key_auth_repository.py +++ b/api/tests/unit_tests/repositories/data_source/test_api_key_auth_repository.py @@ -4,8 +4,8 @@ from sqlalchemy import select from sqlalchemy.orm import Session, sessionmaker from models.source import DataSourceApiKeyAuthBinding -from repositories.data_source_api_key_auth_repository import SQLAlchemyDataSourceApiKeyAuthBindingRepository -from services.entities.data_source_api_key_auth_entities import DataSourceApiKeyAuthCredentials +from repositories.data_source.api_key_auth_repository import SQLAlchemyDataSourceApiKeyAuthBindingRepository +from services.data_source.entities.api_key_auth import DataSourceApiKeyAuthCredentials def _binding( diff --git a/api/tests/unit_tests/repositories/data_source/test_credential_repository.py b/api/tests/unit_tests/repositories/data_source/test_credential_repository.py new file mode 100644 index 00000000000..23ab5055a5a --- /dev/null +++ b/api/tests/unit_tests/repositories/data_source/test_credential_repository.py @@ -0,0 +1,496 @@ +import json + +import pytest +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker + +from models.credential_permission import CredentialPermission, CredentialType +from models.dataset import Dataset, Document +from models.enums import DataSourceType, DocumentCreatedFrom, PermissionEnum +from models.oauth import DatasourceOauthParamConfig, DatasourceOauthTenantParamConfig, DatasourceProvider +from repositories.data_source.credential_repository import SQLAlchemyDatasourceCredentialRepository +from services.data_source.entities.credential import DatasourceCredentialRecord + + +def test_provider_reads_return_scoped_snapshots_and_preserve_default_selection( + sqlite_session_factory: sessionmaker[Session], +) -> None: + repository = SQLAlchemyDatasourceCredentialRepository(session_factory=sqlite_session_factory) + with sqlite_session_factory.begin() as session: + selected = _provider(is_default=True) + other = _provider() + other.name = "Other credential" + foreign = _provider(workspace_id="workspace-2", is_default=True) + wrong_provider = _provider(provider="other-provider", is_default=True) + wrong_plugin = _provider(plugin_id="other-plugin", is_default=True) + session.add_all([selected, other, foreign, wrong_provider, wrong_plugin]) + selected_id, other_id = selected.id, other.id + forbidden_ids = [row.id for row in (foreign, wrong_provider, wrong_plugin)] + + records = repository.list_provider_credentials( + workspace_id="workspace-1", provider="notion_datasource", plugin_id="langgenius/notion_datasource" + ) + assert [record.id for record in records] == [selected_id, other_id] + assert all(isinstance(record, DatasourceCredentialRecord) for record in records) + assert records[0].encrypted_credentials == {"integration_secret": "encrypted"} + default = repository.get_provider_credential( + workspace_id="workspace-1", + provider="notion_datasource", + plugin_id="langgenius/notion_datasource", + credential_id=None, + ) + assert default is not None + assert default.id == selected_id + for credential_id in forbidden_ids: + assert ( + repository.get_provider_credential( + workspace_id="workspace-1", + provider="notion_datasource", + plugin_id="langgenius/notion_datasource", + credential_id=credential_id, + ) + is None + ) + + +def _provider( + *, + workspace_id: str = "workspace-1", + owner_id: str = "owner-1", + provider: str = "notion_datasource", + plugin_id: str = "langgenius/notion_datasource", + visibility: PermissionEnum = PermissionEnum.ONLY_ME, + is_default: bool = False, +) -> DatasourceProvider: + return DatasourceProvider( + tenant_id=workspace_id, + name="Notion", + provider=provider, + plugin_id=plugin_id, + auth_type="oauth2", + encrypted_credentials={"integration_secret": "encrypted"}, + user_id=owner_id, + visibility=visibility, + is_default=is_default, + expires_at=-1, + ) + + +def _notion_source(*, workspace_id: str, credential_id: str | None) -> tuple[Dataset, Document]: + dataset = Dataset( + id="dataset-1", + tenant_id=workspace_id, + name="Notion dataset", + data_source_type=DataSourceType.NOTION_IMPORT, + created_by="owner-1", + ) + document = Document( + id="document-1", + tenant_id=workspace_id, + dataset_id=dataset.id, + position=1, + data_source_type=DataSourceType.NOTION_IMPORT, + data_source_info=json.dumps( + { + **({"credential_id": credential_id} if credential_id is not None else {}), + "notion_workspace_id": "notion-workspace-1", + "notion_page_id": "notion-page-1", + "type": "page", + } + ), + batch="batch-1", + name="Page", + created_from=DocumentCreatedFrom.API, + created_by="owner-1", + ) + return dataset, document + + +def test_get_visible_scopes_workspace_provider_plugin_and_actor( + sqlite_session_factory: sessionmaker[Session], +) -> None: + repository = SQLAlchemyDatasourceCredentialRepository(session_factory=sqlite_session_factory) + with sqlite_session_factory.begin() as session: + credential = _provider(visibility=PermissionEnum.PARTIAL_TEAM) + session.add(credential) + session.flush() + session.add( + CredentialPermission( + credential_id=credential.id, + credential_type=CredentialType.DATASOURCE_PROVIDER, + account_id="member-1", + tenant_id="workspace-1", + ) + ) + credential_id = credential.id + + record = repository.get_visible( + workspace_id="workspace-1", + actor_id="member-1", + credential_id=credential_id, + provider="notion_datasource", + plugin_id="langgenius/notion_datasource", + ) + + assert record is not None + assert record.encrypted_credentials == {"integration_secret": "encrypted"} + assert ( + repository.get_visible( + workspace_id="workspace-2", + actor_id="member-1", + credential_id=credential_id, + provider="notion_datasource", + plugin_id="langgenius/notion_datasource", + ) + is None + ) + assert ( + repository.get_visible( + workspace_id="workspace-1", + actor_id="member-1", + credential_id=credential_id, + provider="other", + plugin_id="langgenius/notion_datasource", + ) + is None + ) + assert ( + repository.get_visible( + workspace_id="workspace-1", + actor_id="member-1", + credential_id=credential_id, + provider="notion_datasource", + plugin_id="other/plugin", + ) + is None + ) + assert ( + repository.get_visible( + workspace_id="workspace-1", + actor_id="outsider-1", + credential_id=credential_id, + provider="notion_datasource", + plugin_id="langgenius/notion_datasource", + ) + is None + ) + + +def test_get_visible_ignores_revoked_partial_permission( + sqlite_session_factory: sessionmaker[Session], +) -> None: + repository = SQLAlchemyDatasourceCredentialRepository(session_factory=sqlite_session_factory) + with sqlite_session_factory.begin() as session: + credential = _provider(visibility=PermissionEnum.PARTIAL_TEAM) + session.add(credential) + session.flush() + session.add( + CredentialPermission( + credential_id=credential.id, + credential_type=CredentialType.DATASOURCE_PROVIDER, + account_id="member-1", + tenant_id="workspace-1", + has_permission=False, + ) + ) + credential_id = credential.id + + assert ( + repository.get_visible( + workspace_id="workspace-1", + actor_id="member-1", + credential_id=credential_id, + provider="notion_datasource", + plugin_id="langgenius/notion_datasource", + ) + is None + ) + + +def test_get_visible_does_not_apply_stale_share_to_private_credential( + sqlite_session_factory: sessionmaker[Session], +) -> None: + repository = SQLAlchemyDatasourceCredentialRepository(session_factory=sqlite_session_factory) + with sqlite_session_factory.begin() as session: + credential = _provider(visibility=PermissionEnum.ONLY_ME) + session.add(credential) + session.flush() + session.add( + CredentialPermission( + credential_id=credential.id, + credential_type=CredentialType.DATASOURCE_PROVIDER, + account_id="member-1", + tenant_id="workspace-1", + ) + ) + credential_id = credential.id + + assert ( + repository.get_visible( + workspace_id="workspace-1", + actor_id="member-1", + credential_id=credential_id, + provider="notion_datasource", + plugin_id="langgenius/notion_datasource", + ) + is None + ) + + +def test_get_oauth_client_config_returns_detached_tenant_and_system_rows( + sqlite_session_factory: sessionmaker[Session], +) -> None: + repository = SQLAlchemyDatasourceCredentialRepository(session_factory=sqlite_session_factory) + with sqlite_session_factory.begin() as session: + session.add_all( + [ + DatasourceOauthTenantParamConfig( + tenant_id="workspace-1", + provider="notion_datasource", + plugin_id="langgenius/notion_datasource", + client_params={"client_secret": "tenant-cipher"}, + enabled=True, + ), + DatasourceOauthTenantParamConfig( + tenant_id="workspace-2", + provider="notion_datasource", + plugin_id="langgenius/notion_datasource", + client_params={"client_secret": "other-cipher"}, + enabled=True, + ), + DatasourceOauthParamConfig( + provider="notion_datasource", + plugin_id="langgenius/notion_datasource", + system_credentials={"client_secret": "system-secret"}, + ), + ] + ) + + record = repository.get_oauth_client_config( + workspace_id="workspace-1", + provider="notion_datasource", + plugin_id="langgenius/notion_datasource", + ) + + assert record.encrypted_tenant_params == {"client_secret": "tenant-cipher"} + assert record.system_credentials == {"client_secret": "system-secret"} + + +def test_update_if_unchanged_detects_concurrent_change( + sqlite_session_factory: sessionmaker[Session], +) -> None: + repository = SQLAlchemyDatasourceCredentialRepository(session_factory=sqlite_session_factory) + with sqlite_session_factory.begin() as session: + credential = _provider() + session.add(credential) + session.flush() + credential_id = credential.id + + record = repository.get_visible( + workspace_id="workspace-1", + actor_id="owner-1", + credential_id=credential_id, + provider="notion_datasource", + plugin_id="langgenius/notion_datasource", + ) + assert record is not None + + with sqlite_session_factory.begin() as session: + stored = session.scalar(select(DatasourceProvider).where(DatasourceProvider.id == credential_id)) + assert stored is not None + stored.encrypted_credentials = {"integration_secret": "concurrent"} + stored.expires_at = 200 + + assert not repository.update_if_unchanged( + record=record, + encrypted_credentials={"integration_secret": "new"}, + expires_at=100, + ) + + +def test_update_if_unchanged_persists_refreshed_snapshot( + sqlite_session_factory: sessionmaker[Session], +) -> None: + repository = SQLAlchemyDatasourceCredentialRepository(session_factory=sqlite_session_factory) + with sqlite_session_factory.begin() as session: + credential = _provider() + session.add(credential) + session.flush() + credential_id = credential.id + + record = repository.get_visible( + workspace_id="workspace-1", + actor_id="owner-1", + credential_id=credential_id, + provider="notion_datasource", + plugin_id="langgenius/notion_datasource", + ) + assert record is not None + assert repository.update_if_unchanged( + record=record, + encrypted_credentials={"integration_secret": "new"}, + expires_at=100, + ) + + with sqlite_session_factory() as session: + stored = session.get(DatasourceProvider, credential_id) + assert stored is not None + assert stored.encrypted_credentials == {"integration_secret": "new"} + assert stored.expires_at == 100 + assert stored.updated_at >= record.updated_at + + +def test_get_for_stored_document_reconstructs_complete_owner_chain( + sqlite_session_factory: sessionmaker[Session], +) -> None: + repository = SQLAlchemyDatasourceCredentialRepository(session_factory=sqlite_session_factory) + with sqlite_session_factory.begin() as session: + credential = _provider(visibility=PermissionEnum.ONLY_ME) + session.add(credential) + session.flush() + dataset, document = _notion_source(workspace_id="workspace-1", credential_id=credential.id) + session.add_all([dataset, document]) + credential_id = credential.id + + record = repository.get_for_stored_document( + workspace_id="workspace-1", + dataset_id="dataset-1", + document_id="document-1", + credential_id=credential_id, + provider="notion_datasource", + plugin_id="langgenius/notion_datasource", + ) + + assert record is not None + assert record.id == credential_id + + +def test_get_for_legacy_stored_document_uses_tenant_default_credential( + sqlite_session_factory: sessionmaker[Session], +) -> None: + repository = SQLAlchemyDatasourceCredentialRepository(session_factory=sqlite_session_factory) + with sqlite_session_factory.begin() as session: + session.add(_provider(is_default=False)) + default_credential = _provider(is_default=True) + default_credential.name = "Default Notion" + session.add(default_credential) + session.flush() + dataset, document = _notion_source(workspace_id="workspace-1", credential_id=None) + session.add_all([dataset, document]) + default_credential_id = default_credential.id + + record = repository.get_for_stored_document( + workspace_id="workspace-1", + dataset_id="dataset-1", + document_id="document-1", + credential_id=None, + provider="notion_datasource", + plugin_id="langgenius/notion_datasource", + ) + + assert record is not None + assert record.id == default_credential_id + + +def test_get_for_legacy_stored_document_does_not_use_non_default_credential( + sqlite_session_factory: sessionmaker[Session], +) -> None: + repository = SQLAlchemyDatasourceCredentialRepository(session_factory=sqlite_session_factory) + with sqlite_session_factory.begin() as session: + session.add(_provider(is_default=False)) + dataset, document = _notion_source(workspace_id="workspace-1", credential_id=None) + session.add_all([dataset, document]) + + record = repository.get_for_stored_document( + workspace_id="workspace-1", + dataset_id="dataset-1", + document_id="document-1", + credential_id=None, + provider="notion_datasource", + plugin_id="langgenius/notion_datasource", + ) + + assert record is None + + +@pytest.mark.parametrize( + ("workspace_id", "provider", "plugin_id"), + [ + ("workspace-2", "notion_datasource", "langgenius/notion_datasource"), + ("workspace-1", "other", "langgenius/notion_datasource"), + ("workspace-1", "notion_datasource", "other/plugin"), + ], +) +def test_get_for_stored_document_rejects_owner_provider_and_plugin_mismatch( + sqlite_session_factory: sessionmaker[Session], + workspace_id: str, + provider: str, + plugin_id: str, +) -> None: + repository = SQLAlchemyDatasourceCredentialRepository(session_factory=sqlite_session_factory) + with sqlite_session_factory.begin() as session: + credential = _provider() + session.add(credential) + session.flush() + dataset, document = _notion_source(workspace_id="workspace-1", credential_id=credential.id) + session.add_all([dataset, document]) + credential_id = credential.id + + assert ( + repository.get_for_stored_document( + workspace_id=workspace_id, + dataset_id="dataset-1", + document_id="document-1", + credential_id=credential_id, + provider=provider, + plugin_id=plugin_id, + ) + is None + ) + + +def test_get_for_stored_document_rejects_deleted_or_different_stored_credential( + sqlite_session_factory: sessionmaker[Session], +) -> None: + repository = SQLAlchemyDatasourceCredentialRepository(session_factory=sqlite_session_factory) + with sqlite_session_factory.begin() as session: + credential = _provider() + session.add(credential) + session.flush() + dataset, document = _notion_source(workspace_id="workspace-1", credential_id="different-credential") + session.add_all([dataset, document]) + credential_id = credential.id + + assert ( + repository.get_for_stored_document( + workspace_id="workspace-1", + dataset_id="dataset-1", + document_id="document-1", + credential_id=credential_id, + provider="notion_datasource", + plugin_id="langgenius/notion_datasource", + ) + is None + ) + + with sqlite_session_factory.begin() as session: + stored_document = session.get(Document, "document-1") + assert stored_document is not None + source_info = stored_document.data_source_info_dict + source_info["credential_id"] = credential_id + stored_document.data_source_info = json.dumps(source_info) + stored_credential = session.get(DatasourceProvider, credential_id) + assert stored_credential is not None + session.delete(stored_credential) + + assert ( + repository.get_for_stored_document( + workspace_id="workspace-1", + dataset_id="dataset-1", + document_id="document-1", + credential_id=credential_id, + provider="notion_datasource", + plugin_id="langgenius/notion_datasource", + ) + is None + ) diff --git a/api/tests/unit_tests/repositories/data_source/test_oauth_binding_repository.py b/api/tests/unit_tests/repositories/data_source/test_oauth_binding_repository.py new file mode 100644 index 00000000000..9e5b1dced33 --- /dev/null +++ b/api/tests/unit_tests/repositories/data_source/test_oauth_binding_repository.py @@ -0,0 +1,247 @@ +from sqlalchemy import null, select, update +from sqlalchemy.orm import Session, sessionmaker + +from models.source import DataSourceOauthBinding +from repositories.data_source.oauth_binding_repository import SQLAlchemyDataSourceOAuthBindingRepository +from services.data_source.entities.oauth import DataSourceOAuthAuthorization + + +def test_upsert_authorization_inserts_and_reenables_existing_binding( + sqlite_session_factory: sessionmaker[Session], +) -> None: + repository = SQLAlchemyDataSourceOAuthBindingRepository(sqlite_session_factory) + authorization = DataSourceOAuthAuthorization(access_token="token", source_info={"pages": []}) + + repository.upsert_authorization( + workspace_id="workspace-1", + provider="notion", + authorization=authorization, + ) + with sqlite_session_factory.begin() as session: + binding = session.scalar(select(DataSourceOauthBinding)) + assert binding is not None + binding.disabled = True + + updated_authorization = DataSourceOAuthAuthorization( + access_token="token", + source_info={"pages": [{"page_id": "page-1"}]}, + ) + repository.upsert_authorization( + workspace_id="workspace-1", + provider="notion", + authorization=updated_authorization, + ) + + with sqlite_session_factory() as session: + bindings = session.scalars(select(DataSourceOauthBinding)).all() + assert len(bindings) == 1 + assert bindings[0].disabled is False + assert bindings[0].source_info == updated_authorization.source_info + + +def test_get_enabled_scopes_by_workspace_provider_and_state( + sqlite_session_factory: sessionmaker[Session], +) -> None: + repository = SQLAlchemyDataSourceOAuthBindingRepository(sqlite_session_factory) + with sqlite_session_factory.begin() as session: + enabled = DataSourceOauthBinding( + tenant_id="workspace-1", + provider="notion", + access_token="enabled-token", + source_info={"pages": []}, + disabled=False, + ) + disabled = DataSourceOauthBinding( + tenant_id="workspace-1", + provider="notion", + access_token="disabled-token", + source_info={"pages": []}, + disabled=True, + ) + session.add_all([enabled, disabled]) + + record = repository.get_enabled( + workspace_id="workspace-1", + provider="notion", + binding_id=enabled.id, + ) + + assert record is not None + assert record.id == enabled.id + assert record.access_token == "enabled-token" + assert ( + repository.get_enabled( + workspace_id="workspace-2", + provider="notion", + binding_id=enabled.id, + ) + is None + ) + assert ( + repository.get_enabled( + workspace_id="workspace-1", + provider="other", + binding_id=enabled.id, + ) + is None + ) + assert ( + repository.get_enabled( + workspace_id="workspace-1", + provider="notion", + binding_id=disabled.id, + ) + is None + ) + + +def test_update_source_info_only_updates_matching_enabled_binding( + sqlite_session_factory: sessionmaker[Session], +) -> None: + repository = SQLAlchemyDataSourceOAuthBindingRepository(sqlite_session_factory) + with sqlite_session_factory.begin() as session: + binding = DataSourceOauthBinding( + tenant_id="workspace-1", + provider="notion", + access_token="token", + source_info={"pages": []}, + disabled=False, + ) + session.add(binding) + + assert repository.update_source_info( + workspace_id="workspace-1", + provider="notion", + binding_id=binding.id, + source_info={"pages": [{"page_id": "page-1"}]}, + ) + assert not repository.update_source_info( + workspace_id="workspace-2", + provider="notion", + binding_id=binding.id, + source_info={"pages": []}, + ) + assert not repository.update_source_info( + workspace_id="workspace-1", + provider="other", + binding_id=binding.id, + source_info={"pages": []}, + ) + + with sqlite_session_factory() as session: + updated = session.get(DataSourceOauthBinding, binding.id) + assert updated is not None + assert updated.source_info == {"pages": [{"page_id": "page-1"}]} + + with sqlite_session_factory.begin() as session: + stored = session.get(DataSourceOauthBinding, binding.id) + assert stored is not None + stored.disabled = True + + assert not repository.update_source_info( + workspace_id="workspace-1", + provider="notion", + binding_id=binding.id, + source_info={"pages": []}, + ) + + +def test_binding_management_lists_enabled_and_changes_state_atomically( + sqlite_session_factory: sessionmaker[Session], +) -> None: + repository = SQLAlchemyDataSourceOAuthBindingRepository(sqlite_session_factory) + with sqlite_session_factory.begin() as session: + enabled = DataSourceOauthBinding( + tenant_id="workspace-1", + provider="notion", + access_token="enabled-token", + source_info={"workspace_id": "notion-workspace"}, + disabled=False, + ) + disabled = DataSourceOauthBinding( + tenant_id="workspace-1", + provider="notion", + access_token="disabled-token", + source_info={}, + disabled=True, + ) + session.add_all([enabled, disabled]) + + summaries = repository.list_enabled_bindings(workspace_id="workspace-1") + + assert [summary.id for summary in summaries] == [enabled.id] + assert ( + repository.change_disabled_state(workspace_id="workspace-1", binding_id=enabled.id, disabled=True) == "updated" + ) + assert ( + repository.change_disabled_state(workspace_id="workspace-1", binding_id=enabled.id, disabled=True) + == "already_disabled" + ) + assert ( + repository.change_disabled_state(workspace_id="workspace-1", binding_id=disabled.id, disabled=False) + == "updated" + ) + assert ( + repository.change_disabled_state(workspace_id="workspace-2", binding_id=disabled.id, disabled=False) + == "not_found" + ) + + +def test_null_disabled_state_is_treated_as_enabled( + sqlite_session_factory: sessionmaker[Session], +) -> None: + repository = SQLAlchemyDataSourceOAuthBindingRepository(sqlite_session_factory) + with sqlite_session_factory.begin() as session: + binding = DataSourceOauthBinding( + tenant_id="workspace-1", + provider="notion", + access_token="token", + source_info={"workspace_id": "notion-workspace"}, + disabled=False, + ) + session.add(binding) + session.flush() + binding_id = binding.id + session.execute( + update(DataSourceOauthBinding).where(DataSourceOauthBinding.id == binding_id).values(disabled=null()) + ) + + assert ( + repository.get_enabled( + workspace_id="workspace-1", + provider="notion", + binding_id=binding_id, + ) + is not None + ) + assert [summary.id for summary in repository.list_enabled_bindings(workspace_id="workspace-1")] == [binding_id] + assert ( + repository.change_disabled_state( + workspace_id="workspace-1", + binding_id=binding_id, + disabled=False, + ) + == "already_enabled" + ) + assert repository.update_source_info( + workspace_id="workspace-1", + provider="notion", + binding_id=binding_id, + source_info={"workspace_id": "updated-workspace"}, + ) + assert ( + repository.change_disabled_state( + workspace_id="workspace-1", + binding_id=binding_id, + disabled=True, + ) + == "updated" + ) + assert ( + repository.get_enabled( + workspace_id="workspace-1", + provider="notion", + binding_id=binding_id, + ) + is None + ) diff --git a/api/tests/unit_tests/repositories/knowledge/__init__.py b/api/tests/unit_tests/repositories/knowledge/__init__.py index e69de29bb2d..3151bb4b71c 100644 --- a/api/tests/unit_tests/repositories/knowledge/__init__.py +++ b/api/tests/unit_tests/repositories/knowledge/__init__.py @@ -0,0 +1 @@ +"""Tests for knowledge persistence adapters.""" diff --git a/api/tests/unit_tests/repositories/knowledge/test_dataset_read_boundaries.py b/api/tests/unit_tests/repositories/knowledge/test_dataset_read_boundaries.py new file mode 100644 index 00000000000..b2e67607f93 --- /dev/null +++ b/api/tests/unit_tests/repositories/knowledge/test_dataset_read_boundaries.py @@ -0,0 +1,303 @@ +"""Behavioral and architecture regressions for explicit knowledge reads.""" + +import ast +from collections.abc import Generator +from contextlib import contextmanager +from datetime import datetime +from pathlib import Path +from urllib.parse import parse_qs, urlsplit + +import pytest +from sqlalchemy import event +from sqlalchemy.orm import Session + +from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom +from core.app.file_access.scope import FileAccessScope, bind_file_access_scope, grant_upload_file_access +from core.tools.signature import verify_tool_file_signature +from extensions.storage.storage_type import StorageType +from fields.document_fields import DocumentResponse +from graphon.file import helpers as file_helpers +from models import dataset as dataset_module +from models.dataset import ( + Dataset, + DatasetMetadata, + DatasetMetadataBinding, + DatasetQuery, + Document, + DocumentSegment, + ExternalKnowledgeApis, + ExternalKnowledgeBindings, +) +from models.enums import CreatorUserRole +from models.model import UploadFile +from models.tools import ToolFile +from repositories.knowledge.dataset_read_repository import get_external_api_bindings_batch +from repositories.knowledge.segment_read_adapter import sign_segment_content +from services.knowledge.dataset_read_service import ( + get_dataset_queries, + get_document_metadata_details, + load_document_details, +) + + +def _dataset() -> Dataset: + return Dataset( + id="dataset-1", tenant_id="tenant-1", name="Knowledge", created_by="account-1", built_in_field_enabled=False + ) + + +def _document(index: int = 1) -> Document: + return Document( + id=f"document-{index}", + tenant_id="tenant-1", + dataset_id="dataset-1", + position=index, + data_source_type="upload_file", + batch="batch-1", + name="source.txt", + created_from="web", + created_by="account-1", + ) + + +def _segment(content: str) -> DocumentSegment: + return DocumentSegment( + tenant_id="tenant-1", + dataset_id="dataset-1", + document_id="document-1", + position=1, + content=content, + word_count=1, + tokens=1, + created_by="account-1", + ) + + +def _upload(tenant_id: str = "tenant-1") -> UploadFile: + return UploadFile( + tenant_id=tenant_id, + storage_type=StorageType.LOCAL, + key="image.png", + name="image.png", + size=10, + extension="png", + mime_type="image/png", + created_by_role=CreatorUserRole.ACCOUNT, + created_by="account-1", + created_at=datetime(2024, 1, 1), + used=True, + ) + + +@contextmanager +def _queries(session: Session) -> Generator[list[str]]: + statements: list[str] = [] + + def count( + _conn: object, _cursor: object, statement: str, _parameters: object, _context: object, _executemany: bool + ) -> None: + statements.append(statement) + + engine = session.get_bind() + event.listen(engine, "before_cursor_execute", count) + try: + yield statements + finally: + event.remove(engine, "before_cursor_execute", count) + + +def test_mixed_links_are_signed_in_place_and_old_parameters_are_replaced(sqlite_session: Session) -> None: + upload = _upload() + tool = ToolFile( + user_id="account-1", tenant_id="tenant-1", conversation_id=None, file_key="tool.png", mimetype="image/png" + ) + sqlite_session.add_all([upload, tool]) + sqlite_session.flush() + paths = [f"/files/tools/{tool.id}.png", f"/files/{upload.id}/image-preview", f"/files/{upload.id}/file-preview"] + segment = _segment( + " BEFORE " + " MIDDLE ".join(path + "?timestamp=1&nonce=old&sign=old" for path in paths) + " AFTER " + ) + signed = sign_segment_content(segment, session=sqlite_session) + assert signed.startswith(" BEFORE ") + assert signed.endswith(" AFTER ") + urls = signed.removeprefix(" BEFORE ").removesuffix(" AFTER ").split(" MIDDLE ") + assert [urlsplit(url).path for url in urls] == paths + for index, url in enumerate(urls): + params = parse_qs(urlsplit(url).query) + assert all(len(value) == 1 for value in params.values()) + values = {key: value[0] for key, value in params.items()} + assert values["nonce"] != "old" + if index == 0: + assert verify_tool_file_signature(tool.id, **values) + elif index == 1: + assert file_helpers.verify_image_signature(upload_file_id=upload.id, **values) + else: + assert file_helpers.verify_file_signature(upload_file_id=upload.id, **values) + # A second pass must also yield valid URLs rather than appending stale parameters. + segment.content = signed + resigned = sign_segment_content(segment, session=sqlite_session) + assert resigned.count("?timestamp=") == 3 + for url in resigned.removeprefix(" BEFORE ").removesuffix(" AFTER ").split(" MIDDLE "): + params = parse_qs(urlsplit(url).query) + assert all(len(values) == 1 for values in params.values()) + assert params["nonce"] != ["old"] + assert params["sign"] != ["old"] + + +def test_foreign_file_reference_never_receives_a_signature(sqlite_session: Session) -> None: + foreign = _upload("tenant-2") + foreign_tool = ToolFile( + user_id="account-2", tenant_id="tenant-2", conversation_id=None, file_key="tool.png", mimetype="image/png" + ) + sqlite_session.add_all([foreign, foreign_tool]) + sqlite_session.flush() + path = f"/files/{foreign.id}/file-preview" + segment = _segment(path + "?timestamp=1&nonce=old&sign=old") + assert sign_segment_content(segment, session=sqlite_session) == path + tool_path = f"/files/tools/{foreign_tool.id}.png" + assert sign_segment_content(_segment(tool_path + "?nonce=old&sign=old"), session=sqlite_session) == tool_path + + +def test_signing_respects_end_user_ownership_in_same_tenant(sqlite_session: Session) -> None: + upload = _upload() + sqlite_session.add(upload) + sqlite_session.flush() + path = f"/files/{upload.id}/file-preview" + scope = FileAccessScope( + tenant_id="tenant-1", user_id="end-user-1", user_from=UserFrom.END_USER, invoke_from=InvokeFrom.WEB_APP + ) + with bind_file_access_scope(scope): + assert sign_segment_content(_segment(path), session=sqlite_session) == path + grant_upload_file_access([upload.id]) + assert "?timestamp=" in sign_segment_content(_segment(path), session=sqlite_session) + + +def test_signing_handles_uppercase_ids_and_ignores_malformed_links(sqlite_session: Session) -> None: + upload = _upload() + sqlite_session.add(upload) + sqlite_session.flush() + path = f"/files/{upload.id.upper()}/file-preview" + signed = sign_segment_content(_segment(path), session=sqlite_session) + params = {key: values[0] for key, values in parse_qs(urlsplit(signed).query).items()} + assert file_helpers.verify_file_signature(upload_file_id=upload.id, **params) + + malformed = "/files/not-a-uuid/file-preview" + with _queries(sqlite_session) as queries: + assert sign_segment_content(_segment(malformed), session=sqlite_session) == malformed + assert queries == [] + + +@pytest.mark.parametrize("content", ["123", "null", "[1]", '{"a":1}', '"text"', "[]", "plain text"]) +def test_plain_text_query_is_not_confused_with_json_records(sqlite_session: Session, content: str) -> None: + query = DatasetQuery( + dataset_id="dataset-1", + content=content, + source="hit_testing", + source_app_id=None, + created_by_role=CreatorUserRole.ACCOUNT, + created_by="account-1", + ) + assert get_dataset_queries(query, session=sqlite_session) == [ + {"content_type": "text_query", "content": content, "file_info": None} + ] + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_document_metadata_respects_dataset_builtin_flag(sqlite_session: Session, enabled: bool) -> None: + dataset = _dataset() + dataset.built_in_field_enabled = enabled + document = _document() + document.doc_metadata = {"custom": "value"} + metadata = DatasetMetadata( + tenant_id=dataset.tenant_id, dataset_id=dataset.id, type="string", name="custom", created_by="account-1" + ) + binding = DatasetMetadataBinding( + tenant_id=dataset.tenant_id, + dataset_id=dataset.id, + document_id=document.id, + metadata_id=metadata.id, + created_by="account-1", + ) + sqlite_session.add_all([dataset, document, metadata, binding]) + sqlite_session.flush() + details = get_document_metadata_details(document, session=sqlite_session) + assert details is not None + assert details[0]["value"] == "value" + assert len([item for item in details if item["id"] == "built-in"]) == (5 if enabled else 0) + + +def test_waiting_document_reports_paused_until_resumed() -> None: + document = _document() + document.indexing_status = "waiting" + document.is_paused = True + assert document.display_status == "paused" + document.is_paused = False + assert document.display_status == "queuing" + + +def test_document_page_load_is_bounded_and_serialization_performs_no_queries(sqlite_session: Session) -> None: + documents = [_document(i) for i in range(1, 21)] + sqlite_session.add_all([_dataset(), *documents]) + sqlite_session.flush() + with _queries(sqlite_session) as single_queries: + load_document_details(documents[:1], session=sqlite_session) + with _queries(sqlite_session) as page_queries: + sources = load_document_details(documents, session=sqlite_session) + assert len(page_queries) == len(single_queries) + with _queries(sqlite_session) as serialization_queries: + responses = [DocumentResponse.model_validate(source, from_attributes=True).model_dump() for source in sources] + assert len(responses) == 20 + assert serialization_queries == [] + + +def test_external_api_bindings_are_batched_and_tenant_scoped(sqlite_session: Session) -> None: + dataset = _dataset() + apis = [ + ExternalKnowledgeApis( + tenant_id="tenant-1", + name=f"API {i}", + description="", + settings="{}", + created_by="account-1", + updated_by=None, + ) + for i in range(20) + ] + foreign_dataset = Dataset(id="dataset-foreign", tenant_id="tenant-2", name="Private", created_by="account-2") + bindings = [ + ExternalKnowledgeBindings( + tenant_id="tenant-1", + external_knowledge_api_id=api.id, + dataset_id=dataset.id, + external_knowledge_id="knowledge", + created_by="account-1", + ) + for api in apis + ] + bindings.append( + ExternalKnowledgeBindings( + tenant_id="tenant-2", + external_knowledge_api_id=apis[0].id, + dataset_id=foreign_dataset.id, + external_knowledge_id="foreign", + created_by="account-2", + ) + ) + sqlite_session.add_all([dataset, foreign_dataset, *apis, *bindings]) + sqlite_session.flush() + with _queries(sqlite_session) as queries: + result = get_external_api_bindings_batch(apis, session=sqlite_session) + assert len(queries) == 1 + assert result == {api.id: [{"id": dataset.id, "name": dataset.name}] for api in apis} + + +def test_dataset_models_have_no_infrastructure_dependencies_or_relative_imports() -> None: + tree = ast.parse(Path(dataset_module.__file__).read_text()) + forbidden = ("repositories", "services", "extensions", "core.tools", "models.engine") + for node in ast.walk(tree): + if isinstance(node, ast.ImportFrom): + assert node.level == 0 + assert not (node.module or "").startswith(forbidden) + assert not any(alias.name in {"Session", "scoped_session"} for alias in node.names) + if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute): + assert ast.unparse(node.func.value) not in {"session", "db.session", "storage"} diff --git a/api/tests/unit_tests/repositories/knowledge/test_dataset_repository.py b/api/tests/unit_tests/repositories/knowledge/test_dataset_repository.py new file mode 100644 index 00000000000..42564cad1ad --- /dev/null +++ b/api/tests/unit_tests/repositories/knowledge/test_dataset_repository.py @@ -0,0 +1,111 @@ +from sqlalchemy.orm import Session, sessionmaker + +from models.dataset import Dataset, DatasetPermission +from repositories.knowledge.dataset_repository import SQLAlchemyDatasetRepository +from services.knowledge.resource_scope import DatasetRef + + +def _dataset( + dataset_id: str, + workspace_id: str, + *, + permission: str = "only_me", + data_source_type: str = "notion_import", +) -> Dataset: + return Dataset( + id=dataset_id, + tenant_id=workspace_id, + name=f"Dataset {dataset_id}", + description="", + provider="vendor", + permission=permission, + created_by="account-1", + maintainer="account-1", + data_source_type=data_source_type, + indexing_technique="economy", + ) + + +def test_estimate_record_and_access_snapshot_are_separate_and_tenant_scoped( + sqlite_session_factory: sessionmaker[Session], +) -> None: + with sqlite_session_factory.begin() as session: + session.add(_dataset("dataset-1", "workspace-1", permission="partial_members")) + session.add( + DatasetPermission( + tenant_id="workspace-1", + dataset_id="dataset-1", + account_id="account-2", + ) + ) + + repository = SQLAlchemyDatasetRepository(session_factory=sqlite_session_factory) + + estimate_record = repository.get_estimate_record(DatasetRef("workspace-1", "dataset-1")) + snapshot = repository.get_access_snapshot( + workspace_id="workspace-1", + dataset_id="dataset-1", + actor_id="account-2", + ) + + assert estimate_record is not None + assert estimate_record.workspace_id == "workspace-1" + assert estimate_record.indexing_technique == "economy" + assert snapshot is not None + assert snapshot.dataset.workspace_id == "workspace-1" + assert snapshot.dataset.permission == "partial_members" + assert snapshot.actor_has_partial_access is True + assert repository.is_notion_dataset(DatasetRef("workspace-1", "dataset-1")) is True + assert repository.is_notion_dataset(DatasetRef("workspace-2", "dataset-1")) is False + assert repository.get_estimate_record(DatasetRef("workspace-2", "dataset-1")) is None + assert ( + repository.get_access_snapshot( + workspace_id="workspace-2", + dataset_id="dataset-1", + actor_id="account-2", + ) + is None + ) + + +def test_access_snapshot_ignores_revoked_permissions( + sqlite_session_factory: sessionmaker[Session], +) -> None: + with sqlite_session_factory.begin() as session: + session.add(_dataset("dataset-1", "workspace-1", permission="partial_members")) + session.add_all( + [ + DatasetPermission( + tenant_id="workspace-1", + dataset_id="dataset-1", + account_id="granted-account", + has_permission=True, + ), + DatasetPermission( + tenant_id="workspace-1", + dataset_id="dataset-1", + account_id="revoked-account", + has_permission=False, + ), + ] + ) + + repository = SQLAlchemyDatasetRepository(session_factory=sqlite_session_factory) + revoked_snapshot = repository.get_access_snapshot( + workspace_id="workspace-1", + dataset_id="dataset-1", + actor_id="revoked-account", + ) + assert revoked_snapshot is not None + assert revoked_snapshot.actor_has_partial_access is False + + +def test_notion_dataset_query_rejects_other_source_types( + sqlite_session_factory: sessionmaker[Session], +) -> None: + with sqlite_session_factory.begin() as session: + session.add(_dataset("dataset-1", "workspace-1", data_source_type="upload_file")) + + repository = SQLAlchemyDatasetRepository(session_factory=sqlite_session_factory) + + assert repository.is_notion_dataset(DatasetRef("workspace-1", "dataset-1")) is False diff --git a/api/tests/unit_tests/repositories/knowledge/test_document_repository.py b/api/tests/unit_tests/repositories/knowledge/test_document_repository.py new file mode 100644 index 00000000000..5820f173e93 --- /dev/null +++ b/api/tests/unit_tests/repositories/knowledge/test_document_repository.py @@ -0,0 +1,273 @@ +import json +from datetime import datetime + +import pytest +from sqlalchemy.orm import Session, sessionmaker + +from core.rag.index_processor.constant.index_type import IndexStructureType +from models.dataset import Dataset, DatasetProcessRule, Document +from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus, ProcessRuleMode +from repositories.knowledge.document_repository import SQLAlchemyDocumentRepository +from services.knowledge.indexing.errors import DocumentIsDeletedPausedError, DocumentIsPausedError +from services.knowledge.resource_scope import DatasetRef + + +def _dataset(dataset_id: str, workspace_id: str) -> Dataset: + return Dataset( + id=dataset_id, + tenant_id=workspace_id, + name=f"Dataset {dataset_id}", + description="", + provider="vendor", + created_by="account-1", + maintainer="account-1", + data_source_type="notion_import", + indexing_technique="economy", + ) + + +def _document( + document_id: str, + *, + workspace_id: str = "workspace-1", + dataset_id: str = "dataset-1", + source_type: str = DataSourceType.NOTION_IMPORT, + enabled: bool = True, + archived: bool = False, + status: str = IndexingStatus.COMPLETED, + batch: str = "batch-1", +) -> Document: + source_info = ( + {"notion_page_id": f"page-{document_id}"} + if source_type == DataSourceType.NOTION_IMPORT + else {"upload_file_id": f"file-{document_id}"} + ) + return Document( + id=document_id, + tenant_id=workspace_id, + dataset_id=dataset_id, + position=1, + data_source_type=source_type, + data_source_info=json.dumps(source_info), + batch=batch, + name=f"Document {document_id}", + created_from=DocumentCreatedFrom.API, + created_by="account-1", + created_at=datetime(2026, 1, 1), + updated_at=datetime(2026, 1, 1), + indexing_status=status, + doc_form=IndexStructureType.PARAGRAPH_INDEX, + word_count=1, + enabled=enabled, + archived=archived, + is_paused=False, + ) + + +def test_document_queries_enforce_the_complete_owner_chain( + sqlite_session_factory: sessionmaker[Session], +) -> None: + with sqlite_session_factory.begin() as session: + session.add_all( + [ + _dataset("dataset-1", "workspace-1"), + _dataset("dataset-2", "workspace-2"), + _document("document-1"), + _document("document-2", workspace_id="workspace-2", dataset_id="dataset-2"), + ] + ) + + repository = SQLAlchemyDocumentRepository(session_factory=sqlite_session_factory) + document_ref = DatasetRef("workspace-1", "dataset-1").document("document-1") + + record = repository.get_estimate_document(document_ref) + + assert record is not None + assert record.workspace_id == "workspace-1" + assert record.data_source_info == {"notion_page_id": "page-document-1"} + assert repository.get_estimate_document(DatasetRef("workspace-2", "dataset-1").document("document-1")) is None + + +def test_pipeline_document_store_owns_sessions_and_scopes_updates( + sqlite_session_factory: sessionmaker[Session], +) -> None: + with sqlite_session_factory.begin() as session: + session.add(_dataset("dataset-1", "workspace-1")) + session.add(_document("document-1")) + + repository = SQLAlchemyDocumentRepository(session_factory=sqlite_session_factory) + assert repository.exists(workspace_id="workspace-1", dataset_id="dataset-1", document_id="document-1") + assert not repository.exists(workspace_id="workspace-2", dataset_id="dataset-1", document_id="document-1") + repository.mark_failed( + workspace_id="workspace-2", dataset_id="dataset-1", document_id="document-1", error="foreign" + ) + with sqlite_session_factory() as session: + document = session.get(Document, "document-1") + assert document is not None + assert document.indexing_status == IndexingStatus.COMPLETED + repository.mark_failed(workspace_id="workspace-1", dataset_id="dataset-1", document_id="document-1", error="failed") + with sqlite_session_factory() as session: + document = session.get(Document, "document-1") + assert document is not None + assert document.indexing_status == IndexingStatus.ERROR + assert document.error == "failed" + + +def test_active_notion_sync_excludes_archived_while_binding_includes_it( + sqlite_session_factory: sessionmaker[Session], +) -> None: + with sqlite_session_factory.begin() as session: + session.add(_dataset("dataset-1", "workspace-1")) + session.add_all( + [ + _document("notion"), + _document("upload", source_type=DataSourceType.UPLOAD_FILE), + _document("disabled", enabled=False), + _document("archived", archived=True), + ] + ) + + repository = SQLAlchemyDocumentRepository(session_factory=sqlite_session_factory) + dataset_ref = DatasetRef("workspace-1", "dataset-1") + + assert repository.list_active_notion_refs(dataset_ref) == (dataset_ref.document("notion"),) + assert repository.list_bound_notion_page_ids(dataset_ref) == frozenset({"page-notion", "page-archived"}) + + +def test_document_records_normalize_invalid_mapping_data_and_ignore_invalid_notion_page_ids( + sqlite_session_factory: sessionmaker[Session], +) -> None: + invalid_json = _document("invalid-json") + invalid_json.data_source_info = "{invalid" + missing_info = _document("missing-info") + missing_info.data_source_info = None + invalid_page_id = _document("invalid-page-id") + invalid_page_id.data_source_info = json.dumps({"notion_page_id": 42}) + with sqlite_session_factory.begin() as session: + session.add(_dataset("dataset-1", "workspace-1")) + session.add_all([invalid_json, missing_info, invalid_page_id]) + + repository = SQLAlchemyDocumentRepository(session_factory=sqlite_session_factory) + dataset_ref = DatasetRef("workspace-1", "dataset-1") + + invalid_record = repository.get_estimate_document(dataset_ref.document("invalid-json")) + missing_record = repository.get_estimate_document(dataset_ref.document("missing-info")) + assert invalid_record is not None + assert invalid_record.data_source_info == {} + assert missing_record is not None + assert missing_record.data_source_info is None + assert repository.list_bound_notion_page_ids(dataset_ref) == frozenset() + + +def test_batch_query_is_scoped_and_keeps_explicit_state_semantics( + sqlite_session_factory: sessionmaker[Session], +) -> None: + with sqlite_session_factory.begin() as session: + session.add_all( + [ + _dataset("dataset-1", "workspace-1"), + _dataset("dataset-2", "workspace-1"), + _dataset("dataset-3", "workspace-2"), + ] + ) + session.add_all( + [ + _document("available"), + _document("waiting", status=IndexingStatus.WAITING), + _document("disabled", enabled=False), + _document("archived", archived=True), + _document("other-batch", batch="batch-2"), + _document("other-dataset", dataset_id="dataset-2"), + _document("other-workspace", workspace_id="workspace-2", dataset_id="dataset-3"), + ] + ) + + repository = SQLAlchemyDocumentRepository(session_factory=sqlite_session_factory) + dataset_ref = DatasetRef("workspace-1", "dataset-1") + + assert {record.id for record in repository.list_estimate_documents_by_batch(dataset_ref, "batch-1")} == { + "available", + "waiting", + "disabled", + "archived", + } + + +@pytest.mark.parametrize( + ("tenant", "dataset", "document"), + [ + ("wrong", "dataset-1", "document-1"), + ("workspace-1", "wrong", "document-1"), + ("workspace-1", "dataset-1", "wrong"), + ], +) +def test_indexing_state_writes_reject_wrong_owner_chain( + sqlite_session_factory: sessionmaker[Session], tenant: str, dataset: str, document: str +) -> None: + with sqlite_session_factory.begin() as session: + session.add_all([_dataset("dataset-1", "workspace-1"), _document("document-1")]) + repository = SQLAlchemyDocumentRepository(session_factory=sqlite_session_factory) + ref = DatasetRef(tenant, dataset).document(document) + for operation in [ + lambda: repository.get_indexing_document(ref), + lambda: repository.mark_splitting(ref), + lambda: repository.complete_indexing(ref, tokens=9, latency=1), + ]: + with pytest.raises(DocumentIsDeletedPausedError): + operation() + repository.fail_indexing(ref, "wrong owner") + with sqlite_session_factory() as session: + persisted_document = session.get(Document, "document-1") + assert persisted_document is not None + assert persisted_document.error is None + + +def test_indexing_load_validates_dataset_owner_and_process_rule(sqlite_session_factory: sessionmaker[Session]) -> None: + with sqlite_session_factory.begin() as session: + document = _document("document-1") + document.dataset_process_rule_id = "other-rule" + rule = DatasetProcessRule( + dataset_id="dataset-2", mode=ProcessRuleMode.AUTOMATIC, rules="{}", created_by="account-1" + ) + rule.id = "other-rule" + session.add_all( + [ + _dataset("dataset-1", "workspace-1"), + _dataset("dataset-2", "workspace-2"), + document, + rule, + ] + ) + repository = SQLAlchemyDocumentRepository(session_factory=sqlite_session_factory) + ref = DatasetRef("workspace-1", "dataset-1").document("document-1") + assert repository.get_indexing_document(ref).processing_rule == {} + with sqlite_session_factory.begin() as session: + persisted_dataset = session.get(Dataset, "dataset-1") + assert persisted_dataset is not None + persisted_dataset.tenant_id = "workspace-2" + with pytest.raises(DocumentIsDeletedPausedError): + repository.get_indexing_document(ref) + + +def test_indexing_transitions_commit_and_preserve_pause(sqlite_session_factory: sessionmaker[Session]) -> None: + with sqlite_session_factory.begin() as session: + session.add_all([_dataset("dataset-1", "workspace-1"), _document("document-1", status=IndexingStatus.PARSING)]) + repository = SQLAlchemyDocumentRepository(session_factory=sqlite_session_factory) + ref = DatasetRef("workspace-1", "dataset-1").document("document-1") + repository.mark_splitting(ref) + with sqlite_session_factory() as session: + persisted_document = session.get(Document, "document-1") + assert persisted_document is not None + assert persisted_document.parsing_completed_at is not None + repository.complete_indexing(ref, tokens=37, latency=2) + with sqlite_session_factory.begin() as session: + row = session.get(Document, "document-1") + assert row is not None + assert (row.indexing_status, row.tokens, row.indexing_latency) == (IndexingStatus.COMPLETED, 37, 2) + row.is_paused = True + with pytest.raises(DocumentIsPausedError): + repository.fail_indexing(ref, "late worker error") + with sqlite_session_factory() as session: + persisted_document = session.get(Document, "document-1") + assert persisted_document is not None + assert persisted_document.error is None diff --git a/api/tests/unit_tests/repositories/knowledge/test_indexing_segments.py b/api/tests/unit_tests/repositories/knowledge/test_indexing_segments.py new file mode 100644 index 00000000000..d4b68d33346 --- /dev/null +++ b/api/tests/unit_tests/repositories/knowledge/test_indexing_segments.py @@ -0,0 +1,119 @@ +from dataclasses import replace + +import pytest +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker + +from core.rag.index_processor.constant.index_type import IndexStructureType +from core.rag.models.document import ChildDocument +from core.rag.models.document import Document as IndexDocument +from models.dataset import ChildChunk, DatasetProcessRule, Document, DocumentSegment +from models.enums import IndexingStatus, ProcessRuleMode, SegmentStatus +from repositories.knowledge.document_repository import SQLAlchemyDocumentRepository +from repositories.knowledge.segment_repository import SQLAlchemySegmentRepository +from services.knowledge.indexing.errors import DocumentIsDeletedPausedError +from services.knowledge.indexing.execution import IndexingDocument +from services.knowledge.resource_scope import DatasetRef +from tests.unit_tests.repositories.knowledge.test_document_repository import _dataset, _document + +IndexingRows = tuple[SQLAlchemySegmentRepository, IndexingDocument, list[IndexDocument]] + + +@pytest.fixture +def indexing_rows(sqlite_session_factory: sessionmaker[Session]) -> IndexingRows: + with sqlite_session_factory.begin() as session: + document = _document("document-1", status=IndexingStatus.SPLITTING) + document.doc_form = IndexStructureType.PARENT_CHILD_INDEX + rule = DatasetProcessRule( + dataset_id="dataset-1", + mode=ProcessRuleMode.HIERARCHICAL, + rules='{"parent_mode": "paragraph"}', + created_by="account-1", + ) + rule.id = "rule-1" + document.dataset_process_rule_id = rule.id + session.add_all([_dataset("dataset-1", "workspace-1"), document, _document("document-2"), rule]) + documents = SQLAlchemyDocumentRepository(session_factory=sqlite_session_factory) + segments = SQLAlchemySegmentRepository(session_factory=sqlite_session_factory) + job = documents.get_indexing_document(DatasetRef("workspace-1", "dataset-1").document("document-1")) + chunks = [ + IndexDocument( + page_content=f"parent-{i}", + metadata={"doc_id": f"node-{i}", "doc_hash": f"hash-{i}"}, + children=[ChildDocument(page_content="child", metadata={"doc_id": f"child-{i}", "doc_hash": "child-hash"})], + ) + for i in range(2) + ] + return segments, job, chunks + + +def test_saved_segments_are_visible_before_workers_and_resume_preserves_total( + sqlite_session_factory: sessionmaker[Session], indexing_rows: IndexingRows +) -> None: + segments, job, chunks = indexing_rows + segments.save_for_indexing(job, chunks, [13, 17]) + segments.complete_indexing_segments(job.ref, ["node-0"]) + resumed, tokens = segments.resume_indexing(job) + assert tokens == 30 + assert [chunk.page_content for chunk in resumed] == ["parent-1"] + assert resumed[0].children is not None + assert resumed[0].children[0].page_content == "child" + with sqlite_session_factory() as session: + persisted_document = session.get(Document, "document-1") + assert persisted_document is not None + assert persisted_document.indexing_status == IndexingStatus.INDEXING + rows = session.scalars(select(DocumentSegment).order_by(DocumentSegment.position)).all() + assert [row.status for row in rows] == [SegmentStatus.COMPLETED, SegmentStatus.INDEXING] + segments.clear_for_indexing(job) + with sqlite_session_factory() as session: + assert session.scalars(select(DocumentSegment)).all() == [] + assert session.scalars(select(ChildChunk)).all() == [] + + +def test_repeated_save_does_not_duplicate_segments_or_children( + sqlite_session_factory: sessionmaker[Session], indexing_rows: IndexingRows +) -> None: + segments, job, chunks = indexing_rows + segments.save_for_indexing(job, chunks, [13, 17]) + segments.save_for_indexing(job, chunks, [19, 23]) + with sqlite_session_factory() as session: + assert len(session.scalars(select(ChildChunk)).all()) == 2 + assert segments.resume_indexing(job)[1] == 42 + + +def test_token_count_mismatch_rolls_back_segments_and_state( + sqlite_session_factory: sessionmaker[Session], indexing_rows: IndexingRows +) -> None: + segments, job, chunks = indexing_rows + with pytest.raises(ValueError): + segments.save_for_indexing(job, chunks, [13]) + with sqlite_session_factory() as session: + assert session.scalars(select(DocumentSegment)).all() == [] + persisted_document = session.get(Document, "document-1") + assert persisted_document is not None + assert persisted_document.indexing_status == IndexingStatus.SPLITTING + + +def test_node_id_collision_cannot_overwrite_another_document(indexing_rows: IndexingRows) -> None: + segments, job, chunks = indexing_rows + other = replace(job, source=job.source._replace(document_ref=job.ref.dataset.document("document-2"))) + segments.save_for_indexing(other, chunks, [3, 5]) + segments.save_for_indexing(job, chunks, [13, 17]) + segments.complete_indexing_segments(job.ref, ["node-0", "node-1"]) + segments.clear_for_indexing(job) + assert segments.resume_indexing(other)[1] == 8 + assert len(segments.resume_indexing(other)[0]) == 2 + + +def test_segment_mutations_reject_foreign_tenant(indexing_rows: IndexingRows) -> None: + segments, job, chunks = indexing_rows + foreign = replace( + job, source=job.source._replace(document_ref=DatasetRef("foreign", "dataset-1").document("document-1")) + ) + for operation in [ + lambda: segments.save_for_indexing(foreign, chunks, [13, 17]), + lambda: segments.clear_for_indexing(foreign), + lambda: segments.complete_indexing_segments(foreign.ref, ["node-0"]), + ]: + with pytest.raises(DocumentIsDeletedPausedError): + operation() diff --git a/api/tests/unit_tests/repositories/knowledge/test_keyword_table_repository.py b/api/tests/unit_tests/repositories/knowledge/test_keyword_table_repository.py new file mode 100644 index 00000000000..e8f85b93eaf --- /dev/null +++ b/api/tests/unit_tests/repositories/knowledge/test_keyword_table_repository.py @@ -0,0 +1,66 @@ +import pytest +from sqlalchemy import select +from sqlalchemy.orm import Session + +from models.dataset import Dataset, DatasetKeywordTable, DocumentSegment +from repositories.knowledge.keyword_table_repository import persist_keyword_table + + +@pytest.mark.parametrize("existing_table", [False, True]) +def test_keyword_writes_roll_back_with_the_callers_transaction(sqlite_session: Session, existing_table: bool) -> None: + sqlite_session.add(Dataset(id="dataset-1", tenant_id="tenant-1", name="Test", created_by="author")) + segment = DocumentSegment( + tenant_id="tenant-1", + dataset_id="dataset-1", + document_id="document-1", + position=1, + content="original", + word_count=8, + tokens=1, + created_by="author", + index_node_id="node-1", + keywords=["original"], + ) + sqlite_session.add(segment) + if existing_table: + sqlite_session.add(DatasetKeywordTable(dataset_id="dataset-1", keyword_table="original")) + sqlite_session.commit() + + persist_keyword_table( + sqlite_session, + tenant_id="tenant-1", + dataset_id="dataset-1", + storage_type="database", + data="replacement", + keywords={"node-1": ["replacement"]}, + ) + assert segment.keywords == ["replacement"] + sqlite_session.rollback() + + assert segment.keywords == ["original"] + row = sqlite_session.scalar(select(DatasetKeywordTable)) + if existing_table: + assert row is not None + assert row.keyword_table == "original" + else: + assert row is None + + +@pytest.mark.parametrize(("tenant_id", "dataset_id"), [("tenant-2", "dataset-1"), ("tenant-1", "missing")]) +def test_keyword_writes_require_a_tenant_owned_dataset( + sqlite_session: Session, tenant_id: str, dataset_id: str +) -> None: + sqlite_session.add(Dataset(id="dataset-1", tenant_id="tenant-1", name="Test", created_by="author")) + sqlite_session.commit() + + with pytest.raises(LookupError, match="Dataset no longer exists"): + persist_keyword_table( + sqlite_session, + tenant_id=tenant_id, + dataset_id=dataset_id, + storage_type="database", + data="replacement", + keywords={}, + ) + + assert sqlite_session.scalar(select(DatasetKeywordTable)) is None diff --git a/api/tests/unit_tests/repositories/test_data_source_oauth_binding_repository.py b/api/tests/unit_tests/repositories/test_data_source_oauth_binding_repository.py deleted file mode 100644 index 468c41f46fd..00000000000 --- a/api/tests/unit_tests/repositories/test_data_source_oauth_binding_repository.py +++ /dev/null @@ -1,120 +0,0 @@ -from sqlalchemy import select -from sqlalchemy.orm import Session, sessionmaker - -from models.source import DataSourceOauthBinding -from repositories.data_source_oauth_binding_repository import SQLAlchemyDataSourceOAuthBindingRepository -from services.entities.data_source_oauth_entities import DataSourceOAuthAuthorization - - -def test_upsert_authorization_inserts_and_reenables_existing_binding( - sqlite_session_factory: sessionmaker[Session], -) -> None: - repository = SQLAlchemyDataSourceOAuthBindingRepository(sqlite_session_factory) - authorization = DataSourceOAuthAuthorization(access_token="token", source_info={"pages": []}) - - repository.upsert_authorization( - workspace_id="workspace-1", - provider="notion", - authorization=authorization, - ) - with sqlite_session_factory.begin() as session: - binding = session.scalar(select(DataSourceOauthBinding)) - assert binding is not None - binding.disabled = True - - updated_authorization = DataSourceOAuthAuthorization( - access_token="token", - source_info={"pages": [{"page_id": "page-1"}]}, - ) - repository.upsert_authorization( - workspace_id="workspace-1", - provider="notion", - authorization=updated_authorization, - ) - - with sqlite_session_factory() as session: - bindings = session.scalars(select(DataSourceOauthBinding)).all() - assert len(bindings) == 1 - assert bindings[0].disabled is False - assert bindings[0].source_info == updated_authorization.source_info - - -def test_get_enabled_scopes_by_workspace_provider_and_state( - sqlite_session_factory: sessionmaker[Session], -) -> None: - repository = SQLAlchemyDataSourceOAuthBindingRepository(sqlite_session_factory) - with sqlite_session_factory.begin() as session: - enabled = DataSourceOauthBinding( - tenant_id="workspace-1", - provider="notion", - access_token="enabled-token", - source_info={"pages": []}, - disabled=False, - ) - disabled = DataSourceOauthBinding( - tenant_id="workspace-1", - provider="notion", - access_token="disabled-token", - source_info={"pages": []}, - disabled=True, - ) - session.add_all([enabled, disabled]) - - record = repository.get_enabled( - workspace_id="workspace-1", - provider="notion", - binding_id=enabled.id, - ) - - assert record is not None - assert record.id == enabled.id - assert record.access_token == "enabled-token" - assert ( - repository.get_enabled( - workspace_id="workspace-2", - provider="notion", - binding_id=enabled.id, - ) - is None - ) - assert ( - repository.get_enabled( - workspace_id="workspace-1", - provider="notion", - binding_id=disabled.id, - ) - is None - ) - - -def test_update_source_info_only_updates_matching_enabled_binding( - sqlite_session_factory: sessionmaker[Session], -) -> None: - repository = SQLAlchemyDataSourceOAuthBindingRepository(sqlite_session_factory) - with sqlite_session_factory.begin() as session: - binding = DataSourceOauthBinding( - tenant_id="workspace-1", - provider="notion", - access_token="token", - source_info={"pages": []}, - disabled=False, - ) - session.add(binding) - - assert repository.update_source_info( - workspace_id="workspace-1", - provider="notion", - binding_id=binding.id, - source_info={"pages": [{"page_id": "page-1"}]}, - ) - assert not repository.update_source_info( - workspace_id="workspace-2", - provider="notion", - binding_id=binding.id, - source_info={"pages": []}, - ) - - with sqlite_session_factory() as session: - updated = session.get(DataSourceOauthBinding, binding.id) - assert updated is not None - assert updated.source_info == {"pages": [{"page_id": "page-1"}]} diff --git a/api/tests/unit_tests/repositories/test_workspace_member_query_repository.py b/api/tests/unit_tests/repositories/test_workspace_member_query_repository.py index 852dc543e67..8aacf2a03b1 100644 --- a/api/tests/unit_tests/repositories/test_workspace_member_query_repository.py +++ b/api/tests/unit_tests/repositories/test_workspace_member_query_repository.py @@ -124,3 +124,34 @@ def test_list_for_workspace_returns_empty_tuple_without_membership( result = WorkspaceMemberQueryRepository(sqlite_session_factory).list_for_workspace("workspace-1") assert result == () + + +def test_get_legacy_role_is_scoped_to_workspace_and_account( + sqlite_session_factory: sessionmaker[Session], +) -> None: + created_at = datetime(2026, 1, 1) + account = make_account("member", status=AccountStatus.ACTIVE, created_at=created_at) + workspace = make_tenant("workspace-1") + other_workspace = make_tenant("workspace-2") + with sqlite_session_factory.begin() as session: + session.add_all( + [ + account, + workspace, + other_workspace, + TenantAccountJoin( + tenant_id=workspace.id, + account_id=account.id, + role=TenantAccountRole.DATASET_OPERATOR, + ), + ] + ) + + repository = WorkspaceMemberQueryRepository(sqlite_session_factory) + + assert ( + repository.get_legacy_role(workspace_id=workspace.id, account_id=account.id) + == TenantAccountRole.DATASET_OPERATOR.value + ) + assert repository.get_legacy_role(workspace_id=other_workspace.id, account_id=account.id) is None + assert repository.get_legacy_role(workspace_id=workspace.id, account_id="other-account") is None diff --git a/api/tests/unit_tests/services/agent/test_agent_services.py b/api/tests/unit_tests/services/agent/test_agent_services.py index dcf3a6e5a3e..e4266eeed99 100644 --- a/api/tests/unit_tests/services/agent/test_agent_services.py +++ b/api/tests/unit_tests/services/agent/test_agent_services.py @@ -6750,7 +6750,7 @@ def test_dataset_rows_filters_malformed_ids(monkeypatch: pytest.MonkeyPatch, sql captured["ids"] = ids return [], 0 - import services.dataset_service as dataset_service_module + import services.knowledge.dataset_service as dataset_service_module from services.agent.knowledge_datasets import get_tenant_knowledge_dataset_rows monkeypatch.setattr(dataset_service_module.DatasetService, "get_datasets_by_ids", fake_get_datasets_by_ids) @@ -6779,7 +6779,7 @@ def test_composer_save_rejects_malformed_knowledge_dataset_ids( captured["tenant_id"] = tenant_id return [], 0 - import services.dataset_service as dataset_service_module + import services.knowledge.dataset_service as dataset_service_module monkeypatch.setattr(dataset_service_module.DatasetService, "get_datasets_by_ids", fake_get_datasets_by_ids) @@ -6818,7 +6818,7 @@ def test_composer_save_rejects_missing_or_out_of_scope_knowledge_datasets( captured["tenant_id"] = tenant_id return [], 0 - import services.dataset_service as dataset_service_module + import services.knowledge.dataset_service as dataset_service_module monkeypatch.setattr(dataset_service_module.DatasetService, "get_datasets_by_ids", fake_get_datasets_by_ids) @@ -6877,7 +6877,7 @@ def test_save_agent_composer_allows_incomplete_knowledge_draft( event.listen(session, "after_flush", count_flush) - import services.dataset_service as dataset_service_module + import services.knowledge.dataset_service as dataset_service_module monkeypatch.setattr( dataset_service_module.DatasetService, diff --git a/api/tests/unit_tests/services/app/test_console_service.py b/api/tests/unit_tests/services/app/test_console_service.py index 48158240e60..75eb694c9e8 100644 --- a/api/tests/unit_tests/services/app/test_console_service.py +++ b/api/tests/unit_tests/services/app/test_console_service.py @@ -62,7 +62,7 @@ from services.entities.dsl_entities import ( Import, ImportStatus, ) -from services.model_provider_service import ModelProviderService +from services.model_provider.service import ModelProviderService CONTEXT = RequestContext("request", "trace", "actor", "workspace") RECORD = AppRecord(id="app", name="Example", mode_compatible_with_agent="chat") diff --git a/api/tests/unit_tests/services/auth/test_data_source_api_key_auth_gateways.py b/api/tests/unit_tests/services/auth/test_data_source_api_key_auth_gateways.py index 8de8e952484..f29296890ff 100644 --- a/api/tests/unit_tests/services/auth/test_data_source_api_key_auth_gateways.py +++ b/api/tests/unit_tests/services/auth/test_data_source_api_key_auth_gateways.py @@ -3,13 +3,13 @@ from unittest.mock import MagicMock, patch import httpx import pytest -from services.auth.data_source_api_key_auth_gateways import ProviderApiKeyAuthCredentialValidator from services.auth.errors import ( DataSourceApiKeyAuthCredentialValidationError, DataSourceApiKeyAuthProviderUnavailableError, UnsupportedDataSourceApiKeyAuthProviderError, ) -from services.entities.data_source_api_key_auth_entities import DataSourceApiKeyAuthCredentials +from services.data_source.auth.api_key_gateways import ProviderApiKeyAuthCredentialValidator +from services.data_source.entities.api_key_auth import DataSourceApiKeyAuthCredentials class TestProviderApiKeyAuthCredentialValidator: @@ -18,17 +18,17 @@ class TestProviderApiKeyAuthCredentialValidator: [ ( "firecrawl", - "services.auth.firecrawl.firecrawl.FirecrawlAuth", + "services.data_source.auth.firecrawl.firecrawl.FirecrawlAuth", DataSourceApiKeyAuthCredentials("bearer", "test_key", {}), ), ( "watercrawl", - "services.auth.watercrawl.watercrawl.WatercrawlAuth", + "services.data_source.auth.watercrawl.watercrawl.WatercrawlAuth", DataSourceApiKeyAuthCredentials("x-api-key", "test_key", {}), ), ( "jinareader", - "services.auth.jina.jina.JinaAuth", + "services.data_source.auth.jina.jina.JinaAuth", DataSourceApiKeyAuthCredentials("bearer", "test_key", {}), ), ], @@ -65,7 +65,7 @@ class TestProviderApiKeyAuthCredentialValidator: auth_instance.validate_credentials.return_value = validation_result with patch( - "services.auth.firecrawl.firecrawl.FirecrawlAuth", + "services.data_source.auth.firecrawl.firecrawl.FirecrawlAuth", autospec=True, return_value=auth_instance, ): @@ -83,7 +83,7 @@ class TestProviderApiKeyAuthCredentialValidator: ) with patch( - "services.auth.firecrawl.firecrawl.FirecrawlAuth", + "services.data_source.auth.firecrawl.firecrawl.FirecrawlAuth", autospec=True, return_value=auth_instance, ): @@ -98,7 +98,7 @@ class TestProviderApiKeyAuthCredentialValidator: auth_instance.validate_credentials.side_effect = httpx.ConnectError("Authentication endpoint unavailable") with patch( - "services.auth.firecrawl.firecrawl.FirecrawlAuth", + "services.data_source.auth.firecrawl.firecrawl.FirecrawlAuth", autospec=True, return_value=auth_instance, ): @@ -119,7 +119,7 @@ class TestProviderApiKeyAuthCredentialValidator: auth_instance.validate_credentials.side_effect = status_error with patch( - "services.auth.firecrawl.firecrawl.FirecrawlAuth", + "services.data_source.auth.firecrawl.firecrawl.FirecrawlAuth", autospec=True, return_value=auth_instance, ): diff --git a/api/tests/unit_tests/services/auth/test_data_source_api_key_auth_service.py b/api/tests/unit_tests/services/auth/test_data_source_api_key_auth_service.py index ee3a31b9004..4ab25292f30 100644 --- a/api/tests/unit_tests/services/auth/test_data_source_api_key_auth_service.py +++ b/api/tests/unit_tests/services/auth/test_data_source_api_key_auth_service.py @@ -4,8 +4,8 @@ from datetime import datetime from unittest.mock import MagicMock from machinery.context import RequestContext -from services.auth.data_source_api_key_auth_service import DataSourceApiKeyAuthService -from services.entities.data_source_api_key_auth_entities import ( +from services.data_source.auth.api_key_service import DataSourceApiKeyAuthService +from services.data_source.entities.api_key_auth import ( DataSourceApiKeyAuthBindingCreate, DataSourceApiKeyAuthBindingRecord, DataSourceApiKeyAuthCredentials, diff --git a/api/tests/unit_tests/services/auth/test_firecrawl_auth.py b/api/tests/unit_tests/services/auth/test_firecrawl_auth.py index 476a685f15b..9514d70ed48 100644 --- a/api/tests/unit_tests/services/auth/test_firecrawl_auth.py +++ b/api/tests/unit_tests/services/auth/test_firecrawl_auth.py @@ -9,8 +9,8 @@ from services.auth.errors import ( DataSourceApiKeyAuthProviderUnavailableError, InvalidDataSourceApiKeyAuthCredentialsError, ) -from services.auth.firecrawl.firecrawl import FirecrawlAuth -from services.entities.data_source_api_key_auth_entities import DataSourceApiKeyAuthCredentials +from services.data_source.auth.firecrawl.firecrawl import FirecrawlAuth +from services.data_source.entities.api_key_auth import DataSourceApiKeyAuthCredentials def _credentials( @@ -64,7 +64,7 @@ class TestFirecrawlAuth: with pytest.raises(InvalidDataSourceApiKeyAuthCredentialsError, match="No API key provided"): FirecrawlAuth(_credentials(api_key="")) - @patch("services.auth.firecrawl.firecrawl.httpx.post", autospec=True) + @patch("services.data_source.auth.firecrawl.firecrawl.httpx.post", autospec=True) def test_should_validate_valid_credentials_successfully(self, mock_post: MagicMock, auth_instance: FirecrawlAuth): """Test successful credential validation""" mock_response = MagicMock() @@ -95,7 +95,7 @@ class TestFirecrawlAuth: (409, "Conflict error"), ], ) - @patch("services.auth.firecrawl.firecrawl.httpx.post", autospec=True) + @patch("services.data_source.auth.firecrawl.firecrawl.httpx.post", autospec=True) def test_should_handle_http_errors( self, mock_post: MagicMock, status_code, error_message, auth_instance: FirecrawlAuth ): @@ -110,7 +110,7 @@ class TestFirecrawlAuth: assert str(exc_info.value) == f"Failed to authorize. Status code: {status_code}. Error: {error_message}" @pytest.mark.parametrize("status_code", [429, 500, 502, 503]) - @patch("services.auth.firecrawl.firecrawl.httpx.post", autospec=True) + @patch("services.data_source.auth.firecrawl.firecrawl.httpx.post", autospec=True) def test_should_map_upstream_failure_to_provider_unavailable( self, mock_post: MagicMock, @@ -136,7 +136,7 @@ class TestFirecrawlAuth: (401, "Not JSON", True, "Failed to authorize. Status code: 401. Error: Not JSON"), ], ) - @patch("services.auth.firecrawl.firecrawl.httpx.post", autospec=True) + @patch("services.data_source.auth.firecrawl.firecrawl.httpx.post", autospec=True) def test_should_handle_unexpected_errors( self, mock_post: MagicMock, @@ -169,7 +169,7 @@ class TestFirecrawlAuth: (httpx.ConnectTimeout, "Connection timeout"), ], ) - @patch("services.auth.firecrawl.firecrawl.httpx.post", autospec=True) + @patch("services.data_source.auth.firecrawl.firecrawl.httpx.post", autospec=True) def test_should_handle_network_errors( self, mock_post: MagicMock, exception_type, exception_message, auth_instance: FirecrawlAuth ): @@ -193,7 +193,7 @@ class TestFirecrawlAuth: FirecrawlAuth(_credentials(auth_type="basic", api_key="super_secret_key_12345")) assert "super_secret_key_12345" not in str(exc_info.value) - @patch("services.auth.firecrawl.firecrawl.httpx.post", autospec=True) + @patch("services.data_source.auth.firecrawl.firecrawl.httpx.post", autospec=True) def test_should_use_custom_base_url_in_validation(self, mock_post): """Test that custom base URL is used in validation and normalized""" mock_response = MagicMock() @@ -208,7 +208,7 @@ class TestFirecrawlAuth: assert result is True assert mock_post.call_args[0][0] == "https://custom.firecrawl.dev/v1/crawl" - @patch("services.auth.firecrawl.firecrawl.httpx.post", autospec=True) + @patch("services.data_source.auth.firecrawl.firecrawl.httpx.post", autospec=True) def test_should_handle_timeout_with_retry_suggestion(self, mock_post: MagicMock, auth_instance: FirecrawlAuth): """Test that timeout errors are handled gracefully with appropriate error message""" mock_post.side_effect = httpx.TimeoutException("The request timed out after 30 seconds") diff --git a/api/tests/unit_tests/services/auth/test_jina_auth.py b/api/tests/unit_tests/services/auth/test_jina_auth.py index a1e41938e9d..507edc466b0 100644 --- a/api/tests/unit_tests/services/auth/test_jina_auth.py +++ b/api/tests/unit_tests/services/auth/test_jina_auth.py @@ -9,8 +9,8 @@ from services.auth.errors import ( DataSourceApiKeyAuthProviderUnavailableError, InvalidDataSourceApiKeyAuthCredentialsError, ) -from services.auth.jina.jina import JinaAuth -from services.entities.data_source_api_key_auth_entities import DataSourceApiKeyAuthCredentials +from services.data_source.auth.jina.jina import JinaAuth +from services.data_source.entities.api_key_auth import DataSourceApiKeyAuthCredentials def _credentials( @@ -41,7 +41,7 @@ class TestJinaAuth: JinaAuth(credentials) assert str(exc_info.value) == "No API key provided" - @patch("services.auth.jina.jina._http_client.post", autospec=True) + @patch("services.data_source.auth.jina.jina._http_client.post", autospec=True) def test_should_validate_valid_credentials_successfully(self, mock_post: MagicMock): """Test successful credential validation""" mock_response = MagicMock() @@ -59,7 +59,7 @@ class TestJinaAuth: json={"url": "https://example.com"}, ) - @patch("services.auth.jina.jina._http_client.post", autospec=True) + @patch("services.data_source.auth.jina.jina._http_client.post", autospec=True) def test_should_handle_http_402_error(self, mock_post: MagicMock): """Test handling of 402 Payment Required error""" mock_response = MagicMock() @@ -74,7 +74,7 @@ class TestJinaAuth: auth.validate_credentials() assert str(exc_info.value) == "Failed to authorize. Status code: 402. Error: Payment required" - @patch("services.auth.jina.jina._http_client.post", autospec=True) + @patch("services.data_source.auth.jina.jina._http_client.post", autospec=True) def test_should_handle_http_error_with_non_json_text_response(self, mock_post): """Test handling of known HTTP errors with non-JSON text response.""" mock_response = MagicMock() @@ -90,7 +90,7 @@ class TestJinaAuth: auth.validate_credentials() assert str(exc_info.value) == "Failed to authorize. Status code: 402. Error: Payment required" - @patch("services.auth.jina.jina._http_client.post", autospec=True) + @patch("services.data_source.auth.jina.jina._http_client.post", autospec=True) def test_should_handle_http_409_error(self, mock_post): """Test handling of 409 Conflict error""" mock_response = MagicMock() @@ -106,7 +106,7 @@ class TestJinaAuth: assert str(exc_info.value) == "Failed to authorize. Status code: 409. Error: Conflict error" @pytest.mark.parametrize("status_code", [429, 500, 502, 503]) - @patch("services.auth.jina.jina._http_client.post", autospec=True) + @patch("services.data_source.auth.jina.jina._http_client.post", autospec=True) def test_should_map_upstream_failure_to_provider_unavailable( self, mock_post: MagicMock, @@ -123,7 +123,7 @@ class TestJinaAuth: assert exc_info.value.provider == "jinareader" assert exc_info.value.status_code == status_code - @patch("services.auth.jina.jina._http_client.post", autospec=True) + @patch("services.data_source.auth.jina.jina._http_client.post", autospec=True) def test_should_handle_unexpected_error_with_text_response(self, mock_post: MagicMock): """Test handling of unexpected errors with text response""" mock_response = MagicMock() @@ -139,7 +139,7 @@ class TestJinaAuth: auth.validate_credentials() assert str(exc_info.value) == "Failed to authorize. Status code: 403. Error: Forbidden" - @patch("services.auth.jina.jina._http_client.post", autospec=True) + @patch("services.data_source.auth.jina.jina._http_client.post", autospec=True) def test_should_handle_unexpected_error_with_non_json_text_response(self, mock_post): """Test handling of unexpected errors with non-JSON text response.""" mock_response = MagicMock() @@ -155,7 +155,7 @@ class TestJinaAuth: auth.validate_credentials() assert str(exc_info.value) == "Failed to authorize. Status code: 403. Error: Forbidden" - @patch("services.auth.jina.jina._http_client.post", autospec=True) + @patch("services.data_source.auth.jina.jina._http_client.post", autospec=True) def test_should_handle_unexpected_error_without_text(self, mock_post): """Test handling of unexpected errors without text response""" mock_response = MagicMock() @@ -171,7 +171,7 @@ class TestJinaAuth: auth.validate_credentials() assert str(exc_info.value) == "Unexpected error occurred while trying to authorize. Status code: 404" - @patch("services.auth.jina.jina._http_client.post", autospec=True) + @patch("services.data_source.auth.jina.jina._http_client.post", autospec=True) def test_should_handle_network_errors(self, mock_post: MagicMock): """Test handling of network connection errors""" mock_post.side_effect = httpx.ConnectError("Network error") diff --git a/api/tests/unit_tests/services/auth/test_watercrawl_auth.py b/api/tests/unit_tests/services/auth/test_watercrawl_auth.py index 09fd6e86a80..e54df0dee9e 100644 --- a/api/tests/unit_tests/services/auth/test_watercrawl_auth.py +++ b/api/tests/unit_tests/services/auth/test_watercrawl_auth.py @@ -9,8 +9,8 @@ from services.auth.errors import ( DataSourceApiKeyAuthProviderUnavailableError, InvalidDataSourceApiKeyAuthCredentialsError, ) -from services.auth.watercrawl.watercrawl import WatercrawlAuth -from services.entities.data_source_api_key_auth_entities import DataSourceApiKeyAuthCredentials +from services.data_source.auth.watercrawl.watercrawl import WatercrawlAuth +from services.data_source.entities.api_key_auth import DataSourceApiKeyAuthCredentials def _credentials( @@ -64,7 +64,7 @@ class TestWatercrawlAuth: with pytest.raises(InvalidDataSourceApiKeyAuthCredentialsError, match="No API key provided"): WatercrawlAuth(_credentials(api_key="")) - @patch("services.auth.watercrawl.watercrawl.httpx.get", autospec=True) + @patch("services.data_source.auth.watercrawl.watercrawl.httpx.get", autospec=True) def test_should_validate_valid_credentials_successfully(self, mock_get, auth_instance): """Test successful credential validation""" mock_response = MagicMock() @@ -87,7 +87,7 @@ class TestWatercrawlAuth: (409, "Conflict error"), ], ) - @patch("services.auth.watercrawl.watercrawl.httpx.get", autospec=True) + @patch("services.data_source.auth.watercrawl.watercrawl.httpx.get", autospec=True) def test_should_handle_http_errors(self, mock_get, status_code, error_message, auth_instance): """Test handling of various HTTP error codes""" mock_response = MagicMock() @@ -100,7 +100,7 @@ class TestWatercrawlAuth: assert str(exc_info.value) == f"Failed to authorize. Status code: {status_code}. Error: {error_message}" @pytest.mark.parametrize("status_code", [429, 500, 502, 503]) - @patch("services.auth.watercrawl.watercrawl.httpx.get", autospec=True) + @patch("services.data_source.auth.watercrawl.watercrawl.httpx.get", autospec=True) def test_should_map_upstream_failure_to_provider_unavailable( self, mock_get: MagicMock, @@ -116,7 +116,7 @@ class TestWatercrawlAuth: assert exc_info.value.provider == "watercrawl" assert exc_info.value.status_code == status_code - @patch("services.auth.watercrawl.watercrawl.httpx.get", autospec=True) + @patch("services.data_source.auth.watercrawl.watercrawl.httpx.get", autospec=True) def test_should_handle_http_error_with_non_json_text_response(self, mock_get, auth_instance): """Test handling of known HTTP errors with non-JSON text response.""" mock_response = MagicMock() @@ -137,7 +137,7 @@ class TestWatercrawlAuth: (401, "Not JSON", True, "Failed to authorize. Status code: 401. Error: Not JSON"), ], ) - @patch("services.auth.watercrawl.watercrawl.httpx.get", autospec=True) + @patch("services.data_source.auth.watercrawl.watercrawl.httpx.get", autospec=True) def test_should_handle_unexpected_errors( self, mock_get, status_code, response_text, has_json_error, expected_error_contains, auth_instance ): @@ -162,7 +162,7 @@ class TestWatercrawlAuth: (httpx.ConnectTimeout, "Connection timeout"), ], ) - @patch("services.auth.watercrawl.watercrawl.httpx.get", autospec=True) + @patch("services.data_source.auth.watercrawl.watercrawl.httpx.get", autospec=True) def test_should_handle_network_errors(self, mock_get, exception_type, exception_message, auth_instance): """Test handling of various network-related errors including timeouts""" mock_get.side_effect = exception_type(exception_message) @@ -184,7 +184,7 @@ class TestWatercrawlAuth: WatercrawlAuth(_credentials(auth_type="bearer", api_key="super_secret_key_12345")) assert "super_secret_key_12345" not in str(exc_info.value) - @patch("services.auth.watercrawl.watercrawl.httpx.get", autospec=True) + @patch("services.data_source.auth.watercrawl.watercrawl.httpx.get", autospec=True) def test_should_use_custom_base_url_in_validation(self, mock_get): """Test that custom base URL is used in validation""" mock_response = MagicMock() @@ -206,7 +206,7 @@ class TestWatercrawlAuth: ("https://app.watercrawl.dev//", "https://app.watercrawl.dev/api/v1/core/crawl-requests/"), ], ) - @patch("services.auth.watercrawl.watercrawl.httpx.get", autospec=True) + @patch("services.data_source.auth.watercrawl.watercrawl.httpx.get", autospec=True) def test_should_use_urljoin_for_url_construction(self, mock_get, base_url, expected_url): """Test that urljoin is used correctly for URL construction with various base URLs""" mock_response = MagicMock() @@ -220,7 +220,7 @@ class TestWatercrawlAuth: # Verify the correct URL was called assert mock_get.call_args[0][0] == expected_url - @patch("services.auth.watercrawl.watercrawl.httpx.get", autospec=True) + @patch("services.data_source.auth.watercrawl.watercrawl.httpx.get", autospec=True) def test_should_handle_timeout_with_retry_suggestion(self, mock_get, auth_instance): """Test that timeout errors are handled gracefully with appropriate error message""" mock_get.side_effect = httpx.TimeoutException("The request timed out after 30 seconds") diff --git a/api/tests/unit_tests/services/data_source/__init__.py b/api/tests/unit_tests/services/data_source/__init__.py new file mode 100644 index 00000000000..2f04e29e200 --- /dev/null +++ b/api/tests/unit_tests/services/data_source/__init__.py @@ -0,0 +1 @@ +"""Tests for data-source application services and adapters.""" diff --git a/api/tests/unit_tests/services/data_source/test_binding_application_service.py b/api/tests/unit_tests/services/data_source/test_binding_application_service.py new file mode 100644 index 00000000000..f8bb0c64bd5 --- /dev/null +++ b/api/tests/unit_tests/services/data_source/test_binding_application_service.py @@ -0,0 +1,89 @@ +from dataclasses import dataclass, field +from datetime import datetime + +import pytest + +from machinery.context import RequestContext +from services.data_source.binding_application_service import ( + BindingAction, + BindingMutationResult, + DataSourceBindingApplicationService, + DataSourceBindingNotFoundError, + DataSourceBindingStateError, +) +from services.data_source.entities.oauth import DataSourceBindingSummary + + +def _context() -> RequestContext: + return RequestContext("request-1", None, "account-1", "workspace-1") + + +@dataclass +class InMemoryBindingStore: + bindings: tuple[DataSourceBindingSummary, ...] = () + mutation_result: BindingMutationResult = "updated" + changes: list[tuple[str, str, bool]] = field(default_factory=list) + + def list_enabled_bindings(self, *, workspace_id: str) -> tuple[DataSourceBindingSummary, ...]: + assert workspace_id == "workspace-1" + return self.bindings + + def change_disabled_state( + self, + *, + workspace_id: str, + binding_id: str, + disabled: bool, + ) -> BindingMutationResult: + self.changes.append((workspace_id, binding_id, disabled)) + return self.mutation_result + + +def test_list_integrations_returns_secret_free_records() -> None: + binding = DataSourceBindingSummary( + id="binding-1", + provider="notion", + created_at=datetime(2026, 1, 1), + disabled=False, + source_info={"workspace_id": "notion-workspace"}, + ) + unsupported = DataSourceBindingSummary( + id="binding-2", + provider="unsupported", + created_at=datetime(2026, 1, 1), + disabled=False, + source_info={}, + ) + service = DataSourceBindingApplicationService(bindings=InMemoryBindingStore(bindings=(binding, unsupported))) + + assert service.list_integrations(_context()) == (binding,) + + +@pytest.mark.parametrize(("action", "disabled"), [(BindingAction.ENABLE, False), (BindingAction.DISABLE, True)]) +def test_change_state_passes_stable_context(action: BindingAction, disabled: bool) -> None: + store = InMemoryBindingStore() + service = DataSourceBindingApplicationService(bindings=store) + + service.change_state(_context(), "binding-1", action) + + assert store.changes == [("workspace-1", "binding-1", disabled)] + + +@pytest.mark.parametrize( + ("action", "result", "error", "message"), + [ + (BindingAction.ENABLE, "not_found", DataSourceBindingNotFoundError, "Data source binding not found"), + (BindingAction.ENABLE, "already_enabled", DataSourceBindingStateError, "Data source is already enabled"), + (BindingAction.DISABLE, "already_disabled", DataSourceBindingStateError, "Data source is already disabled"), + ], +) +def test_change_state_translates_repository_outcomes( + action: BindingAction, + result: BindingMutationResult, + error: type[Exception], + message: str, +) -> None: + service = DataSourceBindingApplicationService(bindings=InMemoryBindingStore(mutation_result=result)) + + with pytest.raises(error, match=message): + service.change_state(_context(), "binding-1", action) diff --git a/api/tests/unit_tests/services/data_source/test_credential_gateway.py b/api/tests/unit_tests/services/data_source/test_credential_gateway.py new file mode 100644 index 00000000000..0bb453f4de1 --- /dev/null +++ b/api/tests/unit_tests/services/data_source/test_credential_gateway.py @@ -0,0 +1,631 @@ +from collections.abc import Mapping +from dataclasses import dataclass, field, replace +from datetime import datetime, timedelta +from typing import cast + +import pytest +from sqlalchemy.orm import Session, sessionmaker + +import services.data_source.credential_adapters as credential_adapters_module +from core.datasource.entities.datasource_entities import ( + DatasourceProviderEntityWithPlugin, + DatasourceProviderIdentity, + DatasourceProviderType, +) +from core.entities.provider_entities import ProviderConfig, ProviderConfigType +from core.plugin.entities.oauth import OAuthSchema +from core.plugin.entities.plugin_daemon import PluginDatasourceProviderEntity, PluginOAuthCredentialsResponse +from core.plugin.impl.datasource import PluginDatasourceManager +from core.plugin.impl.oauth import OAuthHandler +from models.oauth import DatasourceOauthParamConfig, DatasourceOauthTenantParamConfig +from models.provider_ids import DatasourceProviderID +from repositories.data_source.credential_repository import SQLAlchemyDatasourceCredentialRepository +from services.data_source.credential_adapters import ( + OAuthDatasourceCredentialRefresher, + PluginDatasourceCredentialCodec, + PluginDatasourceOAuthClientResolver, +) +from services.data_source.credential_gateway import ( + ActorAwareDatasourceCredentialGateway, + DatasourceCredentialConcurrentUpdateError, + DatasourceCredentialError, + DatasourceCredentialNotFoundError, + DatasourceCredentialRefreshError, + RefreshedDatasourceCredential, + TrustedStoredDatasourceCredentialGateway, +) +from services.data_source.entities.credential import DatasourceCredentialRecord + + +def _record( + *, + auth_type: str = "oauth2", + expires_at: int = -1, + updated_at: datetime | None = None, +) -> DatasourceCredentialRecord: + return DatasourceCredentialRecord( + id="credential-1", + workspace_id="workspace-1", + owner_id="credential-owner-1", + name="Notion", + provider="notion_datasource", + plugin_id="langgenius/notion_datasource", + auth_type=auth_type, + encrypted_credentials={"integration_secret": "encrypted"}, + expires_at=expires_at, + updated_at=updated_at or datetime(2026, 1, 1), + ) + + +@dataclass +class ScriptedCredentialStore: + records: list[DatasourceCredentialRecord | None] + update_result: bool = True + get_calls: list[tuple[str, str, str, str, str]] = field(default_factory=list) + updates: list[tuple[DatasourceCredentialRecord, Mapping[str, object], int]] = field(default_factory=list) + events: list[str] = field(default_factory=list) + + def get_visible( + self, + *, + workspace_id: str, + actor_id: str, + credential_id: str, + provider: str, + plugin_id: str, + ) -> DatasourceCredentialRecord | None: + self.events.append("read") + self.get_calls.append((workspace_id, actor_id, credential_id, provider, plugin_id)) + return self.records.pop(0) + + def update_if_unchanged( + self, + *, + record: DatasourceCredentialRecord, + encrypted_credentials: Mapping[str, object], + expires_at: int, + ) -> bool: + self.events.append("update") + self.updates.append((record, encrypted_credentials, expires_at)) + return self.update_result + + +@dataclass +class ReversibleCredentialCodec: + def decrypt(self, record: DatasourceCredentialRecord) -> dict[str, object]: + return {"integration_secret": f"plain:{record.encrypted_credentials['integration_secret']}"} + + def encrypt(self, record: DatasourceCredentialRecord, credentials: Mapping[str, object]) -> dict[str, object]: + del record + return {"integration_secret": f"cipher:{credentials['integration_secret']}"} + + +@dataclass +class StoredCredentialCatalog: + record: DatasourceCredentialRecord | None + calls: list[dict[str, str | None]] = field(default_factory=list) + updates: list[tuple[DatasourceCredentialRecord, Mapping[str, object], int]] = field(default_factory=list) + + def get_for_stored_document( + self, + *, + workspace_id: str, + dataset_id: str, + document_id: str, + credential_id: str | None, + provider: str, + plugin_id: str, + ) -> DatasourceCredentialRecord | None: + self.calls.append( + { + "workspace_id": workspace_id, + "dataset_id": dataset_id, + "document_id": document_id, + "credential_id": credential_id, + "provider": provider, + "plugin_id": plugin_id, + } + ) + return self.record + + def update_if_unchanged( + self, + *, + record: DatasourceCredentialRecord, + encrypted_credentials: Mapping[str, object], + expires_at: int, + ) -> bool: + self.updates.append((record, encrypted_credentials, expires_at)) + return True + + +@dataclass +class RecordingCredentialRefresher: + calls: list[tuple[str, str, DatasourceCredentialRecord, Mapping[str, object]]] = field(default_factory=list) + events: list[str] = field(default_factory=list) + + def refresh( + self, + *, + workspace_id: str, + actor_id: str, + record: DatasourceCredentialRecord, + credentials: Mapping[str, object], + ) -> RefreshedDatasourceCredential: + self.events.append("refresh") + self.calls.append((workspace_id, actor_id, record, credentials)) + return RefreshedDatasourceCredential({"integration_secret": "refreshed"}, 500) + + +@dataclass +class StaticProviderManager: + provider: PluginDatasourceProviderEntity + calls: list[tuple[str, str]] = field(default_factory=list) + + def fetch_datasource_provider(self, *, tenant_id: str, provider_id: str) -> PluginDatasourceProviderEntity: + self.calls.append((tenant_id, provider_id)) + return self.provider + + +@dataclass +class RecordingOAuthHandler: + response: PluginOAuthCredentialsResponse | None = None + error: Exception | None = None + calls: list[dict[str, object]] = field(default_factory=list) + + def refresh_credentials( + self, + tenant_id: str, + user_id: str, + plugin_id: str, + provider: str, + redirect_uri: str, + system_credentials: Mapping[str, object], + credentials: Mapping[str, object], + ) -> PluginOAuthCredentialsResponse: + self.calls.append( + { + "tenant_id": tenant_id, + "user_id": user_id, + "plugin_id": plugin_id, + "provider": provider, + "redirect_uri": redirect_uri, + "system_credentials": system_credentials, + "credentials": credentials, + } + ) + if self.error is not None: + raise self.error + assert self.response is not None + return self.response + + +@dataclass +class RecordingCredentialEncrypter: + decrypted: Mapping[str, object] + calls: list[Mapping[str, object]] = field(default_factory=list) + + def decrypt(self, credentials: Mapping[str, object]) -> Mapping[str, object]: + self.calls.append(credentials) + return self.decrypted + + +class StaticOAuthClientResolver: + def resolve(self, *, workspace_id: str, provider_id: DatasourceProviderID) -> dict[str, object]: + del workspace_id, provider_id + return {} + + +def _oauth_refresher( + session_factory: sessionmaker[Session], + *, + manager: StaticProviderManager, + handler: RecordingOAuthHandler, +) -> OAuthDatasourceCredentialRefresher: + repository = SQLAlchemyDatasourceCredentialRepository(session_factory=session_factory) + oauth_clients = PluginDatasourceOAuthClientResolver( + configs=repository, + provider_manager=cast(PluginDatasourceManager, manager), + ) + return OAuthDatasourceCredentialRefresher( + oauth_clients=oauth_clients, + oauth_handler=cast(OAuthHandler, handler), + ) + + +def _provider(*, with_oauth_schema: bool = True) -> PluginDatasourceProviderEntity: + oauth_schema = ( + OAuthSchema( + client_schema=[ProviderConfig(type=ProviderConfigType.SECRET_INPUT, name="client_secret")], + credentials_schema=[ + ProviderConfig(type=ProviderConfigType.SECRET_INPUT, name="integration_secret"), + ProviderConfig(type=ProviderConfigType.TEXT_INPUT, name="workspace_name"), + ], + ) + if with_oauth_schema + else None + ) + return PluginDatasourceProviderEntity( + provider="notion_datasource", + plugin_unique_identifier="langgenius/notion_datasource:1.0.0", + plugin_id="langgenius/notion_datasource", + declaration=DatasourceProviderEntityWithPlugin( + identity=DatasourceProviderIdentity( + author="langgenius", + name="notion_datasource", + description={"en_US": "Notion", "zh_Hans": "Notion"}, + icon="notion.svg", + label={"en_US": "Notion", "zh_Hans": "Notion"}, + ), + credentials_schema=[ + ProviderConfig(type=ProviderConfigType.SECRET_INPUT, name="api_secret"), + ProviderConfig(type=ProviderConfigType.TEXT_INPUT, name="label"), + ], + oauth_schema=oauth_schema, + provider_type=DatasourceProviderType.ONLINE_DOCUMENT, + ), + ) + + +def _gateway( + store: ScriptedCredentialStore, refresher: RecordingCredentialRefresher | None = None +) -> tuple[ActorAwareDatasourceCredentialGateway, RecordingCredentialRefresher]: + refresh = refresher or RecordingCredentialRefresher() + return ( + ActorAwareDatasourceCredentialGateway( + credentials=store, + codec=ReversibleCredentialCodec(), + refresher=refresh, + now=lambda: 100, + ), + refresh, + ) + + +def _resolve(gateway: ActorAwareDatasourceCredentialGateway) -> dict[str, object]: + return gateway.resolve( + workspace_id="workspace-1", + actor_id="actor-1", + credential_id="credential-1", + provider="notion_datasource", + plugin_id="langgenius/notion_datasource", + ) + + +def test_resolve_fails_closed_when_credential_is_not_visible() -> None: + gateway, _ = _gateway(ScriptedCredentialStore(records=[None])) + + with pytest.raises(DatasourceCredentialNotFoundError): + _resolve(gateway) + + +def test_trusted_stored_source_resolver_uses_distinct_owner_chain_port() -> None: + catalog = StoredCredentialCatalog(_record()) + gateway = TrustedStoredDatasourceCredentialGateway( + credentials=catalog, + codec=ReversibleCredentialCodec(), + refresher=RecordingCredentialRefresher(), + now=lambda: 100, + ) + + result = gateway.resolve_for_document( + workspace_id="workspace-1", + dataset_id="dataset-1", + document_id="document-1", + credential_id="credential-1", + provider="notion_datasource", + plugin_id="langgenius/notion_datasource", + ) + + assert result == {"integration_secret": "plain:encrypted"} + assert catalog.calls == [ + { + "workspace_id": "workspace-1", + "dataset_id": "dataset-1", + "document_id": "document-1", + "credential_id": "credential-1", + "provider": "notion_datasource", + "plugin_id": "langgenius/notion_datasource", + } + ] + + +def test_trusted_stored_source_resolver_fails_closed_for_missing_credential() -> None: + gateway = TrustedStoredDatasourceCredentialGateway( + credentials=StoredCredentialCatalog(None), + codec=ReversibleCredentialCodec(), + refresher=RecordingCredentialRefresher(), + now=lambda: 100, + ) + + with pytest.raises(DatasourceCredentialNotFoundError): + gateway.resolve_for_document( + workspace_id="workspace-1", + dataset_id="dataset-1", + document_id="document-1", + credential_id="credential-1", + provider="notion_datasource", + plugin_id="langgenius/notion_datasource", + ) + + +def test_trusted_stored_source_refreshes_with_persisted_credential_owner() -> None: + record = _record(expires_at=120) + catalog = StoredCredentialCatalog(record) + refresher = RecordingCredentialRefresher() + gateway = TrustedStoredDatasourceCredentialGateway( + credentials=catalog, + codec=ReversibleCredentialCodec(), + refresher=refresher, + now=lambda: 100, + ) + + result = gateway.resolve_for_document( + workspace_id="workspace-1", + dataset_id="dataset-1", + document_id="document-1", + credential_id="credential-1", + provider="notion_datasource", + plugin_id="langgenius/notion_datasource", + ) + + assert result == {"integration_secret": "refreshed"} + assert refresher.calls == [("workspace-1", "credential-owner-1", record, {"integration_secret": "plain:encrypted"})] + assert catalog.updates == [(record, {"integration_secret": "cipher:refreshed"}, 500)] + + +def test_resolve_decrypts_without_refreshing_non_expiring_credential() -> None: + store = ScriptedCredentialStore(records=[_record()]) + gateway, refresher = _gateway(store) + + assert _resolve(gateway) == {"integration_secret": "plain:encrypted"} + assert refresher.calls == [] + assert store.updates == [] + + +def test_refresh_runs_between_read_and_conditional_update() -> None: + record = _record(expires_at=120) + events: list[str] = [] + store = ScriptedCredentialStore(records=[record], events=events) + gateway, refresher = _gateway(store, RecordingCredentialRefresher(events=events)) + + assert _resolve(gateway) == {"integration_secret": "refreshed"} + assert refresher.calls == [("workspace-1", "actor-1", record, {"integration_secret": "plain:encrypted"})] + assert store.updates == [(record, {"integration_secret": "cipher:refreshed"}, 500)] + assert events == ["read", "refresh", "update"] + + +def test_concurrent_refresh_uses_newer_non_expiring_snapshot() -> None: + stale = _record(expires_at=120) + latest = replace( + stale, + encrypted_credentials={"integration_secret": "latest"}, + expires_at=500, + updated_at=stale.updated_at + timedelta(seconds=1), + ) + store = ScriptedCredentialStore(records=[stale, latest], update_result=False) + gateway, _ = _gateway(store) + + assert _resolve(gateway) == {"integration_secret": "plain:latest"} + assert store.get_calls == [ + ("workspace-1", "actor-1", "credential-1", "notion_datasource", "langgenius/notion_datasource"), + ("workspace-1", "actor-1", "credential-1", "notion_datasource", "langgenius/notion_datasource"), + ] + + +def test_concurrent_refresh_fails_closed_if_latest_snapshot_is_still_expired() -> None: + stale = _record(expires_at=120) + latest = replace(stale, updated_at=stale.updated_at + timedelta(seconds=1)) + gateway, _ = _gateway(ScriptedCredentialStore(records=[stale, latest], update_result=False)) + + with pytest.raises(DatasourceCredentialConcurrentUpdateError): + _resolve(gateway) + + +@pytest.mark.parametrize( + ("auth_type", "secret_name"), + [("api-key", "api_secret"), ("oauth2", "integration_secret")], +) +def test_plugin_codec_encrypts_and_decrypts_only_declared_secret_fields( + monkeypatch: pytest.MonkeyPatch, + auth_type: str, + secret_name: str, +) -> None: + manager = StaticProviderManager(_provider()) + codec = PluginDatasourceCredentialCodec(provider_manager=cast(PluginDatasourceManager, manager)) + record = replace( + _record(auth_type=auth_type), + encrypted_credentials={secret_name: "cipher", "label": "plain"}, + ) + decrypt_calls: list[tuple[str, str]] = [] + encrypt_calls: list[tuple[str, str]] = [] + + def decrypt_token(workspace_id: str, value: str) -> str: + decrypt_calls.append((workspace_id, value)) + return f"decrypted:{value}" + + def encrypt_token(workspace_id: str, value: str) -> str: + encrypt_calls.append((workspace_id, value)) + return f"encrypted:{value}" + + monkeypatch.setattr(credential_adapters_module.encrypter, "decrypt_token", decrypt_token) + monkeypatch.setattr(credential_adapters_module.encrypter, "encrypt_token", encrypt_token) + + assert codec.decrypt(record) == {secret_name: "decrypted:cipher", "label": "plain"} + assert codec.encrypt(record, {secret_name: "secret", "label": "plain"}) == { + secret_name: "encrypted:secret", + "label": "plain", + } + assert decrypt_calls == [("workspace-1", "cipher")] + assert encrypt_calls == [("workspace-1", "secret")] + + non_string_record = replace(record, encrypted_credentials={secret_name: 42}) + assert codec.decrypt(non_string_record) == {secret_name: 42} + assert codec.encrypt(record, {secret_name: None}) == {secret_name: None} + assert decrypt_calls == [("workspace-1", "cipher")] + assert encrypt_calls == [("workspace-1", "secret")] + + +def test_plugin_codec_rejects_unsupported_credential_type() -> None: + codec = PluginDatasourceCredentialCodec( + provider_manager=cast(PluginDatasourceManager, StaticProviderManager(_provider())) + ) + + with pytest.raises(DatasourceCredentialError, match="Unsupported credential type"): + codec.decrypt(_record(auth_type="unauthorized")) + + +def test_oauth_refresher_prefers_enabled_tenant_client_and_returns_refreshed_snapshot( + monkeypatch: pytest.MonkeyPatch, + sqlite_session_factory: sessionmaker[Session], +) -> None: + with sqlite_session_factory.begin() as session: + session.add_all( + [ + DatasourceOauthTenantParamConfig( + tenant_id="workspace-1", + provider="notion_datasource", + plugin_id="langgenius/notion_datasource", + client_params={"client_secret": "tenant-cipher"}, + enabled=True, + ), + DatasourceOauthParamConfig( + provider="notion_datasource", + plugin_id="langgenius/notion_datasource", + system_credentials={"client_secret": "system-secret"}, + ), + ] + ) + manager = StaticProviderManager(_provider()) + handler = RecordingOAuthHandler( + response=PluginOAuthCredentialsResponse( + credentials={"integration_secret": "refreshed-secret"}, + expires_at=500, + ) + ) + encrypter = RecordingCredentialEncrypter({"client_secret": "tenant-secret"}) + + def create_encrypter(**kwargs: object) -> tuple[RecordingCredentialEncrypter, None]: + assert kwargs["tenant_id"] == "workspace-1" + return encrypter, None + + monkeypatch.setattr(credential_adapters_module, "create_provider_encrypter", create_encrypter) + refresher = _oauth_refresher(sqlite_session_factory, manager=manager, handler=handler) + + result = refresher.refresh( + workspace_id="workspace-1", + actor_id="actor-1", + record=_record(), + credentials={"integration_secret": "old-secret"}, + ) + + assert result == RefreshedDatasourceCredential({"integration_secret": "refreshed-secret"}, 500) + assert encrypter.calls == [{"client_secret": "tenant-cipher"}] + assert handler.calls[0]["system_credentials"] == {"client_secret": "tenant-secret"} + assert handler.calls[0]["credentials"] == {"integration_secret": "old-secret"} + assert handler.calls[0]["user_id"] == "actor-1" + + +def test_oauth_refresher_rejects_tenant_client_when_provider_has_no_oauth_schema( + sqlite_session_factory: sessionmaker[Session], +) -> None: + with sqlite_session_factory.begin() as session: + session.add( + DatasourceOauthTenantParamConfig( + tenant_id="workspace-1", + provider="notion_datasource", + plugin_id="langgenius/notion_datasource", + client_params={"client_secret": "tenant-cipher"}, + enabled=True, + ) + ) + provider = _provider(with_oauth_schema=False) + handler = RecordingOAuthHandler() + refresher = _oauth_refresher( + sqlite_session_factory, + manager=StaticProviderManager(provider), + handler=handler, + ) + + with pytest.raises(DatasourceCredentialError, match="oauth schema not found"): + refresher.refresh( + workspace_id="workspace-1", + actor_id="actor-1", + record=_record(), + credentials={"integration_secret": "old"}, + ) + + assert handler.calls == [] + + +@pytest.mark.parametrize("verified", [True, False]) +def test_oauth_refresher_uses_system_client_only_for_verified_plugin( + monkeypatch: pytest.MonkeyPatch, + sqlite_session_factory: sessionmaker[Session], + verified: bool, +) -> None: + with sqlite_session_factory.begin() as session: + session.add_all( + [ + DatasourceOauthTenantParamConfig( + tenant_id="workspace-1", + provider="notion_datasource", + plugin_id="langgenius/notion_datasource", + client_params={"client_secret": "disabled-tenant-secret"}, + enabled=False, + ), + DatasourceOauthParamConfig( + provider="notion_datasource", + plugin_id="langgenius/notion_datasource", + system_credentials={"client_secret": "system-secret"}, + ), + ] + ) + manager = StaticProviderManager(_provider()) + handler = RecordingOAuthHandler( + response=PluginOAuthCredentialsResponse(credentials={"integration_secret": "refreshed"}, expires_at=500) + ) + verification_calls: list[tuple[str, str]] = [] + + def is_plugin_verified(workspace_id: str, plugin_unique_identifier: str) -> bool: + verification_calls.append((workspace_id, plugin_unique_identifier)) + return verified + + monkeypatch.setattr(credential_adapters_module.PluginService, "is_plugin_verified", is_plugin_verified) + refresher = _oauth_refresher(sqlite_session_factory, manager=manager, handler=handler) + + if not verified: + with pytest.raises(DatasourceCredentialError, match="OAuth client is not configured"): + refresher.refresh( + workspace_id="workspace-1", + actor_id="actor-1", + record=_record(), + credentials={"integration_secret": "old"}, + ) + assert handler.calls == [] + else: + refresher.refresh( + workspace_id="workspace-1", + actor_id="actor-1", + record=_record(), + credentials={"integration_secret": "old"}, + ) + assert handler.calls[0]["system_credentials"] == {"client_secret": "system-secret"} + + assert verification_calls == [("workspace-1", "langgenius/notion_datasource:1.0.0")] + + +def test_oauth_refresher_translates_plugin_refresh_failure() -> None: + handler = RecordingOAuthHandler(error=RuntimeError("plugin unavailable")) + refresher = OAuthDatasourceCredentialRefresher( + oauth_clients=StaticOAuthClientResolver(), + oauth_handler=cast(OAuthHandler, handler), + ) + + with pytest.raises(DatasourceCredentialRefreshError, match="credential-1"): + refresher.refresh( + workspace_id="workspace-1", + actor_id="actor-1", + record=_record(), + credentials={"integration_secret": "old"}, + ) diff --git a/api/tests/unit_tests/services/data_source/test_notion_import_adapters.py b/api/tests/unit_tests/services/data_source/test_notion_import_adapters.py new file mode 100644 index 00000000000..da149afe07f --- /dev/null +++ b/api/tests/unit_tests/services/data_source/test_notion_import_adapters.py @@ -0,0 +1,202 @@ +from collections.abc import Sequence +from unittest.mock import Mock, create_autospec + +import pytest + +from core.datasource.__base.datasource_runtime import DatasourceRuntime +from core.datasource.entities.datasource_entities import ( + DatasourceProviderType, + OnlineDocumentInfo, + OnlineDocumentPage, + OnlineDocumentPagesMessage, +) +from core.datasource.online_document.online_document_plugin import OnlineDocumentDatasourcePlugin +from core.rag.extractor.notion_extractor import NotionExtractor +from core.rag.models.document import Document +from services.data_source.credential_gateway import ( + ActorDatasourceCredentialResolver, + DatasourceCredentialError, + DatasourceCredentialNotFoundError, +) +from services.data_source.entities.notion_import import NotionPageType +from services.data_source.notion_import_adapters import PluginNotionSourceGateway +from services.data_source.notion_import_application_service import NotionImportCredentialUnavailableError + + +def _credentials(credentials: dict[str, object], *, error: DatasourceCredentialError | None = None) -> Mock: + resolver = create_autospec(ActorDatasourceCredentialResolver, instance=True, spec_set=True) + resolver.resolve.return_value = credentials + resolver.resolve.side_effect = error + return resolver + + +def _runtime(messages: Sequence[OnlineDocumentPagesMessage] = ()) -> Mock: + runtime = create_autospec(OnlineDocumentDatasourcePlugin, instance=True) + runtime.runtime = DatasourceRuntime(tenant_id="workspace-1") + runtime.get_online_document_pages.return_value = iter(messages) + runtime.datasource_provider_type.return_value = DatasourceProviderType.ONLINE_DOCUMENT + return runtime + + +def _page(page_id: str, *, with_icon: bool = True, page_type: str = "page") -> OnlineDocumentPage: + return OnlineDocumentPage( + page_id=page_id, + page_name=f"Page {page_id}", + page_icon={"type": "emoji", "emoji": "📄"} if with_icon else None, + type=page_type, + last_edited_time="2026-01-01T00:00:00Z", + parent_id=None, + ) + + +def test_list_authorized_pages_groups_paginated_results_by_workspace() -> None: + runtime = _runtime( + messages=( + OnlineDocumentPagesMessage( + result=[ + OnlineDocumentInfo( + workspace_id="w1", + workspace_name="One", + workspace_icon=None, + total=1, + pages=[_page("p1")], + ) + ] + ), + OnlineDocumentPagesMessage( + result=[ + OnlineDocumentInfo( + workspace_id="w2", + workspace_name="Two", + workspace_icon=None, + total=1, + pages=[_page("p2", with_icon=False)], + ), + OnlineDocumentInfo( + workspace_id="w1", + workspace_name="One", + workspace_icon=None, + total=1, + pages=[_page("p3")], + ), + ] + ), + ) + ) + resolver = _credentials({"integration_secret": "secret"}) + loader = Mock(return_value=runtime) + gateway = PluginNotionSourceGateway(credentials=resolver, runtime_loader=loader) + + workspaces = gateway.list_authorized_pages( + workspace_id="workspace-1", actor_id="actor-1", credential_id="credential-1" + ) + + assert [workspace.workspace_id for workspace in workspaces] == ["w1", "w2"] + assert [page.page_id for page in workspaces[0].pages] == ["p1", "p3"] + assert [page.page_id for page in workspaces[1].pages] == ["p2"] + assert workspaces[0].pages[0].page_icon is not None + assert workspaces[0].pages[0].page_icon.emoji == "📄" + assert workspaces[1].pages[0].page_icon is None + assert runtime.runtime.credentials == {"integration_secret": "secret"} + runtime.get_online_document_pages.assert_called_once_with( + user_id="actor-1", + datasource_parameters={}, + provider_type=DatasourceProviderType.ONLINE_DOCUMENT, + ) + resolver.resolve.assert_called_once_with( + workspace_id="workspace-1", + actor_id="actor-1", + credential_id="credential-1", + provider="notion_datasource", + plugin_id="langgenius/notion_datasource", + ) + loader.assert_called_once_with( + provider_id="langgenius/notion_datasource/notion_datasource", + datasource_name="notion_datasource", + tenant_id="workspace-1", + datasource_type=DatasourceProviderType.ONLINE_DOCUMENT, + ) + + +def test_list_authorized_pages_skips_unknown_page_types() -> None: + runtime = _runtime( + messages=( + OnlineDocumentPagesMessage( + result=[ + OnlineDocumentInfo( + workspace_id="w1", + workspace_name="One", + workspace_icon=None, + total=2, + pages=[_page("known"), _page("unknown", page_type="collection")], + ) + ] + ), + ) + ) + gateway = PluginNotionSourceGateway( + credentials=_credentials({"integration_secret": "secret"}), + runtime_loader=Mock(return_value=runtime), + ) + + workspaces = gateway.list_authorized_pages( + workspace_id="workspace-1", actor_id="actor-1", credential_id="credential-1" + ) + + assert [page.page_id for page in workspaces[0].pages] == ["known"] + + +def test_preview_requires_secret_and_passes_it_only_to_extractor() -> None: + factory = create_autospec(NotionExtractor, spec_set=True) + factory.return_value.extract.return_value = [Document(page_content="one"), Document(page_content="two")] + gateway = PluginNotionSourceGateway( + credentials=_credentials({"integration_secret": "secret"}), + runtime_loader=Mock(), + extractor_factory=factory, + ) + + content = gateway.preview_page( + workspace_id="workspace-1", + actor_id="actor-1", + credential_id="credential-1", + page_id="page-1", + page_type=NotionPageType.PAGE, + ) + + assert content == "one\ntwo" + extractor_args = factory.call_args.kwargs + assert extractor_args["notion_obj_id"] == "page-1" + assert extractor_args["notion_page_type"] == "page" + assert extractor_args["notion_access_token"] == "secret" + assert extractor_args["notion_token_loader"]() == "secret" + assert extractor_args["tenant_id"] == "workspace-1" + + +def test_preview_fails_closed_when_secret_is_missing() -> None: + gateway = PluginNotionSourceGateway( + credentials=_credentials({}), + runtime_loader=Mock(), + ) + + with pytest.raises(NotionImportCredentialUnavailableError): + gateway.preview_page( + workspace_id="workspace-1", + actor_id="actor-1", + credential_id="credential-1", + page_id="page-1", + page_type=NotionPageType.PAGE, + ) + + +def test_gateway_translates_credential_infrastructure_failure() -> None: + gateway = PluginNotionSourceGateway( + credentials=_credentials({}, error=DatasourceCredentialNotFoundError()), + runtime_loader=Mock(), + ) + + with pytest.raises(NotionImportCredentialUnavailableError): + gateway.list_authorized_pages( + workspace_id="workspace-1", + actor_id="actor-1", + credential_id="credential-1", + ) diff --git a/api/tests/unit_tests/services/data_source/test_notion_import_application_service.py b/api/tests/unit_tests/services/data_source/test_notion_import_application_service.py new file mode 100644 index 00000000000..e54fb40242a --- /dev/null +++ b/api/tests/unit_tests/services/data_source/test_notion_import_application_service.py @@ -0,0 +1,96 @@ +from unittest.mock import Mock, create_autospec + +import pytest + +from machinery.context import RequestContext +from services.data_source.entities.notion_import import AuthorizedNotionPage, NotionPageType, NotionWorkspace +from services.data_source.notion_import_application_service import ( + DatasetIsNotNotionSourceError, + NotionDatasetReader, + NotionDocumentBindingReader, + NotionImportApplicationService, + NotionSourceGateway, +) +from services.knowledge.dataset_access import AccessibleDataset, DatasetAccess +from services.knowledge.resource_scope import DatasetRef + +CONTEXT = RequestContext("request-1", None, "account-1", "workspace-1") +DATASET = DatasetRef("workspace-1", "dataset-1") +type ServiceFixture = tuple[NotionImportApplicationService, Mock, Mock, Mock, Mock] + + +@pytest.fixture +def fixture() -> ServiceFixture: + access = create_autospec(DatasetAccess, instance=True, spec_set=True) + access.require_accessible.return_value = AccessibleDataset(id="dataset-1", workspace_id="workspace-1") + datasets = create_autospec(NotionDatasetReader, instance=True, spec_set=True) + datasets.is_notion_dataset.return_value = True + documents = create_autospec(NotionDocumentBindingReader, instance=True, spec_set=True) + documents.list_bound_notion_page_ids.return_value = frozenset() + source = create_autospec(NotionSourceGateway, instance=True, spec_set=True) + source.list_authorized_pages.return_value = () + source.preview_page.return_value = "preview" + service = NotionImportApplicationService( + dataset_access=access, datasets=datasets, documents=documents, source=source + ) + return service, access, datasets, documents, source + + +def test_list_pages_preserves_workspaces_and_marks_bound_pages(fixture: ServiceFixture) -> None: + service, access, datasets, documents, source = fixture + page_1 = AuthorizedNotionPage("page-1", "One", None, None, NotionPageType.PAGE) + page_2 = AuthorizedNotionPage("page-2", "Two", None, "page-1", NotionPageType.PAGE) + source.list_authorized_pages.return_value = ( + NotionWorkspace("notion-workspace-1", "First", None, (page_1,)), + NotionWorkspace("notion-workspace-2", "Second", None, (page_2,)), + ) + documents.list_bound_notion_page_ids.return_value = frozenset({"page-2"}) + + result = service.list_pages(CONTEXT, credential_id="credential-1", dataset_id="dataset-1") + + assert [workspace.workspace_id for workspace in result.workspaces] == ["notion-workspace-1", "notion-workspace-2"] + assert result.workspaces[0].pages[0].is_bound is False + assert result.workspaces[1].pages[0].is_bound is True + access.require_accessible.assert_called_once_with(CONTEXT, "dataset-1") + datasets.is_notion_dataset.assert_called_once_with(DATASET) + documents.list_bound_notion_page_ids.assert_called_once_with(DATASET) + source.list_authorized_pages.assert_called_once_with( + workspace_id="workspace-1", actor_id="account-1", credential_id="credential-1" + ) + + +def test_list_pages_without_dataset_does_not_read_document_bindings(fixture: ServiceFixture) -> None: + service, access, _, documents, _ = fixture + + service.list_pages(CONTEXT, credential_id="credential-1") + + access.require_accessible.assert_not_called() + documents.list_bound_notion_page_ids.assert_not_called() + + +def test_list_pages_rejects_non_notion_dataset_before_external_call(fixture: ServiceFixture) -> None: + service, _, datasets, _, source = fixture + datasets.is_notion_dataset.return_value = False + + with pytest.raises(DatasetIsNotNotionSourceError): + service.list_pages(CONTEXT, credential_id="credential-1", dataset_id="dataset-1") + + datasets.is_notion_dataset.assert_called_once_with(DATASET) + source.list_authorized_pages.assert_not_called() + + +def test_preview_passes_explicit_actor_and_workspace(fixture: ServiceFixture) -> None: + service, _, _, _, source = fixture + + content = service.preview_page( + CONTEXT, credential_id="credential-1", page_id="page-1", page_type=NotionPageType.PAGE + ) + + assert content == "preview" + source.preview_page.assert_called_once_with( + workspace_id="workspace-1", + actor_id="account-1", + credential_id="credential-1", + page_id="page-1", + page_type=NotionPageType.PAGE, + ) diff --git a/api/tests/unit_tests/services/data_source/test_notion_oauth_gateway.py b/api/tests/unit_tests/services/data_source/test_notion_oauth_gateway.py new file mode 100644 index 00000000000..fd96f1d9bb7 --- /dev/null +++ b/api/tests/unit_tests/services/data_source/test_notion_oauth_gateway.py @@ -0,0 +1,401 @@ +import json +from dataclasses import dataclass +from typing import Never, cast + +import httpx +import pytest + +from core.helper.ssrf_proxy import SSRFProxy +from core.tools.errors import ToolSSRFError +from services.data_source.notion_oauth_gateway import NotionDataSourceGateway +from services.data_source.oauth_service import DataSourceOAuthError + + +@dataclass(frozen=True) +class RejectingSSRFClient: + error: Exception + + def post(self, *_args: object, **_kwargs: object) -> Never: + raise self.error + + +def _gateway(handler: httpx.MockTransport) -> tuple[NotionDataSourceGateway, httpx.Client]: + client = httpx.Client(transport=handler) + return ( + NotionDataSourceGateway( + client_id="client-id", + client_secret="client-secret", + redirect_uri="https://api.example/oauth/data-source/callback/notion", + http_client=client, + ), + client, + ) + + +def test_get_authorization_url_encodes_configured_values() -> None: + gateway, client = _gateway(httpx.MockTransport(lambda _request: httpx.Response(500))) + try: + result = gateway.get_authorization_url() + finally: + client.close() + + assert result == ( + "https://api.notion.com/v1/oauth/authorize?client_id=client-id&response_type=code&" + "redirect_uri=https%3A%2F%2Fapi.example%2Foauth%2Fdata-source%2Fcallback%2Fnotion&owner=user" + ) + + +def test_authorize_exchanges_code_and_builds_source_info_without_persistence() -> None: + search_types: list[str] = [] + + def handler(request: httpx.Request) -> httpx.Response: + if request.url.path == "/v1/oauth/token": + assert request.headers["authorization"].startswith("Basic ") + return httpx.Response( + 200, + json={ + "access_token": "notion-token", + "workspace_name": "Notion Workspace", + "workspace_icon": "https://example/icon.png", + "workspace_id": "notion-workspace", + }, + ) + if request.url.path == "/v1/search": + payload = json.loads(request.content) + search_types.append(payload["filter"]["value"]) + return httpx.Response(200, json={"results": [], "has_more": False}) + raise AssertionError(f"Unexpected Notion request: {request.url}") + + gateway, client = _gateway(httpx.MockTransport(handler)) + try: + authorization = gateway.authorize("authorization-code") + finally: + client.close() + + assert authorization.access_token == "notion-token" + assert authorization.source_info == { + "workspace_name": "Notion Workspace", + "workspace_icon": "https://example/icon.png", + "workspace_id": "notion-workspace", + "pages": [], + "total": 0, + } + assert search_types == ["page", "database"] + + +@pytest.mark.parametrize("token_payload", [{}, {"access_token": ""}, {"access_token": 42}]) +def test_authorize_rejects_missing_or_invalid_access_token(token_payload: dict[str, object]) -> None: + gateway, client = _gateway(httpx.MockTransport(lambda _request: httpx.Response(200, json=token_payload))) + try: + with pytest.raises(DataSourceOAuthError, match="did not include an access token"): + gateway.authorize("authorization-code") + finally: + client.close() + + +def test_authorize_internal_reads_workspace_name_and_builds_source_info() -> None: + def handler(request: httpx.Request) -> httpx.Response: + if request.url.path == "/v1/users/me": + assert request.headers["authorization"] == "Bearer notion-token" + return httpx.Response( + 200, + json={"object": "user", "type": "bot", "bot": {"workspace_name": "Internal Workspace"}}, + ) + if request.url.path == "/v1/search": + return httpx.Response(200, json={"results": [], "has_more": False}) + raise AssertionError(f"Unexpected Notion request: {request.url}") + + gateway, client = _gateway(httpx.MockTransport(handler)) + try: + authorization = gateway.authorize_internal("notion-token", "notion-workspace") + finally: + client.close() + + assert authorization.access_token == "notion-token" + assert authorization.source_info == { + "workspace_name": "Internal Workspace", + "workspace_icon": None, + "workspace_id": "notion-workspace", + "pages": [], + "total": 0, + } + + +@pytest.mark.parametrize( + ("operation", "expected_message"), + [ + ("authorize", "Notion OAuth response was invalid"), + ("authorize_internal", "Notion integration response was invalid"), + ("refresh", "Notion integration response was invalid"), + ], +) +def test_public_operations_translate_malformed_page_payload( + operation: str, + expected_message: str, +) -> None: + def handler(request: httpx.Request) -> httpx.Response: + if request.url.path == "/v1/oauth/token": + return httpx.Response(200, json={"access_token": "notion-token"}) + if request.url.path == "/v1/users/me": + return httpx.Response(200, json={"object": "user", "type": "bot", "bot": {}}) + if request.url.path == "/v1/search": + payload = json.loads(request.content) + results: list[dict[str, object]] = [{}] if payload["filter"]["value"] == "page" else [] + return httpx.Response(200, json={"results": results, "has_more": False}) + raise AssertionError(f"Unexpected Notion request: {request.url}") + + gateway, client = _gateway(httpx.MockTransport(handler)) + + def invoke() -> None: + if operation == "authorize": + gateway.authorize("authorization-code") + elif operation == "authorize_internal": + gateway.authorize_internal("notion-token", "notion-workspace") + else: + gateway.refresh("notion-token", {}) + + try: + with pytest.raises(DataSourceOAuthError, match=expected_message): + invoke() + finally: + client.close() + + +def test_authorized_pages_map_pagination_icons_and_parent_chains() -> None: + page_requests: list[dict[str, object]] = [] + + def handler(request: httpx.Request) -> httpx.Response: + if request.url.path == "/v1/search": + payload = json.loads(request.content) + object_type = payload["filter"]["value"] + if object_type == "page": + page_requests.append(payload) + if "start_cursor" not in payload: + return httpx.Response( + 200, + json={ + "results": [ + { + "id": "page-1", + "properties": {"Name": {"title": [{"plain_text": "First page"}]}}, + "icon": {"type": "external", "external": {"url": "/page-icon"}}, + "parent": {"type": "block_id", "block_id": "child-block"}, + } + ], + "has_more": True, + "next_cursor": "next-page", + }, + ) + return httpx.Response( + 200, + json={ + "results": [ + { + "id": "page-2", + "properties": {}, + "icon": {"type": "emoji", "emoji": "📄"}, + "parent": {"type": "workspace", "workspace": True}, + } + ], + "has_more": False, + }, + ) + return httpx.Response( + 200, + json={ + "results": [ + { + "id": "database-1", + "title": [{"plain_text": "Database"}], + "icon": {"type": "emoji", "emoji": "🗃️"}, + "parent": {"type": "page_id", "page_id": "page-1"}, + }, + { + "id": "database-2", + "title": [], + "icon": None, + "parent": {"type": "workspace", "workspace": True}, + }, + ], + "has_more": False, + }, + ) + if request.url.path == "/v1/blocks/child-block": + return httpx.Response(200, json={"parent": {"type": "block_id", "block_id": "root-block"}}) + if request.url.path == "/v1/blocks/root-block": + return httpx.Response(200, json={"parent": {"type": "page_id", "page_id": "parent-page"}}) + raise AssertionError(f"Unexpected Notion request: {request.url}") + + gateway, client = _gateway(httpx.MockTransport(handler)) + try: + pages = gateway.get_authorized_pages("notion-token") + finally: + client.close() + + assert pages == [ + { + "page_id": "page-1", + "page_name": "First page", + "page_icon": {"type": "url", "url": "https://www.notion.so/page-icon"}, + "parent_id": "parent-page", + "type": "page", + }, + { + "page_id": "page-2", + "page_name": "Untitled", + "page_icon": {"type": "emoji", "emoji": "📄"}, + "parent_id": "root", + "type": "page", + }, + { + "page_id": "database-1", + "page_name": "Database", + "page_icon": {"type": "emoji", "emoji": "🗃️"}, + "parent_id": "page-1", + "type": "database", + }, + { + "page_id": "database-2", + "page_name": "Untitled", + "page_icon": None, + "parent_id": "root", + "type": "database", + }, + ] + assert page_requests == [ + {"filter": {"value": "page", "property": "object"}}, + {"filter": {"value": "page", "property": "object"}, "start_cursor": "next-page"}, + ] + + +def test_refresh_preserves_workspace_metadata() -> None: + gateway, client = _gateway( + httpx.MockTransport( + lambda request: ( + httpx.Response(200, json={"results": [], "has_more": False}) + if request.url.path == "/v1/search" + else httpx.Response(404) + ) + ) + ) + try: + source_info = gateway.refresh( + "notion-token", + { + "workspace_name": "Workspace", + "workspace_icon": None, + "workspace_id": "notion-workspace", + "pages": [ + { + "page_id": "old", + "page_name": "Old page", + "page_icon": None, + "parent_id": "root", + "type": "page", + } + ], + "total": 1, + }, + ) + finally: + client.close() + + assert source_info == { + "workspace_name": "Workspace", + "workspace_icon": None, + "workspace_id": "notion-workspace", + "pages": [], + "total": 0, + } + + +def test_search_rejects_missing_cursor_when_more_results_are_declared() -> None: + requests = 0 + + def handler(_request: httpx.Request) -> httpx.Response: + nonlocal requests + requests += 1 + return httpx.Response(200, json={"results": [], "has_more": True, "next_cursor": None}) + + gateway, client = _gateway(httpx.MockTransport(handler)) + try: + with pytest.raises(DataSourceOAuthError, match="valid next_cursor"): + gateway.notion_page_search("notion-token") + finally: + client.close() + + assert requests == 1 + + +def test_search_rejects_non_list_results() -> None: + gateway, client = _gateway( + httpx.MockTransport(lambda _request: httpx.Response(200, json={"results": {}, "has_more": False})) + ) + try: + with pytest.raises(DataSourceOAuthError, match="search response was invalid"): + gateway.notion_page_search("notion-token") + finally: + client.close() + + +@pytest.mark.parametrize( + ("payload", "expected_message"), + [ + ({"parent": []}, "block response was invalid"), + ({"parent": {"type": "page_id"}}, "block parent was invalid"), + ({"parent": {"type": "page_id", "page_id": 42}}, "parent identifier was invalid"), + ], +) +def test_block_parent_rejects_invalid_provider_payload( + payload: dict[str, object], + expected_message: str, +) -> None: + gateway, client = _gateway(httpx.MockTransport(lambda _request: httpx.Response(200, json=payload))) + try: + with pytest.raises(DataSourceOAuthError, match=expected_message): + gateway.notion_block_parent_page_id("notion-token", "block-1") + finally: + client.close() + + +def test_provider_http_failure_is_translated_to_application_boundary_error() -> None: + gateway, client = _gateway(httpx.MockTransport(lambda _request: httpx.Response(503, text="unavailable"))) + try: + with pytest.raises(DataSourceOAuthError, match="unsuccessful response"): + gateway.authorize("authorization-code") + finally: + client.close() + + +def test_provider_non_object_json_is_rejected() -> None: + gateway, client = _gateway(httpx.MockTransport(lambda _request: httpx.Response(200, json=[]))) + try: + with pytest.raises(DataSourceOAuthError, match="response was invalid"): + gateway.authorize("authorization-code") + finally: + client.close() + + +def test_provider_get_failure_is_translated_to_application_boundary_error() -> None: + def handler(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("offline", request=request) + + gateway, client = _gateway(httpx.MockTransport(handler)) + try: + with pytest.raises(DataSourceOAuthError, match="Notion request failed"): + gateway.authorize_internal("notion-token", "notion-workspace") + finally: + client.close() + + +def test_ssrf_rejection_is_translated_to_application_boundary_error() -> None: + client = RejectingSSRFClient(ToolSSRFError("blocked")) + gateway = NotionDataSourceGateway( + client_id="client-id", + client_secret="client-secret", + redirect_uri="https://api.example/oauth/data-source/callback/notion", + http_client=cast(SSRFProxy, client), + ) + + with pytest.raises(DataSourceOAuthError, match="Notion request failed"): + gateway.authorize("authorization-code") diff --git a/api/tests/unit_tests/services/data_source/test_oauth_service.py b/api/tests/unit_tests/services/data_source/test_oauth_service.py new file mode 100644 index 00000000000..66ebb3dfc67 --- /dev/null +++ b/api/tests/unit_tests/services/data_source/test_oauth_service.py @@ -0,0 +1,236 @@ +from collections.abc import Mapping +from dataclasses import dataclass, field + +import pytest + +from machinery.context import RequestContext +from services.data_source.entities.oauth import ( + DataSourceOAuthAuthorization, + DataSourceOAuthBindingRecord, + DataSourceOAuthCallback, +) +from services.data_source.oauth_service import ( + DataSourceOAuthConfigurationError, + DataSourceOAuthError, + DataSourceOAuthService, + InvalidDataSourceOAuthCodeError, +) + + +def _authorization() -> DataSourceOAuthAuthorization: + return DataSourceOAuthAuthorization(access_token="token", source_info={"pages": []}) + + +@dataclass +class RecordingProviderGateway: + authorization_url: str = "https://notion.example/authorize" + authorization: DataSourceOAuthAuthorization = field(default_factory=_authorization) + refreshed_source_info: Mapping[str, object] = field(default_factory=lambda: {"pages": [{"page_id": "page-1"}]}) + events: list[tuple[str, object]] = field(default_factory=list) + + def get_authorization_url(self) -> str: + self.events.append(("get_authorization_url", ())) + return self.authorization_url + + def authorize(self, code: str) -> DataSourceOAuthAuthorization: + self.events.append(("authorize", code)) + return self.authorization + + def authorize_internal(self, access_token: str, workspace_id: str) -> DataSourceOAuthAuthorization: + self.events.append(("authorize_internal", (access_token, workspace_id))) + return self.authorization + + def refresh(self, access_token: str, source_info: Mapping[str, object]) -> Mapping[str, object]: + self.events.append(("refresh", (access_token, source_info))) + return self.refreshed_source_info + + +@dataclass +class InMemoryBindingRepository: + binding: DataSourceOAuthBindingRecord | None = None + update_result: bool = True + events: list[tuple[str, object]] = field(default_factory=list) + + def upsert_authorization( + self, + *, + workspace_id: str, + provider: str, + authorization: DataSourceOAuthAuthorization, + ) -> None: + self.events.append(("persist", (workspace_id, provider, authorization))) + + def get_enabled( + self, + *, + workspace_id: str, + provider: str, + binding_id: str, + ) -> DataSourceOAuthBindingRecord | None: + self.events.append(("read", (workspace_id, provider, binding_id))) + return self.binding + + def update_source_info( + self, + *, + workspace_id: str, + provider: str, + binding_id: str, + source_info: Mapping[str, object], + ) -> bool: + self.events.append(("update", (workspace_id, provider, binding_id, source_info))) + return self.update_result + + +@pytest.fixture +def request_context() -> RequestContext: + return RequestContext( + request_id="request-1", + trace_id="trace-1", + account_id="account-1", + active_workspace_id="workspace-1", + ) + + +def test_start_authorization_delegates_to_external_provider(request_context: RequestContext) -> None: + events: list[tuple[str, object]] = [] + provider = RecordingProviderGateway(events=events) + bindings = InMemoryBindingRepository(events=events) + service = DataSourceOAuthService(provider_name="notion", provider_gateway=provider, bindings=bindings) + + result = service.start_authorization(request_context) + + assert result == "https://notion.example/authorize" + assert events == [("get_authorization_url", ())] + + +def test_start_authorization_persists_internal_authorization(request_context: RequestContext) -> None: + events: list[tuple[str, object]] = [] + authorization = DataSourceOAuthAuthorization(access_token="secret", source_info={"pages": []}) + provider = RecordingProviderGateway(authorization=authorization, events=events) + bindings = InMemoryBindingRepository(events=events) + service = DataSourceOAuthService( + provider_name="notion", + provider_gateway=provider, + bindings=bindings, + is_internal_provider=True, + internal_access_token="secret", + ) + + result = service.start_authorization(request_context) + + assert result == "internal" + assert events == [ + ("authorize_internal", ("secret", "workspace-1")), + ("persist", ("workspace-1", "notion", authorization)), + ] + + +def test_start_authorization_rejects_missing_internal_secret(request_context: RequestContext) -> None: + service = DataSourceOAuthService( + provider_name="notion", + provider_gateway=RecordingProviderGateway(), + bindings=InMemoryBindingRepository(), + is_internal_provider=True, + ) + + with pytest.raises(DataSourceOAuthConfigurationError, match="Internal secret is not set"): + service.start_authorization(request_context) + + +@pytest.mark.parametrize( + ("code", "error", "expected_error"), + [ + ("auth/code", None, None), + (None, "access denied", "access denied"), + (None, None, "Access denied"), + ], +) +def test_complete_callback_returns_framework_neutral_result( + code: str | None, + error: str | None, + expected_error: str | None, +) -> None: + service = DataSourceOAuthService( + provider_name="notion", + provider_gateway=RecordingProviderGateway(), + bindings=InMemoryBindingRepository(), + ) + + assert service.complete_callback(code=code, error=error) == DataSourceOAuthCallback( + provider="notion", + code=code, + error=expected_error, + ) + + +def test_bind_authorizes_before_persisting(request_context: RequestContext) -> None: + events: list[tuple[str, object]] = [] + authorization = _authorization() + provider = RecordingProviderGateway(authorization=authorization, events=events) + bindings = InMemoryBindingRepository(events=events) + service = DataSourceOAuthService(provider_name="notion", provider_gateway=provider, bindings=bindings) + + service.bind(request_context, code="code-1") + + assert events == [ + ("authorize", "code-1"), + ("persist", ("workspace-1", "notion", authorization)), + ] + + +def test_bind_rejects_empty_code(request_context: RequestContext) -> None: + service = DataSourceOAuthService( + provider_name="notion", + provider_gateway=RecordingProviderGateway(), + bindings=InMemoryBindingRepository(), + ) + + with pytest.raises(InvalidDataSourceOAuthCodeError): + service.bind(request_context, code="") + + +def test_sync_reads_before_provider_io_and_persists_afterward(request_context: RequestContext) -> None: + events: list[tuple[str, object]] = [] + binding = DataSourceOAuthBindingRecord(id="binding-1", access_token="token", source_info={"pages": []}) + refreshed_source_info = {"pages": [{"page_id": "page-1"}]} + bindings = InMemoryBindingRepository(binding=binding, events=events) + provider = RecordingProviderGateway(refreshed_source_info=refreshed_source_info, events=events) + service = DataSourceOAuthService(provider_name="notion", provider_gateway=provider, bindings=bindings) + + service.sync(request_context, binding_id="binding-1") + + assert events == [ + ("read", ("workspace-1", "notion", "binding-1")), + ("refresh", ("token", {"pages": []})), + ("update", ("workspace-1", "notion", "binding-1", refreshed_source_info)), + ] + + +def test_sync_rejects_missing_binding_without_provider_io(request_context: RequestContext) -> None: + events: list[tuple[str, object]] = [] + service = DataSourceOAuthService( + provider_name="notion", + provider_gateway=RecordingProviderGateway(events=events), + bindings=InMemoryBindingRepository(events=events), + ) + + with pytest.raises(DataSourceOAuthError, match="Data source binding not found"): + service.sync(request_context, binding_id="missing") + + assert events == [("read", ("workspace-1", "notion", "missing"))] + + +def test_sync_rejects_binding_removed_before_update(request_context: RequestContext) -> None: + events: list[tuple[str, object]] = [] + binding = DataSourceOAuthBindingRecord(id="binding-1", access_token="token", source_info={"pages": []}) + service = DataSourceOAuthService( + provider_name="notion", + provider_gateway=RecordingProviderGateway(events=events), + bindings=InMemoryBindingRepository(binding=binding, update_result=False, events=events), + ) + + with pytest.raises(DataSourceOAuthError, match="Data source binding not found"): + service.sync(request_context, binding_id="binding-1") + + assert [event for event, _ in events] == ["read", "refresh", "update"] diff --git a/api/tests/unit_tests/services/dataset_service/test_save_document_with_dataset_id_mysql_snapshot.py b/api/tests/unit_tests/services/dataset_service/test_save_document_with_dataset_id_mysql_snapshot.py index e8af9b89ca1..62b120a0d31 100644 --- a/api/tests/unit_tests/services/dataset_service/test_save_document_with_dataset_id_mysql_snapshot.py +++ b/api/tests/unit_tests/services/dataset_service/test_save_document_with_dataset_id_mysql_snapshot.py @@ -20,7 +20,7 @@ from __future__ import annotations import inspect -from services.dataset_service import DocumentService +from services.knowledge.dataset_service import DocumentService def test_session_committed_before_upload_file_lookup() -> None: @@ -37,13 +37,13 @@ def test_session_committed_before_upload_file_lookup() -> None: invariant the reproduction in the issue depends on, without needing a live MySQL database. """ - source_lines = inspect.getsource(DocumentService.save_document_with_dataset_id).splitlines() + source_lines = inspect.getsource(DocumentService.save_prepared_documents).splitlines() upload_select_line = next( (i for i, line in enumerate(source_lines) if "select(UploadFile)" in line), None, ) - assert upload_select_line is not None, "save_document_with_dataset_id must contain a select(UploadFile) call" + assert upload_select_line is not None, "save_prepared_documents must contain a select(UploadFile) call" # Find every ``session.commit()`` call, then keep the last one # whose line number is BEFORE the UploadFile SELECT. If the fix is @@ -56,7 +56,7 @@ def test_session_committed_before_upload_file_lookup() -> None: last_commit_before_upload = i assert last_commit_before_upload >= 0, ( - "save_document_with_dataset_id must call session.commit() at least once before the UploadFile SELECT" + "save_prepared_documents must call session.commit() at least once before the UploadFile SELECT" ) # The fix is a focused call: one extra ``session.commit()`` between @@ -66,5 +66,5 @@ def test_session_committed_before_upload_file_lookup() -> None: # doesn't accidentally delete the snapshot refresh without # understanding why it's there. assert "#41735" in "\n".join(source_lines), ( - "The fix's invariant comment referencing issue #41735 must be present in save_document_with_dataset_id" + "The fix's invariant comment referencing issue #41735 must be present in save_prepared_documents" ) diff --git a/api/tests/unit_tests/services/dataset_service_test_helpers.py b/api/tests/unit_tests/services/dataset_service_test_helpers.py index 6153afa3396..6b554f85bb8 100644 --- a/api/tests/unit_tests/services/dataset_service_test_helpers.py +++ b/api/tests/unit_tests/services/dataset_service_test_helpers.py @@ -2,7 +2,7 @@ These factories and lightweight builders are reused across the dataset, document, and segment service test modules that exercise -``api/services/dataset_service.py``. +``api/services/knowledge/dataset_service.py``. """ import json @@ -31,28 +31,6 @@ from models.dataset import ( DocumentSegment, ) from models.model import UploadFile -from services.dataset_service import ( - DatasetCollectionBindingService, - DatasetPermissionService, - DatasetService, - DocumentService, - SegmentService, -) -from services.entities.knowledge_entities.knowledge_entities import ( - ChildChunkUpdateArgs, - DataSource, - FileInfo, - InfoList, - KnowledgeConfig, - NotionIcon, - NotionInfo, - NotionPage, - ProcessRule, - RerankingModel, - RetrievalModel, - SegmentUpdateArgs, - WebsiteInfo, -) from services.entities.knowledge_entities.rag_pipeline_entities import ( IconInfo as PipelineIconInfo, ) @@ -71,6 +49,27 @@ from services.errors.chunk import ChildChunkDeleteIndexError, ChildChunkIndexing from services.errors.dataset import DatasetNameDuplicateError from services.errors.document import DocumentIndexingError from services.errors.file import FileNotExistsError +from services.knowledge.dataset_service import ( + DatasetCollectionBindingService, + DatasetPermissionService, + DatasetService, + DocumentService, + SegmentService, +) +from services.knowledge.entities.knowledge_entities import ( + DataSource, + FileInfo, + InfoList, + KnowledgeConfig, + NotionIcon, + NotionInfo, + NotionPage, + ProcessRule, + RerankingModel, + RetrievalModel, + WebsiteInfo, +) +from services.knowledge.entities.segments import ChildChunkUpdateArgs, SegmentUpdateArgs __all__ = [ "Account", diff --git a/api/tests/unit_tests/services/document_service_validation.py b/api/tests/unit_tests/services/document_service_validation.py index ff571eb856c..8ec2b3efc50 100644 --- a/api/tests/unit_tests/services/document_service_validation.py +++ b/api/tests/unit_tests/services/document_service_validation.py @@ -116,8 +116,8 @@ from core.rag.entities import PreProcessingRule, Rule, Segmentation from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType from graphon.model_runtime.entities.model_entities import ModelType from models.dataset import Dataset, DatasetProcessRule, Document -from services.dataset_service import DatasetService, DocumentService -from services.entities.knowledge_entities.knowledge_entities import ( +from services.knowledge.dataset_service import DatasetService, DocumentService +from services.knowledge.entities.knowledge_entities import ( DataSource, FileInfo, InfoList, @@ -419,7 +419,7 @@ class TestDatasetServiceCheckDatasetModelSetting: Provides a mocked ModelManager that can be used to verify model instance retrieval and error handling. """ - with patch("services.dataset_service.ModelManager.for_tenant") as mock_manager: + with patch("services.knowledge.dataset_service.ModelManager.for_tenant") as mock_manager: yield mock_manager def test_check_dataset_model_setting_high_quality_success(self, mock_model_manager): @@ -568,7 +568,7 @@ class TestDatasetServiceCheckEmbeddingModelSetting: Provides a mocked ModelManager that can be used to verify model instance retrieval and error handling. """ - with patch("services.dataset_service.ModelManager.for_tenant") as mock_manager: + with patch("services.knowledge.dataset_service.ModelManager.for_tenant") as mock_manager: yield mock_manager def test_check_embedding_model_setting_success(self, mock_model_manager): @@ -690,7 +690,7 @@ class TestDatasetServiceCheckRerankingModelSetting: Provides a mocked ModelManager that can be used to verify model instance retrieval and error handling. """ - with patch("services.dataset_service.ModelManager.for_tenant") as mock_manager: + with patch("services.knowledge.dataset_service.ModelManager.for_tenant") as mock_manager: yield mock_manager def test_check_reranking_model_setting_success(self, mock_model_manager): diff --git a/api/tests/unit_tests/services/knowledge/__init__.py b/api/tests/unit_tests/services/knowledge/__init__.py index e69de29bb2d..71872abbaa3 100644 --- a/api/tests/unit_tests/services/knowledge/__init__.py +++ b/api/tests/unit_tests/services/knowledge/__init__.py @@ -0,0 +1 @@ +"""Tests for knowledge application services and adapters.""" diff --git a/api/tests/unit_tests/services/knowledge/test_api_key_service.py b/api/tests/unit_tests/services/knowledge/test_api_key_service.py index f725f387de4..ce844001ffe 100644 --- a/api/tests/unit_tests/services/knowledge/test_api_key_service.py +++ b/api/tests/unit_tests/services/knowledge/test_api_key_service.py @@ -7,12 +7,14 @@ import pytest from sqlalchemy import select from sqlalchemy.orm import Session, sessionmaker +from extensions.application_services.datasets import build_dataset_dependencies from extensions.application_services.knowledge import build_dataset_api_key_service from machinery.context import RequestContext from models.account import Tenant, TenantAccountJoin, TenantAccountRole from models.dataset import Dataset, DatasetPermission from models.enums import PermissionEnum from models.model import ApiToken +from repositories.workspace_member_query_repository import WorkspaceMemberQueryRepository from services.api_token_service import ApiTokenCache from services.auth.api_key_contracts import ApiKeyNotFoundError from services.errors.account import NoPermissionError @@ -36,7 +38,14 @@ def service( ] ) sqlite_session.commit() - return build_dataset_api_key_service(database_client=sqlite_session_factory) + dataset_access = build_dataset_dependencies( + database_client=sqlite_session_factory, + workspace_roles=WorkspaceMemberQueryRepository(session_factory=sqlite_session_factory), + ).access + return build_dataset_api_key_service( + database_client=sqlite_session_factory, + dataset_access=dataset_access, + ) @pytest.mark.parametrize("workspace", [False, True]) diff --git a/api/tests/unit_tests/services/knowledge/test_dataset_access_service.py b/api/tests/unit_tests/services/knowledge/test_dataset_access_service.py new file mode 100644 index 00000000000..40a6bf92622 --- /dev/null +++ b/api/tests/unit_tests/services/knowledge/test_dataset_access_service.py @@ -0,0 +1,171 @@ +from dataclasses import dataclass + +import pytest + +from machinery.context import RequestContext +from services.knowledge.dataset_access import ( + DatasetAccessDeniedError, + DatasetAccessRecord, + DatasetAccessService, + DatasetAccessSnapshot, + DatasetNotFoundError, +) + + +def _context(*, actor_id: str = "actor-1", workspace_id: str = "workspace-1") -> RequestContext: + return RequestContext( + request_id="request-1", + trace_id=None, + account_id=actor_id, + active_workspace_id=workspace_id, + ) + + +def _snapshot( + *, + permission: str, + maintainer_id: str = "maintainer-1", + partial_access: bool = False, +) -> DatasetAccessSnapshot: + return DatasetAccessSnapshot( + dataset=DatasetAccessRecord( + id="dataset-1", + workspace_id="workspace-1", + maintainer_id=maintainer_id, + permission=permission, + ), + actor_has_partial_access=partial_access, + ) + + +@dataclass +class DatasetReaderStub: + snapshot: DatasetAccessSnapshot | None + + def get_access_snapshot( + self, + *, + workspace_id: str, + dataset_id: str, + actor_id: str, + ) -> DatasetAccessSnapshot | None: + assert workspace_id == "workspace-1" + assert dataset_id == "dataset-1" + assert actor_id == "actor-1" + return self.snapshot + + +@dataclass +class WorkspaceRoleReaderStub: + role: str | None + calls: int = 0 + + def get_legacy_role(self, *, workspace_id: str, account_id: str) -> str | None: + assert workspace_id == "workspace-1" + assert account_id == "actor-1" + self.calls += 1 + return self.role + + +@pytest.mark.parametrize( + ("permission", "role", "maintainer_id", "partial_access"), + [ + ("only_me", "owner", "maintainer-1", False), + ("only_me", "normal", "actor-1", False), + ("all_team_members", "normal", "maintainer-1", False), + ("partial_members", "normal", "maintainer-1", True), + ], +) +def test_require_accessible_allows_each_legacy_access_path( + permission: str, + role: str, + maintainer_id: str, + partial_access: bool, +) -> None: + roles = WorkspaceRoleReaderStub(role) + service = DatasetAccessService( + datasets=DatasetReaderStub( + _snapshot( + permission=permission, + maintainer_id=maintainer_id, + partial_access=partial_access, + ) + ), + workspace_roles=roles, + legacy_permissions_enabled=True, + ) + + result = service.require_accessible(_context(), "dataset-1") + + assert result.id == "dataset-1" + assert roles.calls == 1 + + +@pytest.mark.parametrize( + ("permission", "partial_access"), + [ + ("only_me", False), + ("partial_members", False), + ("unknown", True), + ], +) +def test_require_accessible_denies_unrecognized_or_ungranted_legacy_access( + permission: str, + partial_access: bool, +) -> None: + service = DatasetAccessService( + datasets=DatasetReaderStub(_snapshot(permission=permission, partial_access=partial_access)), + workspace_roles=WorkspaceRoleReaderStub("normal"), + legacy_permissions_enabled=True, + ) + + with pytest.raises(DatasetAccessDeniedError): + service.require_accessible(_context(), "dataset-1") + + +@pytest.mark.parametrize( + ("permission", "maintainer_id", "partial_access"), + [ + ("only_me", "actor-1", False), + ("all_team_members", "maintainer-1", False), + ("partial_members", "maintainer-1", True), + ], +) +def test_require_accessible_denies_non_members_even_when_dataset_policy_would_allow( + permission: str, + maintainer_id: str, + partial_access: bool, +) -> None: + service = DatasetAccessService( + datasets=DatasetReaderStub( + _snapshot(permission=permission, maintainer_id=maintainer_id, partial_access=partial_access) + ), + workspace_roles=WorkspaceRoleReaderStub(None), + legacy_permissions_enabled=True, + ) + + with pytest.raises(DatasetAccessDeniedError): + service.require_accessible(_context(), "dataset-1") + + +def test_require_accessible_only_enforces_owner_chain_when_legacy_permissions_are_disabled() -> None: + roles = WorkspaceRoleReaderStub("normal") + service = DatasetAccessService( + datasets=DatasetReaderStub(_snapshot(permission="only_me")), + workspace_roles=roles, + legacy_permissions_enabled=False, + ) + + assert service.require_accessible(_context(), "dataset-1").id == "dataset-1" + assert roles.calls == 0 + + +def test_require_accessible_hides_missing_or_cross_workspace_dataset() -> None: + service = DatasetAccessService( + datasets=DatasetReaderStub(None), + workspace_roles=WorkspaceRoleReaderStub("owner"), + legacy_permissions_enabled=True, + ) + + with pytest.raises(DatasetNotFoundError): + service.require_accessible(_context(), "dataset-1") diff --git a/api/tests/unit_tests/services/knowledge/test_dataset_adapters.py b/api/tests/unit_tests/services/knowledge/test_dataset_adapters.py new file mode 100644 index 00000000000..06d1eb39524 --- /dev/null +++ b/api/tests/unit_tests/services/knowledge/test_dataset_adapters.py @@ -0,0 +1,379 @@ +from collections.abc import Callable, Iterator +from unittest.mock import patch + +import pytest +from sqlalchemy import func, select +from sqlalchemy.orm import Session, sessionmaker + +from controllers.console.datasets.datasets import ( + DatasetDetailResponse, + DatasetListResponse, + DatasetQueryListResponse, + RelatedAppListResponse, +) +from extensions.application_services.app import AppServices +from machinery.context import RequestContext +from models import Account, App, Dataset, Document +from models.account import Tenant, TenantAccountJoin, TenantAccountRole +from models.dataset import ( + AppDatasetJoin, + DatasetPermission, + DatasetPermissionEnum, + DatasetQuery, + DocumentSegment, +) +from models.enums import CreatorUserRole, DatasetQuerySource, IndexingStatus, SegmentStatus +from services.enterprise import rbac_service +from services.knowledge.dataset_access import DatasetNotFoundError +from services.knowledge.datasets.adapters import SQLAlchemyDatasetOperations +from services.knowledge.datasets.application import DatasetListFilter +from services.knowledge.documents.adapters import SQLAlchemyDocumentOperations +from services.knowledge.entities.document_creation import DocumentIndexingJobs +from services.knowledge.entities.knowledge_entities import KnowledgeConfig +from services.knowledge.resource_scope import DatasetRef +from services.tag_application_service import CreateTagInput, TagApplicationService, TagBindingInput +from tests.unit_tests.config_override import apply_config_overrides + +CONTEXT = RequestContext("request", None, "actor", "tenant") +REF = DatasetRef("tenant", "dataset") + + +def dataset(**values: object) -> Dataset: + return Dataset( + **{ + "id": "dataset", + "tenant_id": "tenant", + "name": "Dataset", + "created_by": "actor", + "maintainer": "actor", + "indexing_technique": "economy", + "permission": "all_team_members", + **values, + } + ) + + +def document(**values: object) -> Document: + return Document( + **{ + "id": "document", + "tenant_id": "tenant", + "dataset_id": "dataset", + "position": 1, + "data_source_type": "local_file", + "batch": "batch", + "name": "Document", + "created_from": "web", + "created_by": "actor", + "indexing_status": IndexingStatus.ERROR, + **values, + } + ) + + +@pytest.fixture +def operations( + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + application_tags: TagApplicationService, + app_services: AppServices, +) -> Iterator[SQLAlchemyDatasetOperations]: + apply_config_overrides(monkeypatch, RBAC_ENABLED=False) + with sqlite_session_factory.begin() as session: + tenant = Tenant(name="Workspace") + tenant.id = "tenant" + account = Account(name="Actor", email="actor@example.com") + account.id = "actor" + session.add_all( + [ + tenant, + account, + TenantAccountJoin(tenant_id="tenant", account_id="actor", role=TenantAccountRole.OWNER), + dataset(), + dataset(id="foreign", tenant_id="other", name="Foreign"), + ] + ) + with ( + patch.object(rbac_service.RBACService.MyPermissions, "get", return_value=rbac_service.MyPermissionsResponse()), + patch.object( + rbac_service.RBACService.DatasetPermissions, "batch_get", return_value={"dataset": ["dataset.preview"]} + ), + patch.object(rbac_service, "try_sync_creator_access_policy_member_bindings"), + ): + yield SQLAlchemyDatasetOperations( + session_factory=sqlite_session_factory, tags=application_tags, app_queries=app_services.queries + ) + + +def test_listing_materializes_page_and_owner_scoped_partial_members( + operations: SQLAlchemyDatasetOperations, sqlite_session_factory: sessionmaker[Session] +) -> None: + with sqlite_session_factory.begin() as session: + row = session.get(Dataset, "dataset") + assert row is not None + row.permission = DatasetPermissionEnum.PARTIAL_TEAM + session.add_all( + [ + DatasetPermission(tenant_id="tenant", dataset_id="dataset", account_id="member"), + DatasetPermission(tenant_id="other", dataset_id="dataset", account_id="foreign"), + ] + ) + result = operations.list_datasets( + CONTEXT, DatasetListFilter(page=0, limit=1000, ids=["dataset", "foreign"]), None, False + ) + response = DatasetListResponse.model_validate(result).model_dump(mode="json") + assert response["total"] == 1 + assert response["page"] == 1 + assert response["limit"] == 100 + assert response["has_more"] is False + assert response["data"][0]["partial_member_list"] == ["member"] + assert response["data"][0]["retrieval_model_dict"]["top_k"] == 2 + + +@pytest.mark.parametrize("foreign_tag", [False, True]) +def test_listing_filters_tags_within_workspace( + operations: SQLAlchemyDatasetOperations, + application_tags: TagApplicationService, + sqlite_session_factory: sessionmaker[Session], + foreign_tag: bool, +) -> None: + with sqlite_session_factory.begin() as session: + session.add(dataset(id="untagged", name="Untagged")) + context = RequestContext("request", None, "actor", "other") if foreign_tag else CONTEXT + tag = application_tags.create_tag(context, CreateTagInput(name="Selected", type="knowledge")) + application_tags.create_bindings( + context, + TagBindingInput(tag_ids=(tag.id,), target_id="foreign" if foreign_tag else "dataset", type="knowledge"), + ) + + result = operations.list_datasets(CONTEXT, DatasetListFilter(tag_ids=[tag.id]), None, False) + response = DatasetListResponse.model_validate(result) + + assert [row.id for row in response.data] == ([] if foreign_tag else ["dataset"]) + assert response.total == (0 if foreign_tag else 1) + + +@pytest.mark.parametrize(("ids", "own"), [([], False), (["dataset"], False)]) +def test_list_visibility_restricts_even_requested_ids( + operations: SQLAlchemyDatasetOperations, ids: list[str], own: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + apply_config_overrides(monkeypatch, RBAC_ENABLED=True) + result = operations.list_datasets(CONTEXT, DatasetListFilter(ids=["dataset"]), ids, own) + assert result["total"] == len(ids) + + +@pytest.mark.parametrize( + "method", + [ + "get_dataset", + "is_in_use", + "queries", + "related_apps", + "indexing_status", + "error_documents", + "partial_members", + "update_dataset", + "delete_dataset", + "set_api_enabled", + ], +) +def test_wrong_tenant_ref_cannot_read_or_write(operations: SQLAlchemyDatasetOperations, method: str) -> None: + args: list[object] = ( + [CONTEXT] if method in {"get_dataset", "update_dataset", "delete_dataset", "set_api_enabled"} else [] + ) + args.append(DatasetRef("other", "dataset")) + if method == "update_dataset": + args.append({"name": "changed"}) + if method == "set_api_enabled": + args.append(True) + extra: dict[str, int] = {"page": 1, "limit": 20} if method == "queries" else {} + methods: dict[str, Callable[..., object]] = { + "get_dataset": operations.get_dataset, + "is_in_use": operations.is_in_use, + "queries": operations.queries, + "related_apps": operations.related_apps, + "indexing_status": operations.indexing_status, + "error_documents": operations.error_documents, + "partial_members": operations.partial_members, + "update_dataset": operations.update_dataset, + "delete_dataset": operations.delete_dataset, + "set_api_enabled": operations.set_api_enabled, + } + with pytest.raises(DatasetNotFoundError): + methods[method](*args, **extra) + + +def test_create_update_and_api_status_commit_owned_changes( + operations: SQLAlchemyDatasetOperations, sqlite_session_factory: sessionmaker[Session] +) -> None: + created = operations.create_dataset( + CONTEXT, {"name": "Created", "description": "desc", "indexing_technique": "economy", "permission": "only_me"} + ) + assert DatasetDetailResponse.model_validate(created).name == "Created" + result = operations.update_dataset(CONTEXT, REF, {"name": "Renamed", "permission": "all_team_members"}) + assert result["name"] == "Renamed" + operations.set_api_enabled(CONTEXT, REF, True) + with sqlite_session_factory() as session: + row = session.get(Dataset, "dataset") + assert row is not None + assert row.name == "Renamed" + assert row.enable_api is True + assert session.get(Dataset, created["id"]) is not None + + +@pytest.mark.parametrize("entry_point", ["empty", "documents"]) +@pytest.mark.parametrize("rbac_enabled", [False, True]) +def test_created_dataset_initializes_rbac_access( + operations: SQLAlchemyDatasetOperations, + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + entry_point: str, + rbac_enabled: bool, +) -> None: + apply_config_overrides(monkeypatch, RBAC_ENABLED=rbac_enabled) + + def save_documents( + created_dataset: Dataset, _config: KnowledgeConfig, _account: Account, *, session: Session + ) -> tuple[list[Document], str, DocumentIndexingJobs]: + row = document(id="created-document", dataset_id=created_dataset.id, data_source_type="upload_file") + session.add(row) + session.flush() + return [row], "batch", DocumentIndexingJobs(DatasetRef(created_dataset.tenant_id, created_dataset.id)) + + with ( + patch.object(rbac_service.RBACService.DatasetAccess, "replace_whitelist") as replace_whitelist, + patch.object(rbac_service, "try_sync_creator_access_policy_member_bindings") as sync_creator, + patch( + "tasks.initialize_created_app_rbac_access_task.initialize_created_app_rbac_access_task.delay" + ) as initialize, + patch( + "services.knowledge.documents.adapters.DocumentService.save_prepared_documents", + side_effect=save_documents, + ), + ): + if entry_point == "empty": + result = operations.create_dataset( + CONTEXT, + {"name": "Created", "description": "desc", "indexing_technique": "economy", "permission": "only_me"}, + ) + dataset_id = result["id"] + else: + result = SQLAlchemyDocumentOperations(session_factory=sqlite_session_factory).initialize_dataset( + CONTEXT, + { + "indexing_technique": "economy", + "data_source": { + "info_list": {"data_source_type": "upload_file", "file_info_list": {"file_ids": ["file"]}} + }, + "process_rule": {"mode": "automatic"}, + }, + ) + dataset_id = result["dataset"]["id"] + if rbac_enabled: + replace_whitelist.assert_called_once_with( + CONTEXT.active_workspace_id, + CONTEXT.account_id, + dataset_id, + rbac_service.ReplaceMemberBindings(automatic_include_workspace_members=False), + ) + else: + replace_whitelist.assert_not_called() + if entry_point == "documents": + if rbac_enabled: + sync_creator.assert_called_once_with( + CONTEXT.active_workspace_id, + CONTEXT.account_id, + rbac_service.RBACResourceType.DATASET, + dataset_id, + ) + else: + sync_creator.assert_not_called() + initialize.assert_not_called() + + with sqlite_session_factory() as session: + assert session.get(Dataset, dataset_id) is not None + + +def test_partial_members_update_and_clear_are_committed( + operations: SQLAlchemyDatasetOperations, sqlite_session_factory: sessionmaker[Session] +) -> None: + result = operations.update_dataset( + CONTEXT, REF, {"permission": "partial_members", "partial_member_list": [{"user_id": "actor"}]} + ) + assert result["partial_member_list"] == ["actor"] + result = operations.update_dataset(CONTEXT, REF, {"permission": "all_team_members"}) + assert result["partial_member_list"] == [] + with sqlite_session_factory() as session: + assert session.scalar(select(func.count()).select_from(DatasetPermission)) == 0 + + +def test_status_counts_and_errors_exclude_foreign_documents( + operations: SQLAlchemyDatasetOperations, sqlite_session_factory: sessionmaker[Session] +) -> None: + with sqlite_session_factory.begin() as session: + session.add_all( + [ + document(), + document(id="foreign-doc", tenant_id="other"), + DocumentSegment( + tenant_id="tenant", + dataset_id="dataset", + document_id="document", + position=1, + content="content", + word_count=1, + tokens=1, + created_by="actor", + status=SegmentStatus.WAITING, + ), + DocumentSegment( + tenant_id="other", + dataset_id="dataset", + document_id="document", + position=1, + content="content", + word_count=1, + tokens=1, + created_by="actor", + ), + ] + ) + result = operations.indexing_status(REF) + assert len(result["data"]) == 1 + assert result["data"][0]["total_segments"] == 1 + assert result["data"][0]["completed_segments"] == 0 + assert operations.error_documents(REF)["total"] == 1 + + +def test_queries_and_related_apps_materialize_after_session_close( + operations: SQLAlchemyDatasetOperations, sqlite_session_factory: sessionmaker[Session] +) -> None: + with sqlite_session_factory.begin() as session: + session.add( + DatasetQuery( + dataset_id="dataset", + content="query", + source=DatasetQuerySource.HIT_TESTING, + source_app_id=None, + created_by_role=CreatorUserRole.ACCOUNT, + created_by="actor", + ) + ) + session.add_all( + [ + App(id="app", tenant_id="tenant", name="App", mode="chat", enable_api=False, enable_site=False), + App( + id="foreign-app", tenant_id="other", name="Other", mode="chat", enable_api=False, enable_site=False + ), + AppDatasetJoin(app_id="app", dataset_id="dataset"), + AppDatasetJoin(app_id="foreign-app", dataset_id="dataset"), + ] + ) + response = DatasetQueryListResponse.model_validate(operations.queries(REF, page=0, limit=0)).model_dump(mode="json") + assert response["limit"] == 1 + assert response["has_more"] is False + assert response["data"][0]["queries"][0]["content"] == "query" + related = RelatedAppListResponse.model_validate(operations.related_apps(REF)).model_dump(mode="json") + assert related["total"] == 1 + assert related["data"][0]["id"] == "app" + assert related["data"][0]["mode"] == "chat" diff --git a/api/tests/unit_tests/services/knowledge/test_dataset_application_service.py b/api/tests/unit_tests/services/knowledge/test_dataset_application_service.py new file mode 100644 index 00000000000..391c12abde7 --- /dev/null +++ b/api/tests/unit_tests/services/knowledge/test_dataset_application_service.py @@ -0,0 +1,160 @@ +from collections.abc import Callable +from unittest.mock import MagicMock, create_autospec + +import pytest + +from machinery.context import RequestContext +from services.knowledge.dataset_access import AccessibleDataset, DatasetAccess, DatasetAccessDeniedError +from services.knowledge.datasets.application import ( + DatasetApplicationService, + DatasetListFilter, + DatasetOperations, + DatasetVisibility, +) +from services.knowledge.datasets.retrieval import retrieval_methods + +CONTEXT = RequestContext("request", None, "actor", "tenant") + + +@pytest.fixture +def operations() -> MagicMock: + return create_autospec(DatasetOperations, instance=True, spec_set=True) + + +@pytest.fixture +def access() -> MagicMock: + result = create_autospec(DatasetAccess, instance=True, spec_set=True) + result.require_accessible.return_value = AccessibleDataset("dataset", "tenant") + return result + + +@pytest.fixture +def service(operations: MagicMock, access: MagicMock) -> DatasetApplicationService: + return DatasetApplicationService( + dataset_access=access, + operations=operations, + rbac_enabled=True, + service_api_url="", + vector_store="milvus", + tidb_fulltext=False, + ) + + +@pytest.mark.parametrize( + ("defaults", "workspace", "overrides"), + [ + (["dataset.preview"], [], {}), + ([], ["dataset.full_access", "dataset.create_and_management"], {}), + ([], [], {"hidden": ["dataset.preview"]}), + ], +) +def test_restricted_whitelist_wins_over_all_other_grants( + defaults: list[str], workspace: list[str], overrides: dict[str, list[str]] +) -> None: + visibility = DatasetVisibility(defaults, workspace, overrides, False, ["allowed"]) + assert visibility.list_scope(rbac_enabled=True) == (["allowed"], False) + + +@pytest.mark.parametrize( + ("visibility", "expected"), + [ + (DatasetVisibility(default_permissions=["dataset.preview"]), (None, False)), + (DatasetVisibility(workspace_permissions=["dataset.acl.preview"]), (None, False)), + (DatasetVisibility(overrides={"a": ["dataset.preview"], "b": ["dataset.edit"]}), (["a"], False)), + (DatasetVisibility(workspace_permissions=["dataset.create_and_management"]), ([], True)), + (DatasetVisibility(unrestricted=False), ([], False)), + ], +) +def test_visibility_scope(visibility: DatasetVisibility, expected: tuple[list[str] | None, bool]) -> None: + assert visibility.list_scope(rbac_enabled=True) == expected + assert visibility.list_scope(rbac_enabled=False) == (None, False) + + +def test_list_applies_visibility_and_reports_model_availability( + service: DatasetApplicationService, operations: MagicMock +) -> None: + operations.visibility.return_value = DatasetVisibility( + default_permissions=["dataset.preview"], + overrides={"a": ["dataset.edit"]}, + unrestricted=False, + whitelist_ids=["a"], + ) + operations.list_datasets.return_value = { + "data": [ + { + "id": "a", + "indexing_technique": "high_quality", + "embedding_model": "missing", + "embedding_model_provider": "provider", + }, + {"id": "b", "indexing_technique": "economy", "embedding_model": None, "embedding_model_provider": None}, + ] + } + operations.embedding_models.return_value = set() + result = service.list_datasets(CONTEXT, DatasetListFilter(ids=["a"])) + operations.list_datasets.assert_called_once_with(CONTEXT, DatasetListFilter(ids=["a"]), ["a"], False) + assert result["data"][0]["embedding_available"] is False + assert result["data"][0]["permission_keys"] == ["dataset.edit"] + assert result["data"][1]["embedding_available"] is True + + +@pytest.mark.parametrize( + ("method", "extra"), + [ + (DatasetApplicationService.get_dataset, {}), + (DatasetApplicationService.update_dataset, {"values": {}}), + (DatasetApplicationService.delete_dataset, {}), + (DatasetApplicationService.is_in_use, {}), + (DatasetApplicationService.queries, {"page": 1, "limit": 20}), + (DatasetApplicationService.related_apps, {}), + (DatasetApplicationService.indexing_status, {}), + (DatasetApplicationService.error_documents, {}), + (DatasetApplicationService.partial_members, {}), + (DatasetApplicationService.auto_disable_logs, {}), + (DatasetApplicationService.set_api_enabled, {"status": "enable"}), + ], +) +def test_access_denial_precedes_all_owned_operations( + service: DatasetApplicationService, + access: MagicMock, + operations: MagicMock, + method: Callable[..., object], + extra: dict[str, object], +) -> None: + access.require_accessible.side_effect = DatasetAccessDeniedError() + with pytest.raises(DatasetAccessDeniedError): + method(service, CONTEXT, dataset_id="dataset", **extra) + assert operations.mock_calls == [] + + +def test_rbac_create_normalizes_permission_without_mutating_input( + service: DatasetApplicationService, operations: MagicMock +) -> None: + values = {"name": "name", "permission": "only_me"} + service.create_dataset(CONTEXT, values=values) + operations.create_dataset.assert_called_once_with(CONTEXT, {"name": "name", "permission": "all_team_members"}) + assert values["permission"] == "only_me" + + +def test_settings_are_injected_and_url_is_normalized(service: DatasetApplicationService) -> None: + assert service.api_base_url(CONTEXT, request_base_url="https://api.example/v1/") == "https://api.example/v1" + assert service.retrieval_settings(CONTEXT)["retrieval_method"] == [ + "semantic_search", + "full_text_search", + "hybrid_search", + ] + assert service.retrieval_settings(CONTEXT, vector_type="milvus", is_mock=True) == { + "retrieval_method": ["semantic_search"] + } + + +@pytest.mark.parametrize("fulltext", [True, False]) +def test_tidb_retrieval_setting(fulltext: bool) -> None: + result = retrieval_methods("tidb_vector", tidb_fulltext=fulltext) + assert ("full_text_search" in result["retrieval_method"]) is fulltext + + +@pytest.mark.parametrize("store", [None, "unsupported"]) +def test_invalid_vector_store(store: str | None) -> None: + with pytest.raises(ValueError): + retrieval_methods(store) diff --git a/api/tests/unit_tests/services/knowledge/test_document_adapters.py b/api/tests/unit_tests/services/knowledge/test_document_adapters.py new file mode 100644 index 00000000000..eb85952f5b0 --- /dev/null +++ b/api/tests/unit_tests/services/knowledge/test_document_adapters.py @@ -0,0 +1,293 @@ +import json +from datetime import datetime +from unittest.mock import patch + +import pytest +from sqlalchemy import Engine, event +from sqlalchemy.orm import Session, sessionmaker + +from controllers.console.datasets.datasets_document import DocumentDetailResponse, DocumentWithSegmentsListResponse +from core.rag.index_processor.constant.index_type import IndexTechniqueType +from models.dataset import Dataset, DatasetProcessRule, Document, DocumentPipelineExecutionLog, DocumentSegment +from models.enums import IndexingStatus, ProcessRuleMode, SegmentStatus +from services.knowledge.documents.adapters import SQLAlchemyDocumentOperations +from services.knowledge.documents.application import DocumentListFilter, DocumentNotFoundError +from services.knowledge.resource_scope import DatasetRef + +REF = DatasetRef("tenant-1", "dataset-1") + + +def _dataset(**values: object) -> Dataset: + return Dataset( + **{ + "id": "dataset-1", + "tenant_id": "tenant-1", + "name": "Dataset", + "created_by": "actor-1", + "maintainer": "actor-1", + "indexing_technique": "economy", + "chunk_structure": "text_model", + **values, + } + ) + + +def _document(**values: object) -> Document: + return Document( + **{ + "id": "document-1", + "tenant_id": "tenant-1", + "dataset_id": "dataset-1", + "position": 1, + "data_source_type": "local_file", + "data_source_info": json.dumps({"file_path": "/path/to/file"}), + "batch": "batch-1", + "name": "Document", + "created_from": "web", + "created_by": "actor-1", + "indexing_status": IndexingStatus.COMPLETED, + "doc_form": "text_model", + "enabled": True, + "archived": False, + "word_count": 10, + "created_at": datetime(2024, 1, 1), + **values, + } + ) + + +def _segment( + segment_id: str, + *, + tenant_id: str = "tenant-1", + dataset_id: str = "dataset-1", + document_id: str = "document-1", + status: SegmentStatus = SegmentStatus.COMPLETED, + completed_at: datetime | None = datetime(2024, 1, 1), + hit_count: int = 0, +) -> DocumentSegment: + segment = DocumentSegment( + tenant_id=tenant_id, + dataset_id=dataset_id, + document_id=document_id, + position=1, + content="content", + word_count=7, + tokens=2, + created_by="actor-1", + status=status, + completed_at=completed_at, + hit_count=hit_count, + ) + segment.id = segment_id + return segment + + +@pytest.fixture +def operations(sqlite_session_factory: sessionmaker[Session]) -> SQLAlchemyDocumentOperations: + with sqlite_session_factory.begin() as session: + session.add_all([_dataset(), _document()]) + return SQLAlchemyDocumentOperations(session_factory=sqlite_session_factory) + + +@pytest.mark.parametrize("foreign", [{"tenant_id": "other"}, {"dataset_id": "other"}]) +def test_document_reads_and_writes_reject_incorrect_owner( + operations: SQLAlchemyDocumentOperations, sqlite_session_factory: sessionmaker[Session], foreign: dict[str, str] +) -> None: + with sqlite_session_factory.begin() as session: + session.add(_document(id="foreign", **foreign)) + ref = REF.document("foreign") + assert operations.get_state(ref) is None + with pytest.raises(DocumentNotFoundError): + operations.get_detail(ref, metadata_only=False) + with pytest.raises(DocumentNotFoundError): + operations.update_document(ref, {"name": "changed"}) + with pytest.raises(DocumentNotFoundError): + operations.rename_document(ref, "changed") + with sqlite_session_factory() as session: + document = session.get(Document, "foreign") + assert document is not None + assert document.name == "Document" + + +def test_document_lookup_uses_active_workspace(operations: SQLAlchemyDocumentOperations) -> None: + assert operations.find_document(workspace_id="other", document_id="document-1") is None + assert operations.find_document(workspace_id="tenant-1", document_id="document-1") == REF.document("document-1") + + +def test_listing_filters_scope_and_materializes_values( + operations: SQLAlchemyDocumentOperations, sqlite_session_factory: sessionmaker[Session] +) -> None: + with sqlite_session_factory.begin() as session: + session.add_all( + [ + _document(id="matching", name="Needle", position=2), + _document(id="disabled", name="Needle", enabled=False), + _document(id="foreign", name="Needle", tenant_id="other"), + _document(id="other-dataset", name="Needle", dataset_id="other"), + ] + ) + result = operations.list_documents(REF, DocumentListFilter(search="Needle", status="available", fetch=True)) + response = DocumentWithSegmentsListResponse.model_validate(result).model_dump(mode="json") + assert response["total"] == 1 + assert response["data"][0]["id"] == "matching" + assert response["data"][0]["data_source_info"] == {"file_path": "/path/to/file"} + assert response["data"][0]["completed_segments"] == 0 + assert response["data"][0]["total_segments"] == 0 + + +def test_listing_hit_count_sort_and_exact_last_page( + operations: SQLAlchemyDocumentOperations, sqlite_session_factory: sessionmaker[Session] +) -> None: + with sqlite_session_factory.begin() as session: + session.add_all( + [ + _document(id="popular", position=2), + _segment("segment-1", hit_count=1), + _segment("segment-2", document_id="popular", hit_count=10), + _segment("foreign", tenant_id="other", hit_count=1000), + ] + ) + first = operations.list_documents(REF, DocumentListFilter(sort="-hit_count", limit=1)) + last = operations.list_documents(REF, DocumentListFilter(sort="-hit_count", limit=1, page=2)) + assert [row["id"] for row in first["data"]] == ["popular"] + assert first["has_more"] is True + assert last["has_more"] is False + assert last["total"] == 2 + + +def test_status_counts_use_one_aggregate_and_complete_owner_chain( + operations: SQLAlchemyDocumentOperations, sqlite_session_factory: sessionmaker[Session], sqlite_engine: Engine +) -> None: + with sqlite_session_factory.begin() as session: + session.add_all( + [ + _segment("complete"), + _segment("waiting", completed_at=None), + _segment("excluded", status=SegmentStatus.RE_SEGMENT), + _segment("foreign", tenant_id="other"), + _segment("wrong-dataset", dataset_id="other"), + _document(id="second", is_paused=True, indexing_status=IndexingStatus.INDEXING), + ] + ) + statements: list[str] = [] + + def record( + _conn: object, _cursor: object, statement: str, _parameters: object, _context: object, _many: bool + ) -> None: + statements.append(statement) + + event.listen(sqlite_engine, "before_cursor_execute", record) + try: + result = operations.get_batch_indexing_status(REF, "batch-1") + finally: + event.remove(sqlite_engine, "before_cursor_execute", record) + rows = {row["id"]: row for row in result["data"]} + assert rows["document-1"]["completed_segments"] == 1 + assert rows["document-1"]["total_segments"] == 2 + assert rows["second"]["indexing_status"] == "paused" + assert sum("FROM document_segments" in statement for statement in statements) == 1 + + +def test_unknown_batch_is_not_found(operations: SQLAlchemyDocumentOperations) -> None: + with pytest.raises(DocumentNotFoundError): + operations.get_batch_indexing_status(REF, "missing") + + +@pytest.mark.parametrize("rules", [None, {"segmentation": {"delimiter": "\\n", "max_tokens": 500}}]) +def test_process_rule_preserves_null_and_legacy_delimiter( + operations: SQLAlchemyDocumentOperations, + sqlite_session_factory: sessionmaker[Session], + rules: dict[str, object] | None, +) -> None: + with sqlite_session_factory.begin() as session: + session.add( + DatasetProcessRule( + dataset_id="dataset-1", + mode=ProcessRuleMode.CUSTOM, + rules=json.dumps(rules) if rules else None, + created_by="actor-1", + ) + ) + result = operations.get_process_rule(REF) + assert result["mode"] == "custom" + assert result["rules"] == rules + + +def test_detail_is_serializable_after_session_closes(operations: SQLAlchemyDocumentOperations) -> None: + result = operations.get_detail(REF.document("document-1"), metadata_only=False) + response = DocumentDetailResponse.model_validate(result).model_dump(mode="json") + assert response["id"] == "document-1" + assert response["data_source_info"] == {"file_path": "/path/to/file"} + assert response["data_source_detail_dict"] == {} + assert response["created_at"] == int(datetime(2024, 1, 1).timestamp()) + + +def test_metadata_only_avoids_loading_full_details(operations: SQLAlchemyDocumentOperations) -> None: + with patch("services.knowledge.documents.adapters.get_document_source_detail") as source: + result = operations.get_detail(REF.document("document-1"), metadata_only=True) + assert set(result) == {"id", "doc_type", "doc_metadata"} + source.assert_not_called() + + +def test_update_is_committed_before_returning( + operations: SQLAlchemyDocumentOperations, sqlite_session_factory: sessionmaker[Session] +) -> None: + operations.update_document(REF.document("document-1"), {"is_paused": True, "paused_by": "actor-1"}) + with sqlite_session_factory() as session: + document = session.get(Document, "document-1") + assert document is not None + assert document.is_paused is True + assert document.paused_by == "actor-1" + + +def test_latest_pipeline_log_and_empty_result( + operations: SQLAlchemyDocumentOperations, sqlite_session_factory: sessionmaker[Session] +) -> None: + assert operations.get_execution_log(REF.document("document-1")) == {} + with sqlite_session_factory.begin() as session: + for year in [2024, 2025]: + row = DocumentPipelineExecutionLog( + pipeline_id="pipeline-1", + document_id="document-1", + datasource_type="local_file", + datasource_info=json.dumps({"year": year}), + datasource_node_id="node-1", + input_data={"year": year}, + created_by="actor-1", + ) + row.created_at = datetime(year, 1, 1) + session.add(row) + result = operations.get_execution_log(REF.document("document-1")) + assert result["datasource_info"] == {"year": 2025} + assert result["input_data"] == {"year": 2025} + + +@pytest.mark.parametrize("settings", [None, {"enable": False}]) +def test_summary_guard_rejects_disabled_summary( + operations: SQLAlchemyDocumentOperations, + sqlite_session_factory: sessionmaker[Session], + settings: dict[str, object] | None, +) -> None: + with sqlite_session_factory.begin() as session: + dataset = session.get(Dataset, "dataset-1") + assert dataset is not None + dataset.indexing_technique = IndexTechniqueType.HIGH_QUALITY + dataset.summary_index_setting = settings + with pytest.raises(ValueError, match="not enabled"): + operations.require_summary_enabled(REF) + + +def test_enable_summary_only_updates_requested_owner( + operations: SQLAlchemyDocumentOperations, sqlite_session_factory: sessionmaker[Session] +) -> None: + with sqlite_session_factory.begin() as session: + session.add(_document(id="foreign", tenant_id="other", need_summary=False)) + operations.enable_summary(REF, ["document-1", "foreign"]) + with sqlite_session_factory() as session: + document = session.get(Document, "document-1") + foreign = session.get(Document, "foreign") + assert document is not None + assert foreign is not None + assert document.need_summary is True + assert foreign.need_summary is False diff --git a/api/tests/unit_tests/services/knowledge/test_document_application_service.py b/api/tests/unit_tests/services/knowledge/test_document_application_service.py new file mode 100644 index 00000000000..952d7024896 --- /dev/null +++ b/api/tests/unit_tests/services/knowledge/test_document_application_service.py @@ -0,0 +1,252 @@ +from collections.abc import Callable +from dataclasses import replace +from typing import Literal +from unittest.mock import MagicMock, create_autospec + +import pytest + +from machinery.context import RequestContext +from services.knowledge.dataset_access import AccessibleDataset, DatasetAccess, DatasetAccessDeniedError +from services.knowledge.documents.application import ( + DatasetDocumentApplicationService, + DocumentArchivedError, + DocumentIndexingStateError, + DocumentInvalidActionError, + DocumentNotFoundError, + DocumentOperations, + DocumentState, +) +from services.knowledge.resource_scope import DatasetRef + +CONTEXT = RequestContext("request-1", None, "actor-1", "tenant-1") +DATASET = DatasetRef("tenant-1", "dataset-1") +DOCUMENT = DATASET.document("document-1") +STATE = DocumentState("document-1", False, "indexing", False, "website_crawl", "text_model", False) + + +@pytest.fixture +def operations() -> MagicMock: + result = create_autospec(DocumentOperations, instance=True, spec_set=True) + result.get_state.return_value = STATE + return result + + +@pytest.fixture +def access() -> MagicMock: + result = create_autospec(DatasetAccess, instance=True, spec_set=True) + result.require_accessible.return_value = AccessibleDataset("dataset-1", "tenant-1") + return result + + +@pytest.fixture +def service(access: MagicMock, operations: MagicMock) -> DatasetDocumentApplicationService: + return DatasetDocumentApplicationService( + dataset_access=access, operations=operations, metadata_schema={"book": {"title": str}, "others": {}} + ) + + +@pytest.mark.parametrize( + "operation", + [ + DatasetDocumentApplicationService.get_document, + DatasetDocumentApplicationService.delete_document, + DatasetDocumentApplicationService.pause_document, + DatasetDocumentApplicationService.recover_document, + DatasetDocumentApplicationService.sync_website, + ], +) +def test_denied_dataset_never_loads_or_mutates_documents( + service: DatasetDocumentApplicationService, + access: MagicMock, + operations: MagicMock, + operation: Callable[..., object], +) -> None: + access.require_accessible.side_effect = DatasetAccessDeniedError() + with pytest.raises(DatasetAccessDeniedError): + operation(service, CONTEXT, dataset_id="dataset-1", document_id="document-1") + assert operations.mock_calls == [] + + +@pytest.mark.parametrize( + "operation", + [ + DatasetDocumentApplicationService.get_document, + DatasetDocumentApplicationService.delete_document, + DatasetDocumentApplicationService.pause_document, + DatasetDocumentApplicationService.recover_document, + DatasetDocumentApplicationService.sync_website, + ], +) +def test_missing_document_stops_before_any_side_effect( + service: DatasetDocumentApplicationService, + access: MagicMock, + operations: MagicMock, + operation: Callable[..., object], +) -> None: + operations.get_state.return_value = None + with pytest.raises(DocumentNotFoundError): + operation(service, CONTEXT, dataset_id="dataset-1", document_id="document-1") + access.require_accessible.assert_called_once_with(CONTEXT, "dataset-1") + operations.get_state.assert_called_once_with(DOCUMENT) + assert len(operations.mock_calls) == 1 + + +@pytest.mark.parametrize( + "operation", + [ + DatasetDocumentApplicationService.pause_document, + DatasetDocumentApplicationService.recover_document, + DatasetDocumentApplicationService.sync_website, + ], +) +def test_archived_documents_cannot_be_reprocessed( + service: DatasetDocumentApplicationService, operations: MagicMock, operation: Callable[..., object] +) -> None: + operations.get_state.return_value = replace(STATE, archived=True) + with pytest.raises(DocumentArchivedError): + operation(service, CONTEXT, dataset_id="dataset-1", document_id="document-1") + assert len(operations.mock_calls) == 1 + + +@pytest.mark.parametrize( + ("action", "status", "paused"), + [("pause", "indexing", True), ("resume", "paused", False), ("resume", "error", False)], +) +def test_processing_transition_records_explicit_actor( + service: DatasetDocumentApplicationService, + operations: MagicMock, + action: Literal["pause", "resume"], + status: str, + paused: bool, +) -> None: + operations.get_state.return_value = replace(STATE, indexing_status=status) + service.update_processing(CONTEXT, dataset_id="dataset-1", document_id="document-1", action=action) + ref, values = operations.update_document.call_args.args + assert ref == DOCUMENT + assert values["is_paused"] is paused + assert values["paused_by"] == (CONTEXT.account_id if paused else None) + assert (values["paused_at"] is not None) is paused + + +@pytest.mark.parametrize( + ("action", "status"), [("pause", "completed"), ("resume", "indexing"), ("invalid", "indexing")] +) +def test_invalid_processing_transition_does_not_write( + service: DatasetDocumentApplicationService, operations: MagicMock, action: Literal["pause", "resume"], status: str +) -> None: + operations.get_state.return_value = replace(STATE, indexing_status=status) + with pytest.raises(DocumentInvalidActionError): + service.update_processing(CONTEXT, dataset_id="dataset-1", document_id="document-1", action=action) + operations.update_document.assert_not_called() + + +@pytest.mark.parametrize( + ("doc_type", "expected"), [("book", {"title": "Title"}), ("others", {"title": "Title", "unknown": 12})] +) +def test_metadata_policy_filters_schema_but_preserves_custom_fields( + service: DatasetDocumentApplicationService, operations: MagicMock, doc_type: str | None, expected: dict[str, object] +) -> None: + service.update_metadata( + CONTEXT, + dataset_id="dataset-1", + document_id="document-1", + doc_type=doc_type, + doc_metadata={"title": "Title", "unknown": 12}, + ) + ref, values = operations.update_document.call_args.args + assert ref == DOCUMENT + assert values["doc_metadata"] == expected + assert values["doc_type"] == doc_type + + +@pytest.mark.parametrize(("doc_type", "metadata"), [(None, {}), ("book", None), ("invalid", {})]) +def test_invalid_metadata_never_writes( + service: DatasetDocumentApplicationService, + operations: MagicMock, + doc_type: str | None, + metadata: dict[str, object] | None, +) -> None: + with pytest.raises(ValueError): + service.update_metadata( + CONTEXT, dataset_id="dataset-1", document_id="document-1", doc_type=doc_type, doc_metadata=metadata + ) + operations.update_document.assert_not_called() + + +def test_retry_filters_missing_archived_completed_and_duplicate_documents( + service: DatasetDocumentApplicationService, operations: MagicMock +) -> None: + operations.get_states.return_value = [ + replace(STATE, id="error", indexing_status="error"), + replace(STATE, id="archived", archived=True), + replace(STATE, id="completed", indexing_status="completed"), + ] + service.retry_documents( + CONTEXT, dataset_id="dataset-1", document_ids=["archived", "missing", "error", "error", "completed"] + ) + operations.retry_documents.assert_called_once_with(DATASET, ["error"], actor_id="actor-1") + + +def test_summary_skips_qa_and_enables_summary_before_dispatch( + service: DatasetDocumentApplicationService, operations: MagicMock +) -> None: + operations.get_states.return_value = [ + STATE, + replace(STATE, id="qa", doc_form="qa_model"), + replace(STATE, id="ready", need_summary=True), + ] + service.generate_summary(CONTEXT, dataset_id="dataset-1", document_ids=["document-1", "qa", "ready"]) + operations.enable_summary.assert_called_once_with(DATASET, ["document-1"]) + assert [call.args[0] for call in operations.dispatch_summary.call_args_list] == [ + DOCUMENT, + DATASET.document("ready"), + ] + names = [call[0] for call in operations.mock_calls] + assert names.index("enable_summary") < names.index("dispatch_summary") + + +def test_summary_missing_document_cannot_partially_schedule( + service: DatasetDocumentApplicationService, operations: MagicMock +) -> None: + operations.get_states.return_value = [STATE] + with pytest.raises(DocumentNotFoundError, match="missing"): + service.generate_summary(CONTEXT, dataset_id="dataset-1", document_ids=["document-1", "missing"]) + operations.enable_summary.assert_not_called() + operations.dispatch_summary.assert_not_called() + + +def test_process_rule_resolves_document_inside_active_workspace( + service: DatasetDocumentApplicationService, operations: MagicMock, access: MagicMock +) -> None: + operations.find_document.return_value = DOCUMENT + service.get_process_rule(CONTEXT, document_id="document-1") + operations.find_document.assert_called_once_with(workspace_id="tenant-1", document_id="document-1") + access.require_accessible.assert_called_once_with(CONTEXT, "dataset-1") + operations.get_process_rule.assert_called_once_with(DATASET) + + +def test_default_process_rule_does_not_load_a_dataset( + service: DatasetDocumentApplicationService, operations: MagicMock, access: MagicMock +) -> None: + service.get_process_rule(CONTEXT, document_id=None) + access.require_accessible.assert_not_called() + operations.get_process_rule.assert_called_once_with(None) + + +def test_pause_requires_indexing_state_and_recover_requires_pause( + service: DatasetDocumentApplicationService, operations: MagicMock +) -> None: + operations.get_state.return_value = replace(STATE, indexing_status="completed") + with pytest.raises(DocumentIndexingStateError): + service.pause_document(CONTEXT, dataset_id="dataset-1", document_id="document-1") + with pytest.raises(DocumentIndexingStateError): + service.recover_document(CONTEXT, dataset_id="dataset-1", document_id="document-1") + operations.pause_document.assert_not_called() + operations.recover_document.assert_not_called() + + +def test_website_sync_rejects_other_sources(service: DatasetDocumentApplicationService, operations: MagicMock) -> None: + operations.get_state.return_value = replace(STATE, data_source_type="upload_file") + with pytest.raises(ValueError, match="not a website"): + service.sync_website(CONTEXT, dataset_id="dataset-1", document_id="document-1") + operations.sync_website.assert_not_called() diff --git a/api/tests/unit_tests/services/knowledge/test_document_creation.py b/api/tests/unit_tests/services/knowledge/test_document_creation.py new file mode 100644 index 00000000000..93ab1df3d0a --- /dev/null +++ b/api/tests/unit_tests/services/knowledge/test_document_creation.py @@ -0,0 +1,356 @@ +"""Exercise document creation with real SQL and assertions at external boundaries.""" + +from collections.abc import Callable, Generator, Iterator +from contextlib import contextmanager +from datetime import datetime +from unittest.mock import Mock + +import pytest +from redis.exceptions import LockNotOwnedError +from sqlalchemy import Engine, event, select +from sqlalchemy.orm import Session, sessionmaker + +from core.model_manager import ModelInstance +from core.rag.index_processor.constant.index_type import IndexTechniqueType +from enums import CloudPlan +from machinery.context import RequestContext +from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole +from models.dataset import Dataset, Document, DocumentSegment +from models.enums import IndexingStatus, SegmentStatus +from services.document_indexing_proxy.batch_indexing_base import BatchDocumentIndexingProxy +from services.document_indexing_proxy.document_indexing_task_proxy import DocumentIndexingTaskProxy +from services.document_indexing_proxy.duplicate_document_indexing_task_proxy import DuplicateDocumentIndexingTaskProxy +from services.entities.feature_entities import FeatureModel +from services.errors.file import FileNotExistsError +from services.knowledge.documents.adapters import SQLAlchemyDocumentOperations +from services.knowledge.resource_scope import DatasetRef +from tests.unit_tests.model_factories import make_upload_file + +CONTEXT = RequestContext("request", None, "actor", "tenant") +REF = DatasetRef("tenant", "dataset") + + +def settings(**values: object) -> dict[str, object]: + return { + "indexing_technique": "economy", + "process_rule": {"mode": "automatic"}, + "data_source": {"info_list": {"data_source_type": "upload_file", "file_info_list": {"file_ids": ["file"]}}}, + **values, + } + + +@pytest.fixture +def operations( + sqlite_session_factory: sessionmaker[Session], config_overrides: Callable[..., None] +) -> SQLAlchemyDocumentOperations: + config_overrides(DEPLOYMENT_EDITION="CLOUD", RBAC_ENABLED=False) + with sqlite_session_factory.begin() as session: + actor = Account(name="Actor", email="actor@example.com") + actor.id = "actor" + tenant = Tenant(name="Tenant") + tenant.id = "tenant" + session.add_all( + [ + actor, + tenant, + TenantAccountJoin(tenant_id="tenant", account_id="actor", role=TenantAccountRole.OWNER), + Dataset( + id="dataset", tenant_id="tenant", name="Dataset", created_by="actor", indexing_technique="economy" + ), + make_upload_file(file_id="file", tenant_id="tenant", name="example.txt", created_by="actor"), + make_upload_file(file_id="foreign-file", tenant_id="foreign", name="foreign.txt"), + ] + ) + return SQLAlchemyDocumentOperations(session_factory=sqlite_session_factory) + + +@pytest.fixture +def outside_transaction(sqlite_engine: Engine) -> Iterator[Callable[[], None]]: + active: set[int] = set() + + def begin(connection: object) -> None: + active.add(id(connection)) + + def end(connection: object) -> None: + active.discard(id(connection)) + + def check() -> None: + assert not active, "External I/O must not hold an open database transaction" + + listeners = [("begin", begin), ("commit", end), ("rollback", end)] + for name, callback in listeners: + event.listen(sqlite_engine, name, callback) + yield check + for name, callback in listeners: + event.remove(sqlite_engine, name, callback) + + +@pytest.fixture +def external_calls( + monkeypatch: pytest.MonkeyPatch, + outside_transaction: Callable[[], None], + sqlite_session_factory: sessionmaker[Session], +) -> list[tuple[str, tuple[str, ...]]]: + calls: list[tuple[str, tuple[str, ...]]] = [] + + def features(_tenant_id: str, *, exclude_vector_space: bool) -> FeatureModel: + outside_transaction() + assert exclude_vector_space is True + result = FeatureModel() + result.billing.subscription.plan = CloudPlan.PROFESSIONAL + return result + + @contextmanager + def lock(_name: str, *, timeout: int) -> Generator[None, None, None]: + outside_transaction() + assert timeout == 600 + yield + outside_transaction() + + def dispatched(kind: str, dataset_id: str, document_ids: tuple[str, ...]) -> None: + outside_transaction() + with sqlite_session_factory() as session: + for document_id in document_ids: + document = session.get(Document, document_id) + assert document is not None + assert document.dataset_id == dataset_id + if kind != "clean": + assert document.indexing_status == IndexingStatus.WAITING + calls.append((kind, document_ids)) + + def create(proxy: BatchDocumentIndexingProxy) -> None: + dispatched("create", proxy._dataset_id, tuple(proxy._document_ids)) + + def duplicate(proxy: BatchDocumentIndexingProxy) -> None: + dispatched("duplicate", proxy._dataset_id, tuple(proxy._document_ids)) + + monkeypatch.setattr("services.knowledge.dataset_service.FeatureService.get_features", features) + monkeypatch.setattr("services.knowledge.documents.adapters.redis_client.lock", lock) + monkeypatch.setattr(DocumentIndexingTaskProxy, "delay", create) + monkeypatch.setattr(DuplicateDocumentIndexingTaskProxy, "delay", duplicate) + monkeypatch.setattr( + "services.knowledge.dataset_service.document_indexing_update_task.delay", + lambda dataset_id, document_id: dispatched("update", dataset_id, (document_id,)), + ) + monkeypatch.setattr( + "services.knowledge.dataset_service.clean_notion_document_task.delay", + lambda document_ids, dataset_id: dispatched("clean", dataset_id, tuple(document_ids)), + ) + return calls + + +@pytest.mark.parametrize("initialize", [False, True]) +def test_upload_preserves_response_and_commits_before_dispatch( + operations: SQLAlchemyDocumentOperations, + external_calls: list[tuple[str, tuple[str, ...]]], + sqlite_session_factory: sessionmaker[Session], + initialize: bool, +) -> None: + result = ( + operations.initialize_dataset(CONTEXT, settings()) + if initialize + else operations.create_documents(CONTEXT, REF, settings()) + ) + (document,) = result["documents"] + assert document["name"] == "example.txt" + assert document["data_source_info_dict"] == {"upload_file_id": "file"} + assert document["indexing_status"] == "waiting" + assert result["batch"] + assert external_calls == [("create", (document["id"],))] + with sqlite_session_factory() as session: + dataset = session.get(Dataset, result["dataset"]["id"]) + assert dataset is not None + assert dataset.tenant_id == "tenant" + assert dataset.name == ("example.txt..." if initialize else "Dataset") + assert dataset.permission == "only_me" + + +def test_duplicate_upload_reuses_document_and_dispatches_duplicate_job( + operations: SQLAlchemyDocumentOperations, external_calls: list[tuple[str, tuple[str, ...]]] +) -> None: + first = operations.create_documents(CONTEXT, REF, settings()) + second = operations.create_documents(CONTEXT, REF, settings()) + first_id = first["documents"][0]["id"] + assert second["documents"][0]["id"] == first_id + assert external_calls == [("create", (first_id,)), ("duplicate", (first_id,))] + + +def test_reprocess_checks_model_outside_transaction_and_preserves_segment_reset( + operations: SQLAlchemyDocumentOperations, + external_calls: list[tuple[str, tuple[str, ...]]], + outside_transaction: Callable[[], None], + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, +) -> None: + created = operations.create_documents(CONTEXT, REF, settings()) + document_id = created["documents"][0]["id"] + with sqlite_session_factory.begin() as session: + document = session.get(Document, document_id) + assert document is not None + document.indexing_status = IndexingStatus.COMPLETED + document.completed_at = datetime(2026, 1, 1) + dataset = session.get(Dataset, REF.dataset_id) + assert dataset is not None + dataset.indexing_technique = IndexTechniqueType.HIGH_QUALITY + dataset.embedding_model = "embedding" + dataset.embedding_model_provider = "provider" + session.add( + DocumentSegment( + tenant_id="tenant", + dataset_id="dataset", + document_id=document_id, + position=1, + content="content", + word_count=1, + tokens=1, + created_by="actor", + status=SegmentStatus.COMPLETED, + ) + ) + manager = Mock() + manager.get_model_instance.side_effect = lambda **_kwargs: outside_transaction() + monkeypatch.setattr("services.knowledge.dataset_service.ModelManager.for_tenant", lambda **_kwargs: manager) + result = operations.create_documents(CONTEXT, REF, settings(original_document_id=document_id, name="Renamed")) + assert result["documents"][0]["id"] == document_id + assert result["documents"][0]["name"] == "Renamed" + assert result["batch"] == created["batch"] + assert external_calls[-1] == ("update", (document_id,)) + manager.get_model_instance.assert_called_once() + with sqlite_session_factory() as session: + segment = session.scalar(select(DocumentSegment).where(DocumentSegment.document_id == document_id)) + assert segment is not None + assert segment.status == SegmentStatus.RE_SEGMENT + + +def test_first_high_quality_upload_resolves_default_model_without_transaction( + operations: SQLAlchemyDocumentOperations, + external_calls: list[tuple[str, tuple[str, ...]]], + outside_transaction: Callable[[], None], + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, +) -> None: + with sqlite_session_factory.begin() as session: + dataset = session.get(Dataset, REF.dataset_id) + assert dataset is not None + dataset.indexing_technique = None + + def model(**_kwargs: object) -> ModelInstance: + outside_transaction() + instance = object.__new__(ModelInstance) + instance.model_name = "default-embedding" + instance.provider = "default-provider" + return instance + + manager = Mock() + manager.get_default_model_instance.side_effect = model + monkeypatch.setattr("services.knowledge.dataset_service.ModelManager.for_tenant", lambda **_kwargs: manager) + operations.create_documents(CONTEXT, REF, settings(indexing_technique="high_quality")) + with sqlite_session_factory() as session: + dataset = session.get(Dataset, REF.dataset_id) + assert dataset is not None + assert dataset.embedding_model == "default-embedding" + assert dataset.embedding_model_provider == "default-provider" + assert dataset.collection_binding_id is not None + assert len(external_calls) == 1 + + +def test_foreign_upload_is_rejected_without_dispatch( + operations: SQLAlchemyDocumentOperations, external_calls: list[tuple[str, tuple[str, ...]]] +) -> None: + config = settings( + data_source={"info_list": {"data_source_type": "upload_file", "file_info_list": {"file_ids": ["foreign-file"]}}} + ) + with pytest.raises(FileNotExistsError): + operations.create_documents(CONTEXT, REF, config) + assert external_calls == [] + + +def test_notion_removal_is_dispatched_after_new_documents_commit( + operations: SQLAlchemyDocumentOperations, external_calls: list[tuple[str, tuple[str, ...]]] +) -> None: + def notion(page_id: str) -> dict[str, object]: + return settings( + data_source={ + "info_list": { + "data_source_type": "notion_import", + "notion_info_list": [ + { + "credential_id": "credential", + "workspace_id": "workspace", + "pages": [{"page_id": page_id, "page_name": page_id, "type": "page"}], + } + ], + } + } + ) + + first = operations.create_documents(CONTEXT, REF, notion("old-page")) + old_id = first["documents"][0]["id"] + operations.create_documents(CONTEXT, REF, notion("new-page")) + assert external_calls[-2] == ("clean", (old_id,)) + assert external_calls[-1][0] == "create" + assert old_id not in external_calls[-1][1] + + +def test_cloud_quota_rejection_does_not_create_documents_or_dispatch( + operations: SQLAlchemyDocumentOperations, + external_calls: list[tuple[str, tuple[str, ...]]], + sqlite_session_factory: sessionmaker[Session], + config_overrides: Callable[..., None], +) -> None: + config_overrides(BATCH_UPLOAD_LIMIT=0) + with pytest.raises(ValueError, match="batch upload limit"): + operations.create_documents(CONTEXT, REF, settings()) + assert external_calls == [] + with sqlite_session_factory() as session: + assert session.scalar(select(Document.id)) is None + + +def test_response_read_failure_still_dispatches_already_committed_documents( + operations: SQLAlchemyDocumentOperations, + external_calls: list[tuple[str, tuple[str, ...]]], + monkeypatch: pytest.MonkeyPatch, +) -> None: + response = Mock(side_effect=ValueError("response read failed")) + monkeypatch.setattr("services.knowledge.documents.adapters._created_response", response) + with pytest.raises(ValueError, match="response read failed"): + operations.create_documents(CONTEXT, REF, settings()) + assert len(external_calls) == 1 + assert external_calls[0][0] == "create" + + +def test_lost_creation_lock_keeps_empty_batch_response( + operations: SQLAlchemyDocumentOperations, + external_calls: list[tuple[str, tuple[str, ...]]], + monkeypatch: pytest.MonkeyPatch, +) -> None: + @contextmanager + def lock(_name: str, **_kwargs: object) -> Generator[None, None, None]: + raise LockNotOwnedError("lock lost") + yield # pragma: no cover + + monkeypatch.setattr("services.knowledge.documents.adapters.redis_client.lock", lock) + result = operations.create_documents(CONTEXT, REF, settings()) + assert result["documents"] == [] + assert result["batch"] + assert result["dataset"]["id"] == REF.dataset_id + assert external_calls == [] + + +def test_billing_http_boundary_does_not_hold_a_transaction( + operations: SQLAlchemyDocumentOperations, + outside_transaction: Callable[[], None], + monkeypatch: pytest.MonkeyPatch, +) -> None: + class StopBeforeNetworkError(Exception): + pass + + def request(_method: str, url: str, **_kwargs: object) -> None: + assert url.endswith("/subscription/info") + outside_transaction() + raise StopBeforeNetworkError() + + monkeypatch.setattr("services.billing_service._http_client.request", request) + with pytest.raises(StopBeforeNetworkError): + operations.create_documents(CONTEXT, REF, settings()) + outside_transaction() diff --git a/api/tests/unit_tests/services/knowledge/test_document_index_cleanup.py b/api/tests/unit_tests/services/knowledge/test_document_index_cleanup.py new file mode 100644 index 00000000000..b0885001843 --- /dev/null +++ b/api/tests/unit_tests/services/knowledge/test_document_index_cleanup.py @@ -0,0 +1,279 @@ +"""Document cleanup keeps SQL transactions closed during external index deletion.""" + +import json +from unittest.mock import MagicMock +from uuid import uuid4 + +import pytest +from sqlalchemy.orm import Session, sessionmaker + +from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType +from models.dataset import ChildChunk, Dataset, DatasetKeywordTable, Document, DocumentSegment, DocumentSegmentSummary +from models.enums import DataSourceType, DocumentCreatedFrom +from services.knowledge.indexing.adapters import cleanup +from services.knowledge.summaries import adapters as summaries +from tasks.batch_clean_document_task import batch_clean_document_task +from tasks.clean_document_task import clean_document_task +from tasks.clean_notion_document_task import clean_notion_document_task + + +@pytest.fixture +def graph(sqlite_session: Session) -> tuple[Dataset, Document, DocumentSegment, DocumentSegmentSummary, ChildChunk]: + dataset = Dataset( + id=str(uuid4()), + tenant_id=str(uuid4()), + name="Dataset", + created_by="admin", + indexing_technique="high_quality", + index_struct='{"type":"pgvector"}', + ) + document = Document( + id=str(uuid4()), + tenant_id=dataset.tenant_id, + dataset_id=dataset.id, + position=1, + data_source_type=DataSourceType.NOTION_IMPORT, + batch="batch", + name="Document", + created_from=DocumentCreatedFrom.WEB, + created_by="admin", + doc_form="text_model", + ) + segment = DocumentSegment( + tenant_id=dataset.tenant_id, + dataset_id=dataset.id, + document_id=document.id, + position=1, + content="Text", + word_count=1, + tokens=1, + created_by="admin", + index_node_id="body-node", + ) + summary = DocumentSegmentSummary( + dataset_id=dataset.id, + document_id=document.id, + chunk_id=segment.id, + summary_content="Summary", + summary_index_node_id="summary-node", + ) + child = ChildChunk( + tenant_id=dataset.tenant_id, + dataset_id=dataset.id, + document_id=document.id, + segment_id=segment.id, + position=1, + content="Child", + word_count=1, + index_node_id="child-node", + created_by="admin", + index_node_hash="child-hash", + ) + sqlite_session.add_all([dataset, document, segment, summary, child]) + sqlite_session.commit() + return dataset, document, segment, summary, child + + +@pytest.fixture +def vector(monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session]) -> MagicMock: + sessions: list[Session] = [] + + def new_session() -> Session: + session = sqlite_session_factory() + sessions.append(session) + return session + + def assert_closed(_node_ids: list[str]) -> None: + assert sessions + assert all(not session.in_transaction() for session in sessions) + + backend = MagicMock() + backend.delete_by_ids.side_effect = assert_closed + monkeypatch.setattr(cleanup, "Vector", MagicMock(return_value=backend)) + monkeypatch.setattr("core.db.session_factory.session_factory.create_session", new_session) + return backend + + +@pytest.mark.parametrize("task_kind", ["single", "batch", "notion"]) +@pytest.mark.parametrize("doc_form", [IndexStructureType.PARAGRAPH_INDEX, IndexStructureType.PARENT_CHILD_INDEX]) +def test_tasks_delete_summary_and_body_indexes_without_an_open_transaction( + graph: tuple[Dataset, Document, DocumentSegment, DocumentSegmentSummary, ChildChunk], + sqlite_session: Session, + vector: MagicMock, + monkeypatch: pytest.MonkeyPatch, + task_kind: str, + doc_form: IndexStructureType, +) -> None: + dataset, document, segment, summary, child = graph + document.doc_form = doc_form + dataset.chunk_structure = doc_form + ids = (dataset.id, document.id, segment.id, summary.id, child.id) + if task_kind != "notion": + sqlite_session.delete(document) + sqlite_session.commit() + monkeypatch.setattr("tasks.clean_document_task.schedule_billing_vector_space_refresh", lambda _tenant_id: None) + monkeypatch.setattr( + "tasks.batch_clean_document_task.schedule_billing_vector_space_refresh", lambda _tenant_id: None + ) + + if task_kind == "single": + clean_document_task(ids[1], ids[0], doc_form, None) + elif task_kind == "batch": + batch_clean_document_task([ids[1]], ids[0], doc_form, []) + else: + clean_notion_document_task([ids[1]], ids[0]) + + vector.delete_by_ids.assert_called_once_with( + ["summary-node", "child-node" if doc_form == IndexStructureType.PARENT_CHILD_INDEX else "body-node"] + ) + sqlite_session.expire_all() + assert sqlite_session.get(DocumentSegmentSummary, ids[3]) is None + assert sqlite_session.get(ChildChunk, ids[4]) is None + assert sqlite_session.get(DocumentSegment, ids[2]) is None + + +def test_failed_vector_cleanup_preserves_records_for_retry( + graph: tuple[Dataset, Document, DocumentSegment, DocumentSegmentSummary, ChildChunk], + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + vector: MagicMock, +) -> None: + dataset, document, _, summary, child = graph + vector.delete_by_ids.side_effect = RuntimeError("unavailable") + with pytest.raises(RuntimeError, match="unavailable"): + cleanup.clean_document_indexes( + dataset_id=dataset.id, document_ids=[document.id], doc_form="text_model", new_session=sqlite_session_factory + ) + assert sqlite_session.get(DocumentSegmentSummary, summary.id) is not None + assert sqlite_session.get(ChildChunk, child.id) is not None + + +@pytest.mark.parametrize("task_kind", ["single", "batch", "notion"]) +def test_cleanup_retries_summaries_after_segments_have_been_deleted( + graph: tuple[Dataset, Document, DocumentSegment, DocumentSegmentSummary, ChildChunk], + sqlite_session: Session, + vector: MagicMock, + monkeypatch: pytest.MonkeyPatch, + task_kind: str, +) -> None: + dataset, document, segment, summary, _ = graph + summary_id = summary.id + sqlite_session.delete(segment) + sqlite_session.commit() + monkeypatch.setattr("tasks.clean_document_task.schedule_billing_vector_space_refresh", lambda _tenant_id: None) + monkeypatch.setattr( + "tasks.batch_clean_document_task.schedule_billing_vector_space_refresh", lambda _tenant_id: None + ) + if task_kind == "single": + clean_document_task(document.id, dataset.id, "text_model", None) + elif task_kind == "batch": + batch_clean_document_task([document.id], dataset.id, "text_model", []) + else: + clean_notion_document_task([document.id], dataset.id) + vector.delete_by_ids.assert_called_once_with(["summary-node"]) + sqlite_session.expire_all() + assert sqlite_session.get(DocumentSegmentSummary, summary_id) is None + + +def test_cleanup_keeps_other_documents_and_datasets( + graph: tuple[Dataset, Document, DocumentSegment, DocumentSegmentSummary, ChildChunk], + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, +) -> None: + dataset, document, _, summary, _ = graph + summary_id = summary.id + retained = [ + DocumentSegmentSummary( + dataset_id=owner_dataset, + document_id=owner_document, + chunk_id=str(uuid4()), + summary_content="Retained", + summary_index_node_id=f"retained-{index}", + ) + for index, (owner_dataset, owner_document) in enumerate( + [ + (dataset.id, str(uuid4())), + (str(uuid4()), document.id), + ] + ) + ] + sqlite_session.add_all(retained) + sqlite_session.commit() + backend = MagicMock() + monkeypatch.setattr(cleanup, "Vector", MagicMock(return_value=backend)) + cleanup.clean_document_indexes( + dataset_id=dataset.id, document_ids=[document.id], doc_form="text_model", new_session=sqlite_session_factory + ) + backend.delete_by_ids.assert_called_once_with(["summary-node", "body-node"]) + sqlite_session.expire_all() + assert sqlite_session.get(DocumentSegmentSummary, summary_id) is None + for record in retained: + assert sqlite_session.get(DocumentSegmentSummary, record.id) is not None + + +def test_economy_cleanup_preserves_unrelated_keyword_entries( + graph: tuple[Dataset, Document, DocumentSegment, DocumentSegmentSummary, ChildChunk], + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], +) -> None: + dataset, document, _, summary, _ = graph + summary_id = summary.id + dataset.indexing_technique = IndexTechniqueType.ECONOMY + keyword_table = DatasetKeywordTable( + dataset_id=dataset.id, + keyword_table=json.dumps({"__data__": {"table": {"term": ["body-node", "unrelated-node"]}}}), + ) + sqlite_session.add(keyword_table) + sqlite_session.commit() + cleanup.clean_document_indexes( + dataset_id=dataset.id, document_ids=[document.id], doc_form="text_model", new_session=sqlite_session_factory + ) + sqlite_session.refresh(keyword_table) + assert json.loads(keyword_table.keyword_table)["__data__"]["table"] == {"term": ["unrelated-node"]} + sqlite_session.expire_all() + assert sqlite_session.get(DocumentSegmentSummary, summary_id) is None + + +def test_legacy_summary_delete_does_not_commit_callers_transaction( + graph: tuple[Dataset, Document, DocumentSegment, DocumentSegmentSummary, ChildChunk], + sqlite_session: Session, + monkeypatch: pytest.MonkeyPatch, +) -> None: + dataset, _, segment, summary, _ = graph + dataset_id, summary_id = dataset.id, summary.id + monkeypatch.setattr(summaries, "Vector", MagicMock()) + + def rollback_delete() -> None: + with sqlite_session.begin(): + dataset.name = "Uncommitted name" + summaries.SummaryIndexAdapter.delete_summaries_for_segments(dataset, [segment.id], session=sqlite_session) + assert sqlite_session.in_transaction() + assert sqlite_session.get(DocumentSegmentSummary, summary_id) is None + raise RuntimeError("rollback caller") + + with pytest.raises(RuntimeError, match="rollback caller"): + rollback_delete() + assert sqlite_session.get(DocumentSegmentSummary, summary_id) is not None + stored_dataset = sqlite_session.get(Dataset, dataset_id) + assert stored_dataset is not None + assert stored_dataset.name == "Dataset" + + +def test_empty_document_selection_does_not_clear_the_dataset_index( + graph: tuple[Dataset, Document, DocumentSegment, DocumentSegmentSummary, ChildChunk], + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + vector: MagicMock, +) -> None: + dataset, _, segment, summary, child = graph + assert ( + cleanup.clean_document_indexes( + dataset_id=dataset.id, document_ids=[], doc_form="text_model", new_session=sqlite_session_factory + ) + is None + ) + vector.delete_by_ids.assert_not_called() + assert sqlite_session.get(DocumentSegment, segment.id) is not None + assert sqlite_session.get(DocumentSegmentSummary, summary.id) is not None + assert sqlite_session.get(ChildChunk, child.id) is not None diff --git a/api/tests/unit_tests/services/knowledge/test_document_sync_application_service.py b/api/tests/unit_tests/services/knowledge/test_document_sync_application_service.py new file mode 100644 index 00000000000..e7e2fcd1d8f --- /dev/null +++ b/api/tests/unit_tests/services/knowledge/test_document_sync_application_service.py @@ -0,0 +1,150 @@ +from dataclasses import dataclass, field + +import pytest + +from machinery.context import RequestContext +from services.knowledge.dataset_access import AccessibleDataset, DatasetAccessDeniedError, DatasetAccessSnapshot +from services.knowledge.document_sync import ( + DocumentSyncApplicationService, + SyncDocumentNotFoundError, + SyncDocumentRecord, + SyncDocumentSourceError, +) +from services.knowledge.document_sync_adapters import CeleryDocumentSyncDispatcher +from services.knowledge.resource_scope import DatasetRef, DocumentRef + + +def _context() -> RequestContext: + return RequestContext("request-1", None, "account-1", "workspace-1") + + +def _dataset() -> AccessibleDataset: + return AccessibleDataset( + id="dataset-1", + workspace_id="workspace-1", + ) + + +def _document(*, data_source_type: str = "notion_import") -> SyncDocumentRecord: + return SyncDocumentRecord( + id="document-1", + data_source_type=data_source_type, + ) + + +@dataclass +class DatasetAccessStub: + error: Exception | None = None + calls: list[str] = field(default_factory=list) + + def require_accessible(self, context: RequestContext, dataset_id: str) -> AccessibleDataset: + assert context == _context() + self.calls.append(dataset_id) + if self.error is not None: + raise self.error + return _dataset() + + def check_access(self, context: RequestContext, snapshot: DatasetAccessSnapshot) -> AccessibleDataset: + raise AssertionError(f"Unexpected snapshot access for dataset {snapshot.dataset.id} by {context.account_id}") + + +@dataclass +class DocumentReaderStub: + document: SyncDocumentRecord | None = field(default_factory=_document) + active_refs: tuple[DocumentRef, ...] = () + list_calls: list[DatasetRef] = field(default_factory=list) + get_calls: list[DocumentRef] = field(default_factory=list) + + def list_active_notion_refs(self, dataset_ref: DatasetRef) -> tuple[DocumentRef, ...]: + self.list_calls.append(dataset_ref) + return self.active_refs + + def get_sync_document(self, document_ref: DocumentRef) -> SyncDocumentRecord | None: + self.get_calls.append(document_ref) + return self.document + + +@dataclass +class DispatcherRecorder: + calls: list[tuple[str, str]] = field(default_factory=list) + + def dispatch(self, *, dataset_id: str, document_id: str) -> None: + self.calls.append((dataset_id, document_id)) + + +def test_celery_dispatcher_forwards_owned_identifiers_to_injected_delay() -> None: + calls: list[tuple[str, str]] = [] + dispatcher = CeleryDocumentSyncDispatcher( + delay=lambda dataset_id, document_id: calls.append((dataset_id, document_id)) + ) + + dispatcher.dispatch(dataset_id="dataset-1", document_id="document-1") + + assert calls == [("dataset-1", "document-1")] + + +def test_sync_dataset_dispatches_snapshot_after_tenant_scoped_read() -> None: + dataset_ref = DatasetRef("workspace-1", "dataset-1") + documents = DocumentReaderStub(active_refs=(dataset_ref.document("document-1"), dataset_ref.document("document-2"))) + dispatcher = DispatcherRecorder() + service = DocumentSyncApplicationService( + dataset_access=DatasetAccessStub(), documents=documents, dispatcher=dispatcher + ) + + count = service.sync_dataset(_context(), "dataset-1") + + assert count == 2 + assert documents.list_calls == [dataset_ref] + assert dispatcher.calls == [("dataset-1", "document-1"), ("dataset-1", "document-2")] + + +def test_sync_document_validates_owner_chain_and_source_before_dispatch() -> None: + documents = DocumentReaderStub() + dispatcher = DispatcherRecorder() + service = DocumentSyncApplicationService( + dataset_access=DatasetAccessStub(), documents=documents, dispatcher=dispatcher + ) + + service.sync_document(_context(), "dataset-1", "document-1") + + assert documents.get_calls == [DatasetRef("workspace-1", "dataset-1").document("document-1")] + assert dispatcher.calls == [("dataset-1", "document-1")] + + +@pytest.mark.parametrize( + ("document", "error"), + [ + (None, SyncDocumentNotFoundError), + (_document(data_source_type="upload_file"), SyncDocumentSourceError), + ], +) +def test_sync_document_rejects_invalid_document_before_dispatch( + document: SyncDocumentRecord | None, error: type[Exception] +) -> None: + dispatcher = DispatcherRecorder() + service = DocumentSyncApplicationService( + dataset_access=DatasetAccessStub(), + documents=DocumentReaderStub(document=document), + dispatcher=dispatcher, + ) + + with pytest.raises(error): + service.sync_document(_context(), "dataset-1", "document-1") + + assert dispatcher.calls == [] + + +def test_sync_stops_before_document_read_and_dispatch_when_dataset_access_is_denied() -> None: + documents = DocumentReaderStub() + dispatcher = DispatcherRecorder() + service = DocumentSyncApplicationService( + dataset_access=DatasetAccessStub(error=DatasetAccessDeniedError()), + documents=documents, + dispatcher=dispatcher, + ) + + with pytest.raises(DatasetAccessDeniedError): + service.sync_document(_context(), "dataset-1", "document-1") + + assert documents.get_calls == [] + assert dispatcher.calls == [] diff --git a/api/tests/unit_tests/services/knowledge/test_external_adapters.py b/api/tests/unit_tests/services/knowledge/test_external_adapters.py new file mode 100644 index 00000000000..dc70a8baccd --- /dev/null +++ b/api/tests/unit_tests/services/knowledge/test_external_adapters.py @@ -0,0 +1,203 @@ +import json +from collections.abc import Iterator +from unittest.mock import patch + +import httpx +import pytest +from sqlalchemy import Engine, event, func, select +from sqlalchemy.orm import Session, sessionmaker + +from constants import HIDDEN_VALUE +from controllers.console.datasets.external import ExternalKnowledgeApiListResponse +from machinery.context import RequestContext +from models.dataset import Dataset, DatasetQuery, ExternalKnowledgeApis, ExternalKnowledgeBindings +from services.entities.external_knowledge_entities.external_knowledge_entities import ExternalDatasetCreatePayload +from services.errors.knowledge_retrieval import ExternalKnowledgeRetrievalError +from services.knowledge.dataset_access import DatasetNotFoundError +from services.knowledge.external.adapters import SQLAlchemyExternalKnowledgeOperations +from services.knowledge.external.application import ExternalTemplateNotFoundError +from services.knowledge.resource_scope import DatasetRef + +CONTEXT = RequestContext("request", None, "actor", "tenant") +REF = DatasetRef("tenant", "dataset") + + +@pytest.fixture +def operations(sqlite_session_factory: sessionmaker[Session]) -> SQLAlchemyExternalKnowledgeOperations: + with sqlite_session_factory.begin() as session: + for workspace, template_id, dataset_id in [ + ("tenant", "template", "dataset"), + ("other", "foreign", "foreign-dataset"), + ]: + template = ExternalKnowledgeApis( + tenant_id=workspace, + created_by="actor", + updated_by="actor", + name="API", + description="", + settings=json.dumps({"endpoint": "https://example", "api_key": "secret"}), + ) + template.id = template_id + session.add_all( + [ + template, + Dataset( + id=dataset_id, tenant_id=workspace, created_by="actor", name="Dataset", provider="external" + ), + ExternalKnowledgeBindings( + tenant_id=workspace, + dataset_id=dataset_id, + external_knowledge_api_id=template_id, + external_knowledge_id="knowledge", + created_by="actor", + ), + ] + ) + return SQLAlchemyExternalKnowledgeOperations(session_factory=sqlite_session_factory) + + +@pytest.fixture +def checked_out(sqlite_engine: Engine) -> Iterator[set[object]]: + active: set[object] = set() + + def checkout(_connection: object, record: object, _proxy: object) -> None: + active.add(record) + + def checkin(_connection: object, record: object) -> None: + active.discard(record) + + event.listen(sqlite_engine, "checkout", checkout) + event.listen(sqlite_engine, "checkin", checkin) + try: + yield active + finally: + event.remove(sqlite_engine, "checkout", checkout) + event.remove(sqlite_engine, "checkin", checkin) + + +def test_template_list_materializes_bindings_and_settings( + operations: SQLAlchemyExternalKnowledgeOperations, +) -> None: + result = ExternalKnowledgeApiListResponse.model_validate( + operations.list_templates("tenant", page=0, limit=0, keyword=None) + ).model_dump(mode="json") + assert result["total"] == 1 + assert result["limit"] == 1 + assert result["has_more"] is False + assert result["data"][0]["settings"] == {"endpoint": "https://example", "api_key": "secret"} + assert result["data"][0]["dataset_bindings"] == [{"id": "dataset", "name": "Dataset"}] + + +@pytest.mark.parametrize("method", ["get_template", "update_template", "delete_template"]) +def test_template_read_write_scope(operations: SQLAlchemyExternalKnowledgeOperations, method: str) -> None: + if method == "update_template": + with pytest.raises(ExternalTemplateNotFoundError): + operations.update_template( + CONTEXT, "foreign", name="changed", settings={"api_key": "key", "endpoint": "https://example"} + ) + else: + operation = operations.get_template if method == "get_template" else operations.delete_template + with pytest.raises(ExternalTemplateNotFoundError): + operation("tenant", "foreign") + assert operations.template_usage("tenant", "foreign") == (False, 0) + + +def test_masked_update_preserves_secret_without_mutating_request( + operations: SQLAlchemyExternalKnowledgeOperations, sqlite_session_factory: sessionmaker[Session] +) -> None: + settings = {"endpoint": "https://changed", "api_key": HIDDEN_VALUE} + operations.update_template(CONTEXT, "template", name="Renamed", settings=settings) + assert settings["api_key"] == HIDDEN_VALUE + with sqlite_session_factory() as session: + template = session.get(ExternalKnowledgeApis, "template") + assert template is not None + assert template.name == "Renamed" + settings = template.settings_dict + assert settings is not None + assert settings["api_key"] == "secret" + + +def test_creation_probes_outside_transaction_then_commits( + operations: SQLAlchemyExternalKnowledgeOperations, + checked_out: set[object], + sqlite_session_factory: sessionmaker[Session], +) -> None: + def probe(_settings: dict[str, object]) -> None: + assert not checked_out + + with patch( + "services.knowledge.external.service.ExternalDatasetService.check_endpoint_and_api_key", side_effect=probe + ): + result = operations.create_template( + CONTEXT, name="New", settings={"endpoint": "https://example", "api_key": "secret"} + ) + with sqlite_session_factory() as session: + assert session.get(ExternalKnowledgeApis, result["id"]) is not None + + +def test_external_dataset_cannot_bind_foreign_template(operations: SQLAlchemyExternalKnowledgeOperations) -> None: + with pytest.raises(ExternalTemplateNotFoundError): + operations.create_dataset( + CONTEXT, + ExternalDatasetCreatePayload( + name="New", external_knowledge_api_id="foreign", external_knowledge_id="knowledge" + ), + ) + + +def test_retrieval_releases_read_session_before_http_and_commits_query( + operations: SQLAlchemyExternalKnowledgeOperations, + checked_out: set[object], + sqlite_session_factory: sessionmaker[Session], +) -> None: + def send(_request: object, _files: object) -> httpx.Response: + assert not checked_out + return httpx.Response( + 200, + json={"records": [{"content": "answer", "title": "title", "score": 0.9, "metadata": {"source": "doc"}}]}, + ) + + with patch( + "services.knowledge.external.service.ExternalDatasetService.process_external_api", side_effect=send + ) as request: + result = operations.retrieve( + CONTEXT, REF, query='"question"', retrieval_model={"top_k": 3}, metadata_filters=None + ) + assert result["query"] == {"content": '"question"'} + assert result["records"][0]["content"] == "answer" + prepared = request.call_args.args[0] + assert prepared.params["query"] == '\\"question\\"' + assert prepared.params["knowledge_id"] == "knowledge" + assert prepared.params["retrieval_setting"] == {"top_k": 3, "score_threshold": 0.0} + with sqlite_session_factory() as session: + saved = session.scalar(select(DatasetQuery)) + assert saved is not None + assert saved.created_by == "actor" + assert saved.content == '"question"' + assert saved.dataset_id == "dataset" + + +def test_retrieval_failure_does_not_record_query( + operations: SQLAlchemyExternalKnowledgeOperations, sqlite_session_factory: sessionmaker[Session] +) -> None: + with patch( + "services.knowledge.external.service.ExternalDatasetService.process_external_api", + return_value=httpx.Response(502, text="unavailable"), + ): + with pytest.raises(ExternalKnowledgeRetrievalError): + operations.retrieve(CONTEXT, REF, query="query", retrieval_model=None, metadata_filters=None) + with sqlite_session_factory() as session: + assert session.scalar(select(func.count()).select_from(DatasetQuery)) == 0 + + +def test_retrieval_rejects_foreign_dataset_before_network(operations: SQLAlchemyExternalKnowledgeOperations) -> None: + with patch("services.knowledge.external.service.ExternalDatasetService.process_external_api") as request: + with pytest.raises(DatasetNotFoundError): + operations.retrieve( + CONTEXT, + DatasetRef("tenant", "foreign-dataset"), + query="query", + retrieval_model=None, + metadata_filters=None, + ) + request.assert_not_called() diff --git a/api/tests/unit_tests/services/knowledge/test_external_application_service.py b/api/tests/unit_tests/services/knowledge/test_external_application_service.py new file mode 100644 index 00000000000..5d98b8581b6 --- /dev/null +++ b/api/tests/unit_tests/services/knowledge/test_external_application_service.py @@ -0,0 +1,81 @@ +from collections.abc import Callable +from unittest.mock import MagicMock, create_autospec + +import pytest + +from machinery.context import RequestContext +from services.knowledge.dataset_access import AccessibleDataset, DatasetAccess, DatasetAccessDeniedError +from services.knowledge.external.application import ( + ExternalHitTestingError, + ExternalKnowledgeApplicationService, + ExternalKnowledgeOperations, +) +from services.knowledge.resource_scope import DatasetRef + +CONTEXT = RequestContext("request", None, "actor", "tenant") + + +@pytest.fixture +def operations() -> MagicMock: + return create_autospec(ExternalKnowledgeOperations, instance=True, spec_set=True) + + +@pytest.fixture +def access() -> MagicMock: + result = create_autospec(DatasetAccess, instance=True, spec_set=True) + result.require_accessible.return_value = AccessibleDataset("dataset", "tenant") + return result + + +@pytest.fixture +def service(operations: MagicMock, access: MagicMock) -> ExternalKnowledgeApplicationService: + return ExternalKnowledgeApplicationService(dataset_access=access, operations=operations) + + +@pytest.mark.parametrize("settings", [{}, {"endpoint": "https://example"}, {"api_key": "secret"}]) +@pytest.mark.parametrize( + "method", [ExternalKnowledgeApplicationService.create_template, ExternalKnowledgeApplicationService.update_template] +) +def test_template_validation_prevents_io( + service: ExternalKnowledgeApplicationService, + operations: MagicMock, + settings: dict[str, str], + method: Callable[..., object], +) -> None: + extra: dict[str, str] = ( + {"template_id": "template"} if method is ExternalKnowledgeApplicationService.update_template else {} + ) + with pytest.raises(ValueError): + method(service, CONTEXT, name="API", settings=settings, **extra) + assert operations.mock_calls == [] + + +def test_retrieval_authorization_precedes_validation_and_network( + service: ExternalKnowledgeApplicationService, access: MagicMock, operations: MagicMock +) -> None: + access.require_accessible.side_effect = DatasetAccessDeniedError() + with pytest.raises(DatasetAccessDeniedError): + service.hit_testing(CONTEXT, dataset_id="dataset", query="", retrieval_model=None, metadata_filters=None) + assert operations.mock_calls == [] + + +@pytest.mark.parametrize("query", ["", "a" * 251]) +def test_invalid_query_prevents_retrieval( + service: ExternalKnowledgeApplicationService, operations: MagicMock, query: str +) -> None: + with pytest.raises(ValueError): + service.hit_testing(CONTEXT, dataset_id="dataset", query=query, retrieval_model=None, metadata_filters=None) + operations.retrieve.assert_not_called() + + +def test_retrieval_failure_is_neutral_and_receives_owned_ref( + service: ExternalKnowledgeApplicationService, operations: MagicMock +) -> None: + operations.retrieve.side_effect = RuntimeError("unavailable") + with pytest.raises(ExternalHitTestingError, match="unavailable"): + service.hit_testing( + CONTEXT, dataset_id="dataset", query="query", retrieval_model={"top_k": 2}, metadata_filters=None + ) + operations.retrieve.assert_called_once_with( + CONTEXT, DatasetRef("tenant", "dataset"), query="query", retrieval_model={"top_k": 2}, metadata_filters=None + ) diff --git a/api/tests/unit_tests/services/knowledge/test_indexing_estimate_adapters.py b/api/tests/unit_tests/services/knowledge/test_indexing_estimate_adapters.py new file mode 100644 index 00000000000..16efd0c6953 --- /dev/null +++ b/api/tests/unit_tests/services/knowledge/test_indexing_estimate_adapters.py @@ -0,0 +1,221 @@ +from collections.abc import Callable +from datetime import UTC, datetime +from unittest.mock import MagicMock, patch + +import pytest +from sqlalchemy.orm import Session, sessionmaker + +from core.entities.knowledge_entities import PreviewDetail +from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError +from core.plugin.impl.exc import PluginDaemonClientSideError +from core.rag.extractor.entity.datasource_type import DatasourceType +from core.rag.extractor.entity.extract_setting import ExtractSetting, NotionInfo +from core.rag.index_processor.constant.index_type import IndexTechniqueType +from core.rag.models.document import ChildDocument, Document +from enums import DeploymentEdition +from extensions.storage.storage_type import StorageType +from models.dataset import Dataset, DatasetProcessRule +from models.enums import CreatorUserRole, ProcessRuleMode +from models.model import UploadFile +from services.knowledge.indexing.adapters.estimate import ( + IndexingEstimateAdapter, + SQLAlchemyProcessRuleReader, +) +from services.knowledge.indexing.estimate import ( + IndexingEstimateExecutionError, + IndexingEstimateProviderUnavailableError, +) +from services.knowledge.resource_scope import DatasetRef + + +def _dataset(dataset_id: str, workspace_id: str) -> Dataset: + return Dataset( + id=dataset_id, + tenant_id=workspace_id, + name=f"Dataset {dataset_id}", + description="", + provider="vendor", + permission="only_me", + created_by="account-1", + maintainer="account-1", + data_source_type="upload_file", + indexing_technique="economy", + ) + + +def test_notion_access_token_is_excluded_from_model_serialization_and_repr() -> None: + notion_info = NotionInfo( + credential_id="credential-1", + notion_workspace_id="workspace-1", + notion_obj_id="page-1", + notion_page_type="page", + tenant_id="tenant-1", + notion_access_token="secret", + ) + + assert "notion_access_token" not in notion_info.model_dump() + assert "secret" not in repr(notion_info) + + +def test_process_rule_reader_loads_tenant_scoped_rule(sqlite_session_factory: sessionmaker[Session]) -> None: + with sqlite_session_factory.begin() as session: + session.add(_dataset("dataset-1", "workspace-1")) + rule = DatasetProcessRule( + dataset_id="dataset-1", + mode=ProcessRuleMode.CUSTOM, + rules='{"segmentation": {"separator": "\\n", "max_tokens": 100}}', + created_by="account-1", + ) + rule.id = "rule-1" + session.add(rule) + + reader = SQLAlchemyProcessRuleReader(session_factory=sqlite_session_factory) + + assert reader.get_by_id(dataset_ref=DatasetRef("workspace-1", "dataset-1"), process_rule_id="rule-1") == { + "mode": "custom", + "rules": {"segmentation": {"separator": "\n", "max_tokens": 100}}, + } + assert reader.get_by_id(dataset_ref=DatasetRef("workspace-2", "dataset-1"), process_rule_id="rule-1") is None + + +@pytest.mark.parametrize("doc_form", ["text_model", "hierarchical_model", "qa_model"]) +def test_estimates_segments_and_limits_previews( + sqlite_session_factory: sessionmaker[Session], doc_form: str, config_overrides: Callable[..., None] +) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + processor = MagicMock() + processor.extract.return_value = [Document(page_content="source")] + processor.transform.return_value = [ + Document( + page_content=f"question-{i}", + metadata={"answer": f"answer-{i}"}, + children=[ChildDocument(page_content="child")], + ) + for i in range(11) + ] + with patch("services.knowledge.indexing.adapters.estimate.IndexProcessorFactory") as factory: + factory.return_value.init_index_processor.return_value = processor + result = IndexingEstimateAdapter(session_factory=sqlite_session_factory).run( + "workspace-1", [ExtractSetting(datasource_type=DatasourceType.WEBSITE)], {"mode": "automatic"}, doc_form + ) + if doc_form == "qa_model": + assert result.total_segments == 220 + assert len(result.qa_preview) == 10 + assert result.qa_preview[-1].answer == "answer-9" + else: + assert result.total_segments == 11 + assert len(result.preview) == 10 + assert result.preview[0].child_chunks == ["child"] + + +def test_estimate_enforces_batch_limit( + sqlite_session_factory: sessionmaker[Session], config_overrides: Callable[..., None] +) -> None: + config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, BATCH_UPLOAD_LIMIT=1) + with pytest.raises(IndexingEstimateExecutionError, match="batch upload limit"): + IndexingEstimateAdapter(session_factory=sqlite_session_factory).run( + "workspace-1", [ExtractSetting(datasource_type=DatasourceType.WEBSITE)] * 2, {"mode": "automatic"} + ) + + +@pytest.mark.parametrize("storage_fails", [False, True]) +def test_preview_cleanup_commits_before_summary_and_excludes_foreign_files( + sqlite_session_factory: sessionmaker[Session], storage_fails: bool +) -> None: + image_ids = [] + with sqlite_session_factory.begin() as session: + for tenant in ["workspace-1", "foreign"]: + upload = UploadFile( + tenant_id=tenant, + storage_type=StorageType.LOCAL, + key=f"{tenant}/image.png", + name="image.png", + size=10, + extension="png", + mime_type="image/png", + created_by_role=CreatorUserRole.ACCOUNT, + created_by="user", + created_at=datetime.now(UTC), + used=True, + ) + image_ids.append(upload.id) + session.add(upload) + processor = MagicMock() + processor.transform.return_value = [Document(page_content="preview")] + + def summary( + _tenant, previews: list[PreviewDetail], _settings, _language, *, session: Session + ) -> list[PreviewDetail]: + assert not session.in_transaction() + with sqlite_session_factory() as observer: + assert observer.get(UploadFile, image_ids[0]) is None + assert observer.get(UploadFile, image_ids[1]) is not None + return [PreviewDetail(content=previews[0].content, summary="summary")] + + processor.generate_summary_preview.side_effect = summary + with ( + patch("services.knowledge.indexing.adapters.estimate.IndexProcessorFactory") as factory, + patch("services.knowledge.indexing.adapters.estimate.get_image_upload_file_ids", return_value=image_ids), + patch("services.knowledge.indexing.adapters.estimate.storage") as storage, + ): + factory.return_value.init_index_processor.return_value = processor + if storage_fails: + storage.delete.side_effect = OSError("storage unavailable") + result = IndexingEstimateAdapter(session_factory=sqlite_session_factory).run( + "workspace-1", + [ExtractSetting(datasource_type=DatasourceType.WEBSITE)], + {"mode": "automatic", "summary_index_setting": {"enable": True}}, + "text_model", + ) + storage.delete.assert_called_once_with("workspace-1/image.png") + assert result.preview[0].summary == "summary" + + +@pytest.mark.parametrize("provider", [None, "embedding-provider"]) +def test_estimate_selects_dataset_embedding_and_rejects_foreign_tenant( + sqlite_session_factory: sessionmaker[Session], provider: str | None +) -> None: + with sqlite_session_factory.begin() as session: + dataset = _dataset("dataset-1", "workspace-1") + dataset.indexing_technique = IndexTechniqueType.HIGH_QUALITY + dataset.embedding_model_provider = provider + dataset.embedding_model = "embedding-model" + session.add(dataset) + adapter = IndexingEstimateAdapter(session_factory=sqlite_session_factory) + with ( + patch("services.knowledge.indexing.adapters.estimate.ModelManager.for_tenant") as manager, + patch("services.knowledge.indexing.adapters.estimate.IndexProcessorFactory"), + ): + adapter.run("workspace-1", [], {"mode": "automatic"}, dataset_id="dataset-1") + selected = ( + manager.return_value.get_model_instance if provider else manager.return_value.get_default_model_instance + ) + selected.assert_called_once() + with pytest.raises(IndexingEstimateExecutionError, match="Dataset not found"): + adapter.run("foreign", [], {"mode": "automatic"}, dataset_id="dataset-1") + + +@pytest.mark.parametrize( + ("processor_error", "expected_error"), + [ + (LLMBadRequestError("provider unavailable"), IndexingEstimateProviderUnavailableError), + (ProviderTokenNotInitError("provider unavailable"), IndexingEstimateProviderUnavailableError), + (PluginDaemonClientSideError("provider unavailable"), IndexingEstimateProviderUnavailableError), + (RuntimeError("processor failed"), IndexingEstimateExecutionError), + ], +) +def test_estimate_translates_infrastructure_errors( + sqlite_session_factory: sessionmaker[Session], + processor_error: Exception, + expected_error: type[Exception], +) -> None: + adapter = IndexingEstimateAdapter(session_factory=sqlite_session_factory) + with ( + patch("services.knowledge.indexing.adapters.estimate.IndexProcessorFactory", side_effect=processor_error), + pytest.raises(expected_error, match="provider unavailable|processor failed"), + ): + adapter.run( + "workspace-1", + [ExtractSetting(datasource_type=DatasourceType.WEBSITE)], + {"mode": "automatic", "rules": {}}, + ) diff --git a/api/tests/unit_tests/services/knowledge/test_indexing_estimate_application_service.py b/api/tests/unit_tests/services/knowledge/test_indexing_estimate_application_service.py new file mode 100644 index 00000000000..0ee5acbc6b5 --- /dev/null +++ b/api/tests/unit_tests/services/knowledge/test_indexing_estimate_application_service.py @@ -0,0 +1,376 @@ +from collections.abc import Sequence +from dataclasses import dataclass +from unittest.mock import Mock, create_autospec + +import pytest + +from core.entities.knowledge_entities import IndexingEstimate +from core.rag.extractor.entity.datasource_type import DatasourceType +from core.rag.extractor.entity.extract_setting import ExtractSetting, UploadFileExtractionInput +from machinery.context import RequestContext +from services.data_source.credential_gateway import DatasourceCredentialNotFoundError +from services.data_source.entities.notion_import import NotionPageType +from services.knowledge.dataset_access import AccessibleDataset, DatasetAccess, DatasetAccessDeniedError +from services.knowledge.entities.indexing_estimate import ( + NewSourcesEstimateCommand, + NotionEstimateSource, + UploadFileEstimateSource, + WebsiteEstimateSource, +) +from services.knowledge.indexing.errors import IndexingInputSourceError, UnsupportedStoredSourceError +from services.knowledge.indexing.estimate import ( + DatasetEstimateRecord, + EstimateDatasetReader, + EstimateDocumentAlreadyFinishedError, + EstimateDocumentNotFoundError, + EstimateDocumentReader, + EstimateDocumentRecord, + EstimateSourceNotFoundError, + IndexingEstimateApplicationService, + IndexingEstimateCredentialUnavailableError, + IndexingEstimateExecutionError, + IndexingEstimateRunner, + NotionEstimateSourceResolver, + ProcessRuleReader, + StoredEstimateSourceResolver, + StoredSource, + UnsupportedEstimateSourceError, + UploadEstimateSourceResolver, + WebsiteEstimateSourceResolver, +) +from services.knowledge.resource_scope import DatasetRef + + +def _context() -> RequestContext: + return RequestContext("request-1", None, "account-1", "workspace-1") + + +def _access_dataset() -> AccessibleDataset: + return AccessibleDataset( + id="dataset-1", + workspace_id="workspace-1", + ) + + +def _estimate_dataset() -> DatasetEstimateRecord: + return DatasetEstimateRecord( + id="dataset-1", + workspace_id="workspace-1", + indexing_technique="high_quality", + ) + + +def _document( + document_id: str = "document-1", + *, + status: str = "waiting", + process_rule_id: str | None = "rule-1", +) -> EstimateDocumentRecord: + return EstimateDocumentRecord( + id=document_id, + workspace_id="workspace-1", + dataset_id="dataset-1", + data_source_type="upload_file", + data_source_info={"upload_file_id": f"file-{document_id}"}, + indexing_status=status, + doc_form="text_model", + doc_language="English", + dataset_process_rule_id=process_rule_id, + ) + + +def _files(*, workspace_id: str, file_ids: Sequence[str], document_model: str) -> dict[str, ExtractSetting]: + return { + file_id: ExtractSetting( + datasource_type=DatasourceType.FILE, + document_model=document_model, + upload_file=UploadFileExtractionInput( + id=file_id, tenant_id=workspace_id, key=file_id, created_by="account-1" + ), + ) + for file_id in file_ids + } + + +@dataclass +class ServiceFixture: + service: IndexingEstimateApplicationService + access: Mock + datasets: Mock + documents: Mock + files: Mock + notion: Mock + websites: Mock + stored: Mock + rules: Mock + runner: Mock + + +@pytest.fixture +def fixture() -> ServiceFixture: + access = create_autospec(DatasetAccess, instance=True, spec_set=True) + access.require_accessible.return_value = _access_dataset() + access.check_access.side_effect = AssertionError("Estimate access must use the selected dataset") + datasets = create_autospec(EstimateDatasetReader, instance=True, spec_set=True) + datasets.get_estimate_record.return_value = _estimate_dataset() + documents = create_autospec(EstimateDocumentReader, instance=True, spec_set=True) + documents.get_estimate_document.return_value = _document() + documents.list_estimate_documents_by_batch.return_value = () + files = create_autospec(UploadEstimateSourceResolver, instance=True, spec_set=True) + files.resolve_selections.side_effect = _files + notion = create_autospec(NotionEstimateSourceResolver, instance=True, spec_set=True) + notion.resolve_selection.return_value = ExtractSetting(datasource_type=DatasourceType.NOTION) + websites = create_autospec(WebsiteEstimateSourceResolver, instance=True, spec_set=True) + websites.resolve_selection.return_value = ExtractSetting(datasource_type=DatasourceType.WEBSITE) + stored = create_autospec(StoredEstimateSourceResolver, instance=True, spec_set=True) + stored.resolve.side_effect = lambda source: ExtractSetting( + datasource_type=source.source_type, document_model=source.document_model + ) + rules = create_autospec(ProcessRuleReader, instance=True, spec_set=True) + rules.get_by_id.return_value = {"mode": "custom", "rules": {"segmentation": dict[str, object]()}} + runner = create_autospec(IndexingEstimateRunner, instance=True, spec_set=True) + runner.run.side_effect = lambda **kwargs: IndexingEstimate( + total_segments=len(kwargs["extract_settings"]), preview=[] + ) + service = IndexingEstimateApplicationService( + dataset_access=access, + datasets=datasets, + documents=documents, + files=files, + notion=notion, + websites=websites, + stored_sources=stored, + process_rules=rules, + runner=runner, + ) + return ServiceFixture(service, access, datasets, documents, files, notion, websites, stored, rules, runner) + + +def test_estimate_new_sources_resolves_each_variant_and_normalizes_rule(fixture: ServiceFixture) -> None: + + result = fixture.service.estimate_new_sources( + _context(), + NewSourcesEstimateCommand( + sources=( + UploadFileEstimateSource("file-1"), + NotionEstimateSource("notion-workspace", "page-1", NotionPageType.PAGE, "credential-1"), + WebsiteEstimateSource("firecrawl", "job-1", "https://example.com", only_main_content=True), + ), + process_rule={"mode": "automatic", "summary_index_setting": {"enable": None}}, + dataset_id="dataset-1", + ), + ) + + assert result.total_segments == 3 + fixture.access.require_accessible.assert_called_once_with(_context(), "dataset-1") + fixture.files.resolve_selections.assert_called_once_with( + workspace_id="workspace-1", + file_ids=("file-1",), + document_model="text_model", + ) + fixture.notion.resolve_selection.assert_called_once_with( + workspace_id="workspace-1", + actor_id="account-1", + credential_id="credential-1", + notion_workspace_id="notion-workspace", + page_id="page-1", + page_type=NotionPageType.PAGE, + document_model="text_model", + ) + fixture.websites.resolve_selection.assert_called_once_with( + workspace_id="workspace-1", + provider="firecrawl", + job_id="job-1", + url="https://example.com", + mode="crawl", + only_main_content=True, + document_model="text_model", + ) + assert fixture.runner.run.call_args.kwargs["tmp_processing_rule"] == {"mode": "automatic", "rules": {}} + assert fixture.runner.run.call_args.kwargs["dataset_id"] == "dataset-1" + + +@pytest.mark.parametrize("dataset_id", [None, ""]) +@pytest.mark.parametrize( + "source", + [ + NotionEstimateSource("notion-workspace", "page-1", NotionPageType.PAGE, "credential-1"), + WebsiteEstimateSource("firecrawl", "job-1", "https://example.com"), + ], + ids=["notion", "website"], +) +def test_estimate_new_dataset_accepts_missing_or_empty_dataset_id( + fixture: ServiceFixture, + dataset_id: str | None, + source: NotionEstimateSource | WebsiteEstimateSource, +) -> None: + result = fixture.service.estimate_new_sources( + _context(), + NewSourcesEstimateCommand( + sources=(source,), + process_rule={"mode": "automatic"}, + dataset_id=dataset_id, + ), + ) + + assert result.total_segments == 1 + fixture.access.require_accessible.assert_not_called() + assert fixture.runner.run.call_args.kwargs["dataset_id"] is None + + +def test_estimate_new_sources_stops_when_dataset_access_is_denied(fixture: ServiceFixture) -> None: + fixture.access.require_accessible.side_effect = DatasetAccessDeniedError() + + with pytest.raises(DatasetAccessDeniedError): + fixture.service.estimate_new_sources( + _context(), + NewSourcesEstimateCommand( + sources=(UploadFileEstimateSource("file-1"),), + process_rule={"mode": "automatic"}, + dataset_id="dataset-1", + ), + ) + + fixture.files.resolve_selections.assert_not_called() + fixture.runner.run.assert_not_called() + + +def test_estimate_new_notion_source_maps_actor_credential_failure_without_running(fixture: ServiceFixture) -> None: + fixture.notion.resolve_selection.side_effect = DatasourceCredentialNotFoundError() + + with pytest.raises(IndexingEstimateCredentialUnavailableError): + fixture.service.estimate_new_sources( + _context(), + NewSourcesEstimateCommand( + sources=( + NotionEstimateSource( + "notion-workspace", + "page-1", + NotionPageType.PAGE, + "private-credential", + ), + ), + process_rule={"mode": "automatic"}, + ), + ) + + fixture.runner.run.assert_not_called() + + +def test_estimate_empty_new_sources_returns_zero_without_runner_call(fixture: ServiceFixture) -> None: + + result = fixture.service.estimate_new_sources( + _context(), + NewSourcesEstimateCommand(sources=(), process_rule={"mode": "automatic"}), + ) + + assert result == IndexingEstimate(total_segments=0, preview=[]) + fixture.runner.run.assert_not_called() + + +def test_estimate_batch_preserves_first_document_contract_and_loads_its_rule(fixture: ServiceFixture) -> None: + documents = (_document("document-1"), _document("document-2")) + fixture.documents.list_estimate_documents_by_batch.return_value = documents + + fixture.service.estimate_batch(_context(), dataset_id="dataset-1", batch="batch-1") + + fixture.stored.resolve.assert_called_once_with(StoredSource.from_document(documents[0])) + fixture.rules.get_by_id.assert_called_once_with( + dataset_ref=DatasetRef("workspace-1", "dataset-1"), + process_rule_id="rule-1", + ) + fixture.documents.list_estimate_documents_by_batch.assert_called_once_with( + DatasetRef("workspace-1", "dataset-1"), + "batch-1", + ) + assert fixture.runner.run.call_args.kwargs["indexing_technique"] == "high_quality" + + +def test_estimate_uses_automatic_rule_when_document_has_no_saved_rule(fixture: ServiceFixture) -> None: + fixture.documents.get_estimate_document.return_value = _document(process_rule_id=None) + + fixture.service.estimate_document(_context(), dataset_id="dataset-1", document_id="document-1") + + fixture.datasets.get_estimate_record.assert_called_once_with(DatasetRef("workspace-1", "dataset-1")) + fixture.documents.get_estimate_document.assert_called_once_with( + DatasetRef("workspace-1", "dataset-1").document("document-1") + ) + fixture.rules.get_by_id.assert_not_called() + assert fixture.runner.run.call_args.kwargs["tmp_processing_rule"] == {"mode": "automatic", "rules": {}} + + +def test_estimate_rejects_missing_saved_process_rule(fixture: ServiceFixture) -> None: + fixture.rules.get_by_id.return_value = None + + with pytest.raises(EstimateSourceNotFoundError, match="rule-1"): + fixture.service.estimate_document(_context(), dataset_id="dataset-1", document_id="document-1") + + fixture.runner.run.assert_not_called() + + +def test_estimate_empty_batch_raises_not_found_without_runner_call(fixture: ServiceFixture) -> None: + + with pytest.raises(EstimateDocumentNotFoundError): + fixture.service.estimate_batch(_context(), dataset_id="dataset-1", batch="batch-1") + + fixture.runner.run.assert_not_called() + + +@pytest.mark.parametrize("status", ["completed", "error"]) +def test_estimate_rejects_finished_document(fixture: ServiceFixture, status: str) -> None: + fixture.documents.get_estimate_document.return_value = _document(status=status) + + with pytest.raises(EstimateDocumentAlreadyFinishedError): + fixture.service.estimate_document(_context(), dataset_id="dataset-1", document_id="document-1") + + fixture.runner.run.assert_not_called() + + +def test_estimate_rejects_missing_document(fixture: ServiceFixture) -> None: + fixture.documents.get_estimate_document.return_value = None + + with pytest.raises(EstimateDocumentNotFoundError): + fixture.service.estimate_document(_context(), dataset_id="dataset-1", document_id="document-1") + + +@pytest.mark.parametrize( + ("source_error", "expected_error"), + [ + (DatasourceCredentialNotFoundError(), IndexingEstimateCredentialUnavailableError), + (IndexingInputSourceError("missing source"), EstimateSourceNotFoundError), + (UnsupportedStoredSourceError("unknown"), UnsupportedEstimateSourceError), + (RuntimeError("failed"), IndexingEstimateExecutionError), + ], +) +def test_estimate_translates_source_boundary_errors( + fixture: ServiceFixture, source_error: Exception, expected_error: type[Exception] +) -> None: + fixture.stored.resolve.side_effect = source_error + + with pytest.raises(expected_error): + fixture.service.estimate_document(_context(), dataset_id="dataset-1", document_id="document-1") + + +def test_multi_file_estimate_reads_once_and_preserves_selection_order(fixture: ServiceFixture) -> None: + fixture.service.estimate_new_sources( + _context(), + NewSourcesEstimateCommand( + sources=( + UploadFileEstimateSource("file-2"), + UploadFileEstimateSource("file-1"), + UploadFileEstimateSource("file-2"), + ), + process_rule={"mode": "automatic"}, + ), + ) + fixture.files.resolve_selections.assert_called_once_with( + workspace_id="workspace-1", + file_ids=("file-2", "file-1"), + document_model="text_model", + ) + settings = fixture.runner.run.call_args.kwargs["extract_settings"] + assert [setting.upload_file.id for setting in settings if setting.upload_file is not None] == [ + "file-2", + "file-1", + "file-2", + ] diff --git a/api/tests/unit_tests/services/knowledge/test_indexing_estimate_entities.py b/api/tests/unit_tests/services/knowledge/test_indexing_estimate_entities.py new file mode 100644 index 00000000000..a4b9c13cfbf --- /dev/null +++ b/api/tests/unit_tests/services/knowledge/test_indexing_estimate_entities.py @@ -0,0 +1,126 @@ +import pytest + +from services.knowledge.entities.indexing_estimate import ( + EstimateValidationError, + normalize_process_rule, +) + + +def test_normalize_process_rule_sets_empty_rules_for_automatic_mode() -> None: + normalized = normalize_process_rule({"mode": "automatic", "rules": {"ignored": True}}) + + assert normalized["rules"] == {} + + +@pytest.mark.parametrize("summary_setting", [None, {"enable": None}]) +def test_normalize_process_rule_treats_absent_summary_enable_as_no_summary(summary_setting: object) -> None: + normalized = normalize_process_rule( + { + "mode": "automatic", + "summary_index_setting": summary_setting, + } + ) + + assert "summary_index_setting" not in normalized + + +def test_normalize_process_rule_rejects_unknown_pre_processing_rule_id() -> None: + with pytest.raises(EstimateValidationError): + normalize_process_rule( + { + "mode": "custom", + "rules": { + "pre_processing_rules": [{"id": "unknown", "enabled": True}], + "segmentation": {"separator": "\n", "max_tokens": 128}, + }, + } + ) + + +def test_normalize_process_rule_deduplicates_custom_rules() -> None: + normalized = normalize_process_rule( + { + "mode": "custom", + "rules": { + "pre_processing_rules": [ + {"id": "remove_stopwords", "enabled": True}, + {"id": "remove_stopwords", "enabled": False}, + ], + "segmentation": {"separator": "\n", "max_tokens": 128}, + }, + } + ) + + rules = normalized["rules"] + assert isinstance(rules, dict) + assert rules["pre_processing_rules"] == [{"id": "remove_stopwords", "enabled": False}] + + +def test_normalize_process_rule_drops_hierarchical_fields_from_custom_mode() -> None: + normalized = normalize_process_rule( + { + "mode": "custom", + "rules": { + "pre_processing_rules": [{"id": "remove_stopwords", "enabled": True}], + "segmentation": {"separator": "\n", "max_tokens": 128}, + "parent_mode": "full-doc", + "subchunk_segmentation": {"separator": "###", "max_tokens": 64}, + }, + } + ) + + assert normalized["rules"] == { + "pre_processing_rules": [{"id": "remove_stopwords", "enabled": True}], + "segmentation": {"separator": "\n", "max_tokens": 128}, + } + + +def test_normalize_process_rule_requires_summary_index_provider_name() -> None: + with pytest.raises(EstimateValidationError, match="Field required"): + normalize_process_rule( + { + "mode": "custom", + "rules": { + "pre_processing_rules": [{"id": "remove_stopwords", "enabled": True}], + "segmentation": {"separator": "\n", "max_tokens": 128}, + }, + "summary_index_setting": {"enable": True, "model_name": "summary-model"}, + } + ) + + +def test_normalize_process_rule_preserves_hierarchical_fields() -> None: + normalized = normalize_process_rule( + { + "mode": "hierarchical", + "rules": { + "pre_processing_rules": [{"id": "remove_stopwords", "enabled": True}], + "segmentation": {"separator": "\n", "max_tokens": 512}, + "parent_mode": "full-doc", + "subchunk_segmentation": {"separator": "###", "max_tokens": 128}, + }, + } + ) + + rules = normalized["rules"] + assert isinstance(rules, dict) + assert rules["parent_mode"] == "full-doc" + assert rules["subchunk_segmentation"] == {"separator": "###", "max_tokens": 128} + + +def test_normalize_process_rule_defaults_hierarchical_parent_mode() -> None: + normalized = normalize_process_rule( + { + "mode": "hierarchical", + "rules": { + "pre_processing_rules": [{"id": "remove_stopwords", "enabled": True}], + "segmentation": {"separator": "\n", "max_tokens": 512}, + "subchunk_segmentation": {"separator": "###", "max_tokens": 128}, + }, + } + ) + + rules = normalized["rules"] + assert isinstance(rules, dict) + assert rules["parent_mode"] == "paragraph" + assert rules["subchunk_segmentation"] == {"separator": "###", "max_tokens": 128} diff --git a/api/tests/unit_tests/services/knowledge/test_indexing_execution.py b/api/tests/unit_tests/services/knowledge/test_indexing_execution.py new file mode 100644 index 00000000000..630e109a33b --- /dev/null +++ b/api/tests/unit_tests/services/knowledge/test_indexing_execution.py @@ -0,0 +1,128 @@ +from unittest.mock import MagicMock + +import pytest + +from core.rag.models.document import Document +from services.knowledge.indexing.errors import DocumentIsDeletedPausedError, DocumentIsPausedError +from services.knowledge.indexing.estimate import StoredSource +from services.knowledge.indexing.execution import DocumentIndexingService, IndexingDocument +from services.knowledge.resource_scope import DatasetRef + +ExecutionFixture = tuple[DocumentIndexingService, MagicMock, IndexingDocument, list[Document]] + + +@pytest.fixture +def execution() -> ExecutionFixture: + ref = DatasetRef("tenant", "dataset").document("document") + document = IndexingDocument( + StoredSource(ref, "upload_file", {}, "text_model"), {"mode": "automatic"}, "English", False + ) + ports = MagicMock() + ports.documents.get_indexing_document.return_value = document + chunks = [Document(page_content="one"), Document(page_content="two")] + ports.backend.extract.return_value = chunks + ports.backend.transform.return_value = chunks + ports.backend.count_tokens.return_value = [7, 11] + ports.backend.describe_error.side_effect = str + ports.segments.resume_indexing.return_value = (chunks[1:], 18) + service = DocumentIndexingService( + documents=ports.documents, + segments=ports.segments, + backend=ports.backend, + enforce_vector_space_admission=True, + clock=iter([10, 13]).__next__, + ) + return service, ports, document, chunks + + +@pytest.mark.parametrize("start", ["parsing", "splitting", "indexing"]) +def test_completion_waits_for_persisted_segments_and_all_index_writes(execution: ExecutionFixture, start: str) -> None: + service, ports, document, chunks = execution + if start == "parsing": + service.run([document.ref]) + elif start == "splitting": + service.run_in_splitting_status(document.ref) + else: + service.run_in_indexing_status(document.ref) + calls = [call[0] for call in ports.mock_calls] + assert calls.index("backend.load") < calls.index("documents.complete_indexing") + ports.documents.complete_indexing.assert_called_once_with(document.ref, tokens=18, latency=3) + if start == "indexing": + ports.backend.extract.assert_not_called() + ports.segments.save_for_indexing.assert_not_called() + ports.backend.load.assert_called_once_with(document, chunks[1:]) + else: + assert calls.index("segments.save_for_indexing") < calls.index("backend.load") + ports.segments.save_for_indexing.assert_called_once_with(document, chunks, [7, 11]) + if start == "splitting": + assert calls.index("segments.clear_for_indexing") < calls.index("backend.extract") + assert ports.backend.ensure_admission.call_count == (start == "parsing") + + +@pytest.mark.parametrize("stage", ["extract", "transform", "ensure_admission", "count_tokens", "load"]) +def test_failed_phase_records_error_without_completing(execution: ExecutionFixture, stage: str) -> None: + service, ports, document, _ = execution + operation = { + "extract": ports.backend.extract, + "transform": ports.backend.transform, + "ensure_admission": ports.backend.ensure_admission, + "count_tokens": ports.backend.count_tokens, + "load": ports.backend.load, + }[stage] + operation.side_effect = RuntimeError("index unavailable") + service.run([document.ref]) + ports.documents.complete_indexing.assert_not_called() + ports.documents.fail_indexing.assert_called_once_with(document.ref, "index unavailable") + if stage != "load": + ports.backend.load.assert_not_called() + + +@pytest.mark.parametrize("stage", ["get_indexing_document", "mark_splitting", "complete_indexing"]) +@pytest.mark.parametrize("error", [DocumentIsPausedError, DocumentIsDeletedPausedError]) +def test_pause_and_deletion_do_not_become_indexing_errors( + execution: ExecutionFixture, stage: str, error: type[Exception] +) -> None: + service, ports, document, _ = execution + operation = { + "get_indexing_document": ports.documents.get_indexing_document, + "mark_splitting": ports.documents.mark_splitting, + "complete_indexing": ports.documents.complete_indexing, + }[stage] + operation.side_effect = error() + if error is DocumentIsPausedError: + with pytest.raises(DocumentIsPausedError): + service.run([document.ref]) + else: + service.run([document.ref]) + ports.documents.fail_indexing.assert_not_called() + + +def test_pause_after_worker_completion_prevents_document_completion(execution: ExecutionFixture) -> None: + service, ports, document, _ = execution + ports.backend.check_paused.side_effect = [None, None, DocumentIsPausedError()] + with pytest.raises(DocumentIsPausedError): + service.run([document.ref]) + ports.documents.complete_indexing.assert_not_called() + ports.documents.fail_indexing.assert_not_called() + + +def test_missing_document_is_skipped(execution: ExecutionFixture) -> None: + service, ports, document, _ = execution + ports.documents.get_indexing_document.return_value = None + service.run([document.ref]) + ports.backend.extract.assert_not_called() + ports.documents.fail_indexing.assert_not_called() + + +def test_indexing_persists_the_backend_error_description(execution: ExecutionFixture) -> None: + service, ports, document, _ = execution + error = RuntimeError("backend details") + ports.backend.extract.side_effect = error + ports.backend.describe_error.side_effect = None + ports.backend.describe_error.return_value = "source file missing" + + service.run([document.ref]) + + ports.backend.describe_error.assert_called_once_with(error) + ports.documents.fail_indexing.assert_called_once_with(document.ref, "source file missing") + ports.documents.complete_indexing.assert_not_called() diff --git a/api/tests/unit_tests/services/knowledge/test_indexing_execution_adapters.py b/api/tests/unit_tests/services/knowledge/test_indexing_execution_adapters.py new file mode 100644 index 00000000000..45ec28a4ec8 --- /dev/null +++ b/api/tests/unit_tests/services/knowledge/test_indexing_execution_adapters.py @@ -0,0 +1,528 @@ +import json +import threading +from collections.abc import Generator, Iterator +from contextlib import contextmanager +from dataclasses import replace +from datetime import datetime +from unittest.mock import MagicMock, patch + +import pytest +from flask import Flask +from sqlalchemy import Connection, Engine, event, inspect, select +from sqlalchemy.orm import Session, sessionmaker +from werkzeug.exceptions import NotFound + +from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError, QuotaExceededError +from core.plugin.impl.exc import PluginDaemonClientSideError +from core.rag.datasource.vdb.vector_factory import Vector +from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType +from core.rag.models.document import AttachmentDocument, ChildDocument +from core.rag.models.document import Document as IndexDocument +from extensions.storage.storage_type import StorageType +from models.dataset import Dataset, DatasetKeywordTable, DatasetProcessRule, Document, DocumentSegment +from models.enums import CreatorUserRole, IndexingStatus, ProcessRuleMode, SegmentStatus +from models.model import UploadFile +from repositories.knowledge.document_repository import SQLAlchemyDocumentRepository +from repositories.knowledge.segment_repository import SQLAlchemySegmentRepository +from repositories.knowledge.upload_file_repository import SQLAlchemyKnowledgeUploadRepository +from services.knowledge.indexing.adapters.execution import IndexingExecutionAdapter +from services.knowledge.indexing.errors import DocumentIsPausedError +from services.knowledge.indexing.execution import DocumentIndexingService, IndexingDocument +from services.knowledge.resource_scope import DatasetRef +from tests.unit_tests.config_override import apply_config_overrides +from tests.unit_tests.repositories.knowledge.test_document_repository import _dataset, _document + +MODULE = "services.knowledge.indexing.adapters.execution" + + +BackendFixture = tuple[ + IndexingExecutionAdapter, + SQLAlchemyDocumentRepository, + SQLAlchemySegmentRepository, + MagicMock, + IndexingDocument, + MagicMock, +] + + +@pytest.fixture +def backend(sqlite_session_factory: sessionmaker[Session]) -> Iterator[BackendFixture]: + with sqlite_session_factory.begin() as session: + session.add_all( + [_dataset("dataset-1", "workspace-1"), _document("document-1", status=IndexingStatus.SPLITTING)] + ) + documents = SQLAlchemyDocumentRepository(session_factory=sqlite_session_factory) + segments = SQLAlchemySegmentRepository(session_factory=sqlite_session_factory) + sources = MagicMock() + adapter = IndexingExecutionAdapter( + session_factory=sqlite_session_factory, + documents=documents, + segments=segments, + uploads=SQLAlchemyKnowledgeUploadRepository(session_factory=sqlite_session_factory), + sources=sources, + ) + job = replace( + documents.get_indexing_document(DatasetRef("workspace-1", "dataset-1").document("document-1")), + processing_rule={"mode": "automatic", "rules": {}}, + ) + with patch(f"{MODULE}.redis_client") as redis: + redis.get.return_value = None + yield adapter, documents, segments, sources, job, redis + + +@pytest.fixture +def active_transactions(sqlite_engine: Engine) -> Iterator[set[Connection]]: + active: set[Connection] = set() + begin, finish = active.add, active.discard + event.listen(sqlite_engine, "begin", begin) + event.listen(sqlite_engine, "commit", finish) + event.listen(sqlite_engine, "rollback", finish) + try: + yield active + finally: + event.remove(sqlite_engine, "begin", begin) + event.remove(sqlite_engine, "commit", finish) + event.remove(sqlite_engine, "rollback", finish) + + +@pytest.mark.parametrize("doc_form", list(IndexStructureType)) +def test_vector_io_releases_transactions_and_preserves_chunk_structure( + backend: BackendFixture, + sqlite_session_factory: sessionmaker[Session], + active_transactions: set[Connection], + monkeypatch: pytest.MonkeyPatch, + doc_form: IndexStructureType, +) -> None: + adapter, _, segments, _, job, _ = backend + apply_config_overrides(monkeypatch, VECTOR_STORE="qdrant", VECTOR_STORE_WHITELIST_ENABLE=True) + files = [ + UploadFile( + tenant_id=tenant_id, + storage_type=StorageType.LOCAL, + key=f"image-{tenant_id}", + name="image.png", + size=1, + extension="png", + mime_type="image/png", + created_by_role=CreatorUserRole.ACCOUNT, + created_by="account-1", + created_at=datetime(2026, 1, 1), + used=True, + ) + for tenant_id in ("workspace-1", "workspace-2") + ] + with sqlite_session_factory.begin() as session: + dataset = session.get(Dataset, "dataset-1") + document = session.get(Document, "document-1") + assert dataset is not None + assert document is not None + dataset.indexing_technique = IndexTechniqueType.HIGH_QUALITY + dataset.is_multimodal = True + document.doc_form = doc_form + session.add_all(files) + child = ChildDocument(page_content="child", metadata={"doc_id": "child-node", "doc_hash": "child-hash"}) + chunk = IndexDocument( + page_content="parent", + metadata={"doc_id": "node", "doc_hash": "hash"}, + children=[child] if doc_form == IndexStructureType.PARENT_CHILD_INDEX else None, + attachments=[ + AttachmentDocument(page_content="image", metadata={"doc_id": file.id, "doc_type": "image"}) + for file in files + ], + ) + segments.save_for_indexing(job, [chunk], [5]) + vector_backend = MagicMock() + vector_factory = MagicMock() + vector_factory.return_value.init_vector.return_value = vector_backend + + def external_call(*_args, **_kwargs) -> list[list[float]]: + assert not active_transactions, "Index I/O must not hold a database transaction" + return [[1.0]] + + def get_vector_factory(vector_type: str) -> MagicMock: + assert not active_transactions, "Vector initialization must follow the whitelist read transaction" + assert vector_type == "qdrant" + return vector_factory + + def load_image(key: str) -> bytes: + assert not active_transactions + assert key == files[0].key + return b"image" + + vector_backend.create.side_effect = external_call + with ( + patch.object(Vector, "get_vector_factory", side_effect=get_vector_factory), + patch( + "core.rag.datasource.vdb.vector_factory._LazyEmbeddings.embed_documents", side_effect=external_call + ) as embed, + patch( + "core.rag.datasource.vdb.vector_factory._LazyEmbeddings.embed_multimodal_documents", + side_effect=external_call, + ) as embed_images, + patch("core.rag.datasource.vdb.vector_factory.storage") as storage, + ): + storage.load_once.side_effect = load_image + adapter._process_chunk(Flask(__name__), job.ref, [chunk], False) + embed.assert_called_once_with(["child" if doc_form == IndexStructureType.PARENT_CHILD_INDEX else "parent"]) + assert embed_images.call_args.args[0] == [{"content": "aW1hZ2U=", "content_type": "image", "file_id": files[0].id}] + storage.load_once.assert_called_once_with(files[0].key) + assert vector_backend.create.call_count == 2 + indexed_text = vector_backend.create.call_args_list[0].kwargs["texts"][0] + assert indexed_text.metadata["doc_id"] == ( + "child-node" if doc_form == IndexStructureType.PARENT_CHILD_INDEX else "node" + ) + assert segments.resume_indexing(job) == ([], 5) + + +@pytest.mark.parametrize("storage_type", ["database", "file"]) +def test_keyword_io_releases_transactions_and_commits_before_unlocking( + backend: BackendFixture, + sqlite_session_factory: sessionmaker[Session], + active_transactions: set[Connection], + storage_type: str, +) -> None: + adapter, _, segments, _, job, redis = backend + original = json.dumps({"__type__": "keyword_table", "__data__": {"table": {"existing": ["old-node"]}}}) + with sqlite_session_factory.begin() as session: + session.add(DatasetKeywordTable(dataset_id="dataset-1", data_source_type=storage_type, keyword_table=original)) + chunk = IndexDocument(page_content="hello", metadata={"doc_id": "node", "doc_hash": "hash"}) + segments.save_for_indexing(job, [chunk], [5]) + locked = False + + @contextmanager + def dataset_lock() -> Generator[None]: + nonlocal locked + locked = True + yield + assert not active_transactions + # A separate session must see the write before the next lock holder can read. + with sqlite_session_factory() as session: + segment = session.scalar(select(DocumentSegment)) + assert segment is not None + assert segment.keywords == ["hello"] + locked = False + + redis.lock.side_effect = lambda *_args, **_kwargs: dataset_lock() + + def external_call(*_args, **_kwargs) -> None: + assert locked + assert not active_transactions, "Keyword I/O must not hold a database transaction" + + def load_keywords(*_args) -> bytes: + external_call() + return original.encode() + + def file_exists(*_args) -> bool: + external_call() + return True + + def extract_keywords(*_args) -> set[str]: + external_call() + return {"hello"} + + with ( + patch("core.rag.datasource.keyword.jieba.jieba.redis_client", redis), + patch("core.rag.datasource.keyword.jieba.jieba.storage") as storage, + patch("core.rag.datasource.keyword.jieba.jieba.JiebaKeywordTableHandler") as handler, + ): + storage.load_once.side_effect = load_keywords + storage.exists.side_effect = file_exists + storage.delete.side_effect = external_call + storage.save.side_effect = external_call + handler.return_value.extract_keywords.side_effect = extract_keywords + adapter._process_chunk(Flask(__name__), job.ref, [chunk], True) + assert not locked + assert redis.lock.call_args.args == ("keyword_indexing_lock_dataset-1",) + _, encoded = segments.get_keyword_table(job.ref.dataset) + if storage_type == "file": + storage.save.assert_called_once() + encoded = storage.save.call_args.args[1] + else: + storage.save.assert_not_called() + assert encoded is not None + assert json.loads(encoded)["__data__"]["table"] == {"existing": ["old-node"], "hello": ["node"]} + assert segments.resume_indexing(job) == ([], 5) + + +def test_extract_attaches_document_metadata_and_owns_local_session(backend: BackendFixture) -> None: + adapter, _, _, sources, job, _ = backend + text = IndexDocument(page_content="source") + with patch(f"{MODULE}.IndexProcessorFactory") as factory: + processor = factory.return_value.init_index_processor.return_value + + def extract(setting: object, *, session: Session, process_rule_mode: str) -> list[IndexDocument]: + assert setting is sources.resolve_for_indexing.return_value + assert not session.in_transaction() + assert process_rule_mode == "automatic" + return [text] + + processor.extract.side_effect = extract + assert adapter.extract(job) == [text] + sources.resolve_for_indexing.assert_called_once_with(job.source) + assert text.metadata == {"dataset_id": "dataset-1", "document_id": "document-1"} + + +@pytest.mark.parametrize( + ("error", "message"), + [ + (ProviderTokenNotInitError(), "Provider Token Not Init"), + (ProviderTokenNotInitError("configure an embedding provider"), "configure an embedding provider"), + (LLMBadRequestError("invalid model input"), "invalid model input"), + (QuotaExceededError(), "Quota Exceeded"), + (PluginDaemonClientSideError("plugin unavailable"), "plugin unavailable"), + (NotFound("source file missing"), "source file missing"), + (RuntimeError("index unavailable"), "index unavailable"), + ], +) +def test_indexing_preserves_provider_and_source_error_descriptions( + backend: BackendFixture, sqlite_session_factory: sessionmaker[Session], error: Exception, message: str +) -> None: + adapter, documents, segments, sources, job, _ = backend + sources.resolve_for_indexing.side_effect = error + with sqlite_session_factory.begin() as session: + rule = DatasetProcessRule( + dataset_id=job.ref.dataset.dataset_id, mode=ProcessRuleMode.AUTOMATIC, rules="{}", created_by="account-1" + ) + session.add(rule) + session.flush() + document = session.get(Document, job.ref.document_id) + assert document is not None + document.dataset_process_rule_id = rule.id + + DocumentIndexingService(documents=documents, segments=segments, backend=adapter).run([job.ref]) + + sources.resolve_for_indexing.assert_called_once() + with sqlite_session_factory() as session: + document = session.get(Document, job.ref.document_id) + assert document is not None + assert document.indexing_status == IndexingStatus.ERROR + assert document.error == message + assert document.completed_at is None + + +@pytest.mark.parametrize("provider", [None, "embedding-provider"]) +def test_transform_selects_embedding_model_and_preserves_processor_options( + backend: BackendFixture, sqlite_session_factory: sessionmaker[Session], provider: str | None +) -> None: + adapter, documents, _, _, job, _ = backend + with sqlite_session_factory.begin() as session: + dataset = session.get(Dataset, "dataset-1") + assert dataset is not None + dataset.indexing_technique = IndexTechniqueType.HIGH_QUALITY + dataset.embedding_model_provider = provider + dataset.embedding_model = "embedding-model" + with ( + patch.object(documents, "get_indexing_user", return_value=MagicMock()) as user, + patch(f"{MODULE}.ModelManager.for_tenant") as manager, + patch(f"{MODULE}.IndexProcessorFactory") as factory, + ): + processor = factory.return_value.init_index_processor.return_value + processor.transform.return_value = list[IndexDocument]() + adapter.transform(job, []) + selected = ( + manager.return_value.get_model_instance if provider else manager.return_value.get_default_model_instance + ) + selected.assert_called_once() + assert processor.transform.call_args.kwargs["embedding_model_instance"] is selected.return_value + assert processor.transform.call_args.kwargs["process_rule"] == job.processing_rule + assert processor.transform.call_args.args == ([], user.return_value) + + +def test_hash_groups_load_worker_owned_inputs_before_completion( + backend: BackendFixture, sqlite_session_factory: sessionmaker[Session] +) -> None: + adapter, _, segments, _, job, _ = backend + with sqlite_session_factory.begin() as session: + persisted_dataset = session.get(Dataset, "dataset-1") + assert persisted_dataset is not None + persisted_dataset.indexing_technique = IndexTechniqueType.HIGH_QUALITY + chunks = [ + IndexDocument(page_content=content, metadata={"doc_id": f"node-{i}", "doc_hash": f"hash-{i}"}) + for i, content in enumerate(["same", "other", "same"]) + ] + segments.save_for_indexing(job, chunks, [3, 5, 7]) + calls: list[tuple[Dataset, list[str]]] = [] + main_thread = threading.get_ident() + + def vector(dataset: Dataset, *, session: Session | None, vector_type: str) -> MagicMock: + assert threading.get_ident() != main_thread + assert inspect(dataset).detached + assert session is None + assert vector_type == "qdrant" + + def create(group: list[IndexDocument]) -> None: + with sqlite_session_factory() as reader: + persisted = reader.scalars(select(DocumentSegment)).all() + assert len(persisted) == 3 + assert all( + row.status == SegmentStatus.INDEXING + for row in persisted + if row.index_node_id in {chunk.metadata["doc_id"] for chunk in group} + ) + calls.append((dataset, [chunk.page_content for chunk in group])) + + backend = MagicMock() + backend.create.side_effect = create + return backend + + with Flask(__name__).app_context(), patch(f"{MODULE}.Vector", side_effect=vector) as factory: + factory.resolve_vector_type.return_value = "qdrant" + adapter.load(job, chunks) + assert sum(group.count("same") for _, group in calls) == 2 + assert sum("same" in group for _, group in calls) == 1 + assert len({id(dataset) for dataset, _ in calls}) == len(calls) + assert segments.resume_indexing(job) == ([], 15) + + +@pytest.mark.parametrize("technique", ["economy", "high_quality"]) +def test_worker_failure_marks_document_error_and_keeps_retryable_segments( + backend: BackendFixture, sqlite_session_factory: sessionmaker[Session], technique: str +) -> None: + adapter, documents, segments, _, job, _ = backend + with sqlite_session_factory.begin() as session: + persisted_dataset = session.get(Dataset, "dataset-1") + assert persisted_dataset is not None + persisted_dataset.indexing_technique = IndexTechniqueType(technique) + chunks = [IndexDocument(page_content="text", metadata={"doc_id": "node", "doc_hash": "hash"})] + segments.save_for_indexing(job, chunks, [9]) + with ( + Flask(__name__).app_context(), + patch(f"{MODULE}.Jieba") as keyword, + patch(f"{MODULE}.Vector") as vector, + ): + keyword.return_value.update_texts.side_effect = RuntimeError("worker failed") + vector.return_value.create.side_effect = RuntimeError("worker failed") + DocumentIndexingService(documents=documents, segments=segments, backend=adapter).run_in_indexing_status(job.ref) + with sqlite_session_factory() as session: + row = session.get(Document, job.ref.document_id) + assert row is not None + assert row.indexing_status == IndexingStatus.ERROR + assert row.error == "worker failed" + assert len(segments.resume_indexing(job)[0]) == 1 + + +def test_worker_pause_does_not_mark_segments_complete(backend: BackendFixture) -> None: + adapter, _, segments, _, job, redis = backend + chunks = [IndexDocument(page_content="text", metadata={"doc_id": "node", "doc_hash": "hash"})] + segments.save_for_indexing(job, chunks, [9]) + redis.get.return_value = b"paused" + with Flask(__name__).app_context(), pytest.raises(DocumentIsPausedError): + adapter.load(job, chunks) + assert len(segments.resume_indexing(job)[0]) == 1 + + +@pytest.mark.parametrize("technique", ["economy", "high_quality"]) +@pytest.mark.parametrize("doc_form", ["text_model", "qa_model", "hierarchical_model"]) +def test_full_indexing_commits_phases_through_real_repositories( + backend: BackendFixture, sqlite_session_factory: sessionmaker[Session], technique: str, doc_form: str +) -> None: + adapter, documents, segments, _, job, _ = backend + with sqlite_session_factory.begin() as session: + dataset = session.get(Dataset, "dataset-1") + assert dataset is not None + dataset.indexing_technique = IndexTechniqueType(technique) + row = session.get(Document, "document-1") + assert row is not None + row.doc_form = IndexStructureType(doc_form) + rule = DatasetProcessRule( + dataset_id="dataset-1", mode=ProcessRuleMode.AUTOMATIC, rules="{}", created_by="account-1" + ) + rule.id = "rule-1" + row.dataset_process_rule_id = rule.id + session.add(rule) + chunk = IndexDocument(page_content="text", metadata={"doc_id": "node", "doc_hash": "hash"}) + with ( + Flask(__name__).app_context(), + patch.object(documents, "get_indexing_user", return_value=MagicMock()), + patch(f"{MODULE}.IndexProcessorFactory") as factory, + patch(f"{MODULE}.ModelManager.for_tenant"), + patch(f"{MODULE}.calculate_segment_token_counts", return_value=[19]), + patch(f"{MODULE}.Jieba"), + patch(f"{MODULE}.Vector"), + ): + processor = factory.return_value.init_index_processor.return_value + processor.extract.return_value = [chunk] + processor.transform.return_value = [chunk] + DocumentIndexingService(documents=documents, segments=segments, backend=adapter).run([job.ref]) + with sqlite_session_factory() as session: + row = session.get(Document, "document-1") + assert row is not None + assert row.indexing_status == IndexingStatus.COMPLETED + assert row.tokens == 19 + assert row.parsing_completed_at is not None + assert row.splitting_completed_at is not None + assert row.completed_at is not None + + +@pytest.mark.parametrize("change", ["paused", "deleted"]) +def test_document_changed_during_worker_io_cannot_be_completed( + backend: BackendFixture, sqlite_session_factory: sessionmaker[Session], change: str +) -> None: + adapter, documents, segments, _, job, _ = backend + chunks = [IndexDocument(page_content="text", metadata={"doc_id": "node", "doc_hash": "hash"})] + segments.save_for_indexing(job, chunks, [9]) + + def change_document(*_args, **_kwargs) -> None: + with sqlite_session_factory.begin() as writer: + row = writer.get(Document, "document-1") + assert row is not None + if change == "paused": + row.is_paused = True + else: + writer.delete(row) + + service = DocumentIndexingService(documents=documents, segments=segments, backend=adapter) + with Flask(__name__).app_context(), patch(f"{MODULE}.Jieba") as keyword: + keyword.return_value.update_texts.side_effect = change_document + if change == "paused": + with pytest.raises(DocumentIsPausedError): + service.run_in_indexing_status(job.ref) + else: + service.run_in_indexing_status(job.ref) + with sqlite_session_factory() as session: + segment = session.scalar(select(DocumentSegment)) + assert segment is not None + assert segment.status == SegmentStatus.INDEXING + + +def test_partial_worker_failure_retries_only_incomplete_segments_and_keeps_total( + backend: BackendFixture, sqlite_session_factory: sessionmaker[Session] +) -> None: + adapter, documents, segments, _, job, _ = backend + with sqlite_session_factory.begin() as session: + persisted_dataset = session.get(Dataset, "dataset-1") + assert persisted_dataset is not None + persisted_dataset.indexing_technique = IndexTechniqueType.HIGH_QUALITY + chunks = [ + IndexDocument(page_content=content, metadata={"doc_id": content, "doc_hash": content}) + for content in ["good", "bad"] + ] + segments.save_for_indexing(job, chunks, [13, 17]) + service = DocumentIndexingService(documents=documents, segments=segments, backend=adapter) + + def load(group: list[IndexDocument]) -> None: + if group[0].page_content == "bad": + raise RuntimeError("one shard failed") + + with ( + Flask(__name__).app_context(), + patch(f"{MODULE}.Vector") as factory, + patch(f"{MODULE}.helper.generate_text_hash", side_effect=lambda content: "0" if content == "good" else "1"), + ): + vector = factory.return_value + vector.create.side_effect = load + service.run_in_indexing_status(job.ref) + pending, tokens = segments.resume_indexing(job) + assert [chunk.page_content for chunk in pending] == ["bad"] + assert tokens == 30 + vector.create.reset_mock() + vector.create.side_effect = None + service.run_in_indexing_status(job.ref) + vector.create.assert_called_once() + assert [chunk.page_content for chunk in vector.create.call_args.args[0]] == ["bad"] + with sqlite_session_factory() as session: + row = session.get(Document, "document-1") + assert row is not None + assert row.indexing_status == IndexingStatus.COMPLETED + assert row.tokens == 30 + assert row.error is None diff --git a/api/tests/unit_tests/services/knowledge/test_indexing_source_adapters.py b/api/tests/unit_tests/services/knowledge/test_indexing_source_adapters.py new file mode 100644 index 00000000000..8037a4d037d --- /dev/null +++ b/api/tests/unit_tests/services/knowledge/test_indexing_source_adapters.py @@ -0,0 +1,342 @@ +from collections.abc import Sequence +from dataclasses import dataclass, field +from unittest.mock import Mock + +import pytest + +from core.rag.extractor.entity.datasource_type import DatasourceType +from core.rag.extractor.entity.extract_setting import UploadFileExtractionInput +from core.rag.extractor.notion_extractor import NotionExtractor +from services.data_source.credential_gateway import ( + DatasourceCredentialNotFoundError, + DatasourceCredentialRefreshError, + DatasourceProviderCredentialStore, +) +from services.data_source.entities.notion_import import NotionPageType +from services.data_source.provider_service import DatasourceProviderService +from services.knowledge.indexing.adapters.sources import ( + CompositeStoredSourceResolver, + FileSourceAdapter, + NotionSourceResolver, + WebsiteSourceAdapter, +) +from services.knowledge.indexing.errors import ( + IndexingInputSourceError, + SourceCredentialUnavailableError, + UnsupportedStoredSourceError, +) +from services.knowledge.indexing.estimate import StoredSource +from services.knowledge.resource_scope import DatasetRef +from tests.unit_tests.config_override import apply_config_overrides + + +def _source( + source_type: str, + source_info: dict[str, object] | None, + *, + document_id: str = "document-1", +) -> StoredSource: + return StoredSource( + document_ref=DatasetRef("workspace-1", "dataset-1").document(document_id), + source_type=source_type, + source_info=source_info, + document_model="text_model", + ) + + +@dataclass +class RecordingUploads: + uploads: dict[str, UploadFileExtractionInput] + calls: list[tuple[str, tuple[str, ...]]] = field(default_factory=list) + + def get_by_ids(self, *, workspace_id: str, file_ids: Sequence[str]) -> dict[str, UploadFileExtractionInput]: + self.calls.append((workspace_id, tuple(file_ids))) + return {file_id: self.uploads[file_id] for file_id in file_ids if file_id in self.uploads} + + +@dataclass +class RecordingActorCredentials: + credentials: dict[str, object] | None = None + error: Exception | None = None + calls: list[dict[str, object]] = field(default_factory=list) + + def resolve(self, **kwargs: object) -> dict[str, object]: + self.calls.append(kwargs) + if self.error is not None: + raise self.error + return self.credentials or {} + + +@dataclass +class RecordingStoredCredentials: + credentials: dict[str, object] | None = None + error: Exception | None = None + calls: list[dict[str, object]] = field(default_factory=list) + + def resolve_for_document(self, **kwargs: object) -> dict[str, object]: + self.calls.append(kwargs) + if self.error is not None: + raise self.error + return self.credentials or {} + + +def _notion_resolver( + *, + actor: RecordingActorCredentials | None = None, + stored: RecordingStoredCredentials | None = None, + integration_token: str | None = None, +) -> NotionSourceResolver: + return NotionSourceResolver( + actor_credentials=actor or RecordingActorCredentials({"integration_secret": "actor-secret"}), + stored_credentials=stored or RecordingStoredCredentials({"integration_secret": "stored-secret"}), + integration_token=integration_token, + ) + + +def test_datasource_adapters_resolve_file_website_and_notion_inputs() -> None: + upload = UploadFileExtractionInput( + id="file-1", + tenant_id="workspace-1", + key="file.txt", + created_by="account-1", + ) + uploads = RecordingUploads({upload.id: upload}) + notion = _notion_resolver() + resolver = CompositeStoredSourceResolver( + adapters={ + "upload_file": FileSourceAdapter(uploads=uploads), + "website_crawl": WebsiteSourceAdapter(), + "notion_import": notion, + } + ) + + file_extraction = resolver.resolve( + _source("upload_file", {"upload_file_id": upload.id}), + ) + website_extraction = resolver.resolve( + _source( + "website_crawl", + { + "provider": "firecrawl", + "job_id": "job-1", + "url": "https://example.com", + "mode": "crawl", + "only_main_content": True, + }, + ), + ) + notion_extraction = resolver.resolve( + _source( + "notion_import", + { + "credential_id": "credential-1", + "notion_workspace_id": "notion-workspace-1", + "notion_page_id": "notion-page-1", + "type": "page", + }, + ), + ) + + assert file_extraction.datasource_type == DatasourceType.FILE + assert file_extraction.upload_file == upload + assert uploads.calls == [("workspace-1", ("file-1",))] + assert website_extraction.website_info is not None + assert website_extraction.website_info.tenant_id == "workspace-1" + assert website_extraction.website_info.only_main_content is True + assert notion_extraction.notion_info is not None + assert notion_extraction.notion_info.notion_access_token == "stored-secret" + + +def test_stored_notion_resolution_uses_trusted_document_owner_chain() -> None: + actor = RecordingActorCredentials({"integration_secret": "actor-secret"}) + stored = RecordingStoredCredentials({"integration_secret": "stored-secret"}) + resolver = _notion_resolver(actor=actor, stored=stored) + + extraction = resolver.resolve( + _source( + "notion_import", + { + "credential_id": "private-credential", + "notion_workspace_id": "notion-workspace-1", + "notion_page_id": "notion-page-1", + "type": "database", + }, + ), + ) + + assert actor.calls == [] + assert stored.calls == [ + { + "workspace_id": "workspace-1", + "dataset_id": "dataset-1", + "document_id": "document-1", + "credential_id": "private-credential", + "provider": "notion_datasource", + "plugin_id": "langgenius/notion_datasource", + } + ] + assert extraction.notion_info is not None + assert extraction.notion_info.notion_page_type == NotionPageType.DATABASE.value + + +@pytest.mark.parametrize( + ("credential_id", "credentials", "error", "expected_token"), + [ + (None, {"integration_secret": "default-secret"}, None, "environment-secret"), + ("credential-1", {"integration_secret": "saved-secret"}, None, "saved-secret"), + ("deleted-credential", {}, DatasourceCredentialNotFoundError(), "environment-secret"), + ("credential-1", {}, DatasourceCredentialRefreshError("credential-1"), "environment-secret"), + ("credential-1", {}, None, "environment-secret"), + ("credential-1", {"integration_secret": ""}, None, "environment-secret"), + ("credential-1", {"integration_secret": 42}, None, "environment-secret"), + ], + ids=["legacy", "saved", "deleted", "refresh-failed", "missing-token", "empty-token", "invalid-token"], +) +@pytest.mark.parametrize("integration_token", ["environment-secret", None, ""]) +def test_stored_notion_preview_and_indexing_use_same_integration( + monkeypatch: pytest.MonkeyPatch, + credential_id: str | None, + credentials: dict[str, object], + error: Exception | None, + expected_token: str, + integration_token: str | None, +) -> None: + apply_config_overrides(monkeypatch, NOTION_INTEGRATION_TOKEN=integration_token) + stored = RecordingStoredCredentials(credentials, error=error) + resolver = _notion_resolver(stored=stored) + indexing_credentials = Mock(return_value=credentials, side_effect=error) + providers = DatasourceProviderService(credentials=Mock(spec=DatasourceProviderCredentialStore)) + monkeypatch.setattr( + providers, + "get_datasource_credentials", + indexing_credentials, + ) + source = _source( + "notion_import", + { + "credential_id": credential_id, + "notion_workspace_id": "notion-workspace-1", + "notion_page_id": "notion-page-1", + "type": "page", + }, + ) + + def indexing_extractor() -> NotionExtractor: + return NotionExtractor( + notion_workspace_id="notion-workspace-1", + notion_obj_id="notion-page-1", + notion_page_type="page", + tenant_id="workspace-1", + credential_id=credential_id, + notion_token_loader=lambda: providers.get_stored_notion_access_token( + tenant_id="workspace-1", credential_id=credential_id + ), + ) + + if expected_token == "environment-secret" and not integration_token: + with pytest.raises(SourceCredentialUnavailableError): + resolver.resolve(source) + with pytest.raises(ValueError, match="Must specify `integration_token`"): + indexing_extractor() + else: + preview = resolver.resolve(source) + extractor = indexing_extractor() + assert preview.notion_info is not None + assert preview.notion_info.notion_access_token == extractor._notion_access_token == expected_token + + expected_ids: list[str] = [credential_id] if credential_id else [] + assert [call["credential_id"] for call in stored.calls] == expected_ids + assert [call.kwargs["credential_id"] for call in indexing_credentials.call_args_list] == expected_ids + + +def test_notion_resolution_falls_back_to_configured_integration_token() -> None: + stored = RecordingStoredCredentials(error=DatasourceCredentialNotFoundError()) + resolver = _notion_resolver(stored=stored, integration_token="environment-secret") + + extraction = resolver.resolve( + _source( + "notion_import", + { + "notion_workspace_id": "notion-workspace-1", + "notion_page_id": "notion-page-1", + "type": "page", + }, + ), + ) + + assert extraction.notion_info is not None + assert extraction.notion_info.notion_access_token == "environment-secret" + + +def test_new_notion_selection_uses_actor_scoped_credential() -> None: + actor = RecordingActorCredentials({"integration_secret": "actor-secret"}) + stored = RecordingStoredCredentials({"integration_secret": "stored-secret"}) + resolver = _notion_resolver(actor=actor, stored=stored) + + extraction = resolver.resolve_selection( + workspace_id="workspace-1", + actor_id="account-1", + credential_id="credential-1", + notion_workspace_id="notion-workspace-1", + page_id="notion-page-1", + page_type=NotionPageType.PAGE, + document_model="text_model", + ) + + assert stored.calls == [] + assert actor.calls[0]["actor_id"] == "account-1" + assert extraction.notion_info is not None + assert extraction.notion_info.notion_access_token == "actor-secret" + + +def test_new_notion_selection_does_not_fall_back_when_actor_cannot_access_credential() -> None: + actor = RecordingActorCredentials(error=DatasourceCredentialNotFoundError()) + resolver = _notion_resolver(actor=actor, integration_token="environment-secret") + + with pytest.raises(DatasourceCredentialNotFoundError): + resolver.resolve_selection( + workspace_id="workspace-1", + actor_id="account-1", + credential_id="private-credential", + notion_workspace_id="notion-workspace-1", + page_id="notion-page-1", + page_type=NotionPageType.PAGE, + document_model="text_model", + ) + + +def test_invalid_or_unknown_sources_are_rejected_at_the_adapter_boundary() -> None: + resolver = CompositeStoredSourceResolver(adapters={"website_crawl": WebsiteSourceAdapter()}) + + with pytest.raises(IndexingInputSourceError, match="provider"): + resolver.resolve( + _source("website_crawl", {"job_id": "job", "url": "url", "mode": "crawl"}), + ) + with pytest.raises(UnsupportedStoredSourceError, match="mystery"): + resolver.resolve(_source("mystery", {})) + + +def test_formal_notion_extraction_carries_source_update_scope_but_preview_does_not() -> None: + source = _source( + "notion_import", + { + "credential_id": "credential-1", + "notion_workspace_id": "notion-workspace-1", + "notion_page_id": "notion-page-1", + "type": "page", + }, + ) + resolver = CompositeStoredSourceResolver(adapters={"notion_import": _notion_resolver()}) + preview = resolver.resolve(source) + formal = resolver.resolve_for_indexing(source) + assert preview.notion_info is not None + assert formal.notion_info is not None + assert preview.notion_info.document is None + document = formal.notion_info.document + assert document is not None + assert (document.tenant_id, document.dataset_id, document.id) == ( + source.document_ref.dataset.tenant_id, + source.document_ref.dataset.dataset_id, + source.document_ref.document_id, + ) + assert document.data_source_info_dict == source.source_info diff --git a/api/tests/unit_tests/services/knowledge/test_retrieval_attachment_signing.py b/api/tests/unit_tests/services/knowledge/test_retrieval_attachment_signing.py new file mode 100644 index 00000000000..23ce2b19dd3 --- /dev/null +++ b/api/tests/unit_tests/services/knowledge/test_retrieval_attachment_signing.py @@ -0,0 +1,305 @@ +"""Published applications sign admin-uploaded images only through valid knowledge bindings.""" + +from datetime import UTC, datetime +from unittest.mock import MagicMock +from uuid import uuid4 + +import pytest +from sqlalchemy import Engine, event +from sqlalchemy.orm import Session + +from core.app.app_config.entities import DatasetEntity, DatasetRetrieveConfigEntity +from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom +from core.app.file_access import FileAccessScope, bind_file_access_scope, get_current_file_access_scope +from core.rag.models.document import Document as RagDocument +from core.rag.retrieval import dataset_retrieval as retrieval_module +from core.rag.retrieval.dataset_retrieval import DatasetRetrieval +from core.workflow.nodes.knowledge_retrieval.retrieval import KnowledgeRetrievalRequest +from extensions.storage.storage_type import StorageType +from models.dataset import Dataset, Document, DocumentSegment, SegmentAttachmentBinding +from models.enums import CreatorUserRole, DataSourceType, DocumentCreatedFrom, SegmentStatus +from models.model import UploadFile +from services.knowledge.retrieval.attachments import authorize_retrieved_segment + + +@pytest.fixture +def graph(sqlite_session: Session) -> tuple[Dataset, Document, DocumentSegment, dict[str, str]]: + dataset = Dataset(id=str(uuid4()), tenant_id=str(uuid4()), name="Dataset", created_by="admin") + document = Document( + id=str(uuid4()), + tenant_id=dataset.tenant_id, + dataset_id=dataset.id, + position=1, + data_source_type=DataSourceType.UPLOAD_FILE, + batch="batch", + name="Document", + created_from=DocumentCreatedFrom.WEB, + created_by="admin", + doc_form="text_model", + doc_metadata={}, + ) + segment = DocumentSegment( + tenant_id=dataset.tenant_id, + dataset_id=dataset.id, + document_id=document.id, + position=1, + content="", + word_count=1, + tokens=1, + created_by="admin", + index_node_id=str(uuid4()), + index_node_hash="hash", + status=SegmentStatus.COMPLETED, + ) + file_ids: dict[str, str] = {} + for kind in ("valid", "unbound", "tenant", "file_tenant", "dataset", "document", "segment"): + upload = UploadFile( + tenant_id=str(uuid4()) if kind == "file_tenant" else dataset.tenant_id, + storage_type=StorageType.LOCAL, + key=f"{kind}.png", + name=f"{kind}.png", + size=1, + extension="png", + mime_type="image/png", + created_by_role=CreatorUserRole.ACCOUNT, + created_by="admin", + created_at=datetime.now(UTC), + used=True, + ) + file_ids[kind] = upload.id + sqlite_session.add(upload) + if kind != "unbound": + sqlite_session.add( + SegmentAttachmentBinding( + tenant_id=str(uuid4()) if kind == "tenant" else dataset.tenant_id, + dataset_id=str(uuid4()) if kind == "dataset" else dataset.id, + document_id=str(uuid4()) if kind == "document" else document.id, + segment_id=str(uuid4()) if kind == "segment" else segment.id, + attachment_id=upload.id, + ) + ) + segment.content = "\n".join( + f"![{kind}](/files/{file_id}/image-preview?sign=old)" for kind, file_id in file_ids.items() + ) + sqlite_session.add_all([dataset, document, segment]) + sqlite_session.commit() + return dataset, document, segment, file_ids + + +@pytest.mark.parametrize("entry", ["chat", "workflow"]) +@pytest.mark.parametrize("vision_enabled", [False, True]) +def test_published_retrieval_grants_bound_images_before_signing( + graph: tuple[Dataset, Document, DocumentSegment, dict[str, str]], + sqlite_session: Session, + monkeypatch: pytest.MonkeyPatch, + entry: str, + vision_enabled: bool, +) -> None: + dataset, document, segment, file_ids = graph + retrieval = DatasetRetrieval() + monkeypatch.setattr(retrieval, "_check_knowledge_rate_limit", MagicMock()) + monkeypatch.setattr(retrieval, "_get_available_datasets", MagicMock(return_value=[dataset])) + monkeypatch.setattr(retrieval, "get_metadata_filter_condition", MagicMock(return_value=(None, None))) + monkeypatch.setattr( + retrieval, + "multiple_retrieve", + MagicMock( + return_value=[ + RagDocument( + page_content=segment.content, + provider="dify", + metadata={ + "doc_id": segment.index_node_id, + "document_id": document.id, + "dataset_id": dataset.id, + "score": 0.8, + }, + ) + ] + ), + ) + model = MagicMock() + model.model_type_instance.get_model_schema.return_value.features = None + monkeypatch.setattr( + retrieval_module.ModelManager, + "for_tenant", + MagicMock( + return_value=MagicMock( + get_model_instance=MagicMock(return_value=model), + ) + ), + ) + scope = FileAccessScope( + tenant_id=dataset.tenant_id, + user_id=str(uuid4()), + user_from=UserFrom.END_USER, + invoke_from=InvokeFrom.WEB_APP, + ) + with bind_file_access_scope(scope): + if entry == "chat": + content, files = retrieval.retrieve( + sqlite_session, + app_id=str(uuid4()), + user_id=scope.user_id, + tenant_id=dataset.tenant_id, + model_config=MagicMock(), + config=DatasetEntity( + dataset_ids=[dataset.id], + retrieve_config=DatasetRetrieveConfigEntity(retrieve_strategy="multiple"), + ), + query="picture", + invoke_from=InvokeFrom.WEB_APP, + show_retrieve_source=False, + hit_callback=MagicMock(), + message_id=str(uuid4()), + vision_enabled=vision_enabled, + ) + assert files is not None + assert [file.storage_key for file in files] == (["valid.png"] if vision_enabled else []) + else: + sources = retrieval.knowledge_retrieval( + sqlite_session, + KnowledgeRetrievalRequest( + tenant_id=dataset.tenant_id, + user_id=scope.user_id, + app_id=str(uuid4()), + user_from=UserFrom.END_USER.value, + dataset_ids=[dataset.id], + query="picture", + retrieval_mode="multiple", + ), + ) + assert len(sources) == 1 + content = sources[0].content + current_scope = get_current_file_access_scope() + assert current_scope is not None + assert current_scope.granted_upload_file_ids == frozenset([file_ids["valid"]]) + assert segment.id in current_scope.granted_retriever_segment_ids + assert content is not None + assert f"/files/{file_ids['valid']}/image-preview?timestamp=" in content + assert "sign=old" not in content + for kind, file_id in file_ids.items(): + if kind != "valid": + assert f"/files/{file_id}/image-preview)" in content + + +@pytest.mark.parametrize("broken_owner", ["tenant", "dataset", "document"]) +def test_invalid_segment_owner_chain_grants_no_access( + graph: tuple[Dataset, Document, DocumentSegment, dict[str, str]], + sqlite_session: Session, + broken_owner: str, +) -> None: + dataset, document, segment, _ = graph + if broken_owner == "tenant": + document.tenant_id = str(uuid4()) + elif broken_owner == "dataset": + document.dataset_id = str(uuid4()) + else: + segment.document_id = str(uuid4()) + sqlite_session.commit() + scope = FileAccessScope( + tenant_id=dataset.tenant_id, + user_id=str(uuid4()), + user_from=UserFrom.END_USER, + invoke_from=InvokeFrom.WEB_APP, + ) + with bind_file_access_scope(scope): + assert ( + authorize_retrieved_segment( + segment, tenant_id=dataset.tenant_id, dataset_ids=[dataset.id], session=sqlite_session + ) + is None + ) + current_scope = get_current_file_access_scope() + assert current_scope is not None + assert current_scope.granted_upload_file_ids == frozenset() + assert current_scope.granted_retriever_segment_ids == frozenset() + + +@pytest.mark.parametrize("with_image", [False, True]) +@pytest.mark.parametrize("signing_fails", [False, True]) +def test_workflow_retrieval_releases_owned_connections_and_preserves_caller_transaction( + graph: tuple[Dataset, Document, DocumentSegment, dict[str, str]], + sqlite_session: Session, + monkeypatch: pytest.MonkeyPatch, + with_image: bool, + signing_fails: bool, +) -> None: + dataset, document, segment, _ = graph + if not with_image: + segment.content = "Plain text" + sqlite_session.commit() + retrieval = DatasetRetrieval() + monkeypatch.setattr(retrieval, "_check_knowledge_rate_limit", MagicMock()) + monkeypatch.setattr(retrieval, "_get_available_datasets", MagicMock(return_value=[dataset])) + monkeypatch.setattr( + retrieval, + "multiple_retrieve", + MagicMock( + return_value=[ + RagDocument( + page_content=segment.content, + provider="dify", + metadata={"doc_id": segment.index_node_id, "document_id": document.id, "dataset_id": dataset.id}, + ) + ] + ), + ) + if signing_fails: + monkeypatch.setattr( + retrieval_module, "sign_segment_content", MagicMock(side_effect=RuntimeError("signing failed")) + ) + request = KnowledgeRetrievalRequest( + tenant_id=dataset.tenant_id, + user_id="visitor", + app_id=str(uuid4()), + user_from=UserFrom.END_USER.value, + dataset_ids=[dataset.id], + query="picture", + retrieval_mode="multiple", + ) + scope = FileAccessScope( + tenant_id=dataset.tenant_id, + user_id="visitor", + user_from=UserFrom.END_USER, + invoke_from=InvokeFrom.WEB_APP, + ) + engine = sqlite_session.get_bind() + assert isinstance(engine, Engine) + sessions: set[Session] = set() + pool_events: list[str] = [] + + def track_transaction(session: Session, *_args: object) -> None: + # Hold references so garbage collection cannot hide an unclosed session. + sessions.add(session) + + def track_checkout(*_args: object) -> None: + pool_events.append("checkout") + + def track_checkin(*_args: object) -> None: + pool_events.append("checkin") + + event.listen(Session, "after_begin", track_transaction) + event.listen(engine, "checkout", track_checkout) + event.listen(engine, "checkin", track_checkin) + try: + with bind_file_access_scope(scope), Session(engine) as caller_session, caller_session.begin() as caller_txn: + if signing_fails: + with pytest.raises(RuntimeError, match="signing failed"): + retrieval.knowledge_retrieval(caller_session, request) + else: + sources = retrieval.knowledge_retrieval(caller_session, request) + assert len(sources) == 1 + assert sources[0].metadata.segment_id == segment.id + assert caller_session.get_transaction() is caller_txn + assert caller_txn.is_active + assert sessions + assert all(not session.in_transaction() for session in sessions if session is not caller_session) + assert pool_events.count("checkout") > 0 + assert pool_events.count("checkout") == pool_events.count("checkin") + finally: + event.remove(Session, "after_begin", track_transaction) + event.remove(engine, "checkout", track_checkout) + event.remove(engine, "checkin", track_checkin) + for session in sessions: + session.close() diff --git a/api/tests/unit_tests/services/knowledge/test_segment_adapters.py b/api/tests/unit_tests/services/knowledge/test_segment_adapters.py new file mode 100644 index 00000000000..a3cac9fa35e --- /dev/null +++ b/api/tests/unit_tests/services/knowledge/test_segment_adapters.py @@ -0,0 +1,1729 @@ +from collections.abc import Iterator, Mapping, Sequence +from contextlib import nullcontext +from dataclasses import dataclass, field +from datetime import datetime +from pathlib import Path +from typing import Protocol +from unittest.mock import create_autospec, patch + +import pytest +from sqlalchemy import event, select +from sqlalchemy.dialects import postgresql +from sqlalchemy.engine import Connection +from sqlalchemy.orm import Session, sessionmaker + +from core.credit_usage import CreditUsageCreatedBy +from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError +from core.model_context import get_credit_usage_metadata +from core.model_manager import ModelManager +from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType +from core.rag.index_processor.index_processor_base import SummaryIndexSettingDict +from core.rag.index_processor.processor.paragraph_index_processor import ParagraphIndexProcessor +from core.rag.models.document import ChildDocument +from core.rag.models.document import Document as IndexDocument +from extensions.storage.storage_type import StorageType +from graphon.model_runtime.entities.llm_entities import LLMResult, LLMUsage +from graphon.model_runtime.entities.message_entities import AssistantPromptMessage +from libs.helper import generate_text_hash +from libs.pagination import PaginatedResult +from machinery.context import RequestContext +from models.dataset import ChildChunk, Dataset, Document, DocumentSegment, SegmentAttachmentBinding +from models.enums import CreatorUserRole, DocumentCreatedFrom, SegmentStatus, SummaryStatus +from models.model import UploadFile +from repositories.knowledge.dataset_repository import SQLAlchemyDatasetRepository +from repositories.knowledge.segment_repository import SQLAlchemySegmentRepository +from repositories.knowledge.upload_file_repository import SQLAlchemyKnowledgeUploadRepository +from repositories.workspace_member_query_repository import WorkspaceMemberQueryRepository +from services.knowledge.dataset_access import DatasetAccessService +from services.knowledge.dataset_service import SegmentService +from services.knowledge.entities.segments import ChildChunkUpdateArgs, SegmentUpdateArgs +from services.knowledge.resource_scope import DatasetRef +from services.knowledge.segments.adapters import ( + CelerySegmentBatchImportDispatcher, + ModelManagerSegmentGuard, + RedisSegmentIndexingState, +) +from services.knowledge.segments.application import ( + DatasetSegmentApplicationService, + SegmentBatchImportDispatcher, + SegmentDatasetRecord, + SegmentDetail, + SegmentIndex, + SegmentIndexingState, + SegmentListFilter, + SegmentModelGuard, + SegmentModelProviderError, +) + + +@dataclass +class SegmentLimitsStub: + SINGLE_CHUNK_ATTACHMENT_LIMIT: int = 10 + + +def test_celery_dispatcher_preserves_task_argument_order() -> None: + calls: list[tuple[str, ...]] = [] + + def delay(*args: str) -> None: + calls.append(args) + + dispatcher = CelerySegmentBatchImportDispatcher(delay=delay) + dispatcher.dispatch( + job_id="job-1", + upload_file_id="file-1", + dataset_id="dataset-1", + document_id="document-1", + workspace_id="workspace-1", + actor_id="account-1", + ) + + assert calls == [("job-1", "file-1", "dataset-1", "document-1", "workspace-1", "account-1")] + + +def _dataset(dataset_id: str, workspace_id: str, *, permission: str = "partial_members") -> Dataset: + return Dataset( + id=dataset_id, + tenant_id=workspace_id, + name="Dataset", + description="", + provider="vendor", + permission=permission, + indexing_technique="economy", + created_by="account-1", + maintainer="maintainer-1", + ) + + +def _document(document_id: str, workspace_id: str, dataset_id: str) -> Document: + return Document( + id=document_id, + tenant_id=workspace_id, + dataset_id=dataset_id, + position=1, + data_source_type="upload_file", + batch="batch-1", + name="Document", + created_from=DocumentCreatedFrom.API, + created_by="account-1", + doc_form="text_model", + word_count=0, + ) + + +def _segment(segment_id: str, workspace_id: str, dataset_id: str, document_id: str) -> DocumentSegment: + segment = DocumentSegment( + tenant_id=workspace_id, + dataset_id=dataset_id, + document_id=document_id, + position=1, + content="content", + word_count=7, + tokens=2, + created_by="account-1", + status=SegmentStatus.COMPLETED, + ) + segment.id = segment_id + return segment + + +def _upload(file_id: str, workspace_id: str, name: str) -> UploadFile: + upload = UploadFile( + tenant_id=workspace_id, + storage_type=StorageType.OPENDAL, + key=f"key-{file_id}", + name=name, + size=1, + extension="csv", + mime_type="text/csv", + created_by_role=CreatorUserRole.ACCOUNT, + created_by="account-1", + created_at=datetime(2026, 1, 1), + used=False, + ) + upload.id = file_id + return upload + + +def test_scope_lookup_enforces_owner_chain( + sqlite_session_factory: sessionmaker[Session], +) -> None: + with sqlite_session_factory.begin() as session: + session.add(_dataset("dataset-1", "workspace-1")) + session.add(_document("document-1", "workspace-1", "dataset-1")) + + store = SQLAlchemyDatasetRepository(session_factory=sqlite_session_factory) + scope = store.get_segment_scope( + dataset_ref=DatasetRef("workspace-1", "dataset-1"), + document_id="document-1", + actor_id="account-1", + ) + + assert scope.dataset is not None + assert scope.document is not None + assert ( + store.get_segment_scope( + dataset_ref=DatasetRef("workspace-2", "dataset-1"), + document_id="document-1", + actor_id="account-1", + ).dataset + is None + ) + + +def test_upload_lookup_is_tenant_scoped(sqlite_session_factory: sessionmaker[Session]) -> None: + with sqlite_session_factory.begin() as session: + session.add_all( + [ + _upload("file-1", "workspace-1", "one.csv"), + _upload("file-2", "workspace-2", "two.csv"), + ] + ) + + store = SQLAlchemyKnowledgeUploadRepository(session_factory=sqlite_session_factory) + + assert store.get_file_name(workspace_id="workspace-1", upload_file_id="file-1") == "one.csv" + assert store.get_file_name(workspace_id="workspace-1", upload_file_id="file-2") is None + + +def test_segment_lookup_rejects_segment_from_another_owner_chain( + sqlite_session_factory: sessionmaker[Session], +) -> None: + with sqlite_session_factory.begin() as session: + session.add_all( + [ + _dataset("dataset-1", "workspace-1"), + _dataset("dataset-2", "workspace-2"), + _document("document-1", "workspace-1", "dataset-1"), + _document("document-2", "workspace-2", "dataset-2"), + _segment("segment-1", "workspace-2", "dataset-2", "document-2"), + ] + ) + + store = SQLAlchemySegmentRepository(session_factory=sqlite_session_factory) + assert store.get_segment(DatasetRef("workspace-1", "dataset-1").document("document-1").segment("segment-2")) is None + + +@pytest.mark.parametrize("owner_field", ["tenant_id", "dataset_id", "document_id"]) +def test_segment_update_state_enforces_complete_owner_chain( + sqlite_session_factory: sessionmaker[Session], owner_field: str +) -> None: + with sqlite_session_factory.begin() as session: + segment = _segment("segment-1", "workspace-1", "dataset-1", "document-1") + setattr(segment, owner_field, "foreign-owner") + session.add(segment) + store = SQLAlchemySegmentRepository(session_factory=sqlite_session_factory) + ref = DatasetRef("workspace-1", "dataset-1").document("document-1").segment("segment-1") + + assert store.get_segment_update_state(ref) is None + + +def test_empty_segment_list_returns_materialized_page(sqlite_session_factory: sessionmaker[Session]) -> None: + with sqlite_session_factory.begin() as session: + session.add(_dataset("dataset-1", "workspace-1")) + session.add(_document("document-1", "workspace-1", "dataset-1")) + + store = SQLAlchemySegmentRepository(session_factory=sqlite_session_factory) + with patch("repositories.knowledge.segment_repository.dify_config") as config: + config.SQLALCHEMY_DATABASE_URI_SCHEME = "mysql" + result = store.list_segments( + DatasetRef("workspace-1", "dataset-1").document("document-1"), + SegmentListFilter(page=1, limit=20, keyword="missing"), + ) + + assert result.items == () + assert result.total == 0 + assert result.limit == 20 + + +def test_postgresql_keyword_filter_keeps_json_array_guard( + sqlite_session_factory: sessionmaker[Session], +) -> None: + with sqlite_session_factory.begin() as session: + session.add(_dataset("dataset-1", "workspace-1")) + session.add(_document("document-1", "workspace-1", "dataset-1")) + + page = PaginatedResult[DocumentSegment](items=[], total=0, page=1, per_page=20) + store = SQLAlchemySegmentRepository(session_factory=sqlite_session_factory) + with ( + patch("repositories.knowledge.segment_repository.dify_config") as config, + patch("repositories.knowledge.segment_repository.paginate_query", return_value=page) as paginate, + ): + config.SQLALCHEMY_DATABASE_URI_SCHEME = "postgresql" + store.list_segments( + DatasetRef("workspace-1", "dataset-1").document("document-1"), + SegmentListFilter(keyword="知识"), + ) + + statement = paginate.call_args.args[0] + sql = str(statement.compile(dialect=postgresql.dialect())) + assert "jsonb_typeof" in sql + assert "jsonb_array_elements_text" in sql + assert "CASE WHEN" in sql + + +def _application( + factory: sessionmaker[Session], *, indexing: SegmentIndex | None = None +) -> DatasetSegmentApplicationService: + datasets = SQLAlchemyDatasetRepository(session_factory=factory) + index = create_autospec(SegmentIndex, instance=True, spec_set=True) + index.count_tokens.return_value = 0 + state = create_autospec(SegmentIndexingState, instance=True, spec_set=True) + state.is_segment_indexing.return_value = False + state.is_document_indexing.return_value = False + return DatasetSegmentApplicationService( + dataset_access=DatasetAccessService( + datasets=datasets, + workspace_roles=WorkspaceMemberQueryRepository(session_factory=factory), + legacy_permissions_enabled=False, + ), + scopes=datasets, + store=SQLAlchemySegmentRepository(session_factory=factory), + limits=SegmentLimitsStub(), + text_hash=generate_text_hash, + index=indexing or index, + uploads=SQLAlchemyKnowledgeUploadRepository(session_factory=factory), + model_guard=create_autospec(SegmentModelGuard, instance=True, spec_set=True), + indexing_state=state, + batch_dispatcher=create_autospec(SegmentBatchImportDispatcher, instance=True, spec_set=True), + job_id_factory=lambda: "job-1", + ) + + +def test_mutations_use_request_actor_without_a_flask_request( + sqlite_session_factory: sessionmaker[Session], +) -> None: + + with sqlite_session_factory.begin() as session: + session.add(_dataset("dataset-1", "workspace-1")) + session.add(_document("document-1", "workspace-1", "dataset-1")) + app = _application(sqlite_session_factory) + context = RequestContext("request-1", None, "author-1", "workspace-1") + created = app.create_segment( + context, dataset_id="dataset-1", document_id="document-1", values={"content": "original"} + ) + assert created.data.created_by == "author-1" + updated = app.update_segment( + RequestContext("request-2", None, "editor-2", "workspace-1"), + dataset_id="dataset-1", + document_id="document-1", + segment_id=created.data.id, + values={"content": "updated content"}, + ) + assert updated.data.updated_by == "editor-2" + assert updated.data.content == "updated content" + assert isinstance(updated.data.created_at, datetime) + with sqlite_session_factory() as session: + segment = session.get(DocumentSegment, created.data.id) + assert segment is not None + assert segment.created_by == "author-1" + assert segment.updated_by == "editor-2" + assert segment.status == SegmentStatus.COMPLETED + + +def test_list_reads_dataset_and_document_once(sqlite_session_factory: sessionmaker[Session]) -> None: + with sqlite_session_factory.begin() as session: + session.add(_dataset("dataset-1", "workspace-1")) + session.add(_document("document-1", "workspace-1", "dataset-1")) + statements: list[str] = [] + engine = sqlite_session_factory.kw["bind"] + + def record(_conn, _cursor, statement: str, _parameters, _context, _executemany) -> None: + statements.append(statement) + + event.listen(engine, "before_cursor_execute", record) + try: + page = _application(sqlite_session_factory).list_segments( + RequestContext("request-1", None, "account-1", "workspace-1"), + dataset_id="dataset-1", + document_id="document-1", + query=SegmentListFilter(), + ) + finally: + event.remove(engine, "before_cursor_execute", record) + assert page.items == () + assert sum("FROM datasets" in statement for statement in statements) == 1 + assert sum("JOIN documents" in statement or "FROM documents" in statement for statement in statements) == 1 + + +@pytest.mark.parametrize( + ("error", "kind"), + [ + (LLMBadRequestError(), "bad_request"), + (ProviderTokenNotInitError("token missing"), "token"), + ], +) +def test_model_guard_translates_provider_errors(error: Exception, kind: str) -> None: + from core.model_manager import ModelManager + + model = create_autospec(ModelManager, instance=True, spec_set=True) + model.get_model_instance.side_effect = error + dataset = SegmentDatasetRecord( + id="dataset-1", + workspace_id="workspace-1", + indexing_technique="high_quality", + embedding_model_provider="provider", + embedding_model="model", + ) + + with patch("services.knowledge.segments.adapters.ModelManager.for_tenant", return_value=model): + with pytest.raises(SegmentModelProviderError) as raised: + ModelManagerSegmentGuard().check(dataset) + + assert raised.value.kind == kind + + +def test_redis_state_uses_existing_keys_and_decodes_status() -> None: + from services.knowledge.segments.adapters import RedisSegmentClient + + redis = create_autospec(RedisSegmentClient, instance=True, spec_set=True) + redis.get.side_effect = [b"running", b"completed"] + state = RedisSegmentIndexingState(redis) + + assert state.is_document_indexing("document-1") is True + state.set_batch_waiting("job-1") + assert state.get_batch_status("job-1") == "completed" + redis.get.assert_any_call("document_document-1_indexing") + redis.setnx.assert_called_once_with("segment_batch_import_job-1", "waiting") + + +def test_child_chunk_mutations_preserve_explicit_actors( + sqlite_session_factory: sessionmaker[Session], +) -> None: + + with sqlite_session_factory.begin() as session: + session.add(_dataset("dataset-1", "workspace-1")) + document = _document("document-1", "workspace-1", "dataset-1") + document.doc_form = IndexStructureType.PARENT_CHILD_INDEX + session.add(document) + session.add(_segment("segment-1", "workspace-1", "dataset-1", "document-1")) + app = _application(sqlite_session_factory) + created = app.create_child_chunk( + RequestContext("create", None, "creator", "workspace-1"), + dataset_id="dataset-1", + document_id="document-1", + segment_id="segment-1", + content="first", + ) + app.update_child_chunk( + RequestContext("update", None, "editor", "workspace-1"), + dataset_id="dataset-1", + document_id="document-1", + segment_id="segment-1", + child_chunk_id=created.id, + content="second", + ) + with sqlite_session_factory() as session: + chunk = session.get(ChildChunk, created.id) + assert chunk is not None + assert chunk.created_by == "creator" + assert chunk.updated_by == "editor" + results = app.update_child_chunks( + RequestContext("batch", None, "batch-editor", "workspace-1"), + dataset_id="dataset-1", + document_id="document-1", + segment_id="segment-1", + chunks=[ChildChunkUpdateArgs(id=created.id, content="third"), ChildChunkUpdateArgs(content="new child")], + ) + with sqlite_session_factory() as session: + chunks = [session.get(ChildChunk, result.id) for result in results] + assert len(chunks) == 2 + existing, added = chunks + assert existing is not None + assert added is not None + assert existing.updated_by == "batch-editor" + assert added.created_by == "batch-editor" + + +@dataclass +class IndexRedisStub: + def get(self, name: str | bytes) -> None: + del name + + def setnx(self, name: str | bytes, value: object) -> bool: + del name, value + return True + + def setex(self, name: str, time: int, value: object) -> bool: + del name, time, value + return True + + def lock(self, name: str, timeout: int) -> nullcontext[None]: + del name, timeout + return nullcontext() + + +@dataclass +class VectorProbe: + active_transactions: set[Connection] + fail: bool = False + writes: int = 0 + deleted_ids: list[str] = field(default_factory=list) + tasks: list[tuple[str, tuple[object, ...]]] = field(default_factory=list) + batches: list[tuple[str, ...]] = field(default_factory=list) + + def add_texts(self, documents: Sequence[IndexDocument], **_kwargs: object) -> None: + assert not self.active_transactions, "Index I/O must not hold a DB transaction" + self.writes += len(documents) + if self.fail: + raise RuntimeError("index unavailable") + + def create(self, documents: Sequence[IndexDocument], **kwargs: object) -> None: + self.batches.append(tuple(document.page_content for document in documents)) + self.add_texts(documents, **kwargs) + + def delete_by_ids(self, ids: Sequence[str]) -> None: + assert ids + assert not self.active_transactions + if self.fail: + raise RuntimeError("index unavailable") + self.deleted_ids.extend(ids) + + def create_multimodal(self, documents: Sequence[IndexDocument], **kwargs: object) -> None: + assert not self.active_transactions + assert kwargs["upload_files"] + self.add_texts(documents) + + +type IndexingProbe = tuple[DatasetSegmentApplicationService, SQLAlchemySegmentRepository, VectorProbe] + + +@pytest.fixture +def indexing_probe( + sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch +) -> Iterator[IndexingProbe]: + from services.knowledge.segments import indexing as indexing_module + + with sqlite_session_factory.begin() as session: + dataset = _dataset("dataset-1", "workspace-1") + dataset.indexing_technique = IndexTechniqueType.HIGH_QUALITY + dataset.index_struct = '{"type": "qdrant"}' + session.add(dataset) + session.add(_document("document-1", "workspace-1", "dataset-1")) + engine = sqlite_session_factory.kw["bind"] + active: set[Connection] = set() + begin = active.add + finish = active.discard + event.listen(engine, "begin", begin) + event.listen(engine, "commit", finish) + event.listen(engine, "rollback", finish) + probe = VectorProbe(active) + + def vector(dataset: Dataset, *, session: Session | None, vector_type: str) -> VectorProbe: + assert dataset.tenant_id == "workspace-1" + assert session is None + assert vector_type == "qdrant" + assert not active + return probe + + class TokenModel: + def get_text_embedding_num_tokens(self, texts: Sequence[str]) -> list[int]: + assert not active + return [len(text) for text in texts] + + class ModelLookup: + def get_model_instance(self, **kwargs: object) -> TokenModel: + assert kwargs["tenant_id"] == "workspace-1" + assert not active + return TokenModel() + + monkeypatch.setattr(indexing_module, "Vector", vector) + monkeypatch.setattr(indexing_module.ModelManager, "for_tenant", lambda **_kwargs: ModelLookup()) + repository = SQLAlchemySegmentRepository(session_factory=sqlite_session_factory) + + def dispatch(kind: str, *args: object) -> None: + assert not active + probe.tasks.append((kind, args)) + + gateway = indexing_module.SegmentIndexingGateway( + segments=repository, + uploads=SQLAlchemyKnowledgeUploadRepository(session_factory=sqlite_session_factory), + redis=IndexRedisStub(), + new_session=sqlite_session_factory, + delete_task=lambda *args: dispatch("delete", *args), + enable_task=lambda *args: dispatch("enable", *args), + disable_task=lambda *args: dispatch("disable", *args), + ) + yield _application(sqlite_session_factory, indexing=gateway), repository, probe + event.remove(engine, "begin", begin) + event.remove(engine, "commit", finish) + event.remove(engine, "rollback", finish) + + +class SegmentUpdateEntry(Protocol): + def __call__( + self, + context: RequestContext, + *, + dataset_id: str, + document_id: str, + segment_id: str, + values: Mapping[str, object], + ) -> SegmentDetail: ... + + +@pytest.fixture(params=["console", "service_api"]) +def segment_update( + request: pytest.FixtureRequest, indexing_probe: IndexingProbe, sqlite_session_factory: sessionmaker[Session] +) -> SegmentUpdateEntry: + app, repository, _ = indexing_probe + if request.param == "console": + return app.update_segment + + def update( + context: RequestContext, + *, + dataset_id: str, + document_id: str, + segment_id: str, + values: Mapping[str, object], + ) -> SegmentDetail: + with sqlite_session_factory() as session: + dataset = session.get(Dataset, dataset_id) + document = session.get(Document, document_id) + segment = session.get(DocumentSegment, segment_id) + assert dataset is not None + assert document is not None + assert segment is not None + result = SegmentService.update_segment( + SegmentUpdateArgs.model_validate(values), + segment, + document, + dataset, + session, + actor_id=context.account_id, + mutations=app.mutations, + ) + assert result.content == (values.get("content") or segment.content) + ref = DatasetRef(context.active_workspace_id, dataset_id).document(document_id).segment(segment_id) + detail = repository.get_segment(ref) + assert detail is not None + return detail + + return update + + +@pytest.mark.parametrize("entry", ["console", "service_api"]) +@pytest.mark.parametrize("enabled", [True, False]) +def test_delete_segment_with_legacy_scalar_keywords( + indexing_probe: IndexingProbe, + sqlite_session_factory: sessionmaker[Session], + entry: str, + enabled: bool, +) -> None: + app, _, probe = indexing_probe + with sqlite_session_factory.begin() as session: + segment = _segment("segment-1", "workspace-1", "dataset-1", "document-1") + segment.keywords = "legacy-scalar" + segment.enabled = enabled + segment.index_node_id = "index-1" + session.add(segment) + document = session.get(Document, "document-1") + assert document is not None + document.word_count = segment.word_count + + if entry == "console": + app.delete_segment( + RequestContext("request", None, "author", "workspace-1"), + dataset_id="dataset-1", + document_id="document-1", + segment_id="segment-1", + ) + else: + with sqlite_session_factory() as session: + dataset = session.get(Dataset, "dataset-1") + document = session.get(Document, "document-1") + segment = session.get(DocumentSegment, "segment-1") + assert dataset is not None + assert document is not None + assert segment is not None + SegmentService.delete_segment(segment, document, dataset, session, mutations=app.mutations) + + with sqlite_session_factory() as session: + assert session.get(DocumentSegment, "segment-1") is None + document = session.get(Document, "document-1") + assert document is not None + assert document.word_count == 0 + if enabled: + assert probe.tasks == [("delete", (["index-1"], "dataset-1", "document-1", ["segment-1"], []))] + else: + assert probe.tasks == [] + + +@pytest.mark.parametrize("keywords", ["legacy-scalar", 42, {"invalid": "object"}, [{"invalid": "item"}]]) +def test_update_repairs_legacy_keywords_without_validating_presentation_first( + segment_update: SegmentUpdateEntry, + sqlite_session_factory: sessionmaker[Session], + keywords: object, +) -> None: + with sqlite_session_factory.begin() as session: + segment = _segment("segment-1", "workspace-1", "dataset-1", "document-1") + segment.keywords = keywords + segment.index_node_id = "index-1" + segment.index_node_hash = generate_text_hash(segment.content) + session.add(segment) + + result = segment_update( + RequestContext("request", None, "editor", "workspace-1"), + dataset_id="dataset-1", + document_id="document-1", + segment_id="segment-1", + values={"content": "content", "keywords": ["replacement"]}, + ) + + assert result.data.status == "completed" + assert result.data.keywords == ("replacement",) + with sqlite_session_factory() as session: + segment = session.get(DocumentSegment, "segment-1") + assert segment is not None + assert segment.keywords == ["replacement"] + + +@pytest.mark.parametrize("content", ["original", "changed text"]) +@pytest.mark.parametrize( + ("attachment_values", "expected_ids", "removed_original"), + [ + ({}, ["image-1"], False), + ({"attachment_ids": None}, ["image-1"], False), + ({"attachment_ids": []}, [], True), + ({"attachment_ids": ["image-2"]}, ["image-2"], True), + ], + ids=["omitted", "null", "clear", "replace"], +) +def test_segment_update_preserves_or_replaces_attachments( + segment_update: SegmentUpdateEntry, + indexing_probe: IndexingProbe, + sqlite_session_factory: sessionmaker[Session], + content: str, + attachment_values: dict[str, object], + expected_ids: list[str], + removed_original: bool, +) -> None: + app, _, probe = indexing_probe + context = RequestContext("request", None, "author", "workspace-1") + with sqlite_session_factory.begin() as session: + session.add_all( + [_upload("image-1", "workspace-1", "chart.png"), _upload("image-2", "workspace-1", "replacement.png")] + ) + dataset = session.get(Dataset, "dataset-1") + assert dataset is not None + dataset.is_multimodal = True + created = app.create_segment( + context, + dataset_id="dataset-1", + document_id="document-1", + values={"content": "original", "attachment_ids": ["image-1"]}, + ) + assert [attachment.id for attachment in created.data.attachments] == ["image-1"] + + updated = segment_update( + context, + dataset_id="dataset-1", + document_id="document-1", + segment_id=created.data.id, + values={"content": content, **attachment_values}, + ) + + assert updated.data.status == "completed" + assert [attachment.id for attachment in updated.data.attachments] == expected_ids + with sqlite_session_factory() as session: + attachment_ids = list( + session.scalars( + select(SegmentAttachmentBinding.attachment_id).where( + SegmentAttachmentBinding.segment_id == created.data.id + ) + ) + ) + assert attachment_ids == expected_ids + assert ("image-1" in probe.deleted_ids) is removed_original + + +def test_segment_indexing_and_summary_use_closed_read_sessions( + segment_update: SegmentUpdateEntry, indexing_probe: IndexingProbe, sqlite_session_factory: sessionmaker[Session] +) -> None: + from models.dataset import DocumentSegmentSummary + + app, repository, probe = indexing_probe + context = RequestContext("request", None, "author", "workspace-1") + created = app.create_segment(context, dataset_id="dataset-1", document_id="document-1", values={"content": "first"}) + updated = segment_update( + context, + dataset_id="dataset-1", + document_id="document-1", + segment_id=created.data.id, + values={"content": "second", "summary": "manual summary"}, + ) + assert updated.data.status == "completed" + assert updated.data.summary == "manual summary" + with sqlite_session_factory() as session: + summary = session.scalar(select(DocumentSegmentSummary)) + assert summary is not None + assert summary.status == "completed" + document = session.get(Document, "document-1") + assert document is not None + assert document.word_count == len("second") + assert probe.writes >= 3 + probe.fail = True + failed = segment_update( + context, + dataset_id="dataset-1", + document_id="document-1", + segment_id=created.data.id, + values={"content": "failed change"}, + ) + assert failed.data.status == "error" + assert failed.data.enabled is False + assert failed.data.error == "index unavailable" + + +@pytest.mark.parametrize("operation", ["create", "update", "replace", "delete"]) +def test_child_index_failure_leaves_database_unchanged(indexing_probe: IndexingProbe, operation: str) -> None: + from services.knowledge.segments.application import ( + ChildChunkDeleteIndexApplicationError, + ChildChunkIndexingApplicationError, + ) + + app, repository, probe = indexing_probe + context = RequestContext("request", None, "author", "workspace-1") + parent = app.create_segment(context, dataset_id="dataset-1", document_id="document-1", values={"content": "parent"}) + child = app.create_child_chunk( + context, + dataset_id="dataset-1", + document_id="document-1", + segment_id=parent.data.id, + content="original", + ) + ref = DatasetRef("workspace-1", "dataset-1").document("document-1").segment(parent.data.id) + before = repository.get_children(ref) + probe.fail = True + + def mutate() -> None: + match operation: + case "create": + app.create_child_chunk( + context, + dataset_id="dataset-1", + document_id="document-1", + segment_id=parent.data.id, + content="new", + ) + case "update": + app.update_child_chunk( + context, + dataset_id="dataset-1", + document_id="document-1", + segment_id=parent.data.id, + child_chunk_id=child.id, + content="changed", + ) + case "replace": + app.update_child_chunks( + context, + dataset_id="dataset-1", + document_id="document-1", + segment_id=parent.data.id, + chunks=[ChildChunkUpdateArgs(content="replacement")], + ) + case "delete": + app.delete_child_chunk( + context, + dataset_id="dataset-1", + document_id="document-1", + segment_id=parent.data.id, + child_chunk_id=child.id, + ) + + with pytest.raises((ChildChunkIndexingApplicationError, ChildChunkDeleteIndexApplicationError)): + mutate() + assert repository.get_children(ref) == before + + +def test_upload_batch_uses_one_scoped_query(sqlite_session_factory: sessionmaker[Session]) -> None: + with sqlite_session_factory.begin() as session: + session.add_all( + [ + _upload("file-1", "workspace-1", "one.csv"), + _upload("file-2", "workspace-1", "two.csv"), + _upload("foreign", "workspace-2", "foreign.csv"), + ] + ) + engine = sqlite_session_factory.kw["bind"] + statements = [] + + def record_sql(_conn, _cursor, statement: str, _parameters, _context, _executemany) -> None: + statements.append(statement) + + event.listen(engine, "before_cursor_execute", record_sql) + try: + uploads = SQLAlchemyKnowledgeUploadRepository(session_factory=sqlite_session_factory).get_by_ids( + workspace_id="workspace-1", file_ids=("file-2", "file-1", "foreign", "missing") + ) + finally: + event.remove(engine, "before_cursor_execute", record_sql) + assert set(uploads) == {"file-1", "file-2"} + assert len(statements) == 1 + + +def test_child_query_escapes_wildcards_and_enforces_complete_scope( + sqlite_session_factory: sessionmaker[Session], +) -> None: + from services.knowledge.segments.application import ChildChunkListFilter + + with sqlite_session_factory.begin() as session: + session.add(_dataset("dataset-1", "workspace-1")) + session.add(_document("document-1", "workspace-1", "dataset-1")) + session.add(_segment("segment-1", "workspace-1", "dataset-1", "document-1")) + for position, (workspace, text) in enumerate( + [("workspace-1", "50%_literal"), ("workspace-1", "5000literal"), ("workspace-2", "50%_literal")], start=1 + ): + session.add( + ChildChunk( + tenant_id=workspace, + dataset_id="dataset-1", + document_id="document-1", + segment_id="segment-1", + position=position, + content=text, + word_count=len(text), + created_by="author", + ) + ) + repo = SQLAlchemySegmentRepository(session_factory=sqlite_session_factory) + page = repo.list_child_chunks( + DatasetRef("workspace-1", "dataset-1").document("document-1").segment("segment-1"), + ChildChunkListFilter(keyword="%_", limit=1), + ) + assert page is not None + assert page.total == 1 + assert [child.content for child in page.items] == ["50%_literal"] + + +def test_batch_status_and_delete_only_modify_scoped_segments( + indexing_probe: IndexingProbe, sqlite_session_factory: sessionmaker[Session] +) -> None: + app, repository, probe = indexing_probe + context = RequestContext("request", None, "author", "workspace-1") + first = app.create_segment(context, dataset_id="dataset-1", document_id="document-1", values={"content": "first"}) + second = app.create_segment(context, dataset_id="dataset-1", document_id="document-1", values={"content": "second"}) + with sqlite_session_factory.begin() as session: + session.add(_segment("foreign", "workspace-2", "dataset-1", "document-1")) + segment_ids = [first.data.id, second.data.id, "foreign"] + app.change_segment_status( + context, dataset_id="dataset-1", document_id="document-1", segment_ids=segment_ids, action="disable" + ) + kind, task_args = probe.tasks[-1] + assert kind == "disable" + assert isinstance(task_args[0], list) + assert set(task_args[0]) == {first.data.id, second.data.id} + assert task_args[1:] == ("dataset-1", "document-1") + with sqlite_session_factory() as session: + first_row = session.get(DocumentSegment, first.data.id) + assert first_row is not None + assert first_row.enabled is False + foreign_row = session.get(DocumentSegment, "foreign") + assert foreign_row is not None + assert foreign_row.enabled is True + app.delete_segments(context, dataset_id="dataset-1", document_id="document-1", segment_ids=segment_ids) + with sqlite_session_factory() as session: + assert session.get(DocumentSegment, first.data.id) is None + assert session.get(DocumentSegment, second.data.id) is None + assert session.get(DocumentSegment, "foreign") is not None + document = session.get(Document, "document-1") + assert document is not None + assert document.word_count == 0 + + +def test_child_regeneration_reuses_splitter_outside_transaction( + segment_update: SegmentUpdateEntry, + indexing_probe: IndexingProbe, + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, +) -> None: + import json + + from core.rag.index_processor.processor.parent_child_index_processor import ParentChildIndexProcessor + from models.dataset import DatasetProcessRule + from models.enums import ProcessRuleMode + + app, repository, probe = indexing_probe + with sqlite_session_factory.begin() as session: + dataset = session.get(Dataset, "dataset-1") + assert dataset is not None + dataset.embedding_model_provider = "provider" + rule = DatasetProcessRule( + dataset_id="dataset-1", + mode=ProcessRuleMode.HIERARCHICAL, + rules=json.dumps({"parent_mode": "full-doc", "subchunk_segmentation": {"max_tokens": 100}}), + created_by="author", + ) + session.add(rule) + document = session.get(Document, "document-1") + assert document is not None + document.doc_form = IndexStructureType.PARENT_CHILD_INDEX + document.dataset_process_rule_id = rule.id + + def split(_self, document: IndexDocument, _rule, _mode, _model) -> list[ChildDocument]: + assert not probe.active_transactions + return [ + ChildDocument( + page_content=document.page_content, metadata={"doc_id": document.page_content, "doc_hash": "hash"} + ) + ] + + monkeypatch.setattr(ParentChildIndexProcessor, "split_child_nodes", split) + context = RequestContext("request", None, "author", "workspace-1") + created = app.create_segment(context, dataset_id="dataset-1", document_id="document-1", values={"content": "first"}) + assert created.data.status == "completed" + updated = segment_update( + context, + dataset_id="dataset-1", + document_id="document-1", + segment_id=created.data.id, + values={"content": "second", "regenerate_child_chunks": True}, + ) + assert updated.data.status == "completed" + ref = DatasetRef("workspace-1", "dataset-1").document("document-1").segment(created.data.id) + children = repository.get_children(ref) + assert children is not None + assert [child.data.content for child in children] == ["second"] + + +@pytest.mark.parametrize( + ("content", "submitted_summary", "summary_enabled", "expected_summary", "generated"), + [ + ("new text", "old summary", True, "new summary", True), + ("new text", None, True, "new summary", True), + ("old text", "old summary", True, "old summary", False), + ("new text", "manual summary", True, "manual summary", False), + ("new text", "", True, None, False), + ("new text", " \n\t ", True, None, False), + ("new text", "old summary", False, "old summary", False), + ("new text", None, False, "old summary", False), + ], +) +def test_summary_updates_preserve_generation_and_manual_edit_rules( + segment_update: SegmentUpdateEntry, + indexing_probe: IndexingProbe, + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + content: str, + submitted_summary: str | None, + summary_enabled: bool, + expected_summary: str | None, + generated: bool, +) -> None: + app, repository, _ = indexing_probe + context = RequestContext("request", None, "author", "workspace-1") + created = app.create_segment( + context, dataset_id="dataset-1", document_id="document-1", values={"content": "old text"} + ) + reference = DatasetRef("workspace-1", "dataset-1").document("document-1").segment(created.data.id) + repository.save_summary(reference, "old summary") + with sqlite_session_factory.begin() as session: + dataset = session.get(Dataset, "dataset-1") + assert dataset is not None + dataset.summary_index_setting = { + "enable": summary_enabled, + "model_name": "llm-model", + "model_provider_name": "provider", + } + calls: list[str] = [] + + def generate(_tenant_id: str, text: str, _setting: object, **_kwargs: object) -> tuple[str, LLMUsage]: + calls.append(text) + return "new summary", LLMUsage.empty_usage() + + monkeypatch.setattr(ParagraphIndexProcessor, "generate_summary_from_inputs", staticmethod(generate)) + updated = segment_update( + context, + dataset_id="dataset-1", + document_id="document-1", + segment_id=created.data.id, + values={"content": content, "summary": submitted_summary}, + ) + assert calls == ([content] if generated else []) + assert updated.data.summary == expected_summary + + +def test_automatic_summary_keeps_knowledge_credit_attribution( + indexing_probe: IndexingProbe, + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, +) -> None: + app, repository, _ = indexing_probe + context = RequestContext("request", None, "author", "workspace-1") + created = app.create_segment( + context, dataset_id="dataset-1", document_id="document-1", values={"content": "old text"} + ) + reference = DatasetRef("workspace-1", "dataset-1").document("document-1").segment(created.data.id) + repository.save_summary(reference, "old summary") + setting: SummaryIndexSettingDict = {"enable": True, "model_name": "llm-model", "model_provider_name": "provider"} + with sqlite_session_factory.begin() as session: + dataset = session.get(Dataset, "dataset-1") + assert dataset is not None + dataset.summary_index_setting = setting + metadata: list[Mapping[str, object] | None] = [] + + class ModelProbe: + credentials: dict[str, str] = {} + + @property + def model_type_instance(self) -> "ModelProbe": + return self + + def get_model_schema(self, _name: str, _credentials: object) -> None: + return None + + def get_text_embedding_num_tokens(self, texts: Sequence[str]) -> list[int]: + return [len(text) for text in texts] + + def invoke_llm(self, **_kwargs: object) -> LLMResult: + metadata.append(get_credit_usage_metadata()) + return LLMResult( + model="llm-model", message=AssistantPromptMessage(content="new summary"), usage=LLMUsage.empty_usage() + ) + + def get_model_instance(self, **_kwargs: object) -> "ModelProbe": + return self + + monkeypatch.setattr(ModelManager, "for_tenant", lambda **_kwargs: ModelProbe()) + # Both public summary entry points must preserve the credit-usage context. + with sqlite_session_factory() as session: + ParagraphIndexProcessor.generate_summary("workspace-1", "old text", setting, session=session) + assert metadata.pop() == {"created_by": CreditUsageCreatedBy.KNOWLEDGE_INDEXING} + updated = app.update_segment( + context, + dataset_id="dataset-1", + document_id="document-1", + segment_id=created.data.id, + values={"content": "new text"}, + ) + assert updated.data.summary == "new summary" + assert metadata == [{"created_by": CreditUsageCreatedBy.KNOWLEDGE_INDEXING}] + assert get_credit_usage_metadata() is None + + +@pytest.mark.parametrize("storage_type", ["database", "file"]) +def test_economy_index_shares_persistence_with_existing_keyword_calls( + indexing_probe: IndexingProbe, + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + storage_type: str, +) -> None: + from core.rag.datasource.keyword.jieba import jieba as jieba_module + from core.rag.datasource.keyword.keyword_factory import Keyword + from repositories.knowledge import keyword_table_repository as table_module + from tests.unit_tests.config_override import apply_config_overrides + + app, repository, probe = indexing_probe + reads: list[str] = [] + writes: list[str] = [] + + class FileStorageProbe: + def load_once(self, key: str) -> bytes: + assert not probe.active_transactions + reads.append(key) + return (tmp_path / key).read_bytes() + + def exists(self, key: str) -> bool: + assert not probe.active_transactions + return (tmp_path / key).exists() + + def save(self, key: str, data: bytes) -> None: + assert not probe.active_transactions + writes.append(key) + path = tmp_path / key + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(data) + + def delete(self, key: str) -> None: + assert not probe.active_transactions + (tmp_path / key).unlink() + + class KeywordExtractor: + def extract_keywords(self, text: str, _keyword_number: int = 10) -> set[str]: + assert not probe.active_transactions + return set(text.split()) + + storage = FileStorageProbe() + monkeypatch.setattr(jieba_module, "storage", storage) + monkeypatch.setattr(jieba_module, "JiebaKeywordTableHandler", KeywordExtractor) + apply_config_overrides(monkeypatch, KEYWORD_DATA_SOURCE_TYPE=storage_type) + with sqlite_session_factory.begin() as session: + dataset = session.get(Dataset, "dataset-1") + assert dataset is not None + dataset.indexing_technique = IndexTechniqueType.ECONOMY + context = RequestContext("request", None, "author", "workspace-1") + first = app.create_segment( + context, dataset_id="dataset-1", document_id="document-1", values={"content": "initial 中文"} + ) + second = app.create_segment( + context, dataset_id="dataset-1", document_id="document-1", values={"content": "other", "keywords": ["manual"]} + ) + updated = app.update_segment( + context, + dataset_id="dataset-1", + document_id="document-1", + segment_id=first.data.id, + values={"content": "replacement"}, + ) + assert first.data.status == second.data.status == updated.data.status == "completed" + assert first.data.index_node_id is not None + assert second.data.index_node_id is not None + assert set(first.data.keywords or ()) == {"initial", "中文"} + assert second.data.keywords == ("manual",) + assert updated.data.keywords == ("replacement",) + dataset_ref = DatasetRef("workspace-1", "dataset-1") + source_type, data = repository.get_keyword_table(dataset_ref) + payload = table_module.load_keyword_table( + storage=storage, tenant_id="workspace-1", dataset_id="dataset-1", storage_type=source_type, data=data + ) + assert payload is not None + assert payload["__data__"]["table"] == { + "replacement": {first.data.index_node_id}, + "manual": {second.data.index_node_id}, + } + if storage_type == "database": + assert reads == writes == [] + else: + assert writes == ["keyword_files/workspace-1/dataset-1.txt"] * 3 + + # The existing session-taking API must consume the exact payload written by + # the repository-backed path. Only that existing API retains caller sessions. + class SessionFileStorage: + def load_once(self, key: str) -> bytes: + return (tmp_path / key).read_bytes() + + def exists(self, key: str) -> bool: + return (tmp_path / key).exists() + + def save(self, key: str, data: bytes) -> None: + (tmp_path / key).write_bytes(data) + + def delete(self, key: str) -> None: + (tmp_path / key).unlink() + + session_storage = SessionFileStorage() + monkeypatch.setattr(jieba_module, "storage", session_storage) + monkeypatch.setattr(jieba_module, "redis_client", IndexRedisStub()) + with sqlite_session_factory() as session: + dataset = session.get(Dataset, "dataset-1") + assert dataset is not None + keyword = Keyword(dataset) + assert keyword.text_exists(first.data.index_node_id, session=session) + assert keyword.text_exists(second.data.index_node_id, session=session) + keyword.delete_by_ids([first.data.index_node_id], session) + session.commit() + source_type, data = repository.get_keyword_table(dataset_ref) + payload = table_module.load_keyword_table( + storage=session_storage, tenant_id="workspace-1", dataset_id="dataset-1", storage_type=source_type, data=data + ) + assert payload is not None + assert payload["__data__"]["table"] == {"manual": {second.data.index_node_id}} + + +@pytest.mark.parametrize("generated_summary", [None, "", " \n\t "], ids=["exception", "empty", "whitespace"]) +def test_summary_generation_failure_preserves_the_completed_segment( + segment_update: SegmentUpdateEntry, + indexing_probe: IndexingProbe, + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + generated_summary: str | None, +) -> None: + from models.dataset import DocumentSegmentSummary + + app, repository, probe = indexing_probe + context = RequestContext("request", None, "author", "workspace-1") + created = app.create_segment( + context, dataset_id="dataset-1", document_id="document-1", values={"content": "original"} + ) + reference = DatasetRef("workspace-1", "dataset-1").document("document-1").segment(created.data.id) + repository.save_summary(reference, "original summary") + with sqlite_session_factory.begin() as session: + dataset = session.get(Dataset, "dataset-1") + assert dataset is not None + dataset.summary_index_setting = {"enable": True, "model_name": "model", "model_provider_name": "provider"} + summary = session.scalar(select(DocumentSegmentSummary)) + assert summary is not None + summary.summary_index_node_id = "original-summary-node" + + def generate(*_args: object, **_kwargs: object) -> tuple[str, LLMUsage]: + assert not probe.active_transactions + if generated_summary is not None: + return generated_summary, LLMUsage.empty_usage() + raise RuntimeError("LLM unavailable") + + monkeypatch.setattr(ParagraphIndexProcessor, "generate_summary_from_inputs", staticmethod(generate)) + updated = segment_update( + context, + dataset_id="dataset-1", + document_id="document-1", + segment_id=created.data.id, + values={"content": "replacement", "summary": "original summary"}, + ) + assert updated.data.content == "replacement" + assert updated.data.enabled is True + assert updated.data.status == "completed" + assert updated.data.summary == "original summary" + assert "original-summary-node" not in probe.deleted_ids + with sqlite_session_factory() as session: + summary = session.scalar(select(DocumentSegmentSummary)) + assert summary is not None + assert summary.status == "error" + assert summary.summary_index_node_id == "original-summary-node" + assert summary.error == ("Generated summary is empty" if generated_summary is not None else "LLM unavailable") + + +@pytest.mark.parametrize("failed_commits", [1, 2, 3]) +def test_summary_commit_retries_use_fresh_sessions_and_persist_outcome( + indexing_probe: IndexingProbe, + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + failed_commits: int, +) -> None: + from models.dataset import DocumentSegmentSummary + from services.knowledge.summaries import adapters as summary_module + + app, _, probe = indexing_probe + context = RequestContext("request", None, "author", "workspace-1") + created = app.create_segment( + context, dataset_id="dataset-1", document_id="document-1", values={"content": "original"} + ) + attempts: list[Session] = [] + delays: list[float] = [] + + def fail_summary_commit(session: Session) -> None: + if any( + isinstance(record, DocumentSegmentSummary) and record.status == SummaryStatus.COMPLETED + for record in session.new | session.dirty + ): + attempts.append(session) + if len(attempts) <= failed_commits: + raise ConnectionError("connection lost before commit") + + def wait(seconds: float) -> None: + assert not probe.active_transactions + delays.append(seconds) + + monkeypatch.setattr(summary_module.time, "sleep", wait) + event.listen(sqlite_session_factory.class_, "before_commit", fail_summary_commit) + try: + updated = app.update_segment( + context, + dataset_id="dataset-1", + document_id="document-1", + segment_id=created.data.id, + values={"content": "original", "summary": "manual summary"}, + ) + assert updated.data.status == "completed" + assert updated.data.enabled is True + assert not probe.active_transactions + finally: + event.remove(sqlite_session_factory.class_, "before_commit", fail_summary_commit) + for session in attempts: + session.close() + + assert len(attempts) == min(failed_commits + 1, 3) + assert len({id(session) for session in attempts}) == len(attempts) + assert delays == [2.0, 4.0][: min(failed_commits, 2)] + with sqlite_session_factory() as session: + summary = session.scalar(select(DocumentSegmentSummary)) + assert summary is not None + assert summary.summary_content == "manual summary" + if failed_commits < 3: + assert summary.status == SummaryStatus.COMPLETED + assert summary.summary_index_node_id is not None + assert summary.summary_index_node_hash is not None + assert summary.error is None + else: + assert summary.status == SummaryStatus.ERROR + assert summary.error == "Vectorization failed: connection lost before commit" + + +def test_summary_noop_does_not_open_a_database_transaction( + sqlite_session_factory: sessionmaker[Session], +) -> None: + from services.knowledge.segments.indexing import SegmentIndexingGateway + + repository = SQLAlchemySegmentRepository(session_factory=sqlite_session_factory) + gateway = SegmentIndexingGateway( + segments=repository, + uploads=SQLAlchemyKnowledgeUploadRepository(session_factory=sqlite_session_factory), + redis=IndexRedisStub(), + new_session=sqlite_session_factory, + delete_task=lambda *_args: None, + enable_task=lambda *_args: None, + disable_task=lambda *_args: None, + ) + engine = sqlite_session_factory.kw["bind"] + connections: list[Connection] = [] + record_connection = connections.append + event.listen(engine, "begin", record_connection) + try: + gateway.update_summary( + DatasetRef("workspace-1", "dataset-1").document("document-1").segment("segment-1"), + summary=None, + content_changed=False, + ) + finally: + event.remove(engine, "begin", record_connection) + assert connections == [] + + +def test_service_api_child_mutations_refresh_rows_and_release_read_transactions( + indexing_probe: IndexingProbe, + sqlite_session_factory: sessionmaker[Session], +) -> None: + app, repository, _ = indexing_probe + context = RequestContext("request", None, "author", "workspace-1") + parent = app.create_segment(context, dataset_id="dataset-1", document_id="document-1", values={"content": "parent"}) + ref = DatasetRef("workspace-1", "dataset-1").document("document-1").segment(parent.data.id) + with sqlite_session_factory() as session: + dataset = session.get(Dataset, "dataset-1") + document = session.get(Document, "document-1") + segment = session.get(DocumentSegment, parent.data.id) + assert dataset is not None + assert document is not None + assert segment is not None + child = SegmentService.create_child_chunk( + "first", + segment, + document, + dataset, + session, + actor_id="creator", + mutations=app.mutations, + ) + child_id = child.id + assert child.position == 1 + assert child.created_by == "creator" + updated = SegmentService.update_child_chunk( + "second", + child, + segment, + document, + dataset, + session, + actor_id="editor", + mutations=app.mutations, + ) + assert updated.content == "second" + assert updated.updated_by == "editor" + SegmentService.delete_child_chunk(updated, dataset, session, mutations=app.mutations) + assert session.get(ChildChunk, child_id) is None + SegmentService.delete_segment(segment, document, dataset, session, mutations=app.mutations) + assert document.word_count == 0 + assert repository.get_segment(ref) is None + + +@pytest.mark.parametrize("operation", ["create", "update", "delete"]) +def test_service_api_child_index_failure_preserves_rows_and_legacy_errors( + indexing_probe: IndexingProbe, + sqlite_session_factory: sessionmaker[Session], + operation: str, +) -> None: + from services.errors.chunk import ChildChunkDeleteIndexError, ChildChunkIndexingError + + app, repository, probe = indexing_probe + context = RequestContext("request", None, "author", "workspace-1") + parent = app.create_segment(context, dataset_id="dataset-1", document_id="document-1", values={"content": "parent"}) + created = app.create_child_chunk( + context, + dataset_id="dataset-1", + document_id="document-1", + segment_id=parent.data.id, + content="original", + ) + ref = DatasetRef("workspace-1", "dataset-1").document("document-1").segment(parent.data.id) + before = repository.get_children(ref) + probe.fail = True + expected = ChildChunkDeleteIndexError if operation == "delete" else ChildChunkIndexingError + with sqlite_session_factory() as session: + dataset = session.get(Dataset, "dataset-1") + document = session.get(Document, "document-1") + segment = session.get(DocumentSegment, parent.data.id) + child = session.get(ChildChunk, created.id) + assert dataset is not None + assert document is not None + assert segment is not None + assert child is not None + + def mutate() -> None: + if operation == "create": + SegmentService.create_child_chunk( + "new", + segment, + document, + dataset, + session, + actor_id="editor", + mutations=app.mutations, + ) + elif operation == "update": + SegmentService.update_child_chunk( + "changed", + child, + segment, + document, + dataset, + session, + actor_id="editor", + mutations=app.mutations, + ) + else: + SegmentService.delete_child_chunk(child, dataset, session, mutations=app.mutations) + + with pytest.raises(expected, match="index unavailable"): + mutate() + assert repository.get_children(ref) == before + + +@pytest.mark.parametrize("values", [{"answer": "new answer"}, {"content": "new question", "answer": "new answer"}]) +def test_service_api_partial_qa_update_uses_shared_word_count_and_indexing( + indexing_probe: IndexingProbe, + sqlite_session_factory: sessionmaker[Session], + values: dict[str, str], +) -> None: + app, _, _ = indexing_probe + with sqlite_session_factory.begin() as session: + document = session.get(Document, "document-1") + assert document is not None + document.doc_form = IndexStructureType.QA_INDEX + created = app.create_segment( + RequestContext("request", None, "author", "workspace-1"), + dataset_id="dataset-1", + document_id="document-1", + values={"content": "question", "answer": "answer"}, + ) + with sqlite_session_factory() as session: + dataset = session.get(Dataset, "dataset-1") + document = session.get(Document, "document-1") + segment = session.get(DocumentSegment, created.data.id) + assert dataset is not None + assert document is not None + assert segment is not None + updated = SegmentService.update_segment( + SegmentUpdateArgs(**values), + segment, + document, + dataset, + session, + actor_id="editor", + mutations=app.mutations, + ) + assert updated.content == values.get("content", "question") + assert updated.answer == "new answer" + assert updated.word_count == len(updated.content) + len(updated.answer) + assert document.word_count == updated.word_count + if "content" in values: + assert updated.tokens == updated.word_count + + +@pytest.mark.parametrize("owner_field", ["tenant_id", "dataset_id", "document_id"]) +def test_legacy_segment_mutation_rejects_mismatched_rows_before_dispatch(owner_field: str) -> None: + from services.knowledge.segments.application import SegmentMutationService + + dataset = _dataset("dataset-1", "workspace-1") + document = _document("document-1", "workspace-1", "dataset-1") + segment = _segment("segment-1", "workspace-1", "dataset-1", "document-1") + setattr(segment, owner_field, "foreign-owner") + mutations = create_autospec(SegmentMutationService, instance=True, spec_set=True) + session = create_autospec(Session, instance=True, spec_set=True) + with pytest.raises(ValueError, match="Segment does not belong"): + SegmentService.update_segment( + SegmentUpdateArgs(content="new"), + segment, + document, + dataset, + session, + actor_id="author", + mutations=mutations, + ) + mutations.update_segment.assert_not_called() + session.commit.assert_not_called() + + +def _persist_index_batch( + session_factory: sessionmaker[Session], *, doc_form: IndexStructureType = IndexStructureType.PARAGRAPH_INDEX +) -> None: + with session_factory.begin() as session: + document = session.get(Document, "document-1") + assert document is not None + document.doc_form = doc_form + for position, content in enumerate(("first", "second"), start=1): + segment = _segment(content, "workspace-1", "dataset-1", "document-1") + segment.position = position + segment.content = content + segment.index_node_id = f"node-{content}" + segment.index_node_hash = f"hash-{content}" + session.add(segment) + + +@pytest.mark.parametrize("doc_form", [IndexStructureType.PARAGRAPH_INDEX, IndexStructureType.QA_INDEX]) +def test_batch_index_materializes_ordered_documents_before_one_vector_write( + indexing_probe: IndexingProbe, sqlite_session_factory: sessionmaker[Session], doc_form: IndexStructureType +) -> None: + app, _, probe = indexing_probe + _persist_index_batch(sqlite_session_factory, doc_form=doc_form) + + app.mutations.index_segments( + DatasetRef("workspace-1", "dataset-1").document("document-1"), segment_ids=["second", "first"] + ) + + assert probe.batches == [("second", "first")] + assert probe.writes == 2 + assert not probe.active_transactions + + +def test_batch_index_failure_persists_errors_on_every_committed_segment( + indexing_probe: IndexingProbe, sqlite_session_factory: sessionmaker[Session] +) -> None: + app, _, probe = indexing_probe + _persist_index_batch(sqlite_session_factory) + probe.fail = True + + with pytest.raises(RuntimeError, match="index unavailable"): + app.mutations.index_segments( + DatasetRef("workspace-1", "dataset-1").document("document-1"), segment_ids=["first", "second"] + ) + + with sqlite_session_factory() as session: + segments = session.scalars(select(DocumentSegment)).all() + assert len(segments) == 2 + assert all(segment.status == "error" and not segment.enabled for segment in segments) + assert all(segment.error == "index unavailable" and segment.disabled_at is not None for segment in segments) + + +@pytest.mark.parametrize("owner_field", ["tenant_id", "dataset_id", "document_id"]) +def test_batch_index_validates_every_segment_owner_before_external_io( + indexing_probe: IndexingProbe, sqlite_session_factory: sessionmaker[Session], owner_field: str +) -> None: + app, _, probe = indexing_probe + _persist_index_batch(sqlite_session_factory) + with sqlite_session_factory.begin() as session: + foreign = session.get(DocumentSegment, "second") + assert foreign is not None + setattr(foreign, owner_field, "foreign-owner") + + with pytest.raises(LookupError, match="Segment no longer exists"): + app.mutations.index_segments( + DatasetRef("workspace-1", "dataset-1").document("document-1"), segment_ids=["first", "second"] + ) + + assert not probe.batches + with sqlite_session_factory() as session: + foreign = session.get(DocumentSegment, "second") + assert foreign is not None + assert foreign.status == "completed" + + +def test_batch_index_preserves_owned_attachments( + indexing_probe: IndexingProbe, sqlite_session_factory: sessionmaker[Session] +) -> None: + app, _, probe = indexing_probe + _persist_index_batch(sqlite_session_factory) + with sqlite_session_factory.begin() as session: + dataset = session.get(Dataset, "dataset-1") + assert dataset is not None + dataset.is_multimodal = True + session.add_all([_upload("owned", "workspace-1", "image.png"), _upload("foreign", "workspace-2", "other.png")]) + for file_id in ("owned", "foreign"): + session.add( + SegmentAttachmentBinding( + tenant_id="workspace-1", + dataset_id="dataset-1", + document_id="document-1", + segment_id="first", + attachment_id=file_id, + ) + ) + + app.mutations.index_segments( + DatasetRef("workspace-1", "dataset-1").document("document-1"), segment_ids=["first", "second"] + ) + + assert probe.writes == 3 + with sqlite_session_factory() as session: + assert len(session.scalars(select(SegmentAttachmentBinding)).all()) == 2 + + +def test_batch_index_keeps_manual_keywords_aligned_with_segments( + indexing_probe: IndexingProbe, + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, +) -> None: + from core.rag.datasource.keyword.jieba import jieba as jieba_module + from tests.unit_tests.config_override import apply_config_overrides + + app, _, probe = indexing_probe + _persist_index_batch(sqlite_session_factory) + with sqlite_session_factory.begin() as session: + dataset = session.get(Dataset, "dataset-1") + assert dataset is not None + dataset.indexing_technique = IndexTechniqueType.ECONOMY + + class Extractor: + def extract_keywords(self, text: str, _keyword_number: int = 10) -> set[str]: + assert not probe.active_transactions + return {text} + + monkeypatch.setattr(jieba_module, "JiebaKeywordTableHandler", Extractor) + apply_config_overrides(monkeypatch, KEYWORD_DATA_SOURCE_TYPE="database") + app.mutations.index_segments( + DatasetRef("workspace-1", "dataset-1").document("document-1"), + segment_ids=["second", "first"], + keywords_list=[["manual"], None], + ) + + with sqlite_session_factory() as session: + assert {s.id: s.keywords for s in session.scalars(select(DocumentSegment))} == { + "first": ["first"], + "second": ["manual"], + } + assert not probe.batches + + +def test_batch_parent_child_index_uses_shared_splitter_without_open_transactions( + indexing_probe: IndexingProbe, + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, +) -> None: + from core.rag.index_processor.processor.parent_child_index_processor import ParentChildIndexProcessor + from models.dataset import DatasetProcessRule + from models.enums import ProcessRuleMode + + app, repository, probe = indexing_probe + _persist_index_batch(sqlite_session_factory, doc_form=IndexStructureType.PARENT_CHILD_INDEX) + with sqlite_session_factory.begin() as session: + rule = DatasetProcessRule( + dataset_id="dataset-1", + mode=ProcessRuleMode.HIERARCHICAL, + rules='{"parent_mode":"full-doc","subchunk_segmentation":{"max_tokens":100}}', + created_by="author", + ) + session.add(rule) + document = session.get(Document, "document-1") + dataset = session.get(Dataset, "dataset-1") + assert document is not None + assert dataset is not None + dataset.embedding_model_provider = "provider" + document.dataset_process_rule_id = rule.id + + def split(_self, document: IndexDocument, _rule, _mode, _model) -> list[ChildDocument]: + assert not probe.active_transactions + return [ + ChildDocument( + page_content=document.page_content, metadata={"doc_id": document.page_content, "doc_hash": "hash"} + ) + ] + + monkeypatch.setattr(ParentChildIndexProcessor, "split_child_nodes", split) + document_ref = DatasetRef("workspace-1", "dataset-1").document("document-1") + app.mutations.index_segments(document_ref, segment_ids=["first", "second"]) + + assert probe.batches == [("first",), ("second",)] + for segment_id in ("first", "second"): + children = repository.get_children(document_ref.segment(segment_id)) + assert children is not None + assert len(children) == 1 + assert children[0].data.content == segment_id diff --git a/api/tests/unit_tests/services/knowledge/test_segment_application_service.py b/api/tests/unit_tests/services/knowledge/test_segment_application_service.py new file mode 100644 index 00000000000..9a7860ee7d6 --- /dev/null +++ b/api/tests/unit_tests/services/knowledge/test_segment_application_service.py @@ -0,0 +1,506 @@ +from dataclasses import dataclass +from datetime import datetime +from typing import Literal +from unittest.mock import Mock, create_autospec + +import pytest + +from libs.helper import generate_text_hash +from machinery.context import RequestContext +from services.knowledge.dataset_access import ( + AccessibleDataset, + DatasetAccess, + DatasetAccessDeniedError, + DatasetAccessRecord, + DatasetAccessSnapshot, +) +from services.knowledge.entities.segments import SegmentRecord, SegmentUpdateArgs +from services.knowledge.resource_scope import DatasetRef +from services.knowledge.segments.application import ( + ChildChunkNotFoundError, + DatasetSegmentApplicationService, + SegmentBatchImportDispatcher, + SegmentBatchImportDispatchError, + SegmentBatchImportNotFoundError, + SegmentDatasetModelUnavailableError, + SegmentDatasetNotFoundError, + SegmentDatasetRecord, + SegmentDetail, + SegmentDocumentIndexingError, + SegmentDocumentNotFoundError, + SegmentDocumentRecord, + SegmentEmbeddingModelUnavailableError, + SegmentIndex, + SegmentIndexingState, + SegmentInvalidFileTypeError, + SegmentListFilter, + SegmentModelGuard, + SegmentModelProviderError, + SegmentMutationScope, + SegmentNotFoundError, + SegmentPage, + SegmentPermissionDeniedError, + SegmentScope, + SegmentScopeReader, + SegmentStatusUpdateError, + SegmentStore, + SegmentUpdateState, + SegmentUploadCatalog, + SegmentUploadFileNotFoundError, +) + + +@dataclass +class SegmentLimitsStub: + SINGLE_CHUNK_ATTACHMENT_LIMIT: int = 10 + + +def _context() -> RequestContext: + return RequestContext( + request_id="request-1", + trace_id=None, + account_id="account-1", + active_workspace_id="workspace-1", + ) + + +def _dataset( + *, + indexing_technique: str = "economy", +) -> SegmentDatasetRecord: + return SegmentDatasetRecord( + id="dataset-1", + workspace_id="workspace-1", + indexing_technique=indexing_technique, + embedding_model_provider="provider", + embedding_model="model", + ) + + +def _document() -> SegmentDocumentRecord: + return SegmentDocumentRecord( + id="document-1", + dataset_id="dataset-1", + workspace_id="workspace-1", + doc_form="text_model", + ) + + +def _snapshot() -> DatasetAccessSnapshot: + return DatasetAccessSnapshot( + DatasetAccessRecord("dataset-1", "workspace-1", "account-1", "all_team_members"), False + ) + + +def _scope( + *, + dataset: SegmentDatasetRecord | None = None, + document: SegmentDocumentRecord | None = None, +) -> SegmentScope: + return SegmentScope( + dataset=_dataset() if dataset is None else dataset, + document=_document() if document is None else document, + access_snapshot=_snapshot(), + ) + + +type ServiceFixture = tuple[DatasetSegmentApplicationService, Mock, Mock, Mock, Mock, Mock, Mock] + + +def _service(scope: SegmentScope | None = None, *, uploads: SegmentUploadCatalog | None = None) -> ServiceFixture: + dataset_access = create_autospec(DatasetAccess, instance=True, spec_set=True) + dataset_access.check_access.return_value = AccessibleDataset(id="dataset-1", workspace_id="workspace-1") + store = create_autospec(SegmentStore, instance=True, spec_set=True) + scopes = create_autospec(SegmentScopeReader, instance=True, spec_set=True) + scopes.get_segment_scope.return_value = scope or _scope() + model_guard = create_autospec(SegmentModelGuard, instance=True, spec_set=True) + indexing_state = create_autospec(SegmentIndexingState, instance=True, spec_set=True) + indexing_state.is_document_indexing.return_value = False + indexing_state.is_segment_indexing.return_value = False + dispatcher = create_autospec(SegmentBatchImportDispatcher, instance=True, spec_set=True) + service = DatasetSegmentApplicationService( + dataset_access=dataset_access, + store=store, + limits=SegmentLimitsStub(), + text_hash=generate_text_hash, + index=create_autospec(SegmentIndex, instance=True, spec_set=True), + uploads=uploads or create_autospec(SegmentUploadCatalog, instance=True, spec_set=True), + scopes=scopes, + model_guard=model_guard, + indexing_state=indexing_state, + batch_dispatcher=dispatcher, + job_id_factory=lambda: "job-1", + ) + return service, dataset_access, store, scopes, model_guard, indexing_state, dispatcher + + +def test_list_segments_uses_complete_owner_scope() -> None: + service, dataset_access, store, scopes, _, _, _ = _service() + expected = SegmentPage(items=(), total=0, total_pages=0, page=1, limit=20) + store.list_segments.return_value = expected + + result = service.list_segments( + _context(), + dataset_id="dataset-1", + document_id="document-1", + query=SegmentListFilter(keyword="needle"), + ) + + assert result is expected + assert scopes.get_segment_scope.call_args.kwargs == { + "dataset_ref": DatasetRef("workspace-1", "dataset-1"), + "document_id": "document-1", + "actor_id": "account-1", + } + assert store.list_segments.call_args.args[0] == _document().ref + dataset_access.check_access.assert_called_once_with(_context(), _snapshot()) + + +@pytest.mark.parametrize( + ("scope", "expected"), + [ + (SegmentScope(None, None), SegmentDatasetNotFoundError), + (SegmentScope(_dataset(), None, _snapshot()), SegmentDocumentNotFoundError), + ], +) +def test_mutations_reject_invalid_scope_before_store_operation( + scope: SegmentScope, + expected: type[Exception], +) -> None: + service, _, store, scopes, _, _, _ = _service(scope) + + with pytest.raises(expected): + service.delete_segment( + _context(), + dataset_id="dataset-1", + document_id="document-1", + segment_id="segment-1", + ) + + store.delete_segments.assert_not_called() + + +@pytest.mark.parametrize( + ("access_error", "expected"), + [ + (DatasetAccessDeniedError(), SegmentPermissionDeniedError), + ], +) +def test_dataset_access_errors_are_translated( + access_error: Exception, + expected: type[Exception], +) -> None: + service, dataset_access, store, scopes, _, _, _ = _service() + dataset_access.check_access.side_effect = access_error + + with pytest.raises(expected): + service.delete_segment( + _context(), + dataset_id="dataset-1", + document_id="document-1", + segment_id="segment-1", + ) + + store.delete_segments.assert_not_called() + + +@pytest.mark.parametrize( + ("kind", "method", "expected", "message"), + [ + ("bad_request", "create_segment", SegmentEmbeddingModelUnavailableError, "No Embedding Model"), + ("token", "create_segment", SegmentEmbeddingModelUnavailableError, "token missing"), + ("bad_request", "delete_segment", SegmentDatasetModelUnavailableError, "No Embedding Model"), + ("token", "delete_segment", SegmentDatasetModelUnavailableError, "dataset is unavailable"), + ], +) +def test_model_failures_are_translated_per_use_case( + kind: Literal["bad_request", "token"], + method: str, + expected: type[Exception], + message: str, +) -> None: + service, _, store, scopes, model_guard, _, _ = _service(_scope(dataset=_dataset(indexing_technique="high_quality"))) + model_guard.check.side_effect = SegmentModelProviderError(kind=kind, description="token missing") + store.get_segment.return_value = None + + if method == "create_segment": + + def invoke() -> object: + return service.create_segment( + _context(), + dataset_id="dataset-1", + document_id="document-1", + values={"content": "content"}, + ) + + else: + + def invoke() -> object: + return service.delete_segment( + _context(), + dataset_id="dataset-1", + document_id="document-1", + segment_id="segment-1", + ) + + with pytest.raises(expected, match=message): + invoke() + + +def test_status_change_is_fenced_by_document_indexing_state() -> None: + service, _, store, scopes, _, state, _ = _service() + state.is_document_indexing.return_value = True + + with pytest.raises(SegmentDocumentIndexingError): + service.change_segment_status( + _context(), + dataset_id="dataset-1", + document_id="document-1", + segment_ids=["segment-1"], + action="disable", + ) + + store.save_segment.assert_not_called() + + +def test_status_change_translates_operation_failure() -> None: + service, _, store, scopes, _, _, _ = _service() + store.get_segments.side_effect = RuntimeError("failed") + + with pytest.raises(SegmentStatusUpdateError, match="failed"): + service.change_segment_status( + _context(), + dataset_id="dataset-1", + document_id="document-1", + segment_ids=["segment-1"], + action="enable", + ) + + +def test_update_and_delete_segment_map_missing_segment() -> None: + service, _, store, scopes, _, _, _ = _service() + store.get_segment_update_state.return_value = None + store.get_segments.return_value = () + + with pytest.raises(SegmentNotFoundError): + service.update_segment( + _context(), + dataset_id="dataset-1", + document_id="document-1", + segment_id="segment-1", + values={"content": "content"}, + ) + with pytest.raises(SegmentNotFoundError): + service.delete_segment( + _context(), + dataset_id="dataset-1", + document_id="document-1", + segment_id="segment-1", + ) + + +@pytest.mark.parametrize( + ("filename", "expected"), + [ + (None, SegmentUploadFileNotFoundError), + ("segments.txt", SegmentInvalidFileTypeError), + ], +) +def test_batch_import_validates_tenant_scoped_upload(filename: str | None, expected: type[Exception]) -> None: + uploads = create_autospec(SegmentUploadCatalog, instance=True, spec_set=True) + uploads.get_file_name.return_value = filename + service, _, store, scopes, _, state, dispatcher = _service(uploads=uploads) + + with pytest.raises(expected): + service.start_batch_import( + _context(), + dataset_id="dataset-1", + document_id="document-1", + upload_file_id="file-1", + ) + + assert uploads.get_file_name.call_args.kwargs["workspace_id"] == "workspace-1" + state.set_batch_waiting.assert_not_called() + dispatcher.dispatch.assert_not_called() + + +def test_batch_import_records_status_before_dispatch() -> None: + uploads = create_autospec(SegmentUploadCatalog, instance=True, spec_set=True) + uploads.get_file_name.return_value = "segments.CSV" + service, _, store, scopes, _, state, dispatcher = _service(uploads=uploads) + calls: list[str] = [] + state.set_batch_waiting.side_effect = lambda _job_id: calls.append("status") + dispatcher.dispatch.side_effect = lambda **_kwargs: calls.append("dispatch") + + result = service.start_batch_import( + _context(), + dataset_id="dataset-1", + document_id="document-1", + upload_file_id="file-1", + ) + + assert result.job_id == "job-1" + assert calls == ["status", "dispatch"] + dispatcher.dispatch.assert_called_once_with( + job_id="job-1", + upload_file_id="file-1", + dataset_id="dataset-1", + document_id="document-1", + workspace_id="workspace-1", + actor_id="account-1", + ) + + +def test_batch_import_translates_dispatch_failure() -> None: + uploads = create_autospec(SegmentUploadCatalog, instance=True, spec_set=True) + uploads.get_file_name.return_value = "segments.csv" + service, _, store, scopes, _, _, dispatcher = _service(uploads=uploads) + dispatcher.dispatch.side_effect = RuntimeError("queue down") + + with pytest.raises(SegmentBatchImportDispatchError, match="queue down"): + service.start_batch_import( + _context(), + dataset_id="dataset-1", + document_id="document-1", + upload_file_id="file-1", + ) + + +def test_batch_import_status_distinguishes_unknown_job() -> None: + service, _, _, scopes, _, state, _ = _service() + state.get_batch_status.return_value = None + with pytest.raises(SegmentBatchImportNotFoundError): + service.get_batch_import_status("missing") + + state.get_batch_status.return_value = "completed" + assert service.get_batch_import_status("job-1").job_status == "completed" + + +def test_child_chunk_update_distinguishes_missing_parent_and_child() -> None: + service, _, store, scopes, _, _, _ = _service() + store.get_children.return_value = None + with pytest.raises(SegmentNotFoundError): + service.update_child_chunk( + _context(), + dataset_id="dataset-1", + document_id="document-1", + segment_id="segment-1", + child_chunk_id="child-1", + content="content", + ) + + store.get_children.return_value = () + with pytest.raises(ChildChunkNotFoundError): + service.update_child_chunk( + _context(), + dataset_id="dataset-1", + document_id="document-1", + segment_id="segment-1", + child_chunk_id="child-1", + content="content", + ) + + +def _segment_detail(*, enabled: bool = True) -> SegmentDetail: + return SegmentDetail( + data=SegmentRecord( + id="segment-1", + position=1, + document_id="document-1", + content="content", + sign_content="content", + answer=None, + word_count=7, + tokens=0, + keywords=None, + index_node_id=None, + index_node_hash=None, + hit_count=0, + enabled=enabled, + disabled_at=None, + disabled_by=None, + status="completed", + created_by="account-1", + created_at=datetime(2026, 1, 1), + updated_at=datetime(2026, 1, 1), + updated_by=None, + indexing_at=None, + completed_at=None, + error=None, + stopped_at=None, + child_chunks=(), + attachments=(), + summary=None, + ), + doc_form="text_model", + ) + + +def test_create_segment_returns_store_detail() -> None: + service, _, store, scopes, _, _, _ = _service() + detail = _segment_detail() + store.get_segment.return_value = detail + + result = service.create_segment( + _context(), + dataset_id="dataset-1", + document_id="document-1", + values={"content": "content"}, + ) + + assert result is detail + assert store.save_segment.call_args.args[1]["created_by"] == _context().account_id + + +@pytest.mark.parametrize( + ("enabled", "indexing", "message"), + [(False, False, "Can't update disabled segment"), (True, True, "Segment is indexing")], +) +def test_shared_update_rejects_disabled_or_indexing_segments(enabled: bool, indexing: bool, message: str) -> None: + service, _, store, _, _, state, _ = _service() + store.get_segment_update_state.return_value = SegmentUpdateState("content", enabled, None) + state.is_segment_indexing.return_value = indexing + with pytest.raises(ValueError, match=message): + service.mutations.update_segment( + SegmentMutationScope(_dataset(), _document()), + actor_id="editor", + segment_id="segment-1", + args=SegmentUpdateArgs(content="changed"), + ) + store.save_segment.assert_not_called() + + +def test_shared_disable_patch_does_not_replace_content() -> None: + service, _, store, _, _, state, _ = _service() + store.get_segment_update_state.return_value = SegmentUpdateState("content", True, None) + service.mutations.update_segment( + SegmentMutationScope(_dataset(), _document()), + actor_id="editor", + segment_id="segment-1", + args=SegmentUpdateArgs(enabled=False), + ) + changes = store.save_segment.call_args.args[1] + assert changes["enabled"] is False + assert changes["disabled_by"] == "editor" + assert "content" not in changes + state.mark_segment_indexing.assert_called_once_with("segment-1") + store.get_segment.assert_not_called() + + +def test_shared_writes_do_not_build_segment_detail() -> None: + service, _, store, _, _, _, _ = _service() + store.get_segment.side_effect = RuntimeError("attachment URL signing unavailable") + store.get_segment_update_state.return_value = SegmentUpdateState("content", True, None) + scope = SegmentMutationScope(_dataset(), _document()) + + created = service.mutations.create_segment(scope, actor_id="author", args=SegmentUpdateArgs(content="content")) + service.mutations.update_segment( + scope, + actor_id="editor", + segment_id=created.segment_id, + args=SegmentUpdateArgs(content="updated"), + ) + + assert store.save_segment.call_count == 2 + assert store.save_segment.call_args.args[1]["content"] == "updated" + store.get_segment.assert_not_called() diff --git a/api/tests/unit_tests/services/rag_pipeline/test_pipeline_generate_service.py b/api/tests/unit_tests/services/rag_pipeline/test_pipeline_generate_service.py index fabc093c9b0..939b54b1559 100644 --- a/api/tests/unit_tests/services/rag_pipeline/test_pipeline_generate_service.py +++ b/api/tests/unit_tests/services/rag_pipeline/test_pipeline_generate_service.py @@ -4,16 +4,22 @@ import pytest from pytest_mock import MockerFixture from sqlalchemy.orm import Session +from core.app.apps.pipeline.pipeline_generator import PipelineGenerator from core.app.entities.app_invoke_entities import InvokeFrom from core.rag.index_processor.constant.index_type import IndexStructureType from models.dataset import Dataset, Document, Pipeline from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus from models.model import Account, App, AppMode, EndUser from models.workflow import Workflow, WorkflowType -from services.dataset_ref_service import DatasetRefService +from services.knowledge.resource_scope import DatasetRef from services.rag_pipeline.pipeline_generate_service import PipelineGenerateService +@pytest.fixture +def pipeline_generator(mocker: MockerFixture): + return mocker.create_autospec(PipelineGenerator, instance=True, spec_set=True) + + def _make_pipeline(*, tenant_id: str = "tenant-1") -> Pipeline: pipeline = Pipeline(tenant_id=tenant_id, name="Pipeline", description="") pipeline.id = "pipeline-1" @@ -128,7 +134,7 @@ def test_get_workflow(mocker: MockerFixture, invoke_from, workflow, expected_err def test_generate_updates_document_status_and_returns_event_stream( - mocker: MockerFixture, sqlite_session: Session + mocker: MockerFixture, sqlite_session: Session, pipeline_generator ) -> None: dataset = _make_dataset() pipeline = _make_pipeline() @@ -141,7 +147,7 @@ def test_generate_updates_document_status_and_returns_event_stream( update_status_mock = mocker.patch.object(PipelineGenerateService, "update_document_status") generator_cls = mocker.patch("services.rag_pipeline.pipeline_generate_service.PipelineGenerator") - generator_instance = generator_cls.return_value + generator_instance = pipeline_generator generator_instance.generate.return_value = "raw-events" generator_cls.convert_to_event_stream.return_value = "stream-events" @@ -152,6 +158,7 @@ def test_generate_updates_document_status_and_returns_event_stream( invoke_from=InvokeFrom.WEB_APP, streaming=True, session=sqlite_session, + generator=pipeline_generator, ) assert result == "stream-events" @@ -163,7 +170,9 @@ def test_generate_updates_document_status_and_returns_event_stream( assert generator_instance.generate.call_args.kwargs["session"] is sqlite_session -def test_generate_rejects_pipeline_dataset_from_another_tenant(mocker: MockerFixture, sqlite_session: Session) -> None: +def test_generate_rejects_pipeline_dataset_from_another_tenant( + mocker: MockerFixture, sqlite_session: Session, pipeline_generator +) -> None: dataset = _make_dataset(tenant_id="tenant-2") pipeline = _make_pipeline() sqlite_session.add_all([dataset, pipeline]) @@ -178,6 +187,7 @@ def test_generate_rejects_pipeline_dataset_from_another_tenant(mocker: MockerFix args={"original_document_id": "doc-1"}, invoke_from=InvokeFrom.WEB_APP, session=sqlite_session, + generator=pipeline_generator, ) update_status_mock.assert_not_called() @@ -186,6 +196,7 @@ def test_generate_rejects_pipeline_dataset_from_another_tenant(mocker: MockerFix def test_generate_rejects_original_document_outside_pipeline_dataset_before_dispatch( mocker: MockerFixture, sqlite_session: Session, + pipeline_generator, ) -> None: dataset = _make_dataset() pipeline = _make_pipeline() @@ -202,11 +213,12 @@ def test_generate_rejects_original_document_outside_pipeline_dataset_before_disp args={"original_document_id": "foreign-doc"}, invoke_from=InvokeFrom.PUBLISHED_PIPELINE, session=sqlite_session, + generator=pipeline_generator, ) sqlite_session.refresh(outside_document) assert outside_document.indexing_status == IndexingStatus.COMPLETED - generator_cls.assert_not_called() + pipeline_generator.generate.assert_not_called() def test_update_document_status_updates_existing_document(sqlite_session: Session) -> None: @@ -214,8 +226,8 @@ def test_update_document_status_updates_existing_document(sqlite_session: Sessio dataset = _make_dataset() sqlite_session.add_all([dataset, document]) sqlite_session.commit() - dataset_ref = DatasetRefService.create_dataset_ref(dataset) - document_ref = DatasetRefService.create_document_ref_from_id(dataset_ref, "doc-1") + dataset_ref = DatasetRef(tenant_id=dataset.tenant_id, dataset_id=dataset.id) + document_ref = dataset_ref.document("doc-1") PipelineGenerateService.update_document_status(document_ref, session=sqlite_session) @@ -235,8 +247,8 @@ def test_update_document_status_rejects_document_outside_owner( sqlite_session: Session, ) -> None: dataset = _make_dataset() - dataset_ref = DatasetRefService.create_dataset_ref(dataset) - document_ref = DatasetRefService.create_document_ref_from_id(dataset_ref, "doc-1") + dataset_ref = DatasetRef(tenant_id=dataset.tenant_id, dataset_id=dataset.id) + document_ref = dataset_ref.document("doc-1") outside_document = _make_document(tenant_id=document_tenant_id, dataset_id=document_dataset_id) sqlite_session.add(outside_document) sqlite_session.commit() @@ -251,11 +263,13 @@ def test_update_document_status_rejects_document_outside_owner( # --- generate_single_iteration --- -def test_generate_single_iteration_delegates(mocker: MockerFixture, sqlite_session: Session) -> None: +def test_generate_single_iteration_delegates( + mocker: MockerFixture, sqlite_session: Session, pipeline_generator +) -> None: mocker.patch.object(PipelineGenerateService, "_get_workflow", return_value=_make_workflow()) generator_cls = mocker.patch("services.rag_pipeline.pipeline_generate_service.PipelineGenerator") - generator_instance = generator_cls.return_value + generator_instance = pipeline_generator generator_instance.single_iteration_generate.return_value = "raw-iter" generator_cls.convert_to_event_stream.return_value = "stream-iter" @@ -264,7 +278,9 @@ def test_generate_single_iteration_delegates(mocker: MockerFixture, sqlite_sessi user = _make_account(account_id="u1") session = sqlite_session - result = PipelineGenerateService.generate_single_iteration(pipeline, user, "node-1", {"key": "val"}, session) + result = PipelineGenerateService.generate_single_iteration( + pipeline, user, "node-1", {"key": "val"}, session, generator=pipeline_generator + ) assert result == "stream-iter" generator_instance.single_iteration_generate.assert_called_once() @@ -274,11 +290,11 @@ def test_generate_single_iteration_delegates(mocker: MockerFixture, sqlite_sessi # --- generate_single_loop --- -def test_generate_single_loop_delegates(mocker: MockerFixture, sqlite_session: Session) -> None: +def test_generate_single_loop_delegates(mocker: MockerFixture, sqlite_session: Session, pipeline_generator) -> None: mocker.patch.object(PipelineGenerateService, "_get_workflow", return_value=_make_workflow()) generator_cls = mocker.patch("services.rag_pipeline.pipeline_generate_service.PipelineGenerator") - generator_instance = generator_cls.return_value + generator_instance = pipeline_generator generator_instance.single_loop_generate.return_value = "raw-loop" generator_cls.convert_to_event_stream.return_value = "stream-loop" @@ -287,7 +303,9 @@ def test_generate_single_loop_delegates(mocker: MockerFixture, sqlite_session: S user = _make_account(account_id="u1") session = sqlite_session - result = PipelineGenerateService.generate_single_loop(pipeline, user, "node-1", {"key": "val"}, session) + result = PipelineGenerateService.generate_single_loop( + pipeline, user, "node-1", {"key": "val"}, session, generator=pipeline_generator + ) assert result == "stream-loop" generator_instance.single_loop_generate.assert_called_once() diff --git a/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_manage_service.py b/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_manage_service.py index dbe3955033f..35e7c01dad6 100644 --- a/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_manage_service.py +++ b/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_manage_service.py @@ -40,7 +40,7 @@ def test_list_rag_pipeline_datasources_marks_authorized(mocker: MockerFixture) - None, ] - result = RagPipelineManageService.list_rag_pipeline_datasources("tenant-1") + result = RagPipelineManageService.list_rag_pipeline_datasources("tenant-1", datasource_providers=provider_instance) assert result == [datasource_1, datasource_2] assert datasource_1.is_authorized is True @@ -63,7 +63,7 @@ def test_list_rag_pipeline_datasources_marks_credential_free_providers_authorize provider_instance = provider_cls.return_value provider_instance.get_datasource_credentials.return_value = None - result = RagPipelineManageService.list_rag_pipeline_datasources("tenant-1") + result = RagPipelineManageService.list_rag_pipeline_datasources("tenant-1", datasource_providers=provider_instance) assert result == [local_file, notion] assert local_file.is_authorized is True @@ -85,7 +85,7 @@ def test_oauth_only_provider_still_requires_authorization( provider_instance = provider_cls.return_value provider_instance.get_datasource_credentials.return_value = None - result = RagPipelineManageService.list_rag_pipeline_datasources("tenant-1") + result = RagPipelineManageService.list_rag_pipeline_datasources("tenant-1", datasource_providers=provider_instance) assert result == [oauth_only] assert oauth_only.is_authorized is False diff --git a/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_service.py b/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_service.py index 3687c0a298b..abfcf2e3437 100644 --- a/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_service.py +++ b/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_service.py @@ -5,12 +5,14 @@ from dataclasses import dataclass from datetime import datetime from types import SimpleNamespace from unittest import mock +from unittest.mock import create_autospec import pytest from pytest_mock import MockerFixture from sqlalchemy.dialects import postgresql from sqlalchemy.orm import Session, sessionmaker +from core.app.apps.pipeline.pipeline_generator import PipelineGenerator from core.app.entities.app_invoke_entities import InvokeFrom from core.rag.index_processor.constant.index_type import IndexStructureType from graphon.enums import ( @@ -32,12 +34,21 @@ from models.dataset import ( ) from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus from models.workflow import Workflow, WorkflowRun +from repositories.knowledge.dataset_read_repository import get_pipeline_dataset +from services.credentials.query import CredentialQuery +from services.data_source.credential_gateway import DatasourceProviderCredentialStore from services.entities.knowledge_entities.rag_pipeline_entities import IconInfo, PipelineTemplateInfoEntity from services.errors.rag_pipeline import RagPipelineResourceNotFoundError +from services.rag_pipeline import rag_pipeline as rag_pipeline_module from services.rag_pipeline.rag_pipeline import RagPipelineService from services.workflow_ref_service import WorkflowRef +@pytest.fixture +def pipeline_generator(mocker: MockerFixture): + return mocker.create_autospec(PipelineGenerator, instance=True, spec_set=True) + + @dataclass class RagPipelineServiceTestContext: service: RagPipelineService @@ -611,7 +622,7 @@ def test_publish_workflow_success(mocker: MockerFixture, rag_pipeline_service: R dataset = _make_dataset() _persist(rag_pipeline_service.session, draft_workflow, pipeline, dataset) mocker.patch("services.rag_pipeline.rag_pipeline.KnowledgeConfiguration.model_validate", return_value=mocker.Mock()) - mock_dataset_service_class = mocker.patch("services.dataset_service.DatasetService") + mock_dataset_service_class = mocker.patch("services.knowledge.dataset_service.DatasetService") result = rag_pipeline_service.service.publish_workflow( session=rag_pipeline_service.session, pipeline=pipeline, account=_make_account() @@ -716,6 +727,9 @@ def test_run_datasource_workflow_node_website_crawl( account=_make_account(), datasource_type="website_crawl", is_published=True, + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), ) events = list(gen) @@ -797,6 +811,9 @@ def test_run_datasource_node_preview_online_document( account=_make_account(), datasource_type="online_document", is_published=True, + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), ) # 4. Assertions @@ -916,7 +933,7 @@ def test_publish_customized_pipeline_template_success( # Mock RagPipelineDslService mock_dsl_service = mocker.Mock() mock_dsl_service.export_rag_pipeline_dsl.return_value = "dsl: content" - mocker.patch("services.rag_pipeline.rag_pipeline_dsl_service.RagPipelineDslService", return_value=mock_dsl_service) + mocker.patch.object(rag_pipeline_module, "RagPipelineDslService", return_value=mock_dsl_service) account = _make_account(account_id="user-123") @@ -968,7 +985,15 @@ def test_get_datasource_plugins_success( mocker.patch("services.rag_pipeline.rag_pipeline.DatasourceProviderService", return_value=mock_provider_service) # 2. Run test - result = rag_pipeline_service.service.get_datasource_plugins("t1", "d1", True) + result = rag_pipeline_service.service.get_datasource_plugins( + "t1", + "d1", + True, + credential_query=mock.create_autospec(CredentialQuery, instance=True), + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), + ) # 3. Assertions assert len(result) == 1 @@ -980,7 +1005,7 @@ def test_get_datasource_plugins_success( def test_retry_error_document_success( - mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext, pipeline_generator ) -> None: dataset = _make_dataset() document = _make_document() @@ -1003,12 +1028,11 @@ def test_retry_error_document_success( _persist(rag_pipeline_service.session, log, pipeline, workflow) # Mock PipelineGenerator - mock_gen_instance = mocker.Mock() - mocker.patch("services.rag_pipeline.rag_pipeline.PipelineGenerator", return_value=mock_gen_instance) + mock_gen_instance = pipeline_generator # 2. Run test user = mocker.Mock() - rag_pipeline_service.service.retry_error_document(dataset, document, user) + rag_pipeline_service.service.retry_error_document(dataset, document, user, generator=pipeline_generator) # 3. Assertions mock_gen_instance.generate.assert_called_once() @@ -1266,6 +1290,9 @@ def test_run_datasource_workflow_node_returns_error_when_workflow_missing( account=_make_account(), datasource_type="online_document", is_published=False, + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), ) ) @@ -1314,6 +1341,9 @@ def test_run_datasource_workflow_node_online_document_success( account=_make_account(), datasource_type=DatasourceProviderType.ONLINE_DOCUMENT, is_published=True, + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), ) ) @@ -1366,6 +1396,9 @@ def test_run_datasource_workflow_node_online_drive_success( account=_make_account(), datasource_type=DatasourceProviderType.ONLINE_DRIVE, is_published=True, + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), ) ) @@ -1599,6 +1632,9 @@ def test_run_datasource_workflow_node_returns_error_when_node_missing( account=_make_account(), datasource_type="online_document", is_published=True, + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), ) ) @@ -1641,7 +1677,8 @@ def test_run_datasource_workflow_node_online_document_exception( mocker.patch("core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", return_value=runtime) mocker.patch( - "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", return_value=None + "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", + return_value=None, ) events = list( @@ -1652,6 +1689,9 @@ def test_run_datasource_workflow_node_online_document_exception( account=_make_account(), datasource_type="online_document", is_published=True, + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), ) ) @@ -1696,7 +1736,8 @@ def test_run_datasource_node_preview_raises_for_stream_non_string( mocker.patch("core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", return_value=runtime) mocker.patch( - "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", return_value=None + "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", + return_value=None, ) with pytest.raises(RuntimeError, match="must be a string"): @@ -1707,6 +1748,9 @@ def test_run_datasource_node_preview_raises_for_stream_non_string( account=_make_account(), datasource_type="online_document", is_published=True, + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), ) @@ -1758,9 +1802,12 @@ def test_get_second_step_parameters_filters_first_step_variables( def test_retry_error_document_raises_when_execution_log_not_found( rag_pipeline_service: RagPipelineServiceTestContext, + pipeline_generator, ) -> None: with pytest.raises(ValueError, match="Document pipeline execution log not found"): - rag_pipeline_service.service.retry_error_document(_make_dataset(), _make_document(), _make_account()) + rag_pipeline_service.service.retry_error_document( + _make_dataset(), _make_document(), _make_account(), generator=pipeline_generator + ) def test_get_datasource_plugins_raises_when_workflow_not_found( @@ -1771,7 +1818,15 @@ def test_get_datasource_plugins_raises_when_workflow_not_found( _persist(rag_pipeline_service.session, dataset, pipeline) with pytest.raises(ValueError, match="Pipeline or workflow not found"): - rag_pipeline_service.service.get_datasource_plugins("t1", "d1", True) + rag_pipeline_service.service.get_datasource_plugins( + "t1", + "d1", + True, + credential_query=mock.create_autospec(CredentialQuery, instance=True), + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), + ) def test_handle_node_run_result_raises_when_no_terminal_event( @@ -1864,7 +1919,8 @@ def test_run_datasource_node_preview_raises_for_unsupported_provider( runtime.datasource_provider_type.return_value = "unsupported" mocker.patch("core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", return_value=runtime) mocker.patch( - "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", return_value=None + "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", + return_value=None, ) with pytest.raises(RuntimeError, match="Unsupported datasource provider"): @@ -1875,6 +1931,9 @@ def test_run_datasource_node_preview_raises_for_unsupported_provider( account=_make_account(), datasource_type="website_crawl", is_published=True, + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), ) @@ -1984,7 +2043,7 @@ def test_publish_workflow_skips_dataset_update_for_non_knowledge_nodes( ) -> None: draft = _make_workflow(graph={"nodes": [{"data": {"type": "start"}}]}) _persist(rag_pipeline_service.session, draft) - dataset_service = mocker.patch("services.dataset_service.DatasetService") + dataset_service = mocker.patch("services.knowledge.dataset_service.DatasetService") result = rag_pipeline_service.service.publish_workflow( session=rag_pipeline_service.session, @@ -2045,7 +2104,8 @@ def test_run_datasource_workflow_node_handles_variable_parameter_types( runtime.datasource_provider_type.return_value = DatasourceProviderType.WEBSITE_CRAWL mocker.patch("core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", return_value=runtime) mocker.patch( - "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", return_value=None + "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", + return_value=None, ) events = list( @@ -2056,6 +2116,9 @@ def test_run_datasource_workflow_node_handles_variable_parameter_types( account=_make_account(), datasource_type="website_crawl", is_published=True, + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), ) ) @@ -2093,7 +2156,8 @@ def test_run_datasource_workflow_node_online_drive_branch( runtime.datasource_provider_type.return_value = DatasourceProviderType.ONLINE_DRIVE mocker.patch("core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", return_value=runtime) mocker.patch( - "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", return_value=None + "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", + return_value=None, ) events = list( @@ -2104,6 +2168,9 @@ def test_run_datasource_workflow_node_online_drive_branch( account=_make_account(), datasource_type="online_drive", is_published=True, + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), ) ) @@ -2143,7 +2210,8 @@ def test_run_datasource_node_preview_not_published_uses_draft( runtime.get_online_document_page_content.side_effect = doc_gen mocker.patch("core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", return_value=runtime) mocker.patch( - "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", return_value=None + "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", + return_value=None, ) result = rag_pipeline_service.service.run_datasource_node_preview( @@ -2153,6 +2221,9 @@ def test_run_datasource_node_preview_not_published_uses_draft( account=_make_account(), datasource_type="online_document", is_published=False, + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), ) assert result == {"x": "v"} @@ -2187,7 +2258,7 @@ def test_publish_customized_pipeline_template_rejects_unowned_workflow_before_ex pipeline = _make_pipeline(workflow_id="wf-1") workflow = _make_workflow(workflow_id="wf-1", tenant_id=workflow_tenant_id, app_id=workflow_app_id) _persist(rag_pipeline_service.session, pipeline, workflow) - dsl_service = mocker.patch("services.rag_pipeline.rag_pipeline_dsl_service.RagPipelineDslService") + dsl_service = mocker.patch.object(rag_pipeline_module, "RagPipelineDslService") with pytest.raises(RagPipelineResourceNotFoundError, match="Workflow not found"): rag_pipeline_service.service.publish_customized_pipeline_template( @@ -2209,7 +2280,7 @@ def test_pipeline_retrieve_dataset_rejects_unowned_dataset( other_tenant_dataset = _make_dataset(tenant_id="t2") _persist(rag_pipeline_service.session, pipeline, other_tenant_dataset) - assert pipeline.retrieve_dataset(session=rag_pipeline_service.session) is None + assert get_pipeline_dataset(pipeline, session=rag_pipeline_service.session) is None @pytest.mark.parametrize( @@ -2230,7 +2301,7 @@ def test_publish_customized_pipeline_template_rejects_missing_or_unowned_draft_b if draft_tenant_id and draft_app_id: resources.append(_make_workflow(workflow_id="wf-draft", tenant_id=draft_tenant_id, app_id=draft_app_id)) _persist(session, *resources) - dsl_service = mocker.patch("services.rag_pipeline.rag_pipeline_dsl_service.RagPipelineDslService") + dsl_service = mocker.patch.object(rag_pipeline_module, "RagPipelineDslService") with pytest.raises(RagPipelineResourceNotFoundError, match="Draft workflow not found"): rag_pipeline_service.service.publish_customized_pipeline_template( @@ -2257,6 +2328,7 @@ def test_get_recommended_plugins_skips_manifest_when_missing( def test_retry_error_document_raises_when_pipeline_missing( rag_pipeline_service: RagPipelineServiceTestContext, + pipeline_generator, ) -> None: exec_log = DocumentPipelineExecutionLog( pipeline_id="p1", @@ -2270,11 +2342,14 @@ def test_retry_error_document_raises_when_pipeline_missing( _persist(rag_pipeline_service.session, exec_log) with pytest.raises(ValueError, match="Pipeline not found"): - rag_pipeline_service.service.retry_error_document(_make_dataset(), _make_document(), _make_account()) + rag_pipeline_service.service.retry_error_document( + _make_dataset(), _make_document(), _make_account(), generator=pipeline_generator + ) def test_retry_error_document_raises_when_workflow_missing( rag_pipeline_service: RagPipelineServiceTestContext, + pipeline_generator, ) -> None: exec_log = DocumentPipelineExecutionLog( pipeline_id="p1", @@ -2289,7 +2364,9 @@ def test_retry_error_document_raises_when_workflow_missing( _persist(rag_pipeline_service.session, exec_log, pipeline) with pytest.raises(ValueError, match="Workflow not found"): - rag_pipeline_service.service.retry_error_document(_make_dataset(), _make_document(), _make_account()) + rag_pipeline_service.service.retry_error_document( + _make_dataset(), _make_document(), _make_account(), generator=pipeline_generator + ) def test_get_datasource_plugins_returns_empty_for_non_datasource_nodes( @@ -2300,7 +2377,18 @@ def test_get_datasource_plugins_returns_empty_for_non_datasource_nodes( workflow = _make_workflow(graph={"nodes": [{"id": "n1", "data": {"type": "start"}}]}) _persist(rag_pipeline_service.session, dataset, pipeline, workflow) - assert rag_pipeline_service.service.get_datasource_plugins("t1", "d1", True) == [] + assert ( + rag_pipeline_service.service.get_datasource_plugins( + "t1", + "d1", + True, + credential_query=mock.create_autospec(CredentialQuery, instance=True), + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), + ) + == [] + ) def test_publish_workflow_raises_when_knowledge_index_dataset_missing( @@ -2330,6 +2418,9 @@ def test_run_datasource_node_preview_raises_when_workflow_missing( account=_make_account(), datasource_type="online_document", is_published=True, + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), ) @@ -2348,6 +2439,9 @@ def test_run_datasource_node_preview_raises_when_node_missing( account=_make_account(), datasource_type="online_document", is_published=True, + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), ) @@ -2385,7 +2479,8 @@ def test_run_datasource_node_preview_keeps_existing_user_input( runtime.get_online_document_page_content.side_effect = gen mocker.patch("core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", return_value=runtime) mocker.patch( - "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", return_value=None + "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", + return_value=None, ) result = rag_pipeline_service.service.run_datasource_node_preview( @@ -2395,6 +2490,9 @@ def test_run_datasource_node_preview_keeps_existing_user_input( account=_make_account(), datasource_type="online_document", is_published=True, + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), ) assert result == {"ok": "1"} @@ -2426,7 +2524,8 @@ def test_run_datasource_node_preview_ignores_non_variable_messages( runtime.get_online_document_page_content.side_effect = gen mocker.patch("core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", return_value=runtime) mocker.patch( - "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", return_value=None + "services.rag_pipeline.rag_pipeline.DatasourceProviderService.get_datasource_credentials", + return_value=None, ) result = rag_pipeline_service.service.run_datasource_node_preview( @@ -2436,6 +2535,9 @@ def test_run_datasource_node_preview_ignores_non_variable_messages( account=_make_account(), datasource_type="online_document", is_published=True, + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), ) assert result == {} @@ -2466,10 +2568,19 @@ def test_get_datasource_plugins_handles_empty_datasource_data_and_non_published( ] _persist(rag_pipeline_service.session, dataset, pipeline, workflow) mocker.patch( - "services.rag_pipeline.rag_pipeline.DatasourceProviderService.list_datasource_credentials", return_value=[] + "services.rag_pipeline.rag_pipeline.DatasourceProviderService.list_datasource_credentials", + return_value=[], ) - result = rag_pipeline_service.service.get_datasource_plugins("t1", "d1", False) + result = rag_pipeline_service.service.get_datasource_plugins( + "t1", + "d1", + False, + credential_query=mock.create_autospec(CredentialQuery, instance=True), + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), + ) assert len(result) == 1 @@ -2510,7 +2621,15 @@ def test_get_datasource_plugins_extracts_user_inputs_and_credentials( return_value=[{"id": "c1", "name": "Cred", "type": "api", "is_default": True}], ) - result = rag_pipeline_service.service.get_datasource_plugins("t1", "d1", True) + result = rag_pipeline_service.service.get_datasource_plugins( + "t1", + "d1", + True, + credential_query=mock.create_autospec(CredentialQuery, instance=True), + datasource_providers=rag_pipeline_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ), + ) assert len(result) == 1 assert len(result[0]["user_input_variables"]) == 2 diff --git a/api/tests/unit_tests/services/test_credential_permission_service.py b/api/tests/unit_tests/services/test_credential_permission_service.py deleted file mode 100644 index 6964f702867..00000000000 --- a/api/tests/unit_tests/services/test_credential_permission_service.py +++ /dev/null @@ -1,201 +0,0 @@ -"""Unit tests for CredentialPermissionService. - -Tests the visibility filtering logic, partial-member read path, -and admin bypass behavior. -""" - -from uuid import uuid4 - -import pytest -from sqlalchemy import select -from sqlalchemy.orm import Session - -from core.plugin.entities.plugin_daemon import CredentialType as TriggerCredentialType -from models.account import Account, TenantAccountRole -from models.credential_permission import CredentialPermission, CredentialType -from models.enums import PermissionEnum -from models.trigger import TriggerSubscription -from services.credential_permission_service import CredentialPermissionService - - -@pytest.fixture -def tenant_id() -> str: - return str(uuid4()) - - -@pytest.fixture -def user_id() -> str: - return str(uuid4()) - - -@pytest.fixture -def other_user_id() -> str: - return str(uuid4()) - - -@pytest.fixture -def credential_id() -> str: - return str(uuid4()) - - -def _subscription( - *, - tenant_id: str, - owner_id: str, - name: str, - visibility: PermissionEnum, -) -> TriggerSubscription: - return TriggerSubscription( - name=name, - tenant_id=tenant_id, - user_id=owner_id, - provider_id="test/provider", - endpoint_id=f"{name}-endpoint", - parameters={}, - properties={}, - credentials={}, - credential_type=TriggerCredentialType.API_KEY, - visibility=visibility, - ) - - -def _user(user_id: str, *, is_admin: bool) -> Account: - user = Account(name="Credential User", email=f"{user_id}@example.com") - user.id = user_id - user.role = TenantAccountRole.ADMIN if is_admin else TenantAccountRole.NORMAL - return user - - -class TestGetPartialMemberList: - def test_returns_empty_when_no_permissions( - self, sqlite_session: Session, credential_id: str, tenant_id: str, user_id: str - ) -> None: - unrelated_permission = CredentialPermission( - credential_id=str(uuid4()), - credential_type=CredentialType.TRIGGER_SUBSCRIPTION, - account_id=user_id, - tenant_id=tenant_id, - ) - sqlite_session.add(unrelated_permission) - sqlite_session.commit() - - result = CredentialPermissionService.get_partial_member_list( - credential_id, CredentialType.TRIGGER_SUBSCRIPTION, session=sqlite_session - ) - - assert result == [] - - def test_returns_account_ids( - self, - sqlite_session: Session, - credential_id: str, - user_id: str, - other_user_id: str, - tenant_id: str, - ) -> None: - sqlite_session.add_all( - [ - CredentialPermission( - credential_id=credential_id, - credential_type=CredentialType.TRIGGER_SUBSCRIPTION, - account_id=account_id, - tenant_id=tenant_id, - ) - for account_id in (user_id, other_user_id) - ] - ) - sqlite_session.commit() - - result = CredentialPermissionService.get_partial_member_list( - credential_id, CredentialType.TRIGGER_SUBSCRIPTION, session=sqlite_session - ) - - assert set(result) == {user_id, other_user_id} - - -class TestApplyVisibilityFilter: - def test_admin_does_not_bypass_personal_visibility( - self, - sqlite_session: Session, - tenant_id: str, - user_id: str, - other_user_id: str, - ) -> None: - private_subscription = _subscription( - tenant_id=tenant_id, - owner_id=other_user_id, - name="private", - visibility=PermissionEnum.ONLY_ME, - ) - sqlite_session.add(private_subscription) - sqlite_session.commit() - - query = CredentialPermissionService.apply_visibility_filter( - select(TriggerSubscription).where(TriggerSubscription.tenant_id == tenant_id), - model_id_column=TriggerSubscription.id, - model_user_id_column=TriggerSubscription.user_id, - model_visibility_column=TriggerSubscription.visibility, - credential_type=CredentialType.TRIGGER_SUBSCRIPTION, - user=_user(user_id, is_admin=True), - ) - - assert sqlite_session.scalars(query).all() == [] - - def test_non_admin_sees_team_owned_and_partial_member_subscriptions( - self, - sqlite_session: Session, - tenant_id: str, - user_id: str, - other_user_id: str, - ) -> None: - team_subscription = _subscription( - tenant_id=tenant_id, - owner_id=other_user_id, - name="team", - visibility=PermissionEnum.ALL_TEAM, - ) - owned_subscription = _subscription( - tenant_id=tenant_id, - owner_id=user_id, - name="owned", - visibility=PermissionEnum.ONLY_ME, - ) - shared_subscription = _subscription( - tenant_id=tenant_id, - owner_id=other_user_id, - name="shared", - visibility=PermissionEnum.PARTIAL_TEAM, - ) - private_subscription = _subscription( - tenant_id=tenant_id, - owner_id=other_user_id, - name="private", - visibility=PermissionEnum.ONLY_ME, - ) - sqlite_session.add_all( - [ - team_subscription, - owned_subscription, - shared_subscription, - private_subscription, - CredentialPermission( - credential_id=shared_subscription.id, - credential_type=CredentialType.TRIGGER_SUBSCRIPTION, - account_id=user_id, - tenant_id=tenant_id, - ), - ] - ) - sqlite_session.commit() - - query = CredentialPermissionService.apply_visibility_filter( - select(TriggerSubscription).where(TriggerSubscription.tenant_id == tenant_id), - model_id_column=TriggerSubscription.id, - model_user_id_column=TriggerSubscription.user_id, - model_visibility_column=TriggerSubscription.visibility, - credential_type=CredentialType.TRIGGER_SUBSCRIPTION, - user=_user(user_id, is_admin=False), - ) - - visible_ids = {subscription.id for subscription in sqlite_session.scalars(query)} - assert visible_ids == {team_subscription.id, owned_subscription.id, shared_subscription.id} diff --git a/api/tests/unit_tests/services/test_data_source_oauth_service.py b/api/tests/unit_tests/services/test_data_source_oauth_service.py deleted file mode 100644 index f2f0189da8a..00000000000 --- a/api/tests/unit_tests/services/test_data_source_oauth_service.py +++ /dev/null @@ -1,182 +0,0 @@ -from unittest.mock import MagicMock, call - -import pytest - -from machinery.context import RequestContext -from services.data_source_oauth_service import ( - DataSourceOAuthConfigurationError, - DataSourceOAuthError, - DataSourceOAuthService, - InvalidDataSourceOAuthCodeError, -) -from services.entities.data_source_oauth_entities import ( - DataSourceOAuthAuthorization, - DataSourceOAuthBindingRecord, - DataSourceOAuthCallback, -) - - -@pytest.fixture -def request_context() -> RequestContext: - return RequestContext( - request_id="request-1", - trace_id="trace-1", - account_id="account-1", - active_workspace_id="workspace-1", - ) - - -def test_start_authorization_delegates_to_external_provider(request_context: RequestContext) -> None: - provider = MagicMock() - provider.get_authorization_url.return_value = "https://notion.example/authorize" - bindings = MagicMock() - service = DataSourceOAuthService( - provider_name="notion", - provider_gateway=provider, - bindings=bindings, - ) - - result = service.start_authorization(request_context) - - assert result == "https://notion.example/authorize" - provider.get_authorization_url.assert_called_once_with() - bindings.upsert_authorization.assert_not_called() - - -def test_start_authorization_persists_internal_authorization(request_context: RequestContext) -> None: - authorization = DataSourceOAuthAuthorization(access_token="secret", source_info={"pages": []}) - provider = MagicMock() - provider.authorize_internal.return_value = authorization - bindings = MagicMock() - service = DataSourceOAuthService( - provider_name="notion", - provider_gateway=provider, - bindings=bindings, - is_internal_provider=True, - internal_access_token="secret", - ) - - result = service.start_authorization(request_context) - - assert result == "internal" - provider.authorize_internal.assert_called_once_with("secret", "workspace-1") - bindings.upsert_authorization.assert_called_once_with( - workspace_id="workspace-1", - provider="notion", - authorization=authorization, - ) - - -def test_start_authorization_rejects_missing_internal_secret(request_context: RequestContext) -> None: - service = DataSourceOAuthService( - provider_name="notion", - provider_gateway=MagicMock(), - bindings=MagicMock(), - is_internal_provider=True, - ) - - with pytest.raises(DataSourceOAuthConfigurationError, match="Internal secret is not set"): - service.start_authorization(request_context) - - -@pytest.mark.parametrize( - ("code", "error", "expected_error"), - [ - ("auth/code", None, None), - (None, "access denied", "access denied"), - (None, None, "Access denied"), - ], -) -def test_complete_callback_returns_framework_neutral_result( - code: str | None, - error: str | None, - expected_error: str | None, -) -> None: - service = DataSourceOAuthService( - provider_name="notion", - provider_gateway=MagicMock(), - bindings=MagicMock(), - ) - - assert service.complete_callback(code=code, error=error) == DataSourceOAuthCallback( - provider="notion", - code=code, - error=expected_error, - ) - - -def test_bind_authorizes_before_persisting(request_context: RequestContext) -> None: - events = MagicMock() - authorization = DataSourceOAuthAuthorization(access_token="token", source_info={"pages": []}) - provider = MagicMock() - provider.authorize.side_effect = lambda code: (events.authorize(code), authorization)[1] - bindings = MagicMock() - bindings.upsert_authorization.side_effect = lambda **kwargs: events.persist(**kwargs) - service = DataSourceOAuthService( - provider_name="notion", - provider_gateway=provider, - bindings=bindings, - ) - - service.bind(request_context, code="code-1") - - assert events.mock_calls == [ - call.authorize("code-1"), - call.persist(workspace_id="workspace-1", provider="notion", authorization=authorization), - ] - - -def test_bind_rejects_empty_code(request_context: RequestContext) -> None: - service = DataSourceOAuthService( - provider_name="notion", - provider_gateway=MagicMock(), - bindings=MagicMock(), - ) - - with pytest.raises(InvalidDataSourceOAuthCodeError): - service.bind(request_context, code="") - - -def test_sync_reads_before_provider_io_and_persists_afterward(request_context: RequestContext) -> None: - events = MagicMock() - binding = DataSourceOAuthBindingRecord(id="binding-1", access_token="token", source_info={"pages": []}) - refreshed_source_info = {"pages": [{"page_id": "page-1"}]} - bindings = MagicMock() - bindings.get_enabled.side_effect = lambda **kwargs: (events.read(**kwargs), binding)[1] - bindings.update_source_info.side_effect = lambda **kwargs: (events.update(**kwargs), True)[1] - provider = MagicMock() - provider.refresh.side_effect = lambda *args: (events.refresh(*args), refreshed_source_info)[1] - service = DataSourceOAuthService( - provider_name="notion", - provider_gateway=provider, - bindings=bindings, - ) - - service.sync(request_context, binding_id="binding-1") - - assert events.mock_calls == [ - call.read(workspace_id="workspace-1", provider="notion", binding_id="binding-1"), - call.refresh("token", {"pages": []}), - call.update( - workspace_id="workspace-1", - provider="notion", - binding_id="binding-1", - source_info=refreshed_source_info, - ), - ] - - -def test_sync_rejects_missing_binding_without_provider_io(request_context: RequestContext) -> None: - provider = MagicMock() - bindings = MagicMock() - bindings.get_enabled.return_value = None - service = DataSourceOAuthService( - provider_name="notion", - provider_gateway=provider, - bindings=bindings, - ) - - with pytest.raises(DataSourceOAuthError, match="Data source binding not found"): - service.sync(request_context, binding_id="missing") - - provider.refresh.assert_not_called() diff --git a/api/tests/unit_tests/services/test_dataset_service_dataset.py b/api/tests/unit_tests/services/test_dataset_service_dataset.py index 4f2dd4fc258..28df36f5ced 100644 --- a/api/tests/unit_tests/services/test_dataset_service_dataset.py +++ b/api/tests/unit_tests/services/test_dataset_service_dataset.py @@ -1,4 +1,4 @@ -"""SQLite-backed tests for dataset-level operations in :mod:`services.dataset_service`. +"""SQLite-backed tests for dataset-level operations in :mod:`services.knowledge.dataset_service`. Mapped objects in this module are real SQLAlchemy models. Provider runtimes, RBAC clients, Celery tasks, and model-manager results remain mocked at their @@ -28,13 +28,13 @@ from models.dataset import ( ExternalKnowledgeBindings, Pipeline, ) -from services.dataset_service import DatasetCollectionBindingService, DatasetPermissionService, DatasetService from services.entities.knowledge_entities.rag_pipeline_entities import ( IconInfo, RagPipelineDatasetCreateEntity, ) from services.errors.account import NoPermissionError from services.errors.dataset import DatasetNameDuplicateError +from services.knowledge.dataset_service import DatasetCollectionBindingService, DatasetPermissionService, DatasetService from tests.unit_tests.model_factories import make_account, make_dataset, make_tenant from .dataset_service_test_helpers import ( @@ -177,7 +177,7 @@ class TestDatasetServiceValidation: def test_check_dataset_model_setting_skips_non_high_quality_datasets(self) -> None: dataset = _dataset(indexing_technique=IndexTechniqueType.ECONOMY) - with patch("services.dataset_service.ModelManager") as model_manager_cls: + with patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls: DatasetService.check_dataset_model_setting(dataset) model_manager_cls.assert_not_called() @@ -185,7 +185,7 @@ class TestDatasetServiceValidation: def test_check_dataset_model_setting_validates_high_quality_embedding(self) -> None: dataset = _dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY) - with patch("services.dataset_service.ModelManager") as model_manager_cls: + with patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls: DatasetService.check_dataset_model_setting(dataset) model_manager_cls.for_tenant.return_value.get_model_instance.assert_called_once_with( @@ -205,7 +205,7 @@ class TestDatasetServiceValidation: def test_check_dataset_model_setting_wraps_provider_errors(self, error: Exception, message: str) -> None: dataset = _dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY) - with patch("services.dataset_service.ModelManager") as model_manager_cls: + with patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls: model_manager_cls.for_tenant.return_value.get_model_instance.side_effect = error with pytest.raises(ValueError, match=message): DatasetService.check_dataset_model_setting(dataset) @@ -220,7 +220,7 @@ class TestDatasetServiceValidation: credentials={"api_key": "secret"}, ) - with patch("services.dataset_service.ModelManager") as model_manager_cls: + with patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls: model_manager_cls.for_tenant.return_value.get_model_instance.return_value = model_instance result = DatasetService.check_is_multimodal_model("tenant-1", "provider", "embedding-model") @@ -235,7 +235,7 @@ class TestDatasetServiceValidation: credentials={}, ) - with patch("services.dataset_service.ModelManager") as model_manager_cls: + with patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls: model_manager_cls.for_tenant.return_value.get_model_instance.return_value = model_instance with pytest.raises(ValueError, match="Model schema not found"): DatasetService.check_is_multimodal_model("tenant-1", "provider", "embedding-model") @@ -254,7 +254,7 @@ class TestDatasetServiceValidation: def test_direct_model_setting_checks_wrap_runtime_errors( self, method: Callable[[str, str, str], None], error: Exception, message: str ) -> None: - with patch("services.dataset_service.ModelManager") as model_manager_cls: + with patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls: model_manager_cls.for_tenant.return_value.get_model_instance.side_effect = error with pytest.raises(ValueError, match=message): method("tenant-1", "provider", "model") @@ -285,7 +285,7 @@ class TestDatasetServiceRetrieval: with ( patch( - "services.dataset_service.enterprise_rbac_service.RBACService.MyPermissions.get", + "services.knowledge.dataset_service.enterprise_rbac_service.RBACService.MyPermissions.get", return_value=SimpleNamespace(workspace=SimpleNamespace(permission_keys=[])), ), ): @@ -371,7 +371,7 @@ class TestDatasetServiceRetrieval: config_overrides(RBAC_ENABLED=True) with ( patch( - "services.dataset_service.enterprise_rbac_service.RBACService.MyPermissions.get", + "services.knowledge.dataset_service.enterprise_rbac_service.RBACService.MyPermissions.get", return_value=SimpleNamespace( workspace=SimpleNamespace(permission_keys=["dataset.create_and_management"]) ), @@ -414,8 +414,10 @@ class TestDatasetServiceCreationAndUpdate: embedding_model = SimpleNamespace(provider="provider", model_name="default-embedding") with ( - patch("services.dataset_service.ModelManager") as model_manager_cls, - patch("services.dataset_service.enterprise_rbac_service.try_sync_creator_access_policy_member_bindings"), + patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls, + patch( + "services.knowledge.dataset_service.enterprise_rbac_service.try_sync_creator_access_policy_member_bindings" + ), ): model_manager_cls.for_tenant.return_value.get_default_model_instance.return_value = embedding_model dataset = DatasetService.create_empty_dataset( @@ -438,7 +440,9 @@ class TestDatasetServiceCreationAndUpdate: sqlite_session.add_all([_external_api(), _external_api(api_id="api-foreign", tenant_id="tenant-2")]) sqlite_session.commit() - with patch("services.dataset_service.enterprise_rbac_service.try_sync_creator_access_policy_member_bindings"): + with patch( + "services.knowledge.dataset_service.enterprise_rbac_service.try_sync_creator_access_policy_member_bindings" + ): dataset = DatasetService.create_empty_dataset( "tenant-1", "External", @@ -465,7 +469,9 @@ class TestDatasetServiceCreationAndUpdate: with ( pytest.raises(ValueError, match="api template not found"), - patch("services.dataset_service.enterprise_rbac_service.try_sync_creator_access_policy_member_bindings"), + patch( + "services.knowledge.dataset_service.enterprise_rbac_service.try_sync_creator_access_policy_member_bindings" + ), ): DatasetService.create_empty_dataset( "tenant-1", @@ -489,7 +495,7 @@ class TestDatasetServiceCreationAndUpdate: permission=DatasetPermissionEnum.ALL_TEAM, ) - with patch("services.dataset_service.current_user", _account()): + with patch("services.knowledge.dataset_service.current_user", _account()): dataset = DatasetService.create_empty_rag_pipeline_dataset("tenant-1", entity, sqlite_session) assert dataset.name == "Untitled 2" @@ -508,7 +514,7 @@ class TestDatasetServiceCreationAndUpdate: ) with ( - patch("services.dataset_service.current_user", _account()), + patch("services.knowledge.dataset_service.current_user", _account()), pytest.raises(DatasetNameDuplicateError, match="already exists"), ): DatasetService.create_empty_rag_pipeline_dataset("tenant-1", entity, sqlite_session) @@ -524,7 +530,7 @@ class TestDatasetServiceCreationAndUpdate: ) with ( - patch("services.dataset_service.current_user", account), + patch("services.knowledge.dataset_service.current_user", account), pytest.raises(ValueError, match="Current user or current user id not found"), ): DatasetService.create_empty_rag_pipeline_dataset("tenant-1", entity, sqlite_session) @@ -640,8 +646,8 @@ class TestDatasetServiceCreationAndUpdate: with ( patch.object(DatasetService, "_handle_indexing_technique_change", return_value="update"), patch.object(DatasetService, "_update_pipeline_knowledge_base_node_data"), - patch("services.dataset_service.deal_dataset_vector_index_task.delay") as vector_task, - patch("services.dataset_service.regenerate_summary_index_task.delay") as summary_task, + patch("services.knowledge.dataset_service.deal_dataset_vector_index_task.delay") as vector_task, + patch("services.knowledge.dataset_service.regenerate_summary_index_task.delay") as summary_task, ): updated = DatasetService._update_internal_dataset( dataset, @@ -685,7 +691,7 @@ class TestDatasetServiceCreationAndUpdate: lambda _session, _previous_transaction: transaction_events.append("rollback"), ) - with patch("services.dataset_service.RagPipelineService") as service_cls: + with patch("services.knowledge.dataset_service.RagPipelineService") as service_cls: service_cls.return_value.get_published_workflow.side_effect = RuntimeError("boom") with pytest.raises(RuntimeError, match="boom"): DatasetService._update_pipeline_knowledge_base_node_data(dataset, "user-1", sqlite_session) @@ -765,8 +771,8 @@ class TestDatasetServiceEmbeddingSettings: def test_configure_high_quality_wraps_provider_error(self, unbound_session: Session) -> None: account = _account() with ( - patch("services.dataset_service.current_user", account), - patch("services.dataset_service.ModelManager") as model_manager_cls, + patch("services.knowledge.dataset_service.current_user", account), + patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls, ): model_manager_cls.for_tenant.return_value.get_model_instance.side_effect = LLMBadRequestError() with pytest.raises(ValueError, match="No Embedding Model available"): @@ -814,8 +820,8 @@ class TestDatasetServiceEmbeddingSettings: collection_binding.id = "collection-binding-1" with ( - patch("services.dataset_service.current_user", account), - patch("services.dataset_service.ModelManager") as model_manager_cls, + patch("services.knowledge.dataset_service.current_user", account), + patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls, patch.object( DatasetCollectionBindingService, "get_dataset_collection_binding", @@ -841,8 +847,8 @@ class TestDatasetServiceEmbeddingSettings: filtered_data: dict[str, object] = {} with ( - patch("services.dataset_service.current_user", _account()), - patch("services.dataset_service.ModelManager") as model_manager_cls, + patch("services.knowledge.dataset_service.current_user", _account()), + patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls, ): model_manager_cls.for_tenant.return_value.get_model_instance.side_effect = ProviderTokenNotInitError( "missing" @@ -886,7 +892,7 @@ class TestDatasetServiceRagPipelineSettings: account._current_tenant = None with ( - patch("services.dataset_service.current_user", account), + patch("services.knowledge.dataset_service.current_user", account), pytest.raises(ValueError, match="Current user or current tenant not found"), ): DatasetService.update_rag_pipeline_dataset_settings( @@ -910,8 +916,8 @@ class TestDatasetServiceRagPipelineSettings: collection_binding.id = "collection-binding-2" with ( - patch("services.dataset_service.current_user", account), - patch("services.dataset_service.ModelManager") as model_manager_cls, + patch("services.knowledge.dataset_service.current_user", account), + patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls, patch.object(DatasetService, "check_is_multimodal_model", return_value=True), patch.object( DatasetCollectionBindingService, @@ -942,7 +948,7 @@ class TestDatasetServiceRagPipelineSettings: sqlite_session.add(dataset) sqlite_session.commit() - with patch("services.dataset_service.current_user", _account()): + with patch("services.knowledge.dataset_service.current_user", _account()): DatasetService.update_rag_pipeline_dataset_settings( dataset, _make_knowledge_configuration( @@ -966,8 +972,8 @@ class TestDatasetServiceRagPipelineSettings: sqlite_session.commit() with ( - patch("services.dataset_service.current_user", _account()), - patch("services.dataset_service.deal_dataset_index_update_task.delay") as update_task, + patch("services.knowledge.dataset_service.current_user", _account()), + patch("services.knowledge.dataset_service.deal_dataset_index_update_task.delay") as update_task, ): DatasetService.update_rag_pipeline_dataset_settings( dataset, @@ -1012,7 +1018,7 @@ class TestDatasetServiceRagPipelineSettings: sqlite_session.commit() with ( - patch("services.dataset_service.current_user", _account()), + patch("services.knowledge.dataset_service.current_user", _account()), pytest.raises(ValueError, match=message), ): DatasetService.update_rag_pipeline_dataset_settings( diff --git a/api/tests/unit_tests/services/test_dataset_service_document.py b/api/tests/unit_tests/services/test_dataset_service_document.py index a450560d231..702e7902065 100644 --- a/api/tests/unit_tests/services/test_dataset_service_document.py +++ b/api/tests/unit_tests/services/test_dataset_service_document.py @@ -6,12 +6,13 @@ from datetime import datetime from sqlalchemy import event, select from sqlalchemy.orm import Session +from core.model_manager import ModelInstance from enums import DeploymentEdition from models.dataset import Dataset, DatasetCollectionBinding, DocumentSegment from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus from models.model import UploadFile from models.source import DataSourceOauthBinding -from services.dataset_ref_service import DatasetRefService +from services.knowledge.resource_scope import DatasetRef from tests.unit_tests.config_override import config_overrides_context from tests.unit_tests.model_factories import make_account, make_tenant, make_upload_file @@ -41,7 +42,6 @@ from .dataset_service_test_helpers import ( RetrievalModel, Rule, Segmentation, - SimpleNamespace, WebsiteInfo, _make_features, _make_lock_context, @@ -248,7 +248,7 @@ class TestDocumentServiceRetrieval: sqlite_session.commit() documents = DocumentService.get_document_by_ids( - DatasetRefService.create_dataset_ref(dataset), + DatasetRef(tenant_id=dataset.tenant_id, dataset_id=dataset.id), ["expected", "disabled", "archived", "waiting", "other-dataset", "other-tenant"], sqlite_session, ) @@ -276,8 +276,8 @@ class TestDocumentServiceMutations: sqlite_session.add_all([dataset, document, other_dataset, other_tenant]) sqlite_session.commit() - with patch("services.dataset_service.batch_clean_document_task") as clean_task: - dataset_ref = DatasetRefService.create_dataset_ref(dataset) + with patch("services.knowledge.dataset_service.batch_clean_document_task") as clean_task: + dataset_ref = DatasetRef(tenant_id=dataset.tenant_id, dataset_id=dataset.id) DocumentService.delete_documents( dataset_ref, [document.id, other_dataset.id, other_tenant.id], @@ -300,14 +300,18 @@ class TestDocumentServiceMutations: commits += 1 event.listen(sqlite_session, "after_commit", count_commit) + dataset = _dataset_row() DocumentService.delete_documents( - DatasetRefService.create_dataset_ref(_dataset_row()), [], IndexStructureType.PARAGRAPH_INDEX, sqlite_session + DatasetRef(tenant_id=dataset.tenant_id, dataset_id=dataset.id), + [], + IndexStructureType.PARAGRAPH_INDEX, + sqlite_session, ) event.remove(sqlite_session, "after_commit", count_commit) assert commits == 0 def test_rename_document_raises_when_dataset_is_missing(self, sqlite_session: Session): - with patch("services.dataset_service.current_user", _account()): + with patch("services.knowledge.dataset_service.current_user", _account()): with pytest.raises(ValueError, match="Dataset not found"): DocumentService.rename_document("dataset-1", "doc-1", "New Name", sqlite_session) @@ -315,7 +319,7 @@ class TestDocumentServiceMutations: dataset = _dataset_row() sqlite_session.add(dataset) sqlite_session.commit() - with patch("services.dataset_service.current_user", _account()): + with patch("services.knowledge.dataset_service.current_user", _account()): with pytest.raises(ValueError, match="Document not found"): DocumentService.rename_document(dataset.id, "doc-1", "New Name", sqlite_session) @@ -324,7 +328,7 @@ class TestDocumentServiceMutations: document = _document_row(tenant_id="tenant-other") sqlite_session.add_all([dataset, document]) sqlite_session.commit() - with patch("services.dataset_service.current_user", _account()): + with patch("services.knowledge.dataset_service.current_user", _account()): with pytest.raises(ValueError, match="No permission"): DocumentService.rename_document(dataset.id, document.id, "New Name", sqlite_session) @@ -357,7 +361,7 @@ class TestDocumentServiceMutations: commits += 1 event.listen(sqlite_session, "after_commit", count_commit) - with patch("services.dataset_service.current_user", _account()): + with patch("services.knowledge.dataset_service.current_user", _account()): result = DocumentService.rename_document(dataset.id, document.id, "New Name", sqlite_session) event.remove(sqlite_session, "after_commit", count_commit) @@ -377,8 +381,8 @@ class TestDocumentServiceMutations: sqlite_session.add(document) sqlite_session.commit() with ( - patch("services.dataset_service.redis_client") as redis, - patch("services.dataset_service.recover_document_indexing_task") as task, + patch("services.knowledge.dataset_service.redis_client") as redis, + patch("services.knowledge.dataset_service.recover_document_indexing_task") as task, ): DocumentService.recover_document(document, sqlite_session) @@ -395,11 +399,11 @@ class TestDocumentServiceMutations: sqlite_session.commit() retry_flags = _RetryFlagStore({f"document_{document.id}_is_retried": "other-request"}) with ( - patch("services.dataset_service.current_user", _account()), - patch("services.dataset_service.redis_client", retry_flags), + patch("services.knowledge.dataset_service.current_user", _account()), + patch("services.knowledge.dataset_service.redis_client", retry_flags), ): with pytest.raises(ValueError, match="being retried"): - DocumentService.retry_document("dataset-1", [document], sqlite_session) + DocumentService.retry_document("dataset-1", [document], sqlite_session, actor_id="user-1") def test_retry_document_leaves_batch_unchanged_when_later_document_is_already_being_retried( self, sqlite_session: Session @@ -412,15 +416,13 @@ class TestDocumentServiceMutations: second_retry_key = "document_doc-2_is_retried" retry_flags = _RetryFlagStore({second_retry_key: "other-request"}) with ( - patch("services.dataset_service.current_user", _account()), - patch("services.dataset_service.redis_client", retry_flags), - patch("services.dataset_service.retry_document_indexing_task") as retry_task, + patch("services.knowledge.dataset_service.current_user", _account()), + patch("services.knowledge.dataset_service.redis_client", retry_flags), + patch("services.knowledge.dataset_service.retry_document_indexing_task") as retry_task, ): with pytest.raises(ValueError, match="being retried"): DocumentService.retry_document( - "dataset-1", - [first_document, second_document], - sqlite_session, + "dataset-1", [first_document, second_document], sqlite_session, actor_id="user-1" ) assert first_document.indexing_status == IndexingStatus.ERROR @@ -444,11 +446,11 @@ class TestDocumentServiceMutations: sqlite_session.commit() with ( - patch("services.dataset_service.current_user", _account()), - patch("services.dataset_service.redis_client", retry_flags), + patch("services.knowledge.dataset_service.current_user", _account()), + patch("services.knowledge.dataset_service.redis_client", retry_flags), ): with pytest.raises(ValueError, match="being retried"): - DocumentService.retry_document("dataset-1", documents, sqlite_session) + DocumentService.retry_document("dataset-1", documents, sqlite_session, actor_id="user-1") assert retry_flags.values[first_retry_key] == "new-owner" assert retry_flags.values[second_retry_key] == "other-request" @@ -464,12 +466,12 @@ class TestDocumentServiceMutations: event.listen(sqlite_session, "before_commit", fail_commit) with ( - patch("services.dataset_service.current_user", _account()), - patch("services.dataset_service.redis_client", retry_flags), - patch("services.dataset_service.retry_document_indexing_task") as retry_task, + patch("services.knowledge.dataset_service.current_user", _account()), + patch("services.knowledge.dataset_service.redis_client", retry_flags), + patch("services.knowledge.dataset_service.retry_document_indexing_task") as retry_task, ): with pytest.raises(RuntimeError, match="database unavailable"): - DocumentService.retry_document("dataset-1", [document], sqlite_session) + DocumentService.retry_document("dataset-1", [document], sqlite_session, actor_id="user-1") event.remove(sqlite_session, "before_commit", fail_commit) assert retry_flags.values == {} @@ -484,11 +486,11 @@ class TestDocumentServiceMutations: sqlite_session.commit() retry_flags = _RetryFlagStore() with ( - patch("services.dataset_service.current_user", _account()), - patch("services.dataset_service.redis_client", retry_flags), - patch("services.dataset_service.retry_document_indexing_task") as task, + patch("services.knowledge.dataset_service.current_user", _account()), + patch("services.knowledge.dataset_service.redis_client", retry_flags), + patch("services.knowledge.dataset_service.retry_document_indexing_task") as task, ): - DocumentService.retry_document("dataset-1", documents, sqlite_session) + DocumentService.retry_document("dataset-1", documents, sqlite_session, actor_id="user-1") sqlite_session.expire_all() statuses = sqlite_session.scalars(select(Document.indexing_status).order_by(Document.id)).all() @@ -498,19 +500,25 @@ class TestDocumentServiceMutations: def test_sync_website_document_raises_when_sync_flag_exists(self, sqlite_session: Session): dataset = _dataset_row() document = _document_row() - with patch("services.dataset_service.redis_client") as mock_redis: + with patch("services.knowledge.dataset_service.redis_client") as mock_redis: mock_redis.get.return_value = "1" with pytest.raises(ValueError, match="being synced"): DocumentService.sync_website_document(dataset, document, sqlite_session) - def test_sync_website_document_rejects_document_outside_dataset(self, sqlite_session: Session): + @pytest.mark.parametrize( + ("document_dataset_id", "document_tenant_id"), + [("dataset-2", "tenant-1"), ("dataset-1", "tenant-2")], + ) + def test_sync_website_document_rejects_document_outside_dataset( + self, document_dataset_id: str, document_tenant_id: str, sqlite_session: Session + ): dataset = _dataset_row() - document = _document_row(dataset_id="dataset-2") + document = _document_row(dataset_id=document_dataset_id, tenant_id=document_tenant_id) with ( pytest.raises(ValueError, match="Document not found"), - patch("services.dataset_service.redis_client") as mock_redis, + patch("services.knowledge.dataset_service.redis_client") as mock_redis, ): DocumentService.sync_website_document(dataset, document, sqlite_session) @@ -526,8 +534,8 @@ class TestDocumentServiceMutations: sqlite_session.commit() with ( - patch("services.dataset_service.redis_client") as mock_redis, - patch("services.dataset_service.sync_website_document_indexing_task") as sync_task, + patch("services.knowledge.dataset_service.redis_client") as mock_redis, + patch("services.knowledge.dataset_service.sync_website_document_indexing_task") as sync_task, ): mock_redis.get.return_value = None @@ -549,7 +557,7 @@ class TestDocumentServiceSaveDocumentWithoutDatasetId: def account_context(self): account = _account() - with patch("services.dataset_service.current_user", account): + with patch("services.knowledge.dataset_service.current_user", account): yield account @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) @@ -579,9 +587,9 @@ class TestDocumentServiceSaveDocumentWithoutDatasetId: first_document = _document_row(name="VeryLongDocumentNameForDataset.txt") with ( - patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()), + patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()), patch( - "services.dataset_service.DatasetCollectionBindingService.get_dataset_collection_binding", + "services.knowledge.dataset_service.DatasetCollectionBindingService.get_dataset_collection_binding", return_value=binding, ), patch.object( @@ -643,7 +651,7 @@ class TestDocumentServiceSaveDocumentWithoutDatasetId: first_document = _document_row(name="Doc") with ( - patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()), + patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()), patch.object( DocumentService, "save_document_with_dataset_id", @@ -677,7 +685,7 @@ class TestDocumentServiceSaveDocumentWithoutDatasetId: with ( patch( - "services.dataset_service.FeatureService.get_features", + "services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features(plan=CloudPlan.SANDBOX), ), patch.object(DocumentService, "check_documents_upload_quota") as check_quota, @@ -697,7 +705,7 @@ class TestDocumentServiceUpdateDocumentWithDatasetId: def account_context(self): account = _account() - with patch("services.dataset_service.current_user", account): + with patch("services.knowledge.dataset_service.current_user", account): yield account def test_update_document_with_dataset_id_raises_when_document_is_missing( @@ -805,8 +813,8 @@ class TestDocumentServiceUpdateDocumentWithDatasetId: with ( patch.object(DatasetService, "check_dataset_model_setting"), - patch("services.dataset_service.naive_utc_now", return_value=updated_at), - patch("services.dataset_service.document_indexing_update_task") as update_task, + patch("services.knowledge.dataset_service.naive_utc_now", return_value=updated_at), + patch("services.knowledge.dataset_service.document_indexing_update_task") as update_task, ): result = DocumentService.update_document_with_dataset_id( dataset, @@ -906,8 +914,8 @@ class TestDocumentServiceUpdateDocumentWithDatasetId: with ( patch.object(DatasetService, "check_dataset_model_setting"), - patch("services.dataset_service.naive_utc_now", return_value=datetime(2026, 2, 1)), - patch("services.dataset_service.document_indexing_update_task") as update_task, + patch("services.knowledge.dataset_service.naive_utc_now", return_value=datetime(2026, 2, 1)), + patch("services.knowledge.dataset_service.document_indexing_update_task") as update_task, ): result = DocumentService.update_document_with_dataset_id( dataset, @@ -933,13 +941,22 @@ class TestDocumentServiceCreateValidation: """Unit tests for document creation validation helpers.""" def test_document_create_args_validate_requires_data_source_or_process_rule(self): - knowledge_config = SimpleNamespace(data_source=None, process_rule=None) + knowledge_config = KnowledgeConfig(indexing_technique="economy") with pytest.raises(ValueError, match="Data source or Process rule is required"): DocumentService.document_create_args_validate(knowledge_config) def test_document_create_args_validate_delegates_to_sub_validators(self): - knowledge_config = SimpleNamespace(data_source=object(), process_rule=object()) + knowledge_config = KnowledgeConfig( + indexing_technique="economy", + data_source=DataSource( + info_list=InfoList( + data_source_type="upload_file", + file_info_list=FileInfo(file_ids=["file-1"]), + ) + ), + process_rule=ProcessRule(mode="automatic"), + ) with ( patch.object(DocumentService, "data_source_args_validate") as validate_data_source, @@ -950,39 +967,20 @@ class TestDocumentServiceCreateValidation: validate_data_source.assert_called_once_with(knowledge_config) validate_process_rule.assert_called_once_with(knowledge_config) - def test_data_source_args_validate_rejects_invalid_type(self): - knowledge_config = SimpleNamespace( - data_source=SimpleNamespace( - info_list=SimpleNamespace( - data_source_type="bad-source", - file_info_list=None, - notion_info_list=None, - website_info_list=None, - ) - ) - ) - - with pytest.raises(ValueError, match="Data source type is invalid"): - DocumentService.data_source_args_validate(knowledge_config) - @pytest.mark.parametrize( - ("data_source_type", "field_name", "message"), + ("data_source_type", "message"), [ - ("upload_file", "file_info_list", "File source info is required"), - ("notion_import", "notion_info_list", "Notion source info is required"), - ("website_crawl", "website_info_list", "Website source info is required"), + ("upload_file", "File source info is required"), + ("notion_import", "Notion source info is required"), + ("website_crawl", "Website source info is required"), ], ) - def test_data_source_args_validate_requires_source_specific_info(self, data_source_type, field_name, message): - info_values = { - "data_source_type": data_source_type, - "file_info_list": object(), - "notion_info_list": object(), - "website_info_list": object(), - } - info_values[field_name] = None - info_list = SimpleNamespace(**info_values) - knowledge_config = SimpleNamespace(data_source=SimpleNamespace(info_list=info_list)) + def test_data_source_args_validate_requires_source_specific_info(self, data_source_type, message): + info_list = InfoList.model_validate({"data_source_type": data_source_type}) + knowledge_config = KnowledgeConfig( + indexing_technique="economy", + data_source=DataSource(info_list=info_list), + ) with pytest.raises(ValueError, match=message): DocumentService.data_source_args_validate(knowledge_config) @@ -1075,7 +1073,7 @@ class TestDocumentServiceSaveDocumentWithDatasetId: account = _account() with ( - patch("services.dataset_service.current_user", account), + patch("services.knowledge.dataset_service.current_user", account), patch.object(DatasetService, "check_doc_form"), ): yield account @@ -1087,7 +1085,7 @@ class TestDocumentServiceSaveDocumentWithDatasetId: dataset = _dataset_row() knowledge_config = _make_upload_knowledge_config(file_ids=None) - with patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()): + with patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()): with pytest.raises(ValueError, match="File source info is required"): DocumentService.save_document_with_dataset_id( dataset, @@ -1105,7 +1103,7 @@ class TestDocumentServiceSaveDocumentWithDatasetId: with ( patch( - "services.dataset_service.FeatureService.get_features", + "services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features(plan=CloudPlan.SANDBOX), ), patch.object(DocumentService, "check_documents_upload_quota") as check_quota, @@ -1132,7 +1130,7 @@ class TestDocumentServiceSaveDocumentWithDatasetId: knowledge_config = _make_upload_knowledge_config(file_ids=["file-1", "file-2"]) with ( - patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()), + patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()), patch.object(DocumentService, "check_documents_upload_quota") as check_quota, ): with pytest.raises(ValueError, match="batch upload limit of 1"): @@ -1155,10 +1153,9 @@ class TestDocumentServiceSaveDocumentWithDatasetId: updated_document.batch = "batch-existing" with ( - patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()), - patch.object( - DocumentService, "update_document_with_dataset_id", return_value=updated_document - ) as update_document, + patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()), + patch("services.knowledge.dataset_service.document_indexing_update_task.delay"), + patch.object(DocumentService, "update_prepared_document", return_value=updated_document) as update_document, ): documents, batch = DocumentService.save_document_with_dataset_id( dataset, @@ -1179,7 +1176,7 @@ class TestDocumentServiceSaveDocumentWithDatasetId: dataset = _dataset_row() knowledge_config = _make_upload_knowledge_config(data_source=None) - with patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()): + with patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()): with pytest.raises(ValueError, match="Data source is required when creating new documents"): DocumentService.save_document_with_dataset_id( dataset, @@ -1200,7 +1197,7 @@ class TestDocumentServiceSaveDocumentWithDatasetId: process_rule=ProcessRule(mode="custom"), ) - with patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()): + with patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()): with pytest.raises(ValueError, match="No process rule found"): DocumentService.save_document_with_dataset_id( dataset, @@ -1209,46 +1206,6 @@ class TestDocumentServiceSaveDocumentWithDatasetId: session=sqlite_session, ) - @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - def test_save_document_with_dataset_id_rejects_invalid_indexing_technique( - self, account_context, unbound_session: Session - ): - dataset = _dataset_row(indexing_technique=None) - knowledge_config = SimpleNamespace( - doc_form=IndexStructureType.PARAGRAPH_INDEX, - original_document_id=None, - data_source=None, - indexing_technique="broken-technique", - ) - - with patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()): - with pytest.raises(ValueError, match="Indexing technique is invalid"): - DocumentService.save_document_with_dataset_id( - dataset, - knowledge_config, - account_context, - session=unbound_session, - ) - - @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) - def test_save_document_with_dataset_id_returns_empty_for_invalid_process_rule_mode( - self, account_context, unbound_session: Session - ): - dataset = _dataset_row() - knowledge_config = _make_upload_knowledge_config(file_ids=["file-1"]) - knowledge_config.process_rule = SimpleNamespace(mode="unsupported-mode", rules=None) - - with patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()): - documents, batch = DocumentService.save_document_with_dataset_id( - dataset, - knowledge_config, - account_context, - session=unbound_session, - ) - - assert documents == [] - assert batch == "" - @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) def test_save_document_with_dataset_id_upload_file_creates_and_reindexes_documents( self, account_context, sqlite_session: Session @@ -1263,13 +1220,13 @@ class TestDocumentServiceSaveDocumentWithDatasetId: sqlite_session.commit() with ( - patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()), - patch("services.dataset_service.redis_client") as mock_redis, - patch("services.dataset_service.DocumentIndexingTaskProxy") as document_proxy_cls, - patch("services.dataset_service.DuplicateDocumentIndexingTaskProxy") as duplicate_proxy_cls, - patch("services.dataset_service.naive_utc_now", return_value=datetime(2026, 2, 1)), - patch("services.dataset_service.time.strftime", return_value="20260101010101"), - patch("services.dataset_service.secrets.randbelow", return_value=23), + patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()), + patch("services.knowledge.dataset_service.redis_client") as mock_redis, + patch("services.knowledge.dataset_service.DocumentIndexingTaskProxy") as document_proxy_cls, + patch("services.knowledge.dataset_service.DuplicateDocumentIndexingTaskProxy") as duplicate_proxy_cls, + patch("services.knowledge.dataset_service.naive_utc_now", return_value=datetime(2026, 2, 1)), + patch("services.knowledge.dataset_service.time.strftime", return_value="20260101010101"), + patch("services.knowledge.dataset_service.secrets.randbelow", return_value=23), ): mock_redis.lock.return_value = _make_lock_context() documents, batch = DocumentService.save_document_with_dataset_id( @@ -1340,11 +1297,11 @@ class TestDocumentServiceSaveDocumentWithDatasetId: sqlite_session.commit() with ( - patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()), - patch("services.dataset_service.redis_client") as mock_redis, - patch("services.dataset_service.clean_notion_document_task") as clean_task, - patch("services.dataset_service.DocumentIndexingTaskProxy") as document_proxy_cls, - patch("services.dataset_service.uuid.uuid4", return_value="doc-new"), + patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()), + patch("services.knowledge.dataset_service.redis_client") as mock_redis, + patch("services.knowledge.dataset_service.clean_notion_document_task") as clean_task, + patch("services.knowledge.dataset_service.DocumentIndexingTaskProxy") as document_proxy_cls, + patch("services.knowledge.dataset_service.uuid.uuid4", return_value="doc-new"), ): mock_redis.lock.return_value = _make_lock_context() documents, _ = DocumentService.save_document_with_dataset_id( @@ -1390,9 +1347,9 @@ class TestDocumentServiceSaveDocumentWithDatasetId: doc_language="English", ) with ( - patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()), - patch("services.dataset_service.redis_client") as mock_redis, - patch("services.dataset_service.DocumentIndexingTaskProxy") as document_proxy_cls, + patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()), + patch("services.knowledge.dataset_service.redis_client") as mock_redis, + patch("services.knowledge.dataset_service.DocumentIndexingTaskProxy") as document_proxy_cls, ): mock_redis.lock.return_value = _make_lock_context() @@ -1448,7 +1405,7 @@ class TestDocumentServiceBatchUpdateStatus: sqlite_session.add_all([dataset, document]) sqlite_session.commit() - with patch("services.dataset_service.redis_client") as mock_redis: + with patch("services.knowledge.dataset_service.redis_client") as mock_redis: mock_redis.get.return_value = "1" with pytest.raises(DocumentIndexingError, match="Busy document is being indexed"): @@ -1470,7 +1427,7 @@ class TestDocumentServiceBatchUpdateStatus: event.listen(sqlite_session, "before_commit", fail_commit) with ( - patch("services.dataset_service.redis_client") as mock_redis, + patch("services.knowledge.dataset_service.redis_client") as mock_redis, ): mock_redis.get.return_value = None @@ -1490,8 +1447,8 @@ class TestDocumentServiceBatchUpdateStatus: sqlite_session.commit() with ( - patch("services.dataset_service.redis_client") as mock_redis, - patch("services.dataset_service.add_document_to_index_task") as add_task, + patch("services.knowledge.dataset_service.redis_client") as mock_redis, + patch("services.knowledge.dataset_service.add_document_to_index_task") as add_task, ): mock_redis.get.return_value = None add_task.delay.side_effect = RuntimeError("task failed") @@ -1513,7 +1470,7 @@ class TestDocumentServiceTenantAndUpdateEdges: def account_context(self): account = _account() - with patch("services.dataset_service.current_user", account): + with patch("services.knowledge.dataset_service.current_user", account): yield account def test_get_tenant_documents_count_scopes_state_and_tenant(self, account_context, sqlite_session: Session): @@ -1563,8 +1520,8 @@ class TestDocumentServiceTenantAndUpdateEdges: with ( patch.object(DatasetService, "check_dataset_model_setting"), - patch("services.dataset_service.naive_utc_now", return_value=updated_at), - patch("services.dataset_service.document_indexing_update_task") as update_task, + patch("services.knowledge.dataset_service.naive_utc_now", return_value=updated_at), + patch("services.knowledge.dataset_service.document_indexing_update_task") as update_task, ): result = DocumentService.update_document_with_dataset_id( dataset, @@ -1686,8 +1643,8 @@ class TestDocumentServiceTenantAndUpdateEdges: with ( patch.object(DatasetService, "check_dataset_model_setting"), - patch("services.dataset_service.naive_utc_now", return_value=datetime(2026, 2, 1)), - patch("services.dataset_service.document_indexing_update_task") as update_task, + patch("services.knowledge.dataset_service.naive_utc_now", return_value=datetime(2026, 2, 1)), + patch("services.knowledge.dataset_service.document_indexing_update_task") as update_task, ): result = DocumentService.update_document_with_dataset_id( dataset, @@ -1719,7 +1676,7 @@ class TestDocumentServiceSaveWithoutDatasetBilling: def account_context(self): account = _account() - with patch("services.dataset_service.current_user", account): + with patch("services.knowledge.dataset_service.current_user", account): yield account @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) @@ -1757,7 +1714,7 @@ class TestDocumentServiceSaveWithoutDatasetBilling: document = _document_row(name="Doc") with ( - patch("services.dataset_service.FeatureService.get_features", return_value=features), + patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=features), patch.object(DocumentService, "check_documents_upload_quota") as check_quota, patch.object( DocumentService, @@ -1799,7 +1756,7 @@ class TestDocumentServiceSaveWithoutDatasetBilling: ) with ( - patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()), + patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()), patch.object(DocumentService, "check_documents_upload_quota") as check_quota, ): with pytest.raises(ValueError, match="batch upload limit of 1"): @@ -1810,132 +1767,6 @@ class TestDocumentServiceSaveWithoutDatasetBilling: check_quota.assert_not_called() -class TestDocumentServiceEstimateValidation: - """Unit tests for estimate_args_validate branches.""" - - def test_estimate_args_validate_rejects_missing_info_list(self): - with pytest.raises(ValueError, match="Field required"): - DocumentService.estimate_args_validate({}) - - def test_estimate_args_validate_sets_empty_rules_for_automatic_mode(self): - args = { - "info_list": {"data_source_type": "upload_file"}, - "process_rule": {"mode": "automatic", "rules": {"ignored": True}}, - } - - DocumentService.estimate_args_validate(args) - - assert args["process_rule"]["rules"] == {} - - def test_estimate_args_validate_rejects_unknown_pre_processing_rule_id(self): - args = { - "info_list": {"data_source_type": "upload_file"}, - "process_rule": { - "mode": "custom", - "rules": { - "pre_processing_rules": [{"id": "unknown", "enabled": True}], - "segmentation": {"separator": "\n", "max_tokens": 128}, - }, - }, - } - - with pytest.raises(ValueError, match="pre_processing_rules id is invalid"): - DocumentService.estimate_args_validate(args) - - def test_estimate_args_validate_deduplicates_rules_for_custom_mode(self): - args = { - "info_list": {"data_source_type": "upload_file"}, - "process_rule": { - "mode": "custom", - "rules": { - "pre_processing_rules": [ - {"id": "remove_stopwords", "enabled": True}, - {"id": "remove_stopwords", "enabled": False}, - ], - "segmentation": {"separator": "\n", "max_tokens": 128}, - }, - }, - } - - DocumentService.estimate_args_validate(args) - - assert args["process_rule"]["rules"]["pre_processing_rules"] == [{"id": "remove_stopwords", "enabled": False}] - - def test_estimate_args_validate_custom_mode_drops_hierarchical_fields(self): - args = { - "info_list": {"data_source_type": "upload_file"}, - "process_rule": { - "mode": "custom", - "rules": { - "pre_processing_rules": [{"id": "remove_stopwords", "enabled": True}], - "segmentation": {"separator": "\n", "max_tokens": 128}, - "parent_mode": "full-doc", - "subchunk_segmentation": {"separator": "###", "max_tokens": 64}, - }, - }, - } - - DocumentService.estimate_args_validate(args) - - assert args["process_rule"]["rules"] == { - "pre_processing_rules": [{"id": "remove_stopwords", "enabled": True}], - "segmentation": {"separator": "\n", "max_tokens": 128}, - } - - def test_estimate_args_validate_requires_summary_index_provider_name(self): - args = { - "info_list": {"data_source_type": "upload_file"}, - "process_rule": { - "mode": "custom", - "rules": { - "pre_processing_rules": [{"id": "remove_stopwords", "enabled": True}], - "segmentation": {"separator": "\n", "max_tokens": 128}, - }, - "summary_index_setting": {"enable": True, "model_name": "summary-model"}, - }, - } - - with pytest.raises(ValueError, match="Field required"): - DocumentService.estimate_args_validate(args) - - def test_estimate_args_validate_preserves_hierarchical_fields(self): - args = { - "info_list": {"data_source_type": "upload_file"}, - "process_rule": { - "mode": "hierarchical", - "rules": { - "pre_processing_rules": [{"id": "remove_stopwords", "enabled": True}], - "segmentation": {"separator": "\n", "max_tokens": 512}, - "parent_mode": "full-doc", - "subchunk_segmentation": {"separator": "###", "max_tokens": 128}, - }, - }, - } - - DocumentService.estimate_args_validate(args) - - assert args["process_rule"]["rules"]["parent_mode"] == "full-doc" - assert args["process_rule"]["rules"]["subchunk_segmentation"] == {"separator": "###", "max_tokens": 128} - - def test_estimate_args_validate_hierarchical_defaults_parent_mode_to_paragraph(self): - args = { - "info_list": {"data_source_type": "upload_file"}, - "process_rule": { - "mode": "hierarchical", - "rules": { - "pre_processing_rules": [{"id": "remove_stopwords", "enabled": True}], - "segmentation": {"separator": "\n", "max_tokens": 512}, - "subchunk_segmentation": {"separator": "###", "max_tokens": 128}, - }, - }, - } - - DocumentService.estimate_args_validate(args) - - assert args["process_rule"]["rules"]["parent_mode"] == "paragraph" - assert args["process_rule"]["rules"]["subchunk_segmentation"] == {"separator": "###", "max_tokens": 128} - - class TestDocumentServiceSaveDocumentAdditionalBranches: """Additional unit tests for dataset bootstrap and process-rule branches.""" @@ -1944,7 +1775,7 @@ class TestDocumentServiceSaveDocumentAdditionalBranches: account = _account() with ( - patch("services.dataset_service.current_user", account), + patch("services.knowledge.dataset_service.current_user", account), patch.object(DatasetService, "check_doc_form"), ): yield account @@ -1969,18 +1800,19 @@ class TestDocumentServiceSaveDocumentAdditionalBranches: binding.id = "binding-1" with ( - patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()), - patch("services.dataset_service.ModelManager") as model_manager_cls, + patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()), + patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls, patch( - "services.dataset_service.DatasetCollectionBindingService.get_dataset_collection_binding", + "services.knowledge.dataset_service.DatasetCollectionBindingService.get_dataset_collection_binding", return_value=binding, ) as get_binding, - patch.object(DocumentService, "update_document_with_dataset_id", return_value=updated_document), + patch.object(DocumentService, "update_prepared_document", return_value=updated_document), + patch("services.knowledge.dataset_service.document_indexing_update_task.delay"), ): - model_manager_cls.for_tenant.return_value.get_default_model_instance.return_value = SimpleNamespace( - model_name="default-embedding", - provider="default-provider", - ) + default_model = object.__new__(ModelInstance) + default_model.model_name = "default-embedding" + default_model.provider = "default-provider" + model_manager_cls.for_tenant.return_value.get_default_model_instance.return_value = default_model documents, batch = DocumentService.save_document_with_dataset_id( dataset, @@ -2035,13 +1867,14 @@ class TestDocumentServiceSaveDocumentAdditionalBranches: updated_document = _document_row(document_id="doc-1") with ( - patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()), - patch("services.dataset_service.ModelManager") as model_manager_cls, + patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()), + patch("services.knowledge.dataset_service.ModelManager") as model_manager_cls, patch( - "services.dataset_service.DatasetCollectionBindingService.get_dataset_collection_binding", + "services.knowledge.dataset_service.DatasetCollectionBindingService.get_dataset_collection_binding", return_value=binding, ) as get_binding, - patch.object(DocumentService, "update_document_with_dataset_id", return_value=updated_document), + patch.object(DocumentService, "update_prepared_document", return_value=updated_document), + patch("services.knowledge.dataset_service.document_indexing_update_task.delay"), ): DocumentService.save_document_with_dataset_id( dataset, knowledge_config, account_context, session=sqlite_session @@ -2073,11 +1906,11 @@ class TestDocumentServiceSaveDocumentAdditionalBranches: sqlite_session.commit() with ( - patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()), - patch("services.dataset_service.redis_client") as mock_redis, - patch("services.dataset_service.DocumentIndexingTaskProxy") as document_proxy_cls, - patch("services.dataset_service.time.strftime", return_value="20260101010101"), - patch("services.dataset_service.secrets.randbelow", return_value=23), + patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()), + patch("services.knowledge.dataset_service.redis_client") as mock_redis, + patch("services.knowledge.dataset_service.DocumentIndexingTaskProxy") as document_proxy_cls, + patch("services.knowledge.dataset_service.time.strftime", return_value="20260101010101"), + patch("services.knowledge.dataset_service.secrets.randbelow", return_value=23), ): mock_redis.lock.return_value = _make_lock_context() documents, batch = DocumentService.save_document_with_dataset_id( @@ -2111,11 +1944,11 @@ class TestDocumentServiceSaveDocumentAdditionalBranches: sqlite_session.commit() with ( - patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()), - patch("services.dataset_service.redis_client") as mock_redis, - patch("services.dataset_service.DocumentIndexingTaskProxy"), - patch("services.dataset_service.time.strftime", return_value="20260101010101"), - patch("services.dataset_service.secrets.randbelow", return_value=23), + patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()), + patch("services.knowledge.dataset_service.redis_client") as mock_redis, + patch("services.knowledge.dataset_service.DocumentIndexingTaskProxy"), + patch("services.knowledge.dataset_service.time.strftime", return_value="20260101010101"), + patch("services.knowledge.dataset_service.secrets.randbelow", return_value=23), ): mock_redis.lock.return_value = _make_lock_context() documents, _ = DocumentService.save_document_with_dataset_id( @@ -2141,11 +1974,11 @@ class TestDocumentServiceSaveDocumentAdditionalBranches: sqlite_session.commit() with ( - patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()), - patch("services.dataset_service.redis_client") as mock_redis, - patch("services.dataset_service.DocumentIndexingTaskProxy"), - patch("services.dataset_service.time.strftime", return_value="20260101010101"), - patch("services.dataset_service.secrets.randbelow", return_value=23), + patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()), + patch("services.knowledge.dataset_service.redis_client") as mock_redis, + patch("services.knowledge.dataset_service.DocumentIndexingTaskProxy"), + patch("services.knowledge.dataset_service.time.strftime", return_value="20260101010101"), + patch("services.knowledge.dataset_service.secrets.randbelow", return_value=23), ): mock_redis.lock.return_value = _make_lock_context() documents, _ = DocumentService.save_document_with_dataset_id( @@ -2170,10 +2003,10 @@ class TestDocumentServiceSaveDocumentAdditionalBranches: sqlite_session.commit() with ( - patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()), - patch("services.dataset_service.redis_client") as mock_redis, - patch("services.dataset_service.time.strftime", return_value="20260101010101"), - patch("services.dataset_service.secrets.randbelow", return_value=23), + patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()), + patch("services.knowledge.dataset_service.redis_client") as mock_redis, + patch("services.knowledge.dataset_service.time.strftime", return_value="20260101010101"), + patch("services.knowledge.dataset_service.secrets.randbelow", return_value=23), ): mock_redis.lock.return_value = _make_lock_context() with pytest.raises(FileNotExistsError, match="One or more files not found"): @@ -2200,8 +2033,8 @@ class TestDocumentServiceSaveDocumentAdditionalBranches: ) with ( - patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()), - patch("services.dataset_service.redis_client") as mock_redis, + patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()), + patch("services.knowledge.dataset_service.redis_client") as mock_redis, ): mock_redis.lock.return_value = _make_lock_context() with pytest.raises(ValueError, match="No notion info list found"): @@ -2229,8 +2062,8 @@ class TestDocumentServiceSaveDocumentAdditionalBranches: ) with ( - patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()), - patch("services.dataset_service.redis_client") as mock_redis, + patch("services.knowledge.dataset_service.FeatureService.get_features", return_value=_make_features()), + patch("services.knowledge.dataset_service.redis_client") as mock_redis, ): mock_redis.lock.return_value = _make_lock_context() with pytest.raises(ValueError, match="No website info list found"): diff --git a/api/tests/unit_tests/services/test_dataset_service_lock_not_owned.py b/api/tests/unit_tests/services/test_dataset_service_lock_not_owned.py index 6049e76227a..d7236393504 100644 --- a/api/tests/unit_tests/services/test_dataset_service_lock_not_owned.py +++ b/api/tests/unit_tests/services/test_dataset_service_lock_not_owned.py @@ -1,5 +1,5 @@ import types -from unittest.mock import Mock +from unittest.mock import create_autospec import pytest from redis.exceptions import LockNotOwnedError @@ -10,7 +10,8 @@ from core.rag.index_processor.constant.index_type import IndexStructureType, Ind from models.account import Account, Tenant from models.dataset import Dataset, DatasetProcessRule, Document, DocumentSegment from models.enums import ProcessRuleMode -from services.dataset_service import DocumentService, SegmentService +from services.knowledge.dataset_service import DocumentService, SegmentService +from services.knowledge.segments.application import SegmentMutationService TENANT_ID = "11111111-1111-1111-1111-111111111111" USER_ID = "22222222-2222-2222-2222-222222222222" @@ -36,7 +37,7 @@ def fake_current_user(monkeypatch: pytest.MonkeyPatch): user = Account(name="Test User", email="test@example.com") user.id = USER_ID user._current_tenant = tenant - monkeypatch.setattr("services.dataset_service.current_user", user) + monkeypatch.setattr("services.knowledge.dataset_service.current_user", user) return user @@ -47,7 +48,7 @@ def fake_features(monkeypatch: pytest.MonkeyPatch): documents_upload_quota=types.SimpleNamespace(limit=10_000, size=0), ) monkeypatch.setattr( - "services.dataset_service.FeatureService.get_features", + "services.knowledge.dataset_service.FeatureService.get_features", lambda tenant_id, **_kwargs: features, ) return features @@ -61,7 +62,7 @@ def fake_lock(monkeypatch: pytest.MonkeyPatch): return FakeLock() # DatasetService imports redis_client directly from extensions.ext_redis - monkeypatch.setattr("services.dataset_service.redis_client.lock", _fake_lock) + monkeypatch.setattr("services.knowledge.dataset_service.redis_client.lock", _fake_lock) # --------------------------------------------------------------------------- @@ -115,7 +116,7 @@ def test_save_document_with_dataset_id_ignores_lock_not_owned( account = fake_current_user # Avoid touching real doc_form logic - monkeypatch.setattr("services.dataset_service.DatasetService.check_doc_form", lambda *a, **k: None) + monkeypatch.setattr("services.knowledge.dataset_service.DatasetService.check_doc_form", lambda *a, **k: None) # Act: this would hit the redis lock, whose __enter__ raises LockNotOwnedError. # Our implementation should catch it and still return (documents, batch). @@ -143,60 +144,6 @@ def test_save_document_with_dataset_id_ignores_lock_not_owned( # --------------------------------------------------------------------------- -@pytest.mark.parametrize("sqlite_session", [(Account, Tenant, Dataset, Document, DocumentSegment)], indirect=True) -def test_add_segment_ignores_lock_not_owned( - monkeypatch: pytest.MonkeyPatch, - fake_current_user, - fake_lock, - sqlite_session: Session, -): - # Arrange - dataset = Dataset( - id=DATASET_ID, - tenant_id=TENANT_ID, - name="Test Dataset", - description="", - created_by=USER_ID, - indexing_technique=IndexTechniqueType.ECONOMY, - ) - document = Document( - id=DOCUMENT_ID, - tenant_id=TENANT_ID, - dataset_id=DATASET_ID, - position=1, - data_source_type="upload_file", - data_source_info="{}", - batch="batch-1", - name="Test Document", - created_from="web", - created_by=USER_ID, - word_count=0, - doc_form=IndexStructureType.QA_INDEX, - ) - sqlite_session.add_all([fake_current_user._current_tenant, fake_current_user, dataset, document]) - sqlite_session.commit() - - # Minimal args required by add_segment - args = { - "content": "question text", - "answer": "answer text", - "keywords": ["k1", "k2"], - } - - monkeypatch.setattr("services.dataset_service.VectorService", Mock()) - - # Act - result = SegmentService.create_segment(args=args, document=document, dataset=dataset, session=sqlite_session) - - # Assert - # Under LockNotOwnedError except, add_segment should swallow the error and return None. - assert result is None - assert not sqlite_session.in_transaction() - assert sqlite_session.scalar(select(func.count(DocumentSegment.id))) == 0 - sqlite_session.refresh(document) - assert document.word_count == 0 - - # --------------------------------------------------------------------------- # 3. Multi-segment creation (multi_create_segment) # --------------------------------------------------------------------------- @@ -235,14 +182,17 @@ def test_multi_create_segment_ignores_lock_not_owned( sqlite_session.add_all([fake_current_user._current_tenant, fake_current_user, dataset, document]) sqlite_session.commit() + mutations = create_autospec(SegmentMutationService, instance=True) result = SegmentService.multi_create_segment( segments=[{"content": "question", "answer": "answer", "keywords": ["key"]}], document=document, dataset=dataset, session=sqlite_session, + mutations=mutations, ) assert result is None + mutations.index_segments.assert_not_called() assert not sqlite_session.in_transaction() assert sqlite_session.scalar(select(func.count(DocumentSegment.id))) == 0 sqlite_session.refresh(document) diff --git a/api/tests/unit_tests/services/test_dataset_service_segment.py b/api/tests/unit_tests/services/test_dataset_service_segment.py index c62c9acc680..f848684aee1 100644 --- a/api/tests/unit_tests/services/test_dataset_service_segment.py +++ b/api/tests/unit_tests/services/test_dataset_service_segment.py @@ -5,12 +5,10 @@ from __future__ import annotations from collections.abc import Callable from datetime import datetime from types import SimpleNamespace -from typing import Literal -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock, create_autospec, patch import pytest -from sqlalchemy import event, select -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, sessionmaker from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType from models import Account @@ -18,17 +16,14 @@ from models.account import Tenant from models.dataset import ( ChildChunk, Dataset, - DatasetProcessRule, Document, DocumentSegment, - DocumentSegmentSummary, - SegmentAttachmentBinding, ) -from models.enums import DataSourceType, DocumentCreatedFrom, ProcessRuleMode, SegmentStatus -from services.dataset_ref_service import DatasetRef, DatasetRefService, DocumentRef, SegmentRef -from services.dataset_service import SegmentService -from services.entities.knowledge_entities.knowledge_entities import ChildChunkUpdateArgs, SegmentUpdateArgs -from services.errors.chunk import ChildChunkDeleteIndexError, ChildChunkIndexingError +from models.enums import DataSourceType, DocumentCreatedFrom, SegmentStatus +from repositories.knowledge.segment_repository import SQLAlchemySegmentRepository +from services.knowledge.dataset_service import SegmentService +from services.knowledge.resource_scope import DatasetRef, DocumentRef, SegmentRef +from services.knowledge.segments.application import SegmentIndex, SegmentIndexingState, SegmentMutationService def _account(*, account_id: str = "user-1", tenant_id: str = "tenant-1") -> Account: @@ -146,41 +141,6 @@ def _persist_chain(session: Session) -> tuple[Dataset, Document, DocumentSegment return dataset, document, segment -class TestDatasetRefService: - def test_dataset_ref_is_plain_named_tuple(self) -> None: - assert DatasetRef("tenant-1", "dataset-1") == ("tenant-1", "dataset-1") - - @pytest.mark.parametrize( - ("document_dataset_id", "document_tenant_id"), - [("dataset-2", "tenant-1"), ("dataset-1", "tenant-2")], - ) - def test_create_document_ref_rejects_document_outside_dataset( - self, document_dataset_id: str, document_tenant_id: str - ) -> None: - dataset_ref = DatasetRef("tenant-1", "dataset-1") - document = _document(dataset_id=document_dataset_id, tenant_id=document_tenant_id) - - assert DatasetRefService.create_document_ref(dataset_ref, document) is None - - def test_create_segment_ref_carries_full_parent_chain(self) -> None: - document_ref = DocumentRef(DatasetRef("tenant-1", "dataset-1"), "document-1") - - assert DatasetRefService.create_segment_ref(document_ref, "segment-1") == SegmentRef(document_ref, "segment-1") - - def test_get_document_by_ref_enforces_full_owner_chain(self, sqlite_session: Session) -> None: - owned = _document() - wrong_dataset = _document(document_id="wrong-dataset", dataset_id="dataset-2") - wrong_tenant = _document(document_id="wrong-tenant", tenant_id="tenant-2") - sqlite_session.add_all([owned, wrong_dataset, wrong_tenant]) - sqlite_session.commit() - - owned_ref = DocumentRef(DatasetRef("tenant-1", "dataset-1"), owned.id) - assert DatasetRefService.get_document_by_ref(owned_ref, session=sqlite_session) is owned - for document in (wrong_dataset, wrong_tenant): - spoofed_ref = DocumentRef(DatasetRef("tenant-1", "dataset-1"), document.id) - assert DatasetRefService.get_document_by_ref(spoofed_ref, session=sqlite_session) is None - - class TestSegmentServiceQueries: def test_get_child_chunks_filters_owner_keyword_and_paginates(self, sqlite_session: Session) -> None: _persist_chain(sqlite_session) @@ -194,7 +154,7 @@ class TestSegmentServiceQueries: ) sqlite_session.commit() - with patch("services.dataset_service.current_user", _account()): + with patch("services.knowledge.dataset_service.current_user", _account()): page = SegmentService.get_child_chunks( "segment-1", "document-1", @@ -277,6 +237,7 @@ class TestSegmentServiceQueries: _segment(segment_id="enabled"), _segment(segment_id="disabled", position=2, enabled=False), _segment(segment_id="other", document_id="document-2", position=1), + _segment(segment_id="foreign", tenant_id="tenant-2"), ] ) sqlite_session.commit() @@ -285,6 +246,7 @@ class TestSegmentServiceQueries: "document-1", "dataset-1", sqlite_session, + tenant_id="tenant-1", status=SegmentStatus.COMPLETED, enabled=True, ) @@ -326,262 +288,49 @@ class TestSegmentServiceValidation: ) -class TestSegmentServiceChildChunks: - def test_create_child_chunk_assigns_next_position_and_commits(self, sqlite_session: Session) -> None: - dataset, document, segment = _persist_chain(sqlite_session) - sqlite_session.add(_child(position=1)) - sqlite_session.commit() - - with ( - patch("services.dataset_service.current_user", _account()), - patch("services.dataset_service.VectorService.create_child_chunk_vector"), - ): - child = SegmentService.create_child_chunk("new child", segment, document, dataset, sqlite_session) - - assert child.position == 2 - assert sqlite_session.get(ChildChunk, child.id) is child - - def test_create_child_chunk_rolls_back_on_vector_failure(self, sqlite_session: Session) -> None: - dataset, document, segment = _persist_chain(sqlite_session) - rollback_events: list[str] = [] - event.listen(sqlite_session, "after_rollback", lambda _session: rollback_events.append("rollback")) - - with ( - patch("services.dataset_service.current_user", _account()), - patch( - "services.dataset_service.VectorService.create_child_chunk_vector", - side_effect=RuntimeError("vector failed"), - ), - pytest.raises(ChildChunkIndexingError, match="vector failed"), - ): - SegmentService.create_child_chunk("new child", segment, document, dataset, sqlite_session) - - assert rollback_events == ["rollback"] - assert sqlite_session.scalars(select(ChildChunk)).all() == [] - - def test_update_child_chunks_updates_deletes_and_creates_real_rows(self, sqlite_session: Session) -> None: - dataset, document, segment = _persist_chain(sqlite_session) - keep = _child(child_id="keep", content="old", position=1) - remove = _child(child_id="remove", content="remove", position=2) - sqlite_session.add_all([keep, remove]) - sqlite_session.commit() - - with ( - patch("services.dataset_service.current_user", _account()), - patch("services.dataset_service.VectorService.update_child_chunk_vector"), - ): - chunks = SegmentService.update_child_chunks( - [ - ChildChunkUpdateArgs(id=keep.id, content="updated"), - ChildChunkUpdateArgs(content="created"), - ], - segment, - document, - dataset, - sqlite_session, - ) - - persisted = sqlite_session.scalars(select(ChildChunk).order_by(ChildChunk.position)).all() - assert {chunk.content for chunk in persisted} == {"updated", "created"} - assert sqlite_session.get(ChildChunk, remove.id) is None - assert {chunk.content for chunk in chunks} == {"updated", "created"} - - def test_update_child_chunk_rolls_back_on_vector_failure(self, sqlite_session: Session) -> None: - dataset, document, segment = _persist_chain(sqlite_session) - child = _child() - sqlite_session.add(child) - sqlite_session.commit() - - with ( - patch("services.dataset_service.current_user", _account()), - patch( - "services.dataset_service.VectorService.update_child_chunk_vector", - side_effect=RuntimeError("vector failed"), - ), - pytest.raises(ChildChunkIndexingError, match="vector failed"), - ): - SegmentService.update_child_chunk("changed", child, segment, document, dataset, sqlite_session) - - sqlite_session.refresh(child) - assert child.content == "child content" - - def test_update_child_chunk_persists_and_updates_vector(self, sqlite_session: Session) -> None: - dataset, document, segment = _persist_chain(sqlite_session) - child = _child() - sqlite_session.add(child) - sqlite_session.commit() - - with ( - patch("services.dataset_service.current_user", _account()), - patch("services.dataset_service.VectorService.update_child_chunk_vector") as vector_update, - ): - updated = SegmentService.update_child_chunk( - "changed child", child, segment, document, dataset, sqlite_session - ) - - assert updated is child - assert child.content == "changed child" - assert child.word_count == len("changed child") - assert sqlite_session.get(ChildChunk, child.id) is child - vector_update.assert_called_once_with([], [child], [], dataset, session=sqlite_session) - - def test_update_child_chunks_rolls_back_all_database_changes_on_vector_failure( - self, sqlite_session: Session - ) -> None: - dataset, document, segment = _persist_chain(sqlite_session) - child = _child(content="original") - sqlite_session.add(child) - sqlite_session.commit() - - with ( - patch("services.dataset_service.current_user", _account()), - patch( - "services.dataset_service.VectorService.update_child_chunk_vector", - side_effect=RuntimeError("vector failed"), - ), - pytest.raises(ChildChunkIndexingError, match="vector failed"), - ): - SegmentService.update_child_chunks( - [ChildChunkUpdateArgs(id=child.id, content="changed")], - segment, - document, - dataset, - sqlite_session, - ) - - sqlite_session.refresh(child) - assert child.content == "original" - - def test_delete_child_chunk_commits_after_vector_delete(self, sqlite_session: Session) -> None: - dataset, _, _ = _persist_chain(sqlite_session) - child = _child() - sqlite_session.add(child) - sqlite_session.commit() - - with patch("services.dataset_service.VectorService.delete_child_chunk_vector"): - SegmentService.delete_child_chunk(child, dataset, sqlite_session) - - assert sqlite_session.get(ChildChunk, child.id) is None - - def test_delete_child_chunk_rolls_back_on_vector_failure(self, sqlite_session: Session) -> None: - dataset, _, _ = _persist_chain(sqlite_session) - child = _child() - sqlite_session.add(child) - sqlite_session.commit() - - with ( - patch( - "services.dataset_service.VectorService.delete_child_chunk_vector", - side_effect=RuntimeError("vector failed"), - ), - pytest.raises(ChildChunkDeleteIndexError, match="vector failed"), - ): - SegmentService.delete_child_chunk(child, dataset, sqlite_session) - - assert sqlite_session.get(ChildChunk, child.id) is not None - - class TestSegmentServiceMutations: - def test_create_segment_persists_position_and_updates_document_count(self, sqlite_session: Session) -> None: + @pytest.mark.parametrize( + ("document_dataset_id", "document_tenant_id"), + [("dataset-2", "tenant-1"), ("dataset-1", "tenant-2")], + ) + def test_delete_segment_rejects_document_outside_dataset( + self, document_dataset_id: str, document_tenant_id: str, sqlite_session: Session + ) -> None: dataset = _dataset() - document = _document(word_count=0) - sqlite_session.add_all([dataset, document, _segment(segment_id="existing", content="old")]) - sqlite_session.commit() - embedding_model = SimpleNamespace(get_text_embedding_num_tokens=lambda *, texts: [len(texts) + 2]) + document = _document(dataset_id=document_dataset_id, tenant_id=document_tenant_id) + segment = _segment(dataset_id=document_dataset_id, tenant_id=document_tenant_id) + mutations = create_autospec(SegmentMutationService, instance=True, spec_set=True) - with ( - patch("services.dataset_service.current_user", _account()), - patch("services.dataset_service.ModelManager") as manager_cls, - patch("services.dataset_service.VectorService.create_segments_vector"), - ): - manager_cls.for_tenant.return_value.get_model_instance.return_value = embedding_model - segment = SegmentService.create_segment( - {"content": "new", "attachment_ids": [], "keywords": ["key"]}, - document, - dataset, - sqlite_session, - ) + with pytest.raises(ValueError, match="Document does not belong to the dataset"): + SegmentService.delete_segment(segment, document, dataset, sqlite_session, mutations=mutations) - assert segment is not None - assert segment.position == 2 - assert document.word_count == 3 + mutations.delete_segment.assert_not_called() - def test_create_segment_persists_qa_answer_and_attachment_bindings(self, sqlite_session: Session) -> None: - dataset = _dataset() - document = _document(doc_form=IndexStructureType.QA_INDEX, word_count=0) - sqlite_session.add_all([dataset, document]) - sqlite_session.commit() - embedding_model = MagicMock() - embedding_model.get_text_embedding_num_tokens.return_value = [7] - - with ( - patch("services.dataset_service.current_user", _account()), - patch("services.dataset_service.ModelManager") as manager_cls, - patch("services.dataset_service.VectorService.create_segments_vector"), - ): - manager_cls.for_tenant.return_value.get_model_instance.return_value = embedding_model - segment = SegmentService.create_segment( - { - "content": "question", - "answer": "answer", - "attachment_ids": ["attachment-1", "attachment-2"], - }, - document, - dataset, - sqlite_session, - ) - - assert segment is not None - assert segment.answer == "answer" - assert segment.word_count == len("question") + len("answer") - assert document.word_count == segment.word_count - bindings = sqlite_session.scalars( - select(SegmentAttachmentBinding).where(SegmentAttachmentBinding.segment_id == segment.id) - ).all() - assert {binding.attachment_id for binding in bindings} == {"attachment-1", "attachment-2"} - - def test_create_segment_marks_real_row_error_on_vector_failure(self, sqlite_session: Session) -> None: - dataset = _dataset() - document = _document(word_count=0) - sqlite_session.add_all([dataset, document]) - sqlite_session.commit() - embedding_model = SimpleNamespace(get_text_embedding_num_tokens=lambda *, texts: [len(texts) + 2]) - - with ( - patch("services.dataset_service.current_user", _account()), - patch("services.dataset_service.ModelManager") as manager_cls, - patch( - "services.dataset_service.VectorService.create_segments_vector", - side_effect=RuntimeError("vector failed"), - ), - ): - manager_cls.for_tenant.return_value.get_model_instance.return_value = embedding_model - segment = SegmentService.create_segment( - {"content": "new", "attachment_ids": []}, - document, - dataset, - sqlite_session, - ) - - assert segment is not None - assert segment.status == SegmentStatus.ERROR - assert segment.enabled is False - assert segment.error == "vector failed" - - def test_multi_create_segment_marks_each_real_row_error_on_vector_failure(self, sqlite_session: Session) -> None: + def test_multi_create_segment_marks_each_real_row_error_on_vector_failure( + self, sqlite_session: Session, sqlite_session_factory: sessionmaker[Session] + ) -> None: dataset = _dataset() document = _document(word_count=0) sqlite_session.add_all([dataset, document]) sqlite_session.commit() embedding_model = SimpleNamespace(get_text_embedding_num_tokens=lambda *, texts: [len(texts) + 1]) + index = create_autospec(SegmentIndex, instance=True, spec_set=True) + + def fail(*_args: object, **_kwargs: object) -> None: + assert not sqlite_session.in_transaction() + raise RuntimeError("vector failed") + + index.create_many.side_effect = fail + mutations = SegmentMutationService( + store=SQLAlchemySegmentRepository(session_factory=sqlite_session_factory), + index=index, + text_hash=lambda text: text, + indexing_state=create_autospec(SegmentIndexingState, instance=True, spec_set=True), + ) with ( - patch("services.dataset_service.current_user", _account()), - patch("services.dataset_service.ModelManager") as manager_cls, - patch( - "services.dataset_service.VectorService.create_segments_vector", - side_effect=RuntimeError("vector failed"), - ), + patch("services.knowledge.dataset_service.current_user", _account()), + patch("services.knowledge.dataset_service.ModelManager") as manager_cls, ): manager_cls.for_tenant.return_value.get_model_instance.return_value = embedding_model segments = SegmentService.multi_create_segment( @@ -589,6 +338,7 @@ class TestSegmentServiceMutations: document, dataset, sqlite_session, + mutations=mutations, ) assert segments is not None @@ -603,11 +353,11 @@ class TestSegmentServiceMutations: sqlite_session.commit() embedding_model = MagicMock() embedding_model.get_text_embedding_num_tokens.side_effect = [[11], [13]] + mutations = create_autospec(SegmentMutationService, instance=True, spec_set=True) with ( - patch("services.dataset_service.current_user", _account()), - patch("services.dataset_service.ModelManager") as manager_cls, - patch("services.dataset_service.VectorService.create_segments_vector") as vector_create, + patch("services.knowledge.dataset_service.current_user", _account()), + patch("services.knowledge.dataset_service.ModelManager") as manager_cls, ): manager_cls.for_tenant.return_value.get_model_instance.return_value = embedding_model segments = SegmentService.multi_create_segment( @@ -618,6 +368,7 @@ class TestSegmentServiceMutations: document, dataset, sqlite_session, + mutations=mutations, ) assert segments is not None @@ -632,531 +383,8 @@ class TestSegmentServiceMutations: assert embedding_model.get_text_embedding_num_tokens.call_args_list[1].kwargs == { "texts": ["question-2answer-2"] } - assert vector_create.call_args.args[:2] == ([["key"], None], segments) - - def test_update_segment_disables_and_dispatches_cleanup(self, sqlite_session: Session) -> None: - dataset, document, segment = _persist_chain(sqlite_session) - - with ( - patch("services.dataset_service.current_user", _account()), - patch("services.dataset_service.disable_segment_from_index_task.delay") as cleanup, - ): - updated = SegmentService.update_segment( - SegmentUpdateArgs(enabled=False), - segment, - document, - dataset, - sqlite_session, - ) - - assert updated.enabled is False - cleanup.assert_called_once_with(segment.id) - - def test_update_segment_same_content_persists_keywords(self, sqlite_session: Session) -> None: - dataset, document, segment = _persist_chain(sqlite_session) - - with ( - patch("services.dataset_service.current_user", _account()), - patch("services.dataset_service.VectorService.update_segment_vector") as vector_update, - patch("services.dataset_service.VectorService.update_multimodel_vector"), - ): - updated = SegmentService.update_segment( - SegmentUpdateArgs(content=segment.content, keywords=["new-keyword"]), - segment, - document, - dataset, - sqlite_session, - ) - - assert updated.keywords == ["new-keyword"] - vector_update.assert_called_once() - - def test_update_segment_omitted_attachment_ids_preserves_existing_bindings(self, sqlite_session: Session) -> None: - dataset, document, segment = _persist_chain(sqlite_session) - - with ( - patch("services.dataset_service.current_user", _account()), - patch("services.dataset_service.VectorService.update_multimodel_vector") as multimodel_update, - ): - SegmentService.update_segment( - SegmentUpdateArgs(content=segment.content), - segment, - document, - dataset, - sqlite_session, - ) - - multimodel_update.assert_not_called() - - def test_update_segment_explicit_empty_attachment_ids_clears_bindings(self, sqlite_session: Session) -> None: - dataset, document, segment = _persist_chain(sqlite_session) - - with ( - patch("services.dataset_service.current_user", _account()), - patch("services.dataset_service.VectorService.update_multimodel_vector") as multimodel_update, - ): - SegmentService.update_segment( - SegmentUpdateArgs(content=segment.content, attachment_ids=[]), - segment, - document, - dataset, - sqlite_session, - ) - - multimodel_update.assert_called_once_with(segment, [], dataset, session=sqlite_session) - - def test_update_segment_explicit_attachment_ids_updates_multimodel_vector(self, sqlite_session: Session) -> None: - dataset, document, segment = _persist_chain(sqlite_session) - attachment_ids = ["attachment-1", "attachment-2"] - - with ( - patch("services.dataset_service.current_user", _account()), - patch("services.dataset_service.VectorService.update_multimodel_vector") as multimodel_update, - ): - SegmentService.update_segment( - SegmentUpdateArgs(content=segment.content, attachment_ids=attachment_ids), - segment, - document, - dataset, - sqlite_session, - ) - - multimodel_update.assert_called_once_with(segment, attachment_ids, dataset, session=sqlite_session) - - def test_update_segment_content_change_uses_embedding_tokens_and_updates_document_count( - self, sqlite_session: Session - ) -> None: - dataset, document, segment = _persist_chain(sqlite_session) - original_document_count = document.word_count - original_segment_count = segment.word_count - assert original_document_count is not None - assert original_segment_count is not None - embedding_model = SimpleNamespace(get_text_embedding_num_tokens=lambda *, texts: [len(texts) + 4]) - - with ( - patch("services.dataset_service.current_user", _account()), - patch("services.dataset_service.ModelManager") as manager_cls, - patch("services.dataset_service.VectorService.update_segment_vector"), - patch("services.dataset_service.VectorService.update_multimodel_vector"), - ): - manager_cls.for_tenant.return_value.get_model_instance.return_value = embedding_model - updated = SegmentService.update_segment( - SegmentUpdateArgs(content="changed content"), - segment, - document, - dataset, - sqlite_session, - ) - - assert updated.content == "changed content" - assert updated.tokens == 5 - assert document.word_count == original_document_count + len("changed content") - original_segment_count - - def test_update_segment_same_qa_content_updates_answer_and_word_counts(self, sqlite_session: Session) -> None: - dataset = _dataset() - document = _document(doc_form=IndexStructureType.QA_INDEX, word_count=20) - segment = _segment(content="question") - sqlite_session.add_all([dataset, document, segment]) - sqlite_session.commit() - - with ( - patch("services.dataset_service.current_user", _account()), - patch("services.dataset_service.redis_client.get", return_value=None), - patch("services.dataset_service.VectorService.update_segment_vector") as vector_update, - patch("services.dataset_service.VectorService.update_multimodel_vector") as multimodel_update, - ): - updated = SegmentService.update_segment( - SegmentUpdateArgs(content="question", answer="new answer"), - segment, - document, - dataset, - sqlite_session, - ) - - assert updated.answer == "new answer" - assert updated.word_count == len("question") + len("new answer") - assert document.word_count == 20 + len("new answer") - vector_update.assert_not_called() - multimodel_update.assert_not_called() - - def test_update_segment_changed_qa_content_tokenizes_question_and_answer(self, sqlite_session: Session) -> None: - dataset = _dataset() - document = _document(doc_form=IndexStructureType.QA_INDEX, word_count=10) - segment = _segment(content="old") - sqlite_session.add_all([dataset, document, segment]) - sqlite_session.commit() - embedding_model = MagicMock() - embedding_model.get_text_embedding_num_tokens.return_value = [21] - - with ( - patch("services.dataset_service.current_user", _account()), - patch("services.dataset_service.redis_client.get", return_value=None), - patch("services.dataset_service.ModelManager") as manager_cls, - patch("services.dataset_service.VectorService.update_segment_vector") as vector_update, - patch("services.dataset_service.VectorService.update_multimodel_vector") as multimodel_update, - ): - manager_cls.for_tenant.return_value.get_model_instance.return_value = embedding_model - updated = SegmentService.update_segment( - SegmentUpdateArgs(content="new question", answer="new answer", keywords=["key"]), - segment, - document, - dataset, - sqlite_session, - ) - - embedding_model.get_text_embedding_num_tokens.assert_called_once_with(texts=["new questionnew answer"]) - assert updated.answer == "new answer" - assert updated.tokens == 21 - assert updated.word_count == len("new question") + len("new answer") - assert document.word_count == 10 + updated.word_count - len("old") - vector_update.assert_called_once_with(["key"], segment, dataset, session=sqlite_session) - multimodel_update.assert_not_called() - - @pytest.mark.parametrize("summary", [None, "same summary"]) - def test_update_segment_changed_content_regenerates_existing_summary( - self, sqlite_session: Session, summary: str | None - ) -> None: - dataset = _dataset() - dataset.summary_index_setting = {"enable": True} - document = _document(word_count=10) - segment = _segment(content="old") - summary_record = DocumentSegmentSummary( - dataset_id=dataset.id, - document_id=document.id, - chunk_id=segment.id, - summary_content="same summary", + mutations.index_segments.assert_called_once_with( + DatasetRef("tenant-1", "dataset-1").document("document-1"), + segment_ids=[segment.id for segment in segments], + keywords_list=[["key"], None], ) - sqlite_session.add_all([dataset, document, segment, summary_record]) - sqlite_session.commit() - embedding_model = MagicMock() - embedding_model.get_text_embedding_num_tokens.return_value = [9] - - with ( - patch("services.dataset_service.current_user", _account()), - patch("services.dataset_service.redis_client.get", return_value=None), - patch("services.dataset_service.ModelManager") as manager_cls, - patch("services.dataset_service.VectorService.update_segment_vector"), - patch("services.dataset_service.VectorService.update_multimodel_vector"), - patch( - "services.summary_index_service.SummaryIndexService.generate_and_vectorize_summary" - ) as generate_summary, - patch("services.summary_index_service.SummaryIndexService.update_summary_for_segment") as update_summary, - ): - manager_cls.for_tenant.return_value.get_model_instance.return_value = embedding_model - updated = SegmentService.update_segment( - SegmentUpdateArgs(content="new content", summary=summary), - segment, - document, - dataset, - sqlite_session, - ) - - assert updated.content == "new content" - generate_summary.assert_called_once_with(segment, dataset, {"enable": True}, session=sqlite_session) - update_summary.assert_not_called() - - def test_update_segment_same_parent_content_regenerates_children_and_manual_summary( - self, sqlite_session: Session - ) -> None: - dataset = _dataset() - document = _document(doc_form=IndexStructureType.PARENT_CHILD_INDEX) - segment = _segment(content="same content") - process_rule = DatasetProcessRule( - dataset_id=dataset.id, - mode=ProcessRuleMode.AUTOMATIC, - rules=None, - created_by="user-1", - ) - process_rule.id = "rule-1" - document.dataset_process_rule_id = process_rule.id - summary_record = DocumentSegmentSummary( - dataset_id=dataset.id, - document_id=document.id, - chunk_id=segment.id, - summary_content="old summary", - ) - sqlite_session.add_all([dataset, process_rule, document, segment, summary_record]) - sqlite_session.commit() - embedding_model = object() - - with ( - patch("services.dataset_service.current_user", _account()), - patch("services.dataset_service.redis_client.get", return_value=None), - patch("services.dataset_service.ModelManager") as manager_cls, - patch("services.dataset_service.VectorService") as vector_service, - patch("services.summary_index_service.SummaryIndexService.update_summary_for_segment") as update_summary, - ): - manager_cls.for_tenant.return_value.get_model_instance.return_value = embedding_model - updated = SegmentService.update_segment( - SegmentUpdateArgs( - content="same content", - regenerate_child_chunks=True, - summary="new summary", - ), - segment, - document, - dataset, - sqlite_session, - ) - - assert updated is segment - vector_service.generate_child_chunks.assert_called_once_with( - segment, - document, - dataset, - embedding_model, - process_rule, - True, - session=sqlite_session, - ) - update_summary.assert_called_once_with(segment, dataset, "new summary", session=sqlite_session) - vector_service.update_multimodel_vector.assert_not_called() - - def test_update_segment_changed_parent_content_uses_default_embedding_and_ignores_summary_failure( - self, sqlite_session: Session - ) -> None: - dataset = _dataset() - dataset.embedding_model_provider = None - document = _document(doc_form=IndexStructureType.PARENT_CHILD_INDEX, word_count=10) - segment = _segment(content="old") - process_rule = DatasetProcessRule( - dataset_id=dataset.id, - mode=ProcessRuleMode.AUTOMATIC, - rules=None, - created_by="user-1", - ) - process_rule.id = "rule-1" - document.dataset_process_rule_id = process_rule.id - summary_record = DocumentSegmentSummary( - dataset_id=dataset.id, - document_id=document.id, - chunk_id=segment.id, - summary_content="old summary", - ) - sqlite_session.add_all([dataset, process_rule, document, segment, summary_record]) - sqlite_session.commit() - token_model = MagicMock() - token_model.get_text_embedding_num_tokens.return_value = [9] - default_embedding_model = object() - - with ( - patch("services.dataset_service.current_user", _account()), - patch("services.dataset_service.redis_client.get", return_value=None), - patch("services.dataset_service.ModelManager") as manager_cls, - patch("services.dataset_service.VectorService") as vector_service, - patch( - "services.summary_index_service.SummaryIndexService.update_summary_for_segment", - side_effect=RuntimeError("summary failed"), - ) as update_summary, - ): - manager = manager_cls.for_tenant.return_value - manager.get_model_instance.return_value = token_model - manager.get_default_model_instance.return_value = default_embedding_model - updated = SegmentService.update_segment( - SegmentUpdateArgs( - content="new parent content", - regenerate_child_chunks=True, - summary="new summary", - ), - segment, - document, - dataset, - sqlite_session, - ) - - assert updated.content == "new parent content" - manager.get_default_model_instance.assert_called_once() - vector_service.generate_child_chunks.assert_called_once_with( - segment, - document, - dataset, - default_embedding_model, - process_rule, - True, - session=sqlite_session, - ) - update_summary.assert_called_once_with(segment, dataset, "new summary", session=sqlite_session) - vector_service.update_multimodel_vector.assert_not_called() - - def test_update_segment_parent_regeneration_marks_economy_segment_error(self, sqlite_session: Session) -> None: - dataset = _dataset() - dataset.indexing_technique = IndexTechniqueType.ECONOMY - document = _document(doc_form=IndexStructureType.PARENT_CHILD_INDEX) - segment = _segment(content="same content") - sqlite_session.add_all([dataset, document, segment]) - sqlite_session.commit() - - with ( - patch("services.dataset_service.current_user", _account()), - patch("services.dataset_service.redis_client.get", return_value=None), - patch("services.dataset_service.VectorService.update_multimodel_vector") as multimodel_update, - ): - updated = SegmentService.update_segment( - SegmentUpdateArgs(content="same content", regenerate_child_chunks=True), - segment, - document, - dataset, - sqlite_session, - ) - - assert updated.enabled is False - assert updated.status == SegmentStatus.ERROR - assert updated.error == "The knowledge base index technique is not high quality!" - multimodel_update.assert_not_called() - - def test_update_segment_rejects_disabled_or_indexing_segment(self, sqlite_session: Session) -> None: - dataset, document, segment = _persist_chain(sqlite_session) - segment.enabled = False - sqlite_session.commit() - - with patch("services.dataset_service.current_user", _account()): - with pytest.raises(ValueError, match="Can't update disabled segment"): - SegmentService.update_segment( - SegmentUpdateArgs(content="changed"), segment, document, dataset, sqlite_session - ) - - segment.enabled = True - sqlite_session.commit() - with ( - patch("services.dataset_service.current_user", _account()), - patch("services.dataset_service.redis_client.get", return_value=b"1"), - pytest.raises(ValueError, match="Segment is indexing"), - ): - SegmentService.update_segment( - SegmentUpdateArgs(content="changed"), segment, document, dataset, sqlite_session - ) - - def test_delete_segment_removes_row_and_updates_document_count(self, sqlite_session: Session) -> None: - dataset, document, segment = _persist_chain(sqlite_session) - child = _child() - sqlite_session.add(child) - sqlite_session.commit() - original_count = document.word_count - assert original_count is not None - assert segment.word_count is not None - - with ( - patch("services.dataset_service.redis_client.get", return_value=None), - patch("services.dataset_service.redis_client.setex") as set_cache, - patch("services.dataset_service.delete_segment_from_index_task.delay") as cleanup, - ): - SegmentService.delete_segment(segment, document, dataset, sqlite_session) - - assert sqlite_session.get(DocumentSegment, segment.id) is None - assert document.word_count == original_count - segment.word_count - set_cache.assert_called_once_with(f"segment_{segment.id}_delete_indexing", 600, 1) - cleanup.assert_called_once_with( - [segment.index_node_id], - dataset.id, - document.id, - [segment.id], - [child.index_node_id], - ) - - def test_delete_segment_rejects_when_delete_is_already_in_progress(self, sqlite_session: Session) -> None: - dataset, document, segment = _persist_chain(sqlite_session) - - with ( - patch("services.dataset_service.redis_client.get", return_value=b"1"), - pytest.raises(ValueError, match="Segment is deleting"), - ): - SegmentService.delete_segment(segment, document, dataset, sqlite_session) - - assert sqlite_session.get(DocumentSegment, segment.id) is segment - - def test_delete_segments_scopes_rows_and_clamps_document_count(self, sqlite_session: Session) -> None: - dataset = _dataset() - document = _document(word_count=5) - owned = _segment(segment_id="owned", content="123456789") - wrong_tenant = _segment(segment_id="wrong-tenant", tenant_id="tenant-2", content="foreign") - wrong_dataset = _segment(segment_id="wrong-dataset", dataset_id="dataset-2", content="foreign") - wrong_document = _segment(segment_id="wrong-document", document_id="document-2", content="foreign") - sqlite_session.add_all([dataset, document, owned, wrong_tenant, wrong_dataset, wrong_document]) - sqlite_session.commit() - - with ( - patch("services.dataset_service.current_user", _account()), - patch("services.dataset_service.delete_segment_from_index_task.delay"), - ): - SegmentService.delete_segments( - [owned.id, wrong_tenant.id, wrong_dataset.id, wrong_document.id], - document, - dataset, - sqlite_session, - ) - - assert sqlite_session.get(DocumentSegment, owned.id) is None - assert sqlite_session.get(DocumentSegment, wrong_tenant.id) is wrong_tenant - assert sqlite_session.get(DocumentSegment, wrong_dataset.id) is wrong_dataset - assert sqlite_session.get(DocumentSegment, wrong_document.id) is wrong_document - assert document.word_count == 0 - - @pytest.mark.parametrize("segment_ids", [[], ["missing"]]) - def test_delete_segments_noops_at_empty_selection_boundary( - self, sqlite_session: Session, segment_ids: list[str] - ) -> None: - dataset, document, segment = _persist_chain(sqlite_session) - - with ( - patch("services.dataset_service.current_user", _account()), - patch("services.dataset_service.delete_segment_from_index_task.delay") as cleanup, - ): - SegmentService.delete_segments(segment_ids, document, dataset, sqlite_session) - - assert sqlite_session.get(DocumentSegment, segment.id) is segment - assert document.word_count == 20 - cleanup.assert_not_called() - - @pytest.mark.parametrize(("action", "initial_enabled"), [("enable", False), ("disable", True)]) - def test_update_segments_status_persists_only_owned_rows( - self, - sqlite_session: Session, - action: Literal["enable", "disable"], - initial_enabled: bool, - ) -> None: - dataset = _dataset() - document = _document() - owned = _segment(segment_id="owned", enabled=initial_enabled) - cached = _segment(segment_id="cached", position=2, enabled=initial_enabled) - decoy = _segment(segment_id="decoy", document_id="document-2", enabled=initial_enabled) - sqlite_session.add_all([dataset, document, owned, cached, decoy]) - sqlite_session.commit() - - with ( - patch("services.dataset_service.current_user", _account()), - patch( - "services.dataset_service.redis_client.get", - side_effect=lambda key: b"1" if key == f"segment_{cached.id}_indexing" else None, - ), - patch("services.dataset_service.enable_segments_to_index_task.delay") as enable_task, - patch("services.dataset_service.disable_segments_from_index_task.delay") as disable_task, - ): - SegmentService.update_segments_status( - [owned.id, cached.id, decoy.id], action, dataset, document, sqlite_session - ) - - assert owned.enabled is (action == "enable") - assert cached.enabled is initial_enabled - assert decoy.enabled is initial_enabled - dispatched_task = enable_task if action == "enable" else disable_task - skipped_task = disable_task if action == "enable" else enable_task - dispatched_task.assert_called_once_with([owned.id], dataset.id, document.id) - skipped_task.assert_not_called() - - @pytest.mark.parametrize("action", ["enable", "disable"]) - def test_update_segments_status_noops_at_empty_selection_boundary( - self, sqlite_session: Session, action: Literal["enable", "disable"] - ) -> None: - dataset, document, segment = _persist_chain(sqlite_session) - initial_enabled = segment.enabled - - with ( - patch("services.dataset_service.current_user", _account()), - patch("services.dataset_service.enable_segments_to_index_task.delay") as enable_task, - patch("services.dataset_service.disable_segments_from_index_task.delay") as disable_task, - ): - SegmentService.update_segments_status([], action, dataset, document, sqlite_session) - - assert segment.enabled is initial_enabled - enable_task.assert_not_called() - disable_task.assert_not_called() diff --git a/api/tests/unit_tests/services/test_datasource_provider_service.py b/api/tests/unit_tests/services/test_datasource_provider_service.py index c52e59d8a98..277edf685fa 100644 --- a/api/tests/unit_tests/services/test_datasource_provider_service.py +++ b/api/tests/unit_tests/services/test_datasource_provider_service.py @@ -14,8 +14,10 @@ from models.base import TypeBase from models.model import EndUser from models.oauth import DatasourceOauthParamConfig, DatasourceOauthTenantParamConfig, DatasourceProvider from models.provider_ids import DatasourceProviderID -from services import datasource_provider_service as service_module -from services.datasource_provider_service import DatasourceProviderService, get_current_user +from repositories.credentials.query_repository import CredentialQueryRepository +from repositories.data_source.credential_repository import SQLAlchemyDatasourceCredentialRepository +from services.data_source import provider_service as service_module +from services.data_source.provider_service import DatasourceProviderService, get_current_user # --------------------------------------------------------------------------- # Helpers @@ -65,9 +67,46 @@ def persist(session: Session, *models: TypeBase) -> None: class TestDatasourceProviderService: """Comprehensive tests for DatasourceProviderService targeting >95% coverage.""" + @pytest.mark.parametrize( + ("credential_id", "credentials", "integration_token", "expected"), + [ + ("cred", {"integration_secret": "saved-token"}, "env-token", "saved-token"), + (None, {"integration_secret": "default-token"}, "env-token", "env-token"), + ("cred", {}, "env-token", "env-token"), + ("cred", {}, None, None), + (None, {}, None, None), + ], + ) + def test_stored_notion_access_token_resolves_saved_credential_or_fallback( + self, service, config_overrides, credential_id, credentials, integration_token, expected + ): + config_overrides(NOTION_INTEGRATION_TOKEN=integration_token) + with patch.object(service, "get_datasource_credentials", return_value=credentials) as resolve: + if expected is None: + with pytest.raises(ValueError, match="Must specify `integration_token`"): + service.get_stored_notion_access_token(tenant_id="tenant", credential_id=credential_id) + else: + assert ( + service.get_stored_notion_access_token(tenant_id="tenant", credential_id=credential_id) == expected + ) + + if credential_id is None: + resolve.assert_not_called() + else: + resolve.assert_called_once_with( + tenant_id="tenant", + credential_id=credential_id, + provider="notion_datasource", + plugin_id="langgenius/notion_datasource", + ) + @pytest.fixture - def service(self): - return DatasourceProviderService() + def service(self, sqlite_session: Session): + return DatasourceProviderService( + credentials=SQLAlchemyDatasourceCredentialRepository( + session_factory=sessionmaker(bind=sqlite_session.get_bind(), expire_on_commit=False) + ) + ) @pytest.fixture def sqlite_session(self, sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch) -> Iterator[Session]: @@ -89,10 +128,10 @@ class TestDatasourceProviderService: patch("core.plugin.impl.base._httpx_client.request", side_effect=lambda **kw: httpx.request(**kw)), patch("core.plugin.impl.base._httpx_client.stream", side_effect=lambda **kw: httpx.stream(**kw)), patch("httpx.request") as mock_httpx, - patch("services.datasource_provider_service.encrypter") as mock_enc, - patch("services.datasource_provider_service.redis_client") as mock_redis, - patch("services.datasource_provider_service.generate_incremental_name") as mock_genname, - patch("services.datasource_provider_service.OAuthHandler") as mock_oauth, + patch("services.data_source.provider_service.encrypter") as mock_enc, + patch("services.data_source.provider_service.redis_client") as mock_redis, + patch("services.data_source.provider_service.generate_incremental_name") as mock_genname, + patch("services.data_source.provider_service.OAuthHandler") as mock_oauth, ): mock_enc.encrypt_token.return_value = "enc_tok" mock_enc.decrypt_token.return_value = "dec_tok" @@ -299,7 +338,7 @@ class TestDatasourceProviderService: # ----------------------------------------------------------------------- def test_should_return_empty_dict_when_credential_not_found(self, service, sqlite_session, mock_user): - with patch("services.datasource_provider_service.get_current_user", return_value=mock_user): + with patch("services.data_source.provider_service.get_current_user", return_value=mock_user): assert service.get_datasource_credentials("t1", "prov", "org/plug") == {} def test_should_refresh_oauth_tokens_when_expired(self, service, sqlite_session, mock_user): @@ -307,7 +346,7 @@ class TestDatasourceProviderService: p = make_provider(auth_type="oauth2", expires_at=0, encrypted_credentials={"tok": "x"}) persist(sqlite_session, p) with ( - patch("services.datasource_provider_service.get_current_user", return_value=mock_user), + patch("services.data_source.provider_service.get_current_user", return_value=mock_user), patch.object(service, "get_oauth_client", return_value={"oc": "v"}), patch.object(service, "decrypt_datasource_provider_credentials", return_value={"tok": "plain"}), ): @@ -324,8 +363,8 @@ class TestDatasourceProviderService: ) persist(sqlite_session, p) with ( - patch("services.datasource_provider_service.get_current_user", return_value=mock_user), - patch("services.datasource_provider_service.OAuthHandler") as oauth_handler, + patch("services.data_source.provider_service.get_current_user", return_value=mock_user), + patch("services.data_source.provider_service.OAuthHandler") as oauth_handler, patch.object(service, "get_oauth_client", return_value={"oc": "v"}), patch.object(service, "decrypt_datasource_provider_credentials", return_value={"tok": "plain"}), ): @@ -338,7 +377,7 @@ class TestDatasourceProviderService: p = make_provider(encrypted_credentials={"k": "v"}) persist(sqlite_session, p) with ( - patch("services.datasource_provider_service.get_current_user", return_value=mock_user), + patch("services.data_source.provider_service.get_current_user", return_value=mock_user), patch.object(service, "decrypt_datasource_provider_credentials", return_value={"k": "plain"}), ): result = service.get_datasource_credentials("t1", "prov", "org/plug") @@ -346,28 +385,84 @@ class TestDatasourceProviderService: def test_should_fetch_by_credential_id_when_provided(self, service, sqlite_session, mock_user): """When credential_id is passed, the credential_id filter path (line 113) is taken.""" - p = make_provider(credential_id="cred-id", provider="other-provider", plugin_id="other-plugin") + p = make_provider(credential_id="cred-id") persist(sqlite_session, p) with ( - patch("services.datasource_provider_service.get_current_user", return_value=mock_user), + patch("services.data_source.provider_service.get_current_user", return_value=mock_user), patch.object(service, "decrypt_datasource_provider_credentials", return_value={"k": "v"}), ): result = service.get_datasource_credentials("t1", "prov", "org/plug", credential_id="cred-id") assert result == {"k": "v"} + @pytest.mark.parametrize("concurrent_update", [False, True]) + def test_refresh_closes_read_session_and_preserves_concurrent_credentials( + self, + service: DatasourceProviderService, + sqlite_session: Session, + mock_user: Account, + monkeypatch: pytest.MonkeyPatch, + concurrent_update: bool, + ) -> None: + import time + + from sqlalchemy import update + + row = make_provider(auth_type="oauth2", expires_at=0, encrypted_credentials={"token": "old"}) + persist(sqlite_session, row) + opened: list[Session] = [] + + class ObservedSession(Session): + def __init__(self, **kwargs) -> None: + super().__init__(**kwargs) + opened.append(self) + + def refresh(**_kwargs: object) -> tuple[dict[str, str], int]: + assert opened + assert all(not session.in_transaction() for session in opened) + if concurrent_update: + with Session(sqlite_session.get_bind()) as session, session.begin(): + session.execute( + update(DatasourceProvider) + .where(DatasourceProvider.id == row.id) + .values(encrypted_credentials={"token": "concurrent"}, expires_at=int(time.time()) + 3600) + ) + return {"token": "refreshed"}, int(time.time()) + 3600 + + service = DatasourceProviderService( + credentials=SQLAlchemyDatasourceCredentialRepository( + session_factory=sessionmaker(bind=sqlite_session.get_bind(), class_=ObservedSession) + ) + ) + monkeypatch.setattr(service_module, "get_current_user", lambda: mock_user) + monkeypatch.setattr(service, "_refresh_datasource_credentials", refresh) + monkeypatch.setattr( + service, + "decrypt_datasource_provider_credentials", + lambda **kwargs: dict(kwargs["datasource_provider"].encrypted_credentials), + ) + result = service.get_datasource_credentials("t1", "prov", "org/plug") + assert result == {"token": "concurrent" if concurrent_update else "refreshed"} + assert all(not session.in_transaction() for session in opened) + + def test_selected_credential_must_belong_to_the_requested_provider( + self, service: DatasourceProviderService, sqlite_session: Session + ) -> None: + persist(sqlite_session, make_provider(provider="other-provider", plugin_id="other-plugin")) + assert service.get_datasource_credentials("t1", "prov", "org/plug", credential_id="cred-id") == {} + # ----------------------------------------------------------------------- # get_all_datasource_credentials_by_provider (lines 176-228) # ----------------------------------------------------------------------- def test_should_return_empty_list_when_no_provider_credentials_exist(self, service, sqlite_session, mock_user): - with patch("services.datasource_provider_service.get_current_user", return_value=mock_user): + with patch("services.data_source.provider_service.get_current_user", return_value=mock_user): assert service.get_all_datasource_credentials_by_provider("t1", "prov", "org/plug") == [] def test_should_refresh_and_return_credentials_when_oauth_expired(self, service, sqlite_session, mock_user): p = make_provider(auth_type="oauth2", expires_at=0, encrypted_credentials={"t": "x"}) persist(sqlite_session, p) with ( - patch("services.datasource_provider_service.get_current_user", return_value=mock_user), + patch("services.data_source.provider_service.get_current_user", return_value=mock_user), patch.object(service, "get_oauth_client", return_value={"oc": "v"}), patch.object(service, "decrypt_datasource_provider_credentials", return_value={"t": "plain"}), ): @@ -391,7 +486,7 @@ class TestDatasourceProviderService: ) persist(sqlite_session, failed_provider, working_provider) with ( - patch("services.datasource_provider_service.get_current_user", return_value=mock_user), + patch("services.data_source.provider_service.get_current_user", return_value=mock_user), patch.object( service, "_refresh_datasource_credentials", @@ -410,7 +505,7 @@ class TestDatasourceProviderService: p = make_provider(auth_type="oauth2", encrypted_credentials={"t": "x"}) persist(sqlite_session, p) with ( - patch("services.datasource_provider_service.get_current_user", return_value=mock_user), + patch("services.data_source.provider_service.get_current_user", return_value=mock_user), patch.object(service, "_refresh_datasource_credentials") as refresh_credentials, patch.object(service, "decrypt_datasource_provider_credentials", return_value={"t": "plain"}), ): @@ -514,7 +609,7 @@ class TestDatasourceProviderService: with ( patch.object(service.provider_manager, "fetch_datasource_provider", return_value=pm), patch( - "services.datasource_provider_service.create_provider_encrypter", + "services.data_source.provider_service.create_provider_encrypter", return_value=(MagicMock(), MagicMock()), ), ): @@ -567,7 +662,9 @@ class TestDatasourceProviderService: enabled=True, ), ) - with patch.object(service, "get_oauth_encrypter", return_value=(self._enc, None)): + with patch( + "services.data_source.credential_adapters.create_provider_encrypter", return_value=(self._enc, None) + ): result = service.get_oauth_client("t1", make_id()) assert result == {"k": "dec"} @@ -582,7 +679,7 @@ class TestDatasourceProviderService: ) with ( patch.object(service.provider_manager, "fetch_datasource_provider"), - patch("services.datasource_provider_service.PluginService.is_plugin_verified", return_value=True), + patch("services.data_source.credential_adapters.PluginService.is_plugin_verified", return_value=True), ): result = service.get_oauth_client("t1", make_id()) assert result == {"k": "sys"} @@ -591,7 +688,7 @@ class TestDatasourceProviderService: """Neither tenant nor system credentials → raises ValueError.""" with ( patch.object(service.provider_manager, "fetch_datasource_provider"), - patch("services.datasource_provider_service.PluginService.is_plugin_verified", return_value=False), + patch("services.data_source.credential_adapters.PluginService.is_plugin_verified", return_value=False), ): with pytest.raises(ValueError, match="Please configure oauth client params"): service.get_oauth_client("t1", make_id()) @@ -734,13 +831,13 @@ class TestDatasourceProviderService: sqlite_session, make_provider(name="clash", provider="provider", plugin_id="org/plugin"), ) - with patch("services.datasource_provider_service.get_current_user", return_value=mock_user): + with patch("services.data_source.provider_service.get_current_user", return_value=mock_user): with pytest.raises(ValueError, match="already exists"): service.add_datasource_api_key_provider("clash", "t1", make_id(), {"sk": "v"}) def test_should_raise_value_error_when_credentials_validation_fails(self, service, sqlite_session, mock_user): with ( - patch("services.datasource_provider_service.get_current_user", return_value=mock_user), + patch("services.data_source.provider_service.get_current_user", return_value=mock_user), patch.object(service.provider_manager, "validate_provider_credentials", side_effect=Exception("bad cred")), patch.object(service, "extract_secret_variables", return_value=[]), ): @@ -749,7 +846,7 @@ class TestDatasourceProviderService: def test_should_add_api_key_provider_and_commit_when_valid(self, service, sqlite_session, mock_user): with ( - patch("services.datasource_provider_service.get_current_user", return_value=mock_user), + patch("services.data_source.provider_service.get_current_user", return_value=mock_user), patch.object(service.provider_manager, "validate_provider_credentials"), patch.object(service, "extract_secret_variables", return_value=["sk"]), ): @@ -761,7 +858,7 @@ class TestDatasourceProviderService: def test_should_acquire_redis_lock_when_adding_api_key_provider(self, service, sqlite_session, mock_user): with ( - patch("services.datasource_provider_service.get_current_user", return_value=mock_user), + patch("services.data_source.provider_service.get_current_user", return_value=mock_user), patch.object(service.provider_manager, "validate_provider_credentials"), patch.object(service, "extract_secret_variables", return_value=[]), ): @@ -805,13 +902,30 @@ class TestDatasourceProviderService: # ----------------------------------------------------------------------- def test_should_return_empty_list_when_no_credentials_stored(self, service, sqlite_session): - assert service.list_datasource_credentials("t1", "prov", "org/plug", session=sqlite_session) == [] + assert ( + service.list_datasource_credentials( + "t1", + "prov", + "org/plug", + credential_query=CredentialQueryRepository( + session_factory=sessionmaker(bind=sqlite_session.get_bind(), expire_on_commit=False) + ), + ) + == [] + ) def test_should_return_masked_credentials_list_when_credentials_exist(self, service, sqlite_session): p = make_provider(encrypted_credentials={"sk": "v"}) persist(sqlite_session, p) with patch.object(service, "extract_secret_variables", return_value=["sk"]): - result = service.list_datasource_credentials("t1", "prov", "org/plug", session=sqlite_session) + result = service.list_datasource_credentials( + "t1", + "prov", + "org/plug", + credential_query=CredentialQueryRepository( + session_factory=sessionmaker(bind=sqlite_session.get_bind(), expire_on_commit=False) + ), + ) assert len(result) == 1 assert result[0]["credential"] == {"sk": "obf"} @@ -820,7 +934,7 @@ class TestDatasourceProviderService: # ----------------------------------------------------------------------- def test_should_aggregate_credentials_for_non_hardcoded_plugin(self, service, sqlite_session): - with patch("services.datasource_provider_service.PluginDatasourceManager") as mock_mgr: + with patch("services.data_source.provider_service.PluginDatasourceManager") as mock_mgr: ds = MagicMock() ds.provider = "prov" ds.plugin_id = "org/plug" @@ -828,12 +942,18 @@ class TestDatasourceProviderService: mock_mgr.return_value.fetch_installed_datasource_providers.return_value = [ds] cred = {"credential": {"k": "v"}, "is_default": True} with patch.object(service, "list_datasource_credentials", return_value=[cred]): - results = service.get_all_datasource_credentials("t1", session=sqlite_session) + results = service.get_all_datasource_credentials( + "t1", + session=sqlite_session, + credential_query=CredentialQueryRepository( + session_factory=sessionmaker(bind=sqlite_session.get_bind(), expire_on_commit=False) + ), + ) assert len(results) == 1 def test_should_include_oauth_schema_for_hardcoded_plugin_ids(self, service, sqlite_session): """Lines 819-871: get_all_datasource_credentials covers hardcoded langgenius plugin IDs.""" - with patch("services.datasource_provider_service.PluginDatasourceManager") as mock_mgr: + with patch("services.data_source.provider_service.PluginDatasourceManager") as mock_mgr: ds = MagicMock() ds.plugin_id = "langgenius/firecrawl_datasource" ds.provider = "firecrawl" @@ -853,7 +973,13 @@ class TestDatasourceProviderService: patch.object(service, "is_tenant_oauth_params_enabled", return_value=False), patch.object(service, "is_system_oauth_params_exist", return_value=False), ): - results = service.get_all_datasource_credentials("t1", session=sqlite_session) + results = service.get_all_datasource_credentials( + "t1", + session=sqlite_session, + credential_query=CredentialQueryRepository( + session_factory=sessionmaker(bind=sqlite_session.get_bind(), expire_on_commit=False) + ), + ) assert len(results) == 1 assert results[0]["oauth_schema"] is not None @@ -877,7 +1003,7 @@ class TestDatasourceProviderService: # ----------------------------------------------------------------------- def test_should_raise_value_error_when_credential_not_found_on_update(self, service, sqlite_session, mock_user): - with patch("services.datasource_provider_service.get_current_user", return_value=mock_user): + with patch("services.data_source.provider_service.get_current_user", return_value=mock_user): with pytest.raises(ValueError, match="not found"): service.update_datasource_credentials("t1", "id", "prov", "org/plug", {}, "name") @@ -885,7 +1011,7 @@ class TestDatasourceProviderService: p = make_provider(credential_id="id", name="old_name", encrypted_credentials={"sk": "e"}) conflict = make_provider(credential_id="conflict-id", name="new_name") persist(sqlite_session, p, conflict) - with patch("services.datasource_provider_service.get_current_user", return_value=mock_user): + with patch("services.data_source.provider_service.get_current_user", return_value=mock_user): with pytest.raises(ValueError, match="already exists"): service.update_datasource_credentials("t1", "id", "prov", "org/plug", {}, "new_name") @@ -895,7 +1021,7 @@ class TestDatasourceProviderService: p = make_provider(credential_id="id", name="old_name", encrypted_credentials={"sk": "e"}) persist(sqlite_session, p) with ( - patch("services.datasource_provider_service.get_current_user", return_value=mock_user), + patch("services.data_source.provider_service.get_current_user", return_value=mock_user), patch.object(service, "extract_secret_variables", return_value=["sk"]), patch.object(service.provider_manager, "validate_provider_credentials", side_effect=Exception("bad")), ): @@ -907,7 +1033,7 @@ class TestDatasourceProviderService: p = make_provider(credential_id="id", name="old_name", encrypted_credentials={"sk": "old_enc"}) persist(sqlite_session, p) with ( - patch("services.datasource_provider_service.get_current_user", return_value=mock_user), + patch("services.data_source.provider_service.get_current_user", return_value=mock_user), patch.object(service, "extract_secret_variables", return_value=["sk"]), patch.object(service.provider_manager, "validate_provider_credentials"), ): diff --git a/api/tests/unit_tests/services/test_external_dataset_service.py b/api/tests/unit_tests/services/test_external_dataset_service.py index c6a0032c227..30fdfeb50b5 100644 --- a/api/tests/unit_tests/services/test_external_dataset_service.py +++ b/api/tests/unit_tests/services/test_external_dataset_service.py @@ -25,7 +25,7 @@ from services.entities.external_knowledge_entities.external_knowledge_entities i ) from services.errors.dataset import DatasetNameDuplicateError from services.errors.knowledge_retrieval import ExternalKnowledgeRetrievalError -from services.external_knowledge_service import ExternalDatasetService +from services.knowledge.external.service import ExternalDatasetService class ExternalDatasetServiceTestDataFactory: @@ -415,7 +415,7 @@ class TestExternalDatasetServiceValidateAPIList: class TestExternalDatasetServiceCreateAPI: """Test create_external_knowledge_api operations.""" - @patch("services.external_knowledge_service.ExternalDatasetService.check_endpoint_and_api_key") + @patch("services.knowledge.external.service.ExternalDatasetService.check_endpoint_and_api_key") def test_create_external_knowledge_api_success_full( self, mock_check, factory: ExternalDatasetServiceTestDataFactory, sqlite_session: Session ): @@ -442,7 +442,7 @@ class TestExternalDatasetServiceCreateAPI: persisted_api = sqlite_session.get(ExternalKnowledgeApis, result.id) assert persisted_api is result - @patch("services.external_knowledge_service.ExternalDatasetService.check_endpoint_and_api_key") + @patch("services.knowledge.external.service.ExternalDatasetService.check_endpoint_and_api_key") def test_create_external_knowledge_api_minimal_fields( self, mock_check, factory: ExternalDatasetServiceTestDataFactory, sqlite_session: Session ): @@ -485,7 +485,7 @@ class TestExternalDatasetServiceCreateAPI: with pytest.raises(ValueError, match="settings is required"): ExternalDatasetService.create_external_knowledge_api("tenant-123", "user-123", args, session=sqlite_session) - @patch("services.external_knowledge_service.ExternalDatasetService.check_endpoint_and_api_key") + @patch("services.knowledge.external.service.ExternalDatasetService.check_endpoint_and_api_key") def test_create_external_knowledge_api_settings_json_serialization( self, mock_check, factory: ExternalDatasetServiceTestDataFactory, sqlite_session: Session ): @@ -508,7 +508,7 @@ class TestExternalDatasetServiceCreateAPI: parsed_settings = json.loads(result.settings) assert parsed_settings == settings - @patch("services.external_knowledge_service.ExternalDatasetService.check_endpoint_and_api_key") + @patch("services.knowledge.external.service.ExternalDatasetService.check_endpoint_and_api_key") def test_create_external_knowledge_api_unicode_handling( self, mock_check, factory: ExternalDatasetServiceTestDataFactory, sqlite_session: Session ): @@ -529,7 +529,7 @@ class TestExternalDatasetServiceCreateAPI: assert result.name == "测试API" assert result.description == "テストの説明" - @patch("services.external_knowledge_service.ExternalDatasetService.check_endpoint_and_api_key") + @patch("services.knowledge.external.service.ExternalDatasetService.check_endpoint_and_api_key") def test_create_external_knowledge_api_long_description( self, mock_check, factory: ExternalDatasetServiceTestDataFactory, sqlite_session: Session ): @@ -555,7 +555,7 @@ class TestExternalDatasetServiceCreateAPI: class TestExternalDatasetServiceCheckEndpoint: """Test check_endpoint_and_api_key operations - extensive coverage.""" - @patch("services.external_knowledge_service.ssrf_proxy") + @patch("services.knowledge.external.service.ssrf_proxy") def test_check_endpoint_success_https(self, mock_proxy, factory: ExternalDatasetServiceTestDataFactory): """Test successful validation with HTTPS endpoint.""" # Arrange @@ -569,7 +569,7 @@ class TestExternalDatasetServiceCheckEndpoint: ExternalDatasetService.check_endpoint_and_api_key(settings) mock_proxy.post.assert_called_once() - @patch("services.external_knowledge_service.ssrf_proxy") + @patch("services.knowledge.external.service.ssrf_proxy") def test_check_endpoint_success_http(self, mock_proxy, factory: ExternalDatasetServiceTestDataFactory): """Test successful validation with HTTP endpoint.""" # Arrange @@ -582,7 +582,7 @@ class TestExternalDatasetServiceCheckEndpoint: # Act & Assert - should not raise ExternalDatasetService.check_endpoint_and_api_key(settings) - @patch("services.external_knowledge_service.ssrf_proxy") + @patch("services.knowledge.external.service.ssrf_proxy") def test_check_endpoint_sends_json_body(self, mock_proxy, factory: ExternalDatasetServiceTestDataFactory): """Regression for #39402: the validation probe must POST a JSON body matching the External Knowledge API retrieval contract, not a body-less request that providers @@ -688,7 +688,7 @@ class TestExternalDatasetServiceCheckEndpoint: with pytest.raises(ValueError, match="invalid endpoint"): ExternalDatasetService.check_endpoint_and_api_key(settings) - @patch("services.external_knowledge_service.ssrf_proxy") + @patch("services.knowledge.external.service.ssrf_proxy") def test_check_endpoint_connection_timeout(self, mock_proxy, factory: ExternalDatasetServiceTestDataFactory): """Test validation fails on connection timeout.""" # Arrange @@ -699,7 +699,7 @@ class TestExternalDatasetServiceCheckEndpoint: with pytest.raises(ValueError, match="failed to connect to the endpoint"): ExternalDatasetService.check_endpoint_and_api_key(settings) - @patch("services.external_knowledge_service.ssrf_proxy") + @patch("services.knowledge.external.service.ssrf_proxy") def test_check_endpoint_network_error(self, mock_proxy, factory: ExternalDatasetServiceTestDataFactory): """Test validation fails on network error.""" # Arrange @@ -710,7 +710,7 @@ class TestExternalDatasetServiceCheckEndpoint: with pytest.raises(ValueError, match="failed to connect to the endpoint"): ExternalDatasetService.check_endpoint_and_api_key(settings) - @patch("services.external_knowledge_service.ssrf_proxy") + @patch("services.knowledge.external.service.ssrf_proxy") def test_check_endpoint_502_bad_gateway(self, mock_proxy, factory: ExternalDatasetServiceTestDataFactory): """Test validation fails with 502 Bad Gateway.""" # Arrange @@ -724,7 +724,7 @@ class TestExternalDatasetServiceCheckEndpoint: with pytest.raises(ValueError, match="Bad Gateway.*failed to connect"): ExternalDatasetService.check_endpoint_and_api_key(settings) - @patch("services.external_knowledge_service.ssrf_proxy") + @patch("services.knowledge.external.service.ssrf_proxy") def test_check_endpoint_404_not_found(self, mock_proxy, factory: ExternalDatasetServiceTestDataFactory): """Test validation fails with 404 Not Found.""" # Arrange @@ -738,7 +738,7 @@ class TestExternalDatasetServiceCheckEndpoint: with pytest.raises(ValueError, match="Not Found.*failed to connect"): ExternalDatasetService.check_endpoint_and_api_key(settings) - @patch("services.external_knowledge_service.ssrf_proxy") + @patch("services.knowledge.external.service.ssrf_proxy") def test_check_endpoint_403_forbidden(self, mock_proxy, factory: ExternalDatasetServiceTestDataFactory): """Test validation fails with 403 Forbidden (auth failure).""" # Arrange @@ -752,7 +752,7 @@ class TestExternalDatasetServiceCheckEndpoint: with pytest.raises(ValueError, match="Forbidden.*Authorization failed"): ExternalDatasetService.check_endpoint_and_api_key(settings) - @patch("services.external_knowledge_service.ssrf_proxy") + @patch("services.knowledge.external.service.ssrf_proxy") def test_check_endpoint_403_message_does_not_echo_api_key( self, mock_proxy, factory: ExternalDatasetServiceTestDataFactory ): @@ -789,7 +789,7 @@ class TestExternalDatasetServiceCheckEndpoint: assert "sk-abcdef" not in message assert "CDEF" not in message - @patch("services.external_knowledge_service.ssrf_proxy") + @patch("services.knowledge.external.service.ssrf_proxy") def test_check_endpoint_other_4xx_codes_pass(self, mock_proxy, factory: ExternalDatasetServiceTestDataFactory): """Test that other 4xx codes don't raise exceptions.""" # Arrange @@ -803,7 +803,7 @@ class TestExternalDatasetServiceCheckEndpoint: # Act & Assert - should not raise ExternalDatasetService.check_endpoint_and_api_key(settings) - @patch("services.external_knowledge_service.ssrf_proxy") + @patch("services.knowledge.external.service.ssrf_proxy") def test_check_endpoint_5xx_codes_except_502_pass(self, mock_proxy, factory: ExternalDatasetServiceTestDataFactory): """Test that 5xx codes except 502 don't raise exceptions.""" # Arrange @@ -817,7 +817,7 @@ class TestExternalDatasetServiceCheckEndpoint: # Act & Assert - should not raise ExternalDatasetService.check_endpoint_and_api_key(settings) - @patch("services.external_knowledge_service.ssrf_proxy") + @patch("services.knowledge.external.service.ssrf_proxy") def test_check_endpoint_with_port_number(self, mock_proxy, factory: ExternalDatasetServiceTestDataFactory): """Test validation with endpoint including port number.""" # Arrange @@ -830,7 +830,7 @@ class TestExternalDatasetServiceCheckEndpoint: # Act & Assert - should not raise ExternalDatasetService.check_endpoint_and_api_key(settings) - @patch("services.external_knowledge_service.ssrf_proxy") + @patch("services.knowledge.external.service.ssrf_proxy") def test_check_endpoint_with_path(self, mock_proxy, factory: ExternalDatasetServiceTestDataFactory): """Test validation with endpoint including path.""" # Arrange @@ -846,7 +846,7 @@ class TestExternalDatasetServiceCheckEndpoint: call_args = mock_proxy.post.call_args assert "/retrieval" in call_args[0][0] - @patch("services.external_knowledge_service.ssrf_proxy") + @patch("services.knowledge.external.service.ssrf_proxy") def test_check_endpoint_authorization_header_format( self, mock_proxy, factory: ExternalDatasetServiceTestDataFactory ): @@ -900,7 +900,7 @@ class TestExternalDatasetServiceGetAPI: class TestExternalDatasetServiceUpdateAPI: """Test update_external_knowledge_api operations.""" - @patch("services.external_knowledge_service.naive_utc_now") + @patch("services.knowledge.external.service.naive_utc_now") def test_update_external_knowledge_api_success_all_fields( self, mock_now, factory: ExternalDatasetServiceTestDataFactory, sqlite_session: Session ): @@ -1271,7 +1271,7 @@ class TestExternalDatasetServiceDocumentValidate: class TestExternalDatasetServiceProcessAPI: """Test process_external_api operations - comprehensive HTTP method coverage.""" - @patch("services.external_knowledge_service.ssrf_proxy") + @patch("services.knowledge.external.service.ssrf_proxy") def test_process_external_api_get_request(self, mock_proxy, factory: ExternalDatasetServiceTestDataFactory): """Test processing GET request.""" # Arrange @@ -1287,7 +1287,7 @@ class TestExternalDatasetServiceProcessAPI: assert result == mock_response mock_proxy.get.assert_called_once() - @patch("services.external_knowledge_service.ssrf_proxy") + @patch("services.knowledge.external.service.ssrf_proxy") def test_process_external_api_post_request_with_data( self, mock_proxy, factory: ExternalDatasetServiceTestDataFactory ): @@ -1307,7 +1307,7 @@ class TestExternalDatasetServiceProcessAPI: call_kwargs = mock_proxy.post.call_args.kwargs assert "data" in call_kwargs - @patch("services.external_knowledge_service.ssrf_proxy") + @patch("services.knowledge.external.service.ssrf_proxy") def test_process_external_api_put_request(self, mock_proxy, factory: ExternalDatasetServiceTestDataFactory): """Test processing PUT request.""" # Arrange @@ -1323,7 +1323,7 @@ class TestExternalDatasetServiceProcessAPI: assert result == mock_response mock_proxy.put.assert_called_once() - @patch("services.external_knowledge_service.ssrf_proxy") + @patch("services.knowledge.external.service.ssrf_proxy") def test_process_external_api_delete_request(self, mock_proxy, factory: ExternalDatasetServiceTestDataFactory): """Test processing DELETE request.""" # Arrange @@ -1339,7 +1339,7 @@ class TestExternalDatasetServiceProcessAPI: assert result == mock_response mock_proxy.delete.assert_called_once() - @patch("services.external_knowledge_service.ssrf_proxy") + @patch("services.knowledge.external.service.ssrf_proxy") def test_process_external_api_patch_request(self, mock_proxy, factory: ExternalDatasetServiceTestDataFactory): """Test processing PATCH request.""" # Arrange @@ -1355,7 +1355,7 @@ class TestExternalDatasetServiceProcessAPI: assert result == mock_response mock_proxy.patch.assert_called_once() - @patch("services.external_knowledge_service.ssrf_proxy") + @patch("services.knowledge.external.service.ssrf_proxy") def test_process_external_api_head_request(self, mock_proxy, factory: ExternalDatasetServiceTestDataFactory): """Test processing HEAD request.""" # Arrange @@ -1380,7 +1380,7 @@ class TestExternalDatasetServiceProcessAPI: with pytest.raises(Exception, match="Invalid http method"): ExternalDatasetService.process_external_api(settings, None) - @patch("services.external_knowledge_service.ssrf_proxy") + @patch("services.knowledge.external.service.ssrf_proxy") def test_process_external_api_with_files(self, mock_proxy, factory: ExternalDatasetServiceTestDataFactory): """Test processing request with file uploads.""" # Arrange @@ -1399,7 +1399,7 @@ class TestExternalDatasetServiceProcessAPI: assert "files" in call_kwargs assert call_kwargs["files"] == files - @patch("services.external_knowledge_service.ssrf_proxy") + @patch("services.knowledge.external.service.ssrf_proxy") def test_process_external_api_follow_redirects(self, mock_proxy, factory: ExternalDatasetServiceTestDataFactory): """Test that follow_redirects is enabled.""" # Arrange @@ -1663,7 +1663,7 @@ class TestExternalDatasetServiceCreateDataset: class TestExternalDatasetServiceFetchRetrieval: """Test fetch_external_knowledge_retrieval operations.""" - @patch("services.external_knowledge_service.ExternalDatasetService.process_external_api") + @patch("services.knowledge.external.service.ExternalDatasetService.process_external_api") def test_fetch_external_knowledge_retrieval_success_with_results( self, mock_process, factory: ExternalDatasetServiceTestDataFactory, sqlite_session: Session ): @@ -1722,7 +1722,7 @@ class TestExternalDatasetServiceFetchRetrieval: "tenant-123", "dataset-123", "query", {}, session=sqlite_session ) - @patch("services.external_knowledge_service.ExternalDatasetService.process_external_api") + @patch("services.knowledge.external.service.ExternalDatasetService.process_external_api") def test_fetch_external_knowledge_retrieval_empty_results( self, mock_process, factory: ExternalDatasetServiceTestDataFactory, sqlite_session: Session ): @@ -1743,7 +1743,7 @@ class TestExternalDatasetServiceFetchRetrieval: # Assert assert len(result) == 0 - @patch("services.external_knowledge_service.ExternalDatasetService.process_external_api") + @patch("services.knowledge.external.service.ExternalDatasetService.process_external_api") def test_fetch_external_knowledge_retrieval_with_score_threshold( self, mock_process, factory: ExternalDatasetServiceTestDataFactory, sqlite_session: Session ): @@ -1777,7 +1777,7 @@ class TestExternalDatasetServiceFetchRetrieval: call_args = mock_process.call_args[0][0] assert call_args.params["retrieval_setting"]["score_threshold"] == 0.75 - @patch("services.external_knowledge_service.ExternalDatasetService.process_external_api") + @patch("services.knowledge.external.service.ExternalDatasetService.process_external_api") def test_fetch_external_knowledge_retrieval_non_200_status_raises_exception( self, mock_process, factory: ExternalDatasetServiceTestDataFactory, sqlite_session: Session ): @@ -1809,7 +1809,7 @@ class TestExternalDatasetServiceFetchRetrieval: (503, "Service Unavailable: Maintenance mode"), ], ) - @patch("services.external_knowledge_service.ExternalDatasetService.process_external_api") + @patch("services.knowledge.external.service.ExternalDatasetService.process_external_api") def test_fetch_external_knowledge_retrieval_various_error_status_codes( self, mock_process, @@ -1836,7 +1836,7 @@ class TestExternalDatasetServiceFetchRetrieval: tenant_id, dataset_id, "query", {"top_k": 5}, session=sqlite_session ) - @patch("services.external_knowledge_service.ExternalDatasetService.process_external_api") + @patch("services.knowledge.external.service.ExternalDatasetService.process_external_api") def test_fetch_external_knowledge_retrieval_empty_response_text( self, mock_process, factory: ExternalDatasetServiceTestDataFactory, sqlite_session: Session ): @@ -1855,7 +1855,7 @@ class TestExternalDatasetServiceFetchRetrieval: "tenant-123", "dataset-123", "query", {"top_k": 5}, session=sqlite_session ) - @patch("services.external_knowledge_service.ExternalDatasetService.process_external_api") + @patch("services.knowledge.external.service.ExternalDatasetService.process_external_api") def test_fetch_external_knowledge_retrieval_invalid_json_response( self, mock_process, factory, sqlite_session: Session ): @@ -1872,7 +1872,7 @@ class TestExternalDatasetServiceFetchRetrieval: "tenant-123", "dataset-123", "query", {"top_k": 5}, session=sqlite_session ) - @patch("services.external_knowledge_service.ExternalDatasetService.process_external_api") + @patch("services.knowledge.external.service.ExternalDatasetService.process_external_api") def test_fetch_external_knowledge_retrieval_invalid_success_payload_shape( self, mock_process, factory, sqlite_session: Session ): @@ -1889,7 +1889,7 @@ class TestExternalDatasetServiceFetchRetrieval: "tenant-123", "dataset-123", "query", {"top_k": 5}, session=sqlite_session ) - @patch("services.external_knowledge_service.ExternalDatasetService.process_external_api") + @patch("services.knowledge.external.service.ExternalDatasetService.process_external_api") def test_fetch_external_knowledge_retrieval_invalid_records_shape( self, mock_process, factory, sqlite_session: Session ): @@ -1906,7 +1906,7 @@ class TestExternalDatasetServiceFetchRetrieval: "tenant-123", "dataset-123", "query", {"top_k": 5}, session=sqlite_session ) - @patch("services.external_knowledge_service.ExternalDatasetService.process_external_api") + @patch("services.knowledge.external.service.ExternalDatasetService.process_external_api") def test_fetch_external_knowledge_retrieval_wraps_transport_errors( self, mock_process, factory, sqlite_session: Session ): diff --git a/api/tests/unit_tests/services/test_metadata_bug_complete.py b/api/tests/unit_tests/services/test_metadata_bug_complete.py index 6b22281fb48..7c5f978642d 100644 --- a/api/tests/unit_tests/services/test_metadata_bug_complete.py +++ b/api/tests/unit_tests/services/test_metadata_bug_complete.py @@ -1,13 +1,25 @@ from pathlib import Path from typing import cast -from unittest.mock import Mock +from unittest.mock import Mock, create_autospec import pytest -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, sessionmaker from models import Account, Tenant -from services.entities.knowledge_entities.knowledge_entities import MetadataArgs -from services.metadata_service import MetadataService +from repositories.knowledge.metadata_repository import SQLAlchemyMetadataRepository +from services.knowledge.dataset_access import DatasetAccess +from services.knowledge.entities.knowledge_entities import MetadataArgs +from services.knowledge.metadata.application import MetadataService +from services.knowledge.resource_scope import DatasetRef + + +def _metadata_service(session: Session) -> MetadataService: + return MetadataService( + store=SQLAlchemyMetadataRepository( + session_factory=sessionmaker(bind=session.get_bind(), expire_on_commit=False) + ), + dataset_access=create_autospec(DatasetAccess, instance=True), + ) def _make_account(account_id: str = "user-456", tenant_id: str = "tenant-123") -> Account: @@ -50,15 +62,17 @@ class TestMetadataBugCompleteValidation: account = _make_account() # Should crash with TypeError with pytest.raises(TypeError, match="object of type 'NoneType' has no len"): - MetadataService.create_metadata( - "dataset-123", mock_metadata_args, account, "tenant-123", session=sqlite_session + _metadata_service(sqlite_session).create_metadata( + DatasetRef("tenant-123", "dataset-123"), mock_metadata_args, actor_id=account.id ) # Test update method as well account = _make_account() none_name = cast(str, None) with pytest.raises(TypeError, match="object of type 'NoneType' has no len"): - MetadataService.update_metadata_name(Mock(), "metadata-456", none_name, account, session=sqlite_session) + _metadata_service(sqlite_session).update_metadata_name( + DatasetRef(Mock().tenant_id, Mock().id), "metadata-456", none_name, actor_id=account.id + ) assert not sqlite_session.in_transaction() def test_3_database_constraints_verification(self) -> None: @@ -103,8 +117,8 @@ class TestMetadataBugCompleteValidation: account = _make_account() with pytest.raises(TypeError, match="object of type 'NoneType' has no len"): - MetadataService.create_metadata( - "dataset-123", mock_metadata_args, account, "tenant-123", session=sqlite_session + _metadata_service(sqlite_session).create_metadata( + DatasetRef("tenant-123", "dataset-123"), mock_metadata_args, actor_id=account.id ) assert not sqlite_session.in_transaction() diff --git a/api/tests/unit_tests/services/test_metadata_nullable_bug.py b/api/tests/unit_tests/services/test_metadata_nullable_bug.py index 4285c6f40d9..0347cee0263 100644 --- a/api/tests/unit_tests/services/test_metadata_nullable_bug.py +++ b/api/tests/unit_tests/services/test_metadata_nullable_bug.py @@ -1,12 +1,25 @@ from typing import cast -from unittest.mock import Mock +from unittest.mock import Mock, create_autospec import pytest -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, sessionmaker from models import Account, Tenant -from services.entities.knowledge_entities.knowledge_entities import MetadataArgs -from services.metadata_service import MetadataService +from repositories.knowledge.metadata_repository import SQLAlchemyMetadataRepository +from services.knowledge.dataset_access import DatasetAccess +from services.knowledge.entities.knowledge_entities import MetadataArgs +from services.knowledge.metadata.application import MetadataService +from services.knowledge.resource_scope import DatasetRef + + +def _metadata_service(session: Session) -> MetadataService: + return MetadataService( + store=SQLAlchemyMetadataRepository( + session_factory=sessionmaker(bind=session.get_bind(), expire_on_commit=False) + ), + dataset_access=create_autospec(DatasetAccess, instance=True), + ) + pytestmark = [ pytest.mark.usefixtures("sqlite_session"), @@ -43,8 +56,8 @@ class TestMetadataNullableBug: account = _make_account() # This should crash with TypeError when calling len(None) with pytest.raises(TypeError, match="object of type 'NoneType' has no len"): - MetadataService.create_metadata( - "dataset-123", mock_metadata_args, account, "tenant-123", session=sqlite_session + _metadata_service(sqlite_session).create_metadata( + DatasetRef("tenant-123", "dataset-123"), mock_metadata_args, actor_id=account.id ) assert not sqlite_session.in_transaction() @@ -54,7 +67,9 @@ class TestMetadataNullableBug: none_name = cast(str, None) # This should crash with TypeError when calling len(None) with pytest.raises(TypeError, match="object of type 'NoneType' has no len"): - MetadataService.update_metadata_name(Mock(), "metadata-456", none_name, account, session=sqlite_session) + _metadata_service(sqlite_session).update_metadata_name( + DatasetRef(Mock().tenant_id, Mock().id), "metadata-456", none_name, actor_id=account.id + ) assert not sqlite_session.in_transaction() def test_api_layer_now_uses_pydantic_validation(self) -> None: diff --git a/api/tests/unit_tests/services/test_metadata_service_session_boundary.py b/api/tests/unit_tests/services/test_metadata_service_session_boundary.py index eea773da175..2f3a6951576 100644 --- a/api/tests/unit_tests/services/test_metadata_service_session_boundary.py +++ b/api/tests/unit_tests/services/test_metadata_service_session_boundary.py @@ -1,293 +1,197 @@ +"""Metadata operations persist through scoped, independently owned transactions.""" + +from collections.abc import Iterator from datetime import datetime -from unittest.mock import patch +from unittest.mock import create_autospec import pytest from sqlalchemy import event, select from sqlalchemy.orm import Session, sessionmaker -from core.rag.index_processor.constant.built_in_field import BuiltInField -from models import Account -from models.dataset import Dataset, DatasetMetadata, DatasetMetadataBinding, Document +from models.dataset import DatasetMetadataBinding, Document from models.enums import DocumentCreatedFrom -from services.dataset_service import DocumentService -from services.entities.knowledge_entities.knowledge_entities import ( +from repositories.knowledge.metadata_repository import SQLAlchemyMetadataRepository +from services.errors.metadata import MetadataResourceNotFoundError +from services.knowledge.dataset_access import DatasetAccess +from services.knowledge.entities.knowledge_entities import ( DocumentMetadataOperation, MetadataArgs, MetadataDetail, MetadataOperationData, ) -from services.errors.metadata import MetadataResourceNotFoundError -from services.metadata_service import MetadataService +from services.knowledge.metadata.application import MetadataService +from services.knowledge.resource_scope import DatasetRef from tests.unit_tests.model_factories import make_account, make_dataset, make_document +REF = DatasetRef("tenant-1", "dataset-1") DOCUMENT_ID = "11111111-1111-1111-1111-111111111111" -FOREIGN_DOCUMENT_ID = "22222222-2222-2222-2222-222222222222" -METADATA_ID = "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa" -FOREIGN_METADATA_ID = "bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb" +FOREIGN_ID = "22222222-2222-2222-2222-222222222222" -def _account() -> Account: - return make_account(name="User", email="user@example.com") - - -def test_create_metadata_flushes_without_committing_caller_session(sqlite_session: Session) -> None: - transaction_events: list[str] = [] - event.listen(sqlite_session, "after_commit", lambda _session: transaction_events.append("commit")) - event.listen(sqlite_session, "after_rollback", lambda _session: transaction_events.append("rollback")) - - metadata = MetadataService.create_metadata( - "dataset-1", - MetadataArgs(type="string", name="author"), - _account(), - "tenant-1", - session=sqlite_session, - ) - - assert metadata.name == "author" - assert sqlite_session.get(DatasetMetadata, metadata.id) is metadata - assert transaction_events == [] - - -def _dataset(*, built_in_field_enabled: bool) -> Dataset: - return make_dataset( - description="", - provider="vendor", - built_in_field_enabled=built_in_field_enabled, - ) - - -def _document() -> Document: - return make_document( - document_id=DOCUMENT_ID, - created_from=DocumentCreatedFrom.API, - created_at=datetime(2026, 1, 1), - updated_at=datetime(2026, 1, 2), - doc_metadata={}, - ) - - -def test_enable_built_in_field_uses_caller_session_for_uploader(sqlite_session: Session) -> None: - dataset = _dataset(built_in_field_enabled=False) - document = _document() - sqlite_session.add_all([_account(), dataset, document]) - sqlite_session.commit() - - with ( - patch.object(MetadataService, "knowledge_base_metadata_lock_check"), - patch.object(DocumentService, "get_working_documents_by_dataset_id", return_value=[document]), - patch("services.metadata_service.redis_client.delete"), - ): - MetadataService.enable_built_in_field(dataset, sqlite_session) - - assert document.doc_metadata[BuiltInField.uploader] == "User" - - -def test_update_documents_metadata_uses_caller_session_for_uploader(sqlite_session: Session) -> None: - dataset = _dataset(built_in_field_enabled=True) - document = _document() - sqlite_session.add_all([_account(), dataset, document]) - sqlite_session.commit() - metadata_args = MetadataOperationData( - operation_data=[ - DocumentMetadataOperation(document_id=document.id, metadata_list=[], partial_update=False), - ] - ) - - with ( - patch.object(MetadataService, "knowledge_base_metadata_lock_check"), - patch("services.metadata_service.redis_client.delete"), - ): - MetadataService.update_documents_metadata( - dataset, - metadata_args, - _account(), - session=sqlite_session, - ) - - assert document.doc_metadata[BuiltInField.uploader] == "User" - - -def test_update_documents_metadata_rejects_foreign_metadata_before_writes(sqlite_session: Session) -> None: - dataset = _dataset(built_in_field_enabled=False) - document = _document() - foreign_metadata = DatasetMetadata( - tenant_id="tenant-2", - dataset_id="dataset-2", - type="string", - name="spoofed", - created_by="account-2", - ) - foreign_metadata.id = FOREIGN_METADATA_ID - sqlite_session.add_all([_account(), dataset, document, foreign_metadata]) - sqlite_session.commit() - transaction_events: list[str] = [] - event.listen(sqlite_session, "after_flush", lambda _session, _context: transaction_events.append("flush")) - event.listen(sqlite_session, "after_commit", lambda _session: transaction_events.append("commit")) - metadata_args = MetadataOperationData( - operation_data=[ - DocumentMetadataOperation( - document_id=DOCUMENT_ID, - metadata_list=[MetadataDetail(id=FOREIGN_METADATA_ID, name="spoofed", value="value")], - partial_update=False, +@pytest.fixture +def metadata_service(sqlite_session_factory: sessionmaker[Session]) -> Iterator[MetadataService]: + with sqlite_session_factory.begin() as session: + account = make_account(account_id="actor", name="User", email="metadata@example.com") + session.add(account) + for ref, document_id in [(REF, DOCUMENT_ID), (DatasetRef("tenant-2", "dataset-2"), FOREIGN_ID)]: + session.add( + make_dataset(dataset_id=ref.dataset_id, tenant_id=ref.tenant_id, name="Dataset", created_by="actor") ) - ] + document = make_document( + document_id=document_id, + tenant_id=ref.tenant_id, + dataset_id=ref.dataset_id, + batch="batch", + created_from=DocumentCreatedFrom.API, + created_by="actor", + created_at=datetime(2026, 1, 1), + updated_at=datetime(2026, 1, 2), + ) + document.enabled = True + document.archived = False + document.indexing_status = "completed" + session.add(document) + return MetadataService( + store=SQLAlchemyMetadataRepository(session_factory=sqlite_session_factory), + dataset_access=create_autospec(DatasetAccess, instance=True), ) - with ( - pytest.raises(MetadataResourceNotFoundError, match="Metadata not found"), - patch.object(MetadataService, "knowledge_base_metadata_lock_check") as lock_check, - ): - MetadataService.update_documents_metadata(dataset, metadata_args, _account(), session=sqlite_session) - lock_check.assert_not_called() - assert transaction_events == [] - assert sqlite_session.scalars(select(DatasetMetadataBinding)).all() == [] - assert sqlite_session.get(Document, DOCUMENT_ID).doc_metadata == {} - - -def test_update_documents_metadata_validates_all_documents_before_writes(sqlite_session: Session) -> None: - dataset = _dataset(built_in_field_enabled=False) - document = _document() - metadata = DatasetMetadata( - tenant_id=dataset.tenant_id, - dataset_id=dataset.id, - type="string", - name="canonical", - created_by="account-1", - ) - metadata.id = METADATA_ID - sqlite_session.add_all([_account(), dataset, document, metadata]) - sqlite_session.commit() - transaction_events: list[str] = [] - event.listen(sqlite_session, "after_flush", lambda _session, _context: transaction_events.append("flush")) - event.listen(sqlite_session, "after_commit", lambda _session: transaction_events.append("commit")) - metadata_detail = MetadataDetail(id=metadata.id, name="spoofed", value="value") - metadata_args = MetadataOperationData( - operation_data=[ - DocumentMetadataOperation(document_id=DOCUMENT_ID, metadata_list=[metadata_detail], partial_update=False), - DocumentMetadataOperation( - document_id=FOREIGN_DOCUMENT_ID, metadata_list=[metadata_detail], partial_update=False - ), - ] - ) - - with ( - pytest.raises(MetadataResourceNotFoundError, match="Document not found"), - patch.object(MetadataService, "knowledge_base_metadata_lock_check") as lock_check, - patch("services.metadata_service.redis_client.delete"), - ): - MetadataService.update_documents_metadata(dataset, metadata_args, _account(), session=sqlite_session) - - lock_check.assert_not_called() - assert transaction_events == [] - assert sqlite_session.scalars(select(DatasetMetadataBinding)).all() == [] - assert sqlite_session.get(Document, DOCUMENT_ID).doc_metadata == {} - - -def test_update_documents_metadata_uses_canonical_metadata_name( - sqlite_session: Session, sqlite_session_factory: sessionmaker[Session] +def test_create_commits_through_owned_session( + metadata_service: MetadataService, sqlite_session_factory: sessionmaker[Session] ) -> None: - dataset = _dataset(built_in_field_enabled=False) - document = _document() - metadata = DatasetMetadata( - tenant_id=dataset.tenant_id, - dataset_id=dataset.id, - type="string", - name="canonical", - created_by="account-1", - ) - metadata.id = METADATA_ID - sqlite_session.add_all([_account(), dataset, document, metadata]) - sqlite_session.commit() - metadata_args = MetadataOperationData( + result = metadata_service.create_metadata(REF, MetadataArgs(type="string", name="author"), actor_id="actor") + fields = metadata_service.get_dataset_metadatas(REF) + assert fields["doc_metadata"] == [{"id": result.id, "type": "string", "name": "author", "count": 0}] + with sqlite_session_factory() as session: + assert session.scalar(select(DatasetMetadataBinding)) is None + + +@pytest.mark.parametrize("name", ["x" * 256, "document_name", "uploader", "upload_date", "last_update_date", "source"]) +def test_invalid_name_is_rejected(metadata_service: MetadataService, name: str) -> None: + with pytest.raises(ValueError): + metadata_service.create_metadata(REF, MetadataArgs(type="string", name=name), actor_id="actor") + assert metadata_service.get_dataset_metadatas(REF)["doc_metadata"] == [] + + +def test_duplicate_name_is_rejected(metadata_service: MetadataService) -> None: + metadata_service.create_metadata(REF, MetadataArgs(type="string", name="author"), actor_id="actor") + with pytest.raises(ValueError, match="already exists"): + metadata_service.create_metadata(REF, MetadataArgs(type="string", name="author"), actor_id="actor") + + +def _operation(document_id: str, metadata_id: str, *, partial: bool = False) -> MetadataOperationData: + return MetadataOperationData( operation_data=[ DocumentMetadataOperation( - document_id=document.id, - metadata_list=[MetadataDetail(id=metadata.id, name="spoofed", value="value")], - partial_update=False, + document_id=document_id, + metadata_list=[MetadataDetail(id=metadata_id, name="spoofed", value="Alice")], + partial_update=partial, ) ] ) - with ( - patch.object(MetadataService, "knowledge_base_metadata_lock_check"), - patch("services.metadata_service.redis_client.delete"), - ): - MetadataService.update_documents_metadata(dataset, metadata_args, _account(), session=sqlite_session) - with sqlite_session_factory() as observer: - persisted_document = observer.get(Document, document.id) - assert persisted_document is not None - assert persisted_document.doc_metadata == {"canonical": "value"} - binding = observer.scalar(select(DatasetMetadataBinding)) - assert binding is not None - assert binding.metadata_id == metadata.id - - -def test_metadata_operation_normalizes_uuid_ids() -> None: - operation = DocumentMetadataOperation( - document_id=DOCUMENT_ID.upper(), - metadata_list=[MetadataDetail(id=METADATA_ID.upper(), name="ignored", value="value")], +@pytest.mark.parametrize("partial", [True, False]) +def test_update_uses_canonical_names_and_preserves_partial_values( + metadata_service: MetadataService, sqlite_session_factory: sessionmaker[Session], partial: bool +) -> None: + field = metadata_service.create_metadata(REF, MetadataArgs(type="string", name="author"), actor_id="actor") + with sqlite_session_factory.begin() as session: + document = session.get(Document, DOCUMENT_ID) + assert document is not None + document.doc_metadata = {"existing": "kept"} + metadata_service.update_documents_metadata( + REF, _operation(DOCUMENT_ID, field.id, partial=partial), actor_id="actor" ) - - assert operation.document_id == DOCUMENT_ID - assert operation.metadata_list[0].id == METADATA_ID - - -def test_document_metadata_details_scopes_binding_to_document_owner(sqlite_session: Session) -> None: - dataset = _dataset(built_in_field_enabled=False) - document = _document() - document.doc_metadata = {"canonical": "value"} - foreign_metadata = DatasetMetadata( - tenant_id="tenant-2", - dataset_id="dataset-2", - type="string", - name="canonical", - created_by="account-2", + metadata_service.update_documents_metadata( + REF, _operation(DOCUMENT_ID, field.id, partial=partial), actor_id="actor" ) - foreign_metadata.id = FOREIGN_METADATA_ID - foreign_binding = DatasetMetadataBinding( - tenant_id="tenant-2", - dataset_id="dataset-2", - document_id=document.id, - metadata_id=foreign_metadata.id, - created_by="account-2", + with sqlite_session_factory() as session: + document = session.get(Document, DOCUMENT_ID) + assert document is not None + assert document.doc_metadata == ({"existing": "kept", "author": "Alice"} if partial else {"author": "Alice"}) + assert len(session.scalars(select(DatasetMetadataBinding)).all()) == 1 + + +@pytest.mark.parametrize("foreign", ["dataset", "document", "metadata"]) +def test_owner_chain_is_validated_before_any_batch_write( + metadata_service: MetadataService, sqlite_session_factory: sessionmaker[Session], foreign: str +) -> None: + field = metadata_service.create_metadata(REF, MetadataArgs(type="string", name="author"), actor_id="actor") + other = metadata_service.create_metadata( + DatasetRef("tenant-2", "dataset-2"), MetadataArgs(type="string", name="foreign"), actor_id="actor" ) - sqlite_session.add_all([_account(), dataset, document, foreign_metadata, foreign_binding]) - sqlite_session.commit() - - details = document.get_doc_metadata_details(session=sqlite_session) - - assert details is not None - assert [detail for detail in details if detail["id"] != "built-in"] == [] - - -def test_get_dataset_metadatas_uses_caller_session(monkeypatch, sqlite_session: Session) -> None: - dataset = _dataset(built_in_field_enabled=False) - sqlite_session.add_all( - [ - DatasetMetadataBinding( - tenant_id="tenant-1", - dataset_id="dataset-1", - document_id=f"document-{index}", - metadata_id="metadata-1", - created_by="account-1", - ) - for index in range(2) - ] + args = _operation(DOCUMENT_ID, field.id) + args.operation_data.extend( + _operation( + FOREIGN_ID if foreign == "document" else DOCUMENT_ID, other.id if foreign == "metadata" else field.id + ).operation_data ) - sqlite_session.commit() + ref = DatasetRef("tenant-2", REF.dataset_id) if foreign == "dataset" else REF + with pytest.raises(MetadataResourceNotFoundError): + metadata_service.update_documents_metadata(ref, args, actor_id="actor") + with sqlite_session_factory() as session: + assert all(document.doc_metadata is None for document in session.scalars(select(Document))) + assert not session.scalars(select(DatasetMetadataBinding)).all() - def get_doc_metadata(_dataset: Dataset, *, session: Session) -> list[dict[str, str]]: - assert session is sqlite_session - return [{"id": "metadata-1", "name": "author", "type": "string"}] - monkeypatch.setattr(Dataset, "get_doc_metadata", get_doc_metadata) +def test_rename_delete_and_counts_share_scoped_bindings( + metadata_service: MetadataService, sqlite_session_factory: sessionmaker[Session] +) -> None: + field = metadata_service.create_metadata(REF, MetadataArgs(type="string", name="author"), actor_id="actor") + metadata_service.update_documents_metadata(REF, _operation(DOCUMENT_ID, field.id), actor_id="actor") + renamed = metadata_service.update_metadata_name(REF, field.id, "writer", actor_id="actor") + assert renamed.name == "writer" + assert metadata_service.get_dataset_metadatas(REF)["doc_metadata"] == [ + {"id": field.id, "name": "writer", "type": "string", "count": 1} + ] + with sqlite_session_factory() as session: + document = session.get(Document, DOCUMENT_ID) + assert document is not None + assert document.doc_metadata == {"writer": "Alice"} + metadata_service.delete_metadata(REF, field.id) + with sqlite_session_factory() as session: + document = session.get(Document, DOCUMENT_ID) + assert document is not None + assert document.doc_metadata == {} + assert not session.scalars(select(DatasetMetadataBinding)).all() - result = MetadataService.get_dataset_metadatas(dataset, sqlite_session) - assert result == { - "doc_metadata": [{"id": "metadata-1", "name": "author", "type": "string", "count": 2}], - "built_in_field_enabled": False, - } - assert sqlite_session.scalar(select(DatasetMetadataBinding).limit(1)) is not None +def test_builtin_fields_use_owned_document_uploader( + metadata_service: MetadataService, sqlite_session_factory: sessionmaker[Session] +) -> None: + metadata_service.enable_built_in_field(REF) + with sqlite_session_factory() as session: + document = session.get(Document, DOCUMENT_ID) + foreign = session.get(Document, FOREIGN_ID) + assert document is not None + assert foreign is not None + assert document.doc_metadata["uploader"] == "User" + assert foreign.doc_metadata is None + metadata_service.disable_built_in_field(REF) + with sqlite_session_factory() as session: + document = session.get(Document, DOCUMENT_ID) + assert document is not None + assert document.doc_metadata == {} + + +def test_commit_failure_rolls_back_document_and_bindings( + metadata_service: MetadataService, sqlite_session_factory: sessionmaker[Session] +) -> None: + field = metadata_service.create_metadata(REF, MetadataArgs(type="string", name="author"), actor_id="actor") + + def fail(_session: Session) -> None: + raise RuntimeError("commit failed") + + event.listen(sqlite_session_factory.class_, "before_commit", fail) + try: + with pytest.raises(RuntimeError, match="commit failed"): + metadata_service.update_documents_metadata(REF, _operation(DOCUMENT_ID, field.id), actor_id="actor") + finally: + event.remove(sqlite_session_factory.class_, "before_commit", fail) + with sqlite_session_factory() as session: + document = session.get(Document, DOCUMENT_ID) + assert document is not None + assert document.doc_metadata is None + assert not session.scalars(select(DatasetMetadataBinding)).all() diff --git a/api/tests/unit_tests/services/test_model_provider_service.py b/api/tests/unit_tests/services/test_model_provider_service.py index ce2dd7c1bcb..6659d5ac7d0 100644 --- a/api/tests/unit_tests/services/test_model_provider_service.py +++ b/api/tests/unit_tests/services/test_model_provider_service.py @@ -1,6 +1,6 @@ from types import SimpleNamespace from typing import Any -from unittest.mock import MagicMock +from unittest.mock import MagicMock, create_autospec import pytest from sqlalchemy import event @@ -15,6 +15,7 @@ from enums import DeploymentEdition from graphon.model_runtime.entities.common_entities import I18nObject from graphon.model_runtime.entities.model_entities import FetchFrom, ModelType, ParameterRule, ParameterType from graphon.model_runtime.entities.provider_entities import ConfigurateMethod +from models.account import Account, TenantAccountRole from models.provider import ( Provider, ProviderCredential, @@ -23,9 +24,10 @@ from models.provider import ( TenantDefaultModel, TenantPreferredModelProvider, ) -from services import model_provider_service as service_module +from services.credentials.query import CredentialQuery, ModelCredentialRecord from services.errors.app_model_config import ProviderNotFoundError -from services.model_provider_service import ModelProviderService, _ProviderSummaryState +from services.model_provider import service as service_module +from services.model_provider.service import ModelProviderService, _ProviderSummaryState def _create_service_with_mocked_manager() -> tuple[ModelProviderService, MagicMock]: @@ -35,6 +37,21 @@ def _create_service_with_mocked_manager() -> tuple[ModelProviderService, MagicMo return service, manager +def test_available_credentials_use_injected_query_for_admin_viewer() -> None: + service, manager = _create_service_with_mocked_manager() + user = Account(name="Admin", email="admin@example.com") + user.id = "actor" + user.role = TenantAccountRole.ADMIN + query = create_autospec(CredentialQuery, instance=True, spec_set=True) + query.list_models.return_value = [ModelCredentialRecord(id="credential", name="Shared credential")] + + result = service.get_provider_available_credentials("tenant", "openai", user=user, credential_query=query) + + query.list_models.assert_called_once_with(workspace_id="tenant", provider="openai", actor_id="actor") + assert result == [CredentialConfiguration(credential_id="credential", credential_name="Shared credential")] + assert manager.mock_calls == [] + + def _build_provider_configuration( *, provider_name: str = "openai", diff --git a/api/tests/unit_tests/services/test_model_provider_service_sanitization.py b/api/tests/unit_tests/services/test_model_provider_service_sanitization.py index 97f3bd6f013..1c75b0e059b 100644 --- a/api/tests/unit_tests/services/test_model_provider_service_sanitization.py +++ b/api/tests/unit_tests/services/test_model_provider_service_sanitization.py @@ -7,7 +7,7 @@ from graphon.model_runtime.entities.common_entities import I18nObject from graphon.model_runtime.entities.model_entities import ModelType from graphon.model_runtime.entities.provider_entities import ConfigurateMethod from models.provider import ProviderType -from services.model_provider_service import ModelProviderService +from services.model_provider.service import ModelProviderService class _FakeConfigurations: diff --git a/api/tests/unit_tests/services/test_notion_data_source_gateway.py b/api/tests/unit_tests/services/test_notion_data_source_gateway.py deleted file mode 100644 index b3754eb799b..00000000000 --- a/api/tests/unit_tests/services/test_notion_data_source_gateway.py +++ /dev/null @@ -1,156 +0,0 @@ -import json -from unittest.mock import MagicMock - -import httpx -import pytest - -from core.helper.ssrf_proxy import SSRFProxy -from core.tools.errors import ToolSSRFError -from services.data_source_oauth_service import DataSourceOAuthError -from services.notion_data_source_gateway import NotionDataSourceGateway - - -def _gateway(handler: httpx.MockTransport) -> tuple[NotionDataSourceGateway, httpx.Client]: - client = httpx.Client(transport=handler) - return ( - NotionDataSourceGateway( - client_id="client-id", - client_secret="client-secret", - redirect_uri="https://api.example/oauth/data-source/callback/notion", - http_client=client, - ), - client, - ) - - -def test_get_authorization_url_encodes_configured_values() -> None: - gateway, client = _gateway(httpx.MockTransport(lambda _request: httpx.Response(500))) - try: - result = gateway.get_authorization_url() - finally: - client.close() - - assert result == ( - "https://api.notion.com/v1/oauth/authorize?client_id=client-id&response_type=code&" - "redirect_uri=https%3A%2F%2Fapi.example%2Foauth%2Fdata-source%2Fcallback%2Fnotion&owner=user" - ) - - -def test_authorize_exchanges_code_and_builds_source_info_without_persistence() -> None: - search_types: list[str] = [] - - def handler(request: httpx.Request) -> httpx.Response: - if request.url.path == "/v1/oauth/token": - assert request.headers["authorization"].startswith("Basic ") - return httpx.Response( - 200, - json={ - "access_token": "notion-token", - "workspace_name": "Notion Workspace", - "workspace_icon": "https://example/icon.png", - "workspace_id": "notion-workspace", - }, - ) - if request.url.path == "/v1/search": - payload = json.loads(request.content) - search_types.append(payload["filter"]["value"]) - return httpx.Response(200, json={"results": [], "has_more": False}) - raise AssertionError(f"Unexpected Notion request: {request.url}") - - gateway, client = _gateway(httpx.MockTransport(handler)) - try: - authorization = gateway.authorize("authorization-code") - finally: - client.close() - - assert authorization.access_token == "notion-token" - assert authorization.source_info == { - "workspace_name": "Notion Workspace", - "workspace_icon": "https://example/icon.png", - "workspace_id": "notion-workspace", - "pages": [], - "total": 0, - } - assert search_types == ["page", "database"] - - -def test_refresh_preserves_workspace_metadata() -> None: - gateway, client = _gateway( - httpx.MockTransport( - lambda request: ( - httpx.Response(200, json={"results": [], "has_more": False}) - if request.url.path == "/v1/search" - else httpx.Response(404) - ) - ) - ) - try: - source_info = gateway.refresh( - "notion-token", - { - "workspace_name": "Workspace", - "workspace_icon": None, - "workspace_id": "notion-workspace", - "pages": [ - { - "page_id": "old", - "page_name": "Old page", - "page_icon": None, - "parent_id": "root", - "type": "page", - } - ], - "total": 1, - }, - ) - finally: - client.close() - - assert source_info == { - "workspace_name": "Workspace", - "workspace_icon": None, - "workspace_id": "notion-workspace", - "pages": [], - "total": 0, - } - - -def test_search_rejects_missing_cursor_when_more_results_are_declared() -> None: - requests = 0 - - def handler(_request: httpx.Request) -> httpx.Response: - nonlocal requests - requests += 1 - return httpx.Response(200, json={"results": [], "has_more": True, "next_cursor": None}) - - gateway, client = _gateway(httpx.MockTransport(handler)) - try: - with pytest.raises(DataSourceOAuthError, match="valid next_cursor"): - gateway.notion_page_search("notion-token") - finally: - client.close() - - assert requests == 1 - - -def test_provider_http_failure_is_translated_to_application_boundary_error() -> None: - gateway, client = _gateway(httpx.MockTransport(lambda _request: httpx.Response(503, text="unavailable"))) - try: - with pytest.raises(DataSourceOAuthError, match="unsuccessful response"): - gateway.authorize("authorization-code") - finally: - client.close() - - -def test_ssrf_rejection_is_translated_to_application_boundary_error() -> None: - client = MagicMock(spec=SSRFProxy) - client.post.side_effect = ToolSSRFError("blocked") - gateway = NotionDataSourceGateway( - client_id="client-id", - client_secret="client-secret", - redirect_uri="https://api.example/oauth/data-source/callback/notion", - http_client=client, - ) - - with pytest.raises(DataSourceOAuthError, match="Notion request failed"): - gateway.authorize("authorization-code") diff --git a/api/tests/unit_tests/services/test_rag_pipeline_task_proxy.py b/api/tests/unit_tests/services/test_rag_pipeline_task_proxy.py index 2b73d643387..0905c78deba 100644 --- a/api/tests/unit_tests/services/test_rag_pipeline_task_proxy.py +++ b/api/tests/unit_tests/services/test_rag_pipeline_task_proxy.py @@ -226,14 +226,13 @@ class TestRagPipelineTaskProxy: assert parsed_json[0]["pipeline_id"] == "pipeline-1" assert parsed_json[1]["pipeline_id"] == "pipeline-2" - @patch("services.rag_pipeline.rag_pipeline_task_proxy.rag_pipeline_run_task") - def test_send_to_direct_queue(self, mock_task: MagicMock): + def test_send_to_direct_queue(self): """Test _send_to_direct_queue method.""" # Arrange proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() proxy._tenant_isolated_task_queue = RagPipelineTaskProxyTestDataFactory.create_mock_tenant_queue() upload_file_id = "file-123" - mock_task.delay = Mock() + mock_task = Mock() # Act proxy._send_to_direct_queue(upload_file_id, mock_task) @@ -246,8 +245,7 @@ class TestRagPipelineTaskProxy: rag_pipeline_invoke_entities_file_id=upload_file_id, tenant_id="tenant-123" ) - @patch("services.rag_pipeline.rag_pipeline_task_proxy.rag_pipeline_run_task") - def test_send_to_tenant_queue_with_existing_task_key(self, mock_task: MagicMock): + def test_send_to_tenant_queue_with_existing_task_key(self): """Test _send_to_tenant_queue when task key exists.""" # Arrange proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() @@ -255,7 +253,7 @@ class TestRagPipelineTaskProxy: has_task_key=True ) upload_file_id = "file-123" - mock_task.delay = Mock() + mock_task = Mock() # Act proxy._send_to_tenant_queue(upload_file_id, mock_task) @@ -265,8 +263,7 @@ class TestRagPipelineTaskProxy: # Celery should not be called directly mock_task.delay.assert_not_called() - @patch("services.rag_pipeline.rag_pipeline_task_proxy.rag_pipeline_run_task") - def test_send_to_tenant_queue_without_task_key(self, mock_task: MagicMock): + def test_send_to_tenant_queue_without_task_key(self): """Test _send_to_tenant_queue when no task key exists.""" # Arrange proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() @@ -274,7 +271,7 @@ class TestRagPipelineTaskProxy: has_task_key=False ) upload_file_id = "file-123" - mock_task.delay = Mock() + mock_task = Mock() # Act proxy._send_to_tenant_queue(upload_file_id, mock_task) @@ -288,8 +285,7 @@ class TestRagPipelineTaskProxy: # The first task should be sent to celery directly, so push tasks should not be called proxy._tenant_isolated_task_queue.push_tasks.assert_not_called() - @patch("services.rag_pipeline.rag_pipeline_task_proxy.rag_pipeline_run_task") - def test_send_to_default_tenant_queue(self, mock_task: MagicMock): + def test_send_to_default_tenant_queue(self): """Test _send_to_default_tenant_queue method.""" # Arrange proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() @@ -300,10 +296,13 @@ class TestRagPipelineTaskProxy: proxy._send_to_default_tenant_queue(upload_file_id) # Assert - proxy._send_to_tenant_queue.assert_called_once_with(upload_file_id, mock_task) + proxy._send_to_tenant_queue.assert_called_once() + sent_file_id, task = proxy._send_to_tenant_queue.call_args.args + assert sent_file_id == upload_file_id + assert task.task == "tasks.rag_pipeline.rag_pipeline_run_task.rag_pipeline_run_task" + assert task.options["queue"] == "pipeline" - @patch("services.rag_pipeline.rag_pipeline_task_proxy.priority_rag_pipeline_run_task") - def test_send_to_priority_tenant_queue(self, mock_task: MagicMock): + def test_send_to_priority_tenant_queue(self): """Test _send_to_priority_tenant_queue method.""" # Arrange proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() @@ -314,10 +313,13 @@ class TestRagPipelineTaskProxy: proxy._send_to_priority_tenant_queue(upload_file_id) # Assert - proxy._send_to_tenant_queue.assert_called_once_with(upload_file_id, mock_task) + proxy._send_to_tenant_queue.assert_called_once() + sent_file_id, task = proxy._send_to_tenant_queue.call_args.args + assert sent_file_id == upload_file_id + assert task.task == "tasks.rag_pipeline.priority_rag_pipeline_run_task.priority_rag_pipeline_run_task" + assert task.options["queue"] == "priority_pipeline" - @patch("services.rag_pipeline.rag_pipeline_task_proxy.priority_rag_pipeline_run_task") - def test_send_to_priority_direct_queue(self, mock_task: MagicMock): + def test_send_to_priority_direct_queue(self): """Test _send_to_priority_direct_queue method.""" # Arrange proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() @@ -328,7 +330,11 @@ class TestRagPipelineTaskProxy: proxy._send_to_priority_direct_queue(upload_file_id) # Assert - proxy._send_to_direct_queue.assert_called_once_with(upload_file_id, mock_task) + proxy._send_to_direct_queue.assert_called_once() + sent_file_id, task = proxy._send_to_direct_queue.call_args.args + assert sent_file_id == upload_file_id + assert task.task == "tasks.rag_pipeline.priority_rag_pipeline_run_task.priority_rag_pipeline_run_task" + assert task.options["queue"] == "priority_pipeline" @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) @patch("services.rag_pipeline.rag_pipeline_task_proxy.FeatureService") diff --git a/api/tests/unit_tests/services/test_rbac_resource_service.py b/api/tests/unit_tests/services/test_rbac_resource_service.py index d93834b8eb6..a11eddcbd9b 100644 --- a/api/tests/unit_tests/services/test_rbac_resource_service.py +++ b/api/tests/unit_tests/services/test_rbac_resource_service.py @@ -1,11 +1,13 @@ """RBACResourceService reads through the session its caller passes in (#37403).""" +import pytest from sqlalchemy.orm import Session from models.agent import Agent, AgentKind, AgentScope, AgentSource, AgentStatus from models.dataset import Dataset from models.model import App, AppMode from services.rbac_resource_service import RBACResourceService +from tests.unit_tests.model_factories import make_document TENANT_ID = "tenant-1" OTHER_TENANT_ID = "tenant-2" @@ -133,3 +135,30 @@ class TestGetDatasetIdByPipeline: sqlite_session.flush() assert RBACResourceService.get_dataset_id_by_pipeline(sqlite_session, TENANT_ID, "pipeline-1") is None + + +class TestGetDatasetIdByDocument: + @pytest.mark.parametrize( + ("dataset_tenant", "document_tenant", "expected"), + [ + (TENANT_ID, TENANT_ID, "dataset-1"), + (OTHER_TENANT_ID, OTHER_TENANT_ID, None), + (OTHER_TENANT_ID, TENANT_ID, None), + (TENANT_ID, OTHER_TENANT_ID, None), + ], + ) + def test_scopes_both_dataset_and_document( + self, sqlite_session: Session, dataset_tenant: str, document_tenant: str, expected: str | None + ) -> None: + sqlite_session.add_all( + [ + _dataset(dataset_id="dataset-1", tenant_id=dataset_tenant), + make_document(document_id="doc-1", dataset_id="dataset-1", tenant_id=document_tenant), + ] + ) + sqlite_session.flush() + + assert RBACResourceService.get_dataset_id_by_document(sqlite_session, TENANT_ID, "doc-1") == expected + + def test_missing_document(self, sqlite_session: Session) -> None: + assert RBACResourceService.get_dataset_id_by_document(sqlite_session, TENANT_ID, "missing") is None diff --git a/api/tests/unit_tests/services/test_summary_index_service.py b/api/tests/unit_tests/services/test_summary_index_service.py index 4dab7608964..c8fb1a36e65 100644 --- a/api/tests/unit_tests/services/test_summary_index_service.py +++ b/api/tests/unit_tests/services/test_summary_index_service.py @@ -1,4 +1,4 @@ -"""SQLite-backed tests for :mod:`services.summary_index_service`.""" +"""SQLite-backed tests for :mod:`services.knowledge.summaries.adapters`.""" from __future__ import annotations @@ -11,11 +11,11 @@ import pytest from sqlalchemy import event, func, select from sqlalchemy.orm import Session, sessionmaker -import services.summary_index_service as summary_module +import services.knowledge.summaries.adapters as summary_module from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType from models.dataset import Dataset, Document, DocumentSegment, DocumentSegmentSummary from models.enums import DataSourceType, DocumentCreatedFrom, SegmentStatus, SummaryStatus -from services.summary_index_service import SummaryIndexService +from services.knowledge.summaries.adapters import SummaryIndexAdapter TENANT_ID = "tenant-1" OTHER_TENANT_ID = "tenant-2" @@ -168,7 +168,7 @@ def _install_summary_generator( usage: MagicMock | None = None, ) -> MagicMock: generate = MagicMock(return_value=(content, usage or _usage())) - paragraph_module = SimpleNamespace(ParagraphIndexProcessor=SimpleNamespace(generate_summary=generate)) + paragraph_module = SimpleNamespace(ParagraphIndexProcessor=SimpleNamespace(generate_summary_from_inputs=generate)) monkeypatch.setitem( sys.modules, "core.rag.index_processor.processor.paragraph_index_processor", @@ -196,14 +196,21 @@ class TestGenerateAndCreate: usage = _usage() generate = _install_summary_generator(monkeypatch, usage=usage) - content, result_usage = SummaryIndexService.generate_summary_for_segment( + def call_model(**kwargs): + assert not sqlite_session.in_transaction() + return "generated summary", usage + + generate.side_effect = call_model + + content, result_usage = SummaryIndexAdapter.generate_summary_for_segment( segment, dataset, {"enable": True}, session=sqlite_session ) assert content == "generated summary" assert result_usage is usage assert generate.call_args.kwargs["document_language"] == "en" - assert generate.call_args.kwargs["session"] is sqlite_session + assert callable(generate.call_args.kwargs["image_loader"]) + assert not sqlite_session.in_transaction() def test_generate_allows_missing_document_but_rejects_empty_summary( self, sqlite_session: Session, monkeypatch: pytest.MonkeyPatch @@ -216,7 +223,7 @@ class TestGenerateAndCreate: _install_summary_generator(monkeypatch, content="") with pytest.raises(ValueError, match="Generated summary is empty"): - SummaryIndexService.generate_summary_for_segment(segment, dataset, {"enable": True}, session=sqlite_session) + SummaryIndexAdapter.generate_summary_for_segment(segment, dataset, {"enable": True}, session=sqlite_session) def test_create_updates_only_matching_dataset_and_reenables(self, sqlite_session: Session) -> None: dataset, document, segment = _graph(sqlite_session) @@ -237,7 +244,7 @@ class TestGenerateAndCreate: existing.disabled_by = "account-2" sqlite_session.commit() - result = SummaryIndexService.create_summary_record( + result = SummaryIndexAdapter.create_summary_record( segment, dataset, "new", status=SummaryStatus.GENERATING, session=sqlite_session ) @@ -255,7 +262,7 @@ class TestGenerateAndCreate: def test_create_persists_new_record_in_current_transaction(self, sqlite_session: Session) -> None: dataset, document, segment = _graph(sqlite_session) - result = SummaryIndexService.create_summary_record( + result = SummaryIndexAdapter.create_summary_record( segment, dataset, "new", status=SummaryStatus.NOT_STARTED, session=sqlite_session ) @@ -279,7 +286,7 @@ class TestVectorizeSummary: vector_class = MagicMock() monkeypatch.setattr(summary_module, "Vector", vector_class) - SummaryIndexService.vectorize_summary(summary, segment, dataset, session=sqlite_session) + SummaryIndexAdapter.vectorize_summary(summary, segment, dataset, session=sqlite_session) vector_class.assert_not_called() @@ -288,19 +295,25 @@ class TestVectorizeSummary: summary = _persist_summary(sqlite_session, dataset, document, segment, content=" ") with pytest.raises(ValueError, match="Summary content is empty"): - SummaryIndexService.vectorize_summary(summary, segment, dataset, session=sqlite_session) + SummaryIndexAdapter.vectorize_summary(summary, segment, dataset, session=sqlite_session) - def test_provided_session_retries_and_flushes_mapped_record( + def test_provided_session_releases_transaction_before_every_vector_attempt( self, sqlite_session: Session, monkeypatch: pytest.MonkeyPatch ) -> None: dataset, document, segment = _graph(sqlite_session) summary = _persist_summary(sqlite_session, dataset, document, segment) vector = _install_vector_dependencies(monkeypatch) - vector.add_texts.side_effect = [ConnectionError("connection timeout"), None] + + def write_vector(*args, **kwargs): + assert not sqlite_session.in_transaction() + if vector.add_texts.call_count == 1: + raise ConnectionError("connection timeout") + + vector.add_texts.side_effect = write_vector sleep = MagicMock() monkeypatch.setattr(summary_module.time, "sleep", sleep) - SummaryIndexService.vectorize_summary(summary, segment, dataset, session=sqlite_session) + SummaryIndexAdapter.vectorize_summary(summary, segment, dataset, session=sqlite_session) assert vector.add_texts.call_count == 2 sleep.assert_called_once_with(2.0) @@ -321,7 +334,7 @@ class TestVectorizeSummary: summary = _persist_summary(sqlite_session, dataset, document, segment) _install_vector_dependencies(monkeypatch) - SummaryIndexService.vectorize_summary(summary, segment, dataset) + SummaryIndexAdapter.vectorize_summary(summary, segment, dataset) with sqlite_session_factory() as observer: stored = observer.get(DocumentSegmentSummary, summary.id) @@ -350,7 +363,7 @@ class TestVectorizeSummary: detached.id = "missing-id" _install_vector_dependencies(monkeypatch) - SummaryIndexService.vectorize_summary(detached, segment, dataset) + SummaryIndexAdapter.vectorize_summary(detached, segment, dataset) with sqlite_session_factory() as observer: refreshed = observer.get(DocumentSegmentSummary, stored.id) @@ -375,14 +388,14 @@ class TestVectorizeSummary: detached.id = "new-summary" _install_vector_dependencies(monkeypatch) - SummaryIndexService.vectorize_summary(detached, segment, dataset) + SummaryIndexAdapter.vectorize_summary(detached, segment, dataset) with sqlite_session_factory() as observer: stored = observer.get(DocumentSegmentSummary, detached.id) assert stored is not None assert stored.status == SummaryStatus.COMPLETED - def test_provided_session_failure_records_error_without_opening_another_session( + def test_provided_session_is_committed_before_independent_error_persistence( self, sqlite_session: Session, monkeypatch: pytest.MonkeyPatch, @@ -391,22 +404,60 @@ class TestVectorizeSummary: summary = _persist_summary(sqlite_session, dataset, document, segment) vector = _install_vector_dependencies(monkeypatch) vector.add_texts.side_effect = RuntimeError("fatal vector failure") - create_session = MagicMock() + create_session = MagicMock(wraps=summary_module.session_factory.create_session) monkeypatch.setattr(summary_module.session_factory, "create_session", create_session) with pytest.raises(RuntimeError, match="fatal vector failure"): - SummaryIndexService.vectorize_summary(summary, segment, dataset, session=sqlite_session) + SummaryIndexAdapter.vectorize_summary(summary, segment, dataset, session=sqlite_session) - create_session.assert_not_called() + assert create_session.call_count == 2 + assert not sqlite_session.in_transaction() assert summary.status == SummaryStatus.ERROR assert summary.error == "Vectorization failed: fatal vector failure" class TestBatchAndGeneration: + @pytest.mark.parametrize("foreign_document", [False, True]) + def test_task_resolves_owned_document_and_releases_read_transaction( + self, + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + monkeypatch: pytest.MonkeyPatch, + foreign_document: bool, + ) -> None: + from tasks.generate_summary_index_task import generate_summary_index_task + + dataset, document, _ = _graph(sqlite_session) + dataset.summary_index_setting = {"enable": True} + document.need_summary = True + if foreign_document: + other = _persist_dataset(sqlite_session, dataset_id="other", tenant_id=OTHER_TENANT_ID) + document.dataset_id, document.tenant_id = other.id, other.tenant_id + sqlite_session.commit() + sessions = [] + + def create_session(): + session = sqlite_session_factory() + sessions.append(session) + return session + + def generate(**kwargs): + assert sessions + assert all(not session.in_transaction() for session in sessions) + return [] + + generate_mock = MagicMock(side_effect=generate) + monkeypatch.setattr(summary_module.session_factory, "create_session", create_session) + monkeypatch.setattr(SummaryIndexAdapter, "generate_summaries_for_document", generate_mock) + + generate_summary_index_task(dataset.id, document.id) + + assert generate_mock.call_count == (0 if foreign_document else 1) + def test_batch_create_no_segments_is_noop(self, sqlite_session: Session) -> None: dataset = _persist_dataset(sqlite_session) - SummaryIndexService.batch_create_summary_records([], dataset) + SummaryIndexAdapter.batch_create_summary_records([], dataset) assert sqlite_session.scalar(select(func.count()).select_from(DocumentSegmentSummary)) == 0 @@ -420,7 +471,7 @@ class TestBatchAndGeneration: existing.disabled_by = "account-2" sqlite_session.commit() - SummaryIndexService.batch_create_summary_records([first, second], dataset, status=SummaryStatus.NOT_STARTED) + SummaryIndexAdapter.batch_create_summary_records([first, second], dataset, status=SummaryStatus.NOT_STARTED) with sqlite_session_factory() as observer: rows = observer.scalars( @@ -438,14 +489,14 @@ class TestBatchAndGeneration: dataset, document, segment = _graph(sqlite_session) summary = _persist_summary(sqlite_session, dataset, document, segment) - SummaryIndexService.update_summary_record_error(segment, dataset, "generation failed") + SummaryIndexAdapter.update_summary_record_error(segment, dataset, "generation failed") with sqlite_session_factory() as observer: stored = observer.get(DocumentSegmentSummary, summary.id) assert stored is not None assert (stored.status, stored.error) == (SummaryStatus.ERROR, "generation failed") other = _persist_segment(sqlite_session, dataset, document, segment_id="seg-2", position=2) - SummaryIndexService.update_summary_record_error(other, dataset, "ignored") + SummaryIndexAdapter.update_summary_record_error(other, dataset, "ignored") assert sqlite_session.scalar(select(func.count()).select_from(DocumentSegmentSummary)) == 1 def test_generate_and_vectorize_creates_commits_and_returns_row( @@ -453,7 +504,7 @@ class TestBatchAndGeneration: ) -> None: dataset, _, segment = _graph(sqlite_session) monkeypatch.setattr( - SummaryIndexService, + SummaryIndexAdapter, "generate_summary_for_segment", MagicMock(return_value=("generated", _usage())), ) @@ -468,9 +519,9 @@ class TestBatchAndGeneration: record.summary_index_node_id = "node-1" sqlite_session.flush() - monkeypatch.setattr(SummaryIndexService, "vectorize_summary", vectorize) + monkeypatch.setattr(SummaryIndexAdapter, "vectorize_summary", vectorize) - result = SummaryIndexService.generate_and_vectorize_summary( + result = SummaryIndexAdapter.generate_and_vectorize_summary( segment, dataset, {"enable": True}, session=sqlite_session ) @@ -486,13 +537,13 @@ class TestBatchAndGeneration: dataset, document, segment = _graph(sqlite_session) summary = _persist_summary(sqlite_session, dataset, document, segment, status=SummaryStatus.NOT_STARTED) monkeypatch.setattr( - SummaryIndexService, + SummaryIndexAdapter, "generate_summary_for_segment", MagicMock(side_effect=RuntimeError("LLM failed")), ) with pytest.raises(RuntimeError, match="LLM failed"): - SummaryIndexService.generate_and_vectorize_summary( + SummaryIndexAdapter.generate_and_vectorize_summary( segment, dataset, {"enable": True}, session=sqlite_session ) @@ -517,7 +568,7 @@ class TestBatchAndGeneration: dataset = _persist_dataset(sqlite_session, indexing_technique=technique) document = _persist_document(sqlite_session, dataset, doc_form=doc_form) - assert SummaryIndexService.generate_summaries_for_document(dataset, document, {"enable": enabled}) == [] + assert SummaryIndexAdapter.generate_summaries_for_document(dataset, document, {"enable": enabled}) == [] def test_generate_document_filters_segments_and_continues_after_failure( self, sqlite_session: Session, monkeypatch: pytest.MonkeyPatch @@ -534,10 +585,10 @@ class TestBatchAndGeneration: ) generate = MagicMock(side_effect=[generated, RuntimeError("boom")]) update_error = MagicMock() - monkeypatch.setattr(SummaryIndexService, "generate_and_vectorize_summary", generate) - monkeypatch.setattr(SummaryIndexService, "update_summary_record_error", update_error) + monkeypatch.setattr(SummaryIndexAdapter, "generate_and_vectorize_summary", generate) + monkeypatch.setattr(SummaryIndexAdapter, "update_summary_record_error", update_error) - result = SummaryIndexService.generate_summaries_for_document( + result = SummaryIndexAdapter.generate_summaries_for_document( dataset, document, {"enable": True}, @@ -564,7 +615,7 @@ class TestEnableDisableDelete: vector = MagicMock(name="vector") monkeypatch.setattr(summary_module, "Vector", MagicMock(return_value=vector)) - SummaryIndexService.disable_summaries_for_segments(dataset, segment_ids=[first.id], disabled_by="account-2") + SummaryIndexAdapter.disable_summaries_for_segments(dataset, segment_ids=[first.id], disabled_by="account-2") vector.delete_by_ids.assert_called_once_with(["node-1"]) with sqlite_session_factory() as observer: @@ -586,7 +637,7 @@ class TestEnableDisableDelete: vector.delete_by_ids.side_effect = RuntimeError("vector unavailable") monkeypatch.setattr(summary_module, "Vector", MagicMock(return_value=vector)) - SummaryIndexService.disable_summaries_for_segments(dataset) + SummaryIndexAdapter.disable_summaries_for_segments(dataset) sqlite_session.refresh(summary) assert summary.enabled is False @@ -621,9 +672,9 @@ class TestEnableDisableDelete: sqlite_session.flush() vectorize_mock = MagicMock(side_effect=vectorize) - monkeypatch.setattr(SummaryIndexService, "vectorize_summary", vectorize_mock) + monkeypatch.setattr(SummaryIndexAdapter, "vectorize_summary", vectorize_mock) - SummaryIndexService.enable_summaries_for_segments(dataset) + SummaryIndexAdapter.enable_summaries_for_segments(dataset) sqlite_session.refresh(good_summary) assert good_summary.enabled is True @@ -632,16 +683,16 @@ class TestEnableDisableDelete: def test_enable_skips_economy_dataset(self, sqlite_session: Session) -> None: dataset = _persist_dataset(sqlite_session, indexing_technique=IndexTechniqueType.ECONOMY) - SummaryIndexService.enable_summaries_for_segments(dataset) + SummaryIndexAdapter.enable_summaries_for_segments(dataset) def test_enable_keeps_failed_summary_disabled( self, sqlite_session: Session, monkeypatch: pytest.MonkeyPatch ) -> None: dataset, document, segment = _graph(sqlite_session) summary = _persist_summary(sqlite_session, dataset, document, segment, enabled=False) - monkeypatch.setattr(SummaryIndexService, "vectorize_summary", MagicMock(side_effect=RuntimeError("boom"))) + monkeypatch.setattr(SummaryIndexAdapter, "vectorize_summary", MagicMock(side_effect=RuntimeError("boom"))) - SummaryIndexService.enable_summaries_for_segments(dataset) + SummaryIndexAdapter.enable_summaries_for_segments(dataset) sqlite_session.refresh(summary) assert summary.enabled is False @@ -656,7 +707,7 @@ class TestEnableDisableDelete: vector = MagicMock(name="vector") monkeypatch.setattr(summary_module, "Vector", MagicMock(return_value=vector)) - SummaryIndexService.delete_summaries_for_segments(dataset, segment_ids=[first.id], session=sqlite_session) + SummaryIndexAdapter.delete_summaries_for_segments(dataset, segment_ids=[first.id], session=sqlite_session) assert sqlite_session.get(DocumentSegmentSummary, deleted.id) is None assert sqlite_session.get(DocumentSegmentSummary, kept.id) is not None @@ -668,7 +719,7 @@ class TestManualUpdate: economy = _persist_dataset(sqlite_session, indexing_technique=IndexTechniqueType.ECONOMY) document = _persist_document(sqlite_session, economy) segment = _persist_segment(sqlite_session, economy, document) - assert SummaryIndexService.update_summary_for_segment(segment, economy, "new", session=sqlite_session) is None + assert SummaryIndexAdapter.update_summary_for_segment(segment, economy, "new", session=sqlite_session) is None quality = _persist_dataset(sqlite_session, dataset_id="dataset-2") qa_document = _persist_document( @@ -676,7 +727,7 @@ class TestManualUpdate: ) qa_segment = _persist_segment(sqlite_session, quality, qa_document, segment_id="seg-2") assert ( - SummaryIndexService.update_summary_for_segment(qa_segment, quality, "new", session=sqlite_session) is None + SummaryIndexAdapter.update_summary_for_segment(qa_segment, quality, "new", session=sqlite_session) is None ) def test_empty_content_deletes_record_even_when_vector_delete_fails( @@ -688,14 +739,14 @@ class TestManualUpdate: vector.delete_by_ids.side_effect = RuntimeError("boom") monkeypatch.setattr(summary_module, "Vector", MagicMock(return_value=vector)) - result = SummaryIndexService.update_summary_for_segment(segment, dataset, " ", session=sqlite_session) + result = SummaryIndexAdapter.update_summary_for_segment(segment, dataset, " ", session=sqlite_session) assert result is None assert sqlite_session.get(DocumentSegmentSummary, summary.id) is None def test_empty_content_without_record_is_noop(self, sqlite_session: Session) -> None: dataset, _, segment = _graph(sqlite_session) - assert SummaryIndexService.update_summary_for_segment(segment, dataset, "", session=sqlite_session) is None + assert SummaryIndexAdapter.update_summary_for_segment(segment, dataset, "", session=sqlite_session) is None def test_existing_record_updates_vectorizes_and_commits( self, @@ -705,22 +756,9 @@ class TestManualUpdate: ) -> None: dataset, document, segment = _graph(sqlite_session) summary = _persist_summary(sqlite_session, dataset, document, segment, content="old", node_id="old-node") - vector = MagicMock(name="vector") - monkeypatch.setattr(summary_module, "Vector", MagicMock(return_value=vector)) + vector = _install_vector_dependencies(monkeypatch) - def vectorize( - record: DocumentSegmentSummary, - _segment: DocumentSegment, - _dataset: Dataset, - session: Session | None = None, - ) -> None: - record.status = SummaryStatus.COMPLETED - record.summary_index_node_hash = "new-hash" - sqlite_session.flush() - - monkeypatch.setattr(SummaryIndexService, "vectorize_summary", vectorize) - - result = SummaryIndexService.update_summary_for_segment(segment, dataset, "new", session=sqlite_session) + result = SummaryIndexAdapter.update_summary_for_segment(segment, dataset, "new", session=sqlite_session) assert result is not None vector.delete_by_ids.assert_called_once_with(["old-node"]) @@ -730,7 +768,7 @@ class TestManualUpdate: assert (stored.summary_content, stored.status, stored.summary_index_node_hash) == ( "new", SummaryStatus.COMPLETED, - "new-hash", + "hash-1", ) def test_new_record_is_persisted(self, sqlite_session: Session, monkeypatch: pytest.MonkeyPatch) -> None: @@ -745,9 +783,9 @@ class TestManualUpdate: record.status = SummaryStatus.COMPLETED sqlite_session.flush() - monkeypatch.setattr(SummaryIndexService, "vectorize_summary", vectorize) + monkeypatch.setattr(SummaryIndexAdapter, "vectorize_summary", vectorize) - result = SummaryIndexService.update_summary_for_segment(segment, dataset, "new", session=sqlite_session) + result = SummaryIndexAdapter.update_summary_for_segment(segment, dataset, "new", session=sqlite_session) assert result is not None assert sqlite_session.get(DocumentSegmentSummary, result.id) is result @@ -758,9 +796,9 @@ class TestManualUpdate: ) -> None: dataset, document, segment = _graph(sqlite_session) summary = _persist_summary(sqlite_session, dataset, document, segment, content="old") - monkeypatch.setattr(SummaryIndexService, "vectorize_summary", MagicMock(side_effect=RuntimeError("boom"))) + monkeypatch.setattr(SummaryIndexAdapter, "vectorize_summary", MagicMock(side_effect=RuntimeError("boom"))) - result = SummaryIndexService.update_summary_for_segment(segment, dataset, "new", session=sqlite_session) + result = SummaryIndexAdapter.update_summary_for_segment(segment, dataset, "new", session=sqlite_session) assert result is summary assert (result.summary_content, result.status, result.error) == ( @@ -785,7 +823,7 @@ class TestManualUpdate: event.listen(sqlite_session, "before_flush", fail_first_flush) try: with pytest.raises(RuntimeError, match="flush boom"): - SummaryIndexService.update_summary_for_segment(segment, dataset, "new", session=sqlite_session) + SummaryIndexAdapter.update_summary_for_segment(segment, dataset, "new", session=sqlite_session) finally: event.remove(sqlite_session, "before_flush", fail_first_flush) @@ -804,13 +842,13 @@ class TestReadModels: other_segment = _persist_segment(sqlite_session, other_dataset, other_document, segment_id="seg-3") _persist_summary(sqlite_session, other_dataset, other_document, other_segment, summary_id="sum-3") - assert SummaryIndexService.get_segment_summary(first.id, dataset.id, session=sqlite_session) is first_summary - assert SummaryIndexService.get_segment_summary(second.id, dataset.id, session=sqlite_session) is None - assert SummaryIndexService.get_segments_summaries( + assert SummaryIndexAdapter.get_segment_summary(first.id, dataset.id, session=sqlite_session) is first_summary + assert SummaryIndexAdapter.get_segment_summary(second.id, dataset.id, session=sqlite_session) is None + assert SummaryIndexAdapter.get_segments_summaries( [first.id, second.id, other_segment.id], dataset.id, session=sqlite_session ) == {first.id: first_summary} - assert SummaryIndexService.get_segments_summaries([], dataset.id, session=sqlite_session) == {} - assert SummaryIndexService.get_document_summaries( + assert SummaryIndexAdapter.get_segments_summaries([], dataset.id, session=sqlite_session) == {} + assert SummaryIndexAdapter.get_document_summaries( document.id, dataset.id, [first.id, second.id], session=sqlite_session ) == [first_summary] @@ -819,7 +857,7 @@ class TestReadModels: summary = _persist_summary(sqlite_session, dataset, document, segment, status=SummaryStatus.GENERATING) assert ( - SummaryIndexService.get_document_summary_index_status( + SummaryIndexAdapter.get_document_summary_index_status( document.id, dataset.id, TENANT_ID, session=sqlite_session ) == "SUMMARIZING" @@ -827,13 +865,13 @@ class TestReadModels: summary.status = SummaryStatus.COMPLETED sqlite_session.commit() assert ( - SummaryIndexService.get_document_summary_index_status( + SummaryIndexAdapter.get_document_summary_index_status( document.id, dataset.id, TENANT_ID, session=sqlite_session ) is None ) assert ( - SummaryIndexService.get_document_summary_index_status( + SummaryIndexAdapter.get_document_summary_index_status( document.id, dataset.id, OTHER_TENANT_ID, session=sqlite_session ) is None @@ -855,13 +893,13 @@ class TestReadModels: status=SummaryStatus.COMPLETED, ) - result = SummaryIndexService.get_documents_summary_index_status( + result = SummaryIndexAdapter.get_documents_summary_index_status( [first_doc.id, second_doc.id, "missing"], dataset.id, TENANT_ID, session=sqlite_session ) assert result == {first_doc.id: "SUMMARIZING", second_doc.id: None, "missing": None} assert ( - SummaryIndexService.get_documents_summary_index_status([], dataset.id, TENANT_ID, session=sqlite_session) + SummaryIndexAdapter.get_documents_summary_index_status([], dataset.id, TENANT_ID, session=sqlite_session) == {} ) @@ -879,9 +917,13 @@ class TestReadModels: created_at=created_at, ) segment_service = SimpleNamespace(get_segments_by_document_and_dataset=MagicMock(return_value=[first, second])) - monkeypatch.setitem(sys.modules, "services.dataset_service", SimpleNamespace(SegmentService=segment_service)) + monkeypatch.setitem( + sys.modules, "services.knowledge.dataset_service", SimpleNamespace(SegmentService=segment_service) + ) - detail = SummaryIndexService.get_document_summary_status_detail(document.id, dataset.id, sqlite_session) + detail = SummaryIndexAdapter.get_document_summary_status_detail( + document.id, dataset.id, sqlite_session, tenant_id=dataset.tenant_id + ) assert detail["total_segments"] == 2 assert detail["summary_status"]["completed"] == 1 diff --git a/api/tests/unit_tests/services/test_trigger_provider_service.py b/api/tests/unit_tests/services/test_trigger_provider_service.py index 90abc9671d9..bb0507203bf 100644 --- a/api/tests/unit_tests/services/test_trigger_provider_service.py +++ b/api/tests/unit_tests/services/test_trigger_provider_service.py @@ -32,6 +32,7 @@ from models.trigger import ( TriggerSubscription, WorkflowPluginTrigger, ) +from repositories.credentials.query_repository import CredentialQueryRepository from services.trigger import trigger_provider_service as service_module from services.trigger.trigger_provider_service import TriggerProviderService @@ -183,7 +184,12 @@ def test_provider_manager_entities_are_forwarded(mocker, trigger_db: TriggerData def test_list_subscriptions_empty_state(trigger_db: TriggerDatabase) -> None: assert ( - TriggerProviderService.list_trigger_provider_subscriptions(trigger_db.tenant_id, trigger_db.provider_id) == [] + TriggerProviderService.list_trigger_provider_subscriptions( + trigger_db.tenant_id, + trigger_db.provider_id, + credential_query=CredentialQueryRepository(session_factory=trigger_db.session_maker), + ) + == [] ) @@ -231,7 +237,9 @@ def test_list_subscriptions_masks_and_counts_distinct_apps( ) subscriptions = TriggerProviderService.list_trigger_provider_subscriptions( - trigger_db.tenant_id, trigger_db.provider_id + trigger_db.tenant_id, + trigger_db.provider_id, + credential_query=CredentialQueryRepository(session_factory=trigger_db.session_maker), ) assert [item.id for item in subscriptions] == [target.id] diff --git a/api/tests/unit_tests/services/test_vector_service.py b/api/tests/unit_tests/services/test_vector_service.py deleted file mode 100644 index a8914f41f1c..00000000000 --- a/api/tests/unit_tests/services/test_vector_service.py +++ /dev/null @@ -1,790 +0,0 @@ -"""Unit tests for `api/services/vector_service.py`.""" - -from __future__ import annotations - -import logging -from dataclasses import dataclass -from datetime import datetime -from typing import Any -from unittest.mock import MagicMock - -import pytest -from sqlalchemy import event -from sqlalchemy.orm import Session - -import services.vector_service as vector_service_module -from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType -from models import UploadFile -from models.dataset import ChildChunk, Dataset, DatasetProcessRule, DocumentSegment, SegmentAttachmentBinding -from models.dataset import Document as DatasetDocument -from models.enums import DataSourceType, DocumentCreatedFrom, ProcessRuleMode -from services.vector_service import VectorService -from tests.unit_tests.model_factories import make_upload_file - - -@dataclass(frozen=True) -class _ChildDocStub: - page_content: str - metadata: dict[str, Any] - - -@dataclass -class _ParentDocStub: - children: list[_ChildDocStub] - - -def _make_dataset( - *, - indexing_technique: str = IndexTechniqueType.HIGH_QUALITY, - doc_form: str = IndexStructureType.PARAGRAPH_INDEX, - tenant_id: str = "tenant-1", - dataset_id: str = "dataset-1", - is_multimodal: bool = False, - embedding_model_provider: str | None = "openai", - embedding_model: str = "text-embedding", -) -> Dataset: - dataset = Dataset( - id=dataset_id, - tenant_id=tenant_id, - name="Dataset", - created_by="account-1", - indexing_technique=indexing_technique, - chunk_structure=doc_form, - is_multimodal=is_multimodal, - embedding_model_provider=embedding_model_provider, - embedding_model=embedding_model, - ) - return dataset - - -def _make_segment( - *, - segment_id: str = "seg-1", - tenant_id: str = "tenant-1", - dataset_id: str = "dataset-1", - document_id: str = "doc-1", - content: str = "hello", - index_node_id: str = "node-1", - index_node_hash: str = "hash-1", - session: Session | None = None, - attachments: list[dict[str, str]] | None = None, -) -> DocumentSegment: - segment = DocumentSegment( - tenant_id=tenant_id, - dataset_id=dataset_id, - document_id=document_id, - position=1, - content=content, - word_count=len(content), - tokens=len(content), - created_by="account-1", - index_node_id=index_node_id, - index_node_hash=index_node_hash, - ) - segment.id = segment_id - if attachments: - assert session is not None - for attachment in attachments: - upload_file = _upload_file( - file_id=attachment["id"], - name=attachment.get("name", f"{attachment['id']}.png"), - tenant_id=tenant_id, - ) - session.add_all( - [ - upload_file, - SegmentAttachmentBinding( - tenant_id=tenant_id, - dataset_id=dataset_id, - document_id=document_id, - segment_id=segment_id, - attachment_id=upload_file.id, - ), - ] - ) - session.flush() - return segment - - -def _upload_file(*, file_id: str = "file-1", name: str = "img.png", tenant_id: str = "tenant-1") -> UploadFile: - return make_upload_file( - file_id=file_id, - tenant_id=tenant_id, - key=f"uploads/{file_id}", - name=name, - extension="png", - mime_type="image/png", - created_at=datetime(2026, 1, 1), - ) - - -def _make_child_chunk( - *, - index_node_id: str, - content: str = "child", - index_node_hash: str = "hash", - tenant_id: str = "tenant-1", - dataset_id: str = "dataset-1", - document_id: str = "doc-1", - segment_id: str = "seg-1", -) -> ChildChunk: - return ChildChunk( - tenant_id=tenant_id, - dataset_id=dataset_id, - document_id=document_id, - segment_id=segment_id, - position=1, - content=content, - word_count=len(content), - created_by="account-1", - index_node_id=index_node_id, - index_node_hash=index_node_hash, - ) - - -def test_create_segments_vector_regular_indexing_loads_documents_and_keywords( - monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - dataset = _make_dataset(is_multimodal=False) - segment = _make_segment() - - index_processor = MagicMock(name="index_processor") - factory_instance = MagicMock(name="IndexProcessorFactory-instance") - factory_instance.init_index_processor.return_value = index_processor - monkeypatch.setattr(vector_service_module, "IndexProcessorFactory", MagicMock(return_value=factory_instance)) - - VectorService.create_segments_vector( - [["k1"]], [segment], dataset, IndexStructureType.PARAGRAPH_INDEX, session=sqlite_session - ) - - index_processor.load.assert_called_once() - args, kwargs = index_processor.load.call_args - assert args[0] == dataset - assert len(args[1]) == 1 - assert args[2] is None - assert kwargs["with_keywords"] is True - assert kwargs["keywords_list"] == [["k1"]] - - -def test_create_segments_vector_regular_indexing_loads_multimodal_documents( - monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - dataset = _make_dataset(is_multimodal=True) - segment = _make_segment( - session=sqlite_session, - attachments=[ - {"id": "img-1", "name": "a.png"}, - {"id": "img-2", "name": "b.png"}, - ], - ) - - index_processor = MagicMock(name="index_processor") - factory_instance = MagicMock(name="IndexProcessorFactory-instance") - factory_instance.init_index_processor.return_value = index_processor - monkeypatch.setattr(vector_service_module, "IndexProcessorFactory", MagicMock(return_value=factory_instance)) - - VectorService.create_segments_vector( - [["k1"]], [segment], dataset, IndexStructureType.PARAGRAPH_INDEX, session=sqlite_session - ) - - assert index_processor.load.call_count == 2 - first_args, first_kwargs = index_processor.load.call_args_list[0] - assert first_args[0] == dataset - assert len(first_args[1]) == 1 - assert first_kwargs["with_keywords"] is True - - second_args, second_kwargs = index_processor.load.call_args_list[1] - assert second_args[0] == dataset - assert second_args[1] == [] - assert len(second_args[2]) == 2 - assert second_kwargs["with_keywords"] is False - assert {document.page_content for document in second_args[2]} == {"a.png", "b.png"} - - -def test_create_segments_vector_with_no_segments_does_not_load( - monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - dataset = _make_dataset() - index_processor = MagicMock(name="index_processor") - factory_instance = MagicMock() - factory_instance.init_index_processor.return_value = index_processor - monkeypatch.setattr(vector_service_module, "IndexProcessorFactory", MagicMock(return_value=factory_instance)) - - VectorService.create_segments_vector(None, [], dataset, IndexStructureType.PARAGRAPH_INDEX, session=sqlite_session) - index_processor.load.assert_not_called() - - -def _persist_parent_child_rows( - session: Session, - *, - segment: DocumentSegment, - include_document: bool = True, - include_rule: bool = True, -) -> tuple[DatasetDocument | None, DatasetProcessRule | None]: - document = None - rule = None - if include_document: - document = DatasetDocument( - id=segment.document_id, - tenant_id=segment.tenant_id, - dataset_id=segment.dataset_id, - position=1, - data_source_type=DataSourceType.UPLOAD_FILE, - dataset_process_rule_id="rule-1", - batch="batch-1", - name="Document", - created_from=DocumentCreatedFrom.API, - created_by="user-1", - doc_language="en", - ) - session.add(document) - if include_rule: - rule = DatasetProcessRule( - dataset_id=segment.dataset_id, - mode=ProcessRuleMode.HIERARCHICAL, - rules='{"parent_mode":"full-doc"}', - created_by="user-1", - ) - rule.id = "rule-1" - session.add(rule) - session.flush() - return document, rule - - -def test_create_segments_vector_parent_child_calls_generate_child_chunks_with_explicit_model( - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, -) -> None: - dataset = _make_dataset( - doc_form=vector_service_module.IndexStructureType.PARENT_CHILD_INDEX, - embedding_model_provider="openai", - indexing_technique=IndexTechniqueType.HIGH_QUALITY, - ) - segment = _make_segment() - - dataset_document, processing_rule = _persist_parent_child_rows(sqlite_session, segment=segment) - assert dataset_document is not None - assert processing_rule is not None - - embedding_model_instance = MagicMock(name="embedding_model_instance") - model_manager_instance = MagicMock(name="model_manager_instance") - model_manager_instance.get_model_instance.return_value = embedding_model_instance - monkeypatch.setattr( - vector_service_module.ModelManager, "for_tenant", MagicMock(return_value=model_manager_instance) - ) - - generate_child_chunks_mock = MagicMock() - monkeypatch.setattr(VectorService, "generate_child_chunks", generate_child_chunks_mock) - - index_processor = MagicMock() - factory_instance = MagicMock() - factory_instance.init_index_processor.return_value = index_processor - monkeypatch.setattr(vector_service_module, "IndexProcessorFactory", MagicMock(return_value=factory_instance)) - - VectorService.create_segments_vector( - None, - [segment], - dataset, - vector_service_module.IndexStructureType.PARENT_CHILD_INDEX, - session=sqlite_session, - ) - - model_manager_instance.get_model_instance.assert_called_once() - generate_child_chunks_mock.assert_called_once_with( - segment, dataset_document, dataset, embedding_model_instance, processing_rule, False, session=sqlite_session - ) - index_processor.load.assert_not_called() - - -def test_create_segments_vector_parent_child_uses_default_embedding_model_when_provider_missing( - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, -) -> None: - dataset = _make_dataset( - doc_form=vector_service_module.IndexStructureType.PARENT_CHILD_INDEX, - embedding_model_provider=None, - indexing_technique=IndexTechniqueType.HIGH_QUALITY, - ) - segment = _make_segment() - - _persist_parent_child_rows(sqlite_session, segment=segment) - - embedding_model_instance = MagicMock() - model_manager_instance = MagicMock() - model_manager_instance.get_default_model_instance.return_value = embedding_model_instance - monkeypatch.setattr( - vector_service_module.ModelManager, "for_tenant", MagicMock(return_value=model_manager_instance) - ) - - generate_child_chunks_mock = MagicMock() - monkeypatch.setattr(VectorService, "generate_child_chunks", generate_child_chunks_mock) - - index_processor = MagicMock() - factory_instance = MagicMock() - factory_instance.init_index_processor.return_value = index_processor - monkeypatch.setattr(vector_service_module, "IndexProcessorFactory", MagicMock(return_value=factory_instance)) - - VectorService.create_segments_vector( - None, - [segment], - dataset, - vector_service_module.IndexStructureType.PARENT_CHILD_INDEX, - session=sqlite_session, - ) - - model_manager_instance.get_default_model_instance.assert_called_once() - generate_child_chunks_mock.assert_called_once() - - -def test_create_segments_vector_parent_child_missing_document_logs_warning_and_continues( - monkeypatch: pytest.MonkeyPatch, - caplog: pytest.LogCaptureFixture, - sqlite_session: Session, -) -> None: - dataset = _make_dataset(doc_form=vector_service_module.IndexStructureType.PARENT_CHILD_INDEX) - segment = _make_segment() - - index_processor = MagicMock() - factory_instance = MagicMock() - factory_instance.init_index_processor.return_value = index_processor - monkeypatch.setattr(vector_service_module, "IndexProcessorFactory", MagicMock(return_value=factory_instance)) - - with caplog.at_level(logging.WARNING, logger="services.vector_service"): - VectorService.create_segments_vector( - None, - [segment], - dataset, - vector_service_module.IndexStructureType.PARENT_CHILD_INDEX, - session=sqlite_session, - ) - assert any(r.levelno >= logging.WARNING for r in caplog.records) - index_processor.load.assert_not_called() - - -def test_create_segments_vector_parent_child_missing_processing_rule_raises( - monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - dataset = _make_dataset(doc_form=vector_service_module.IndexStructureType.PARENT_CHILD_INDEX) - segment = _make_segment() - - _persist_parent_child_rows(sqlite_session, segment=segment, include_rule=False) - - with pytest.raises(ValueError, match="No processing rule found"): - VectorService.create_segments_vector( - None, - [segment], - dataset, - vector_service_module.IndexStructureType.PARENT_CHILD_INDEX, - session=sqlite_session, - ) - - -def test_create_segments_vector_parent_child_non_high_quality_raises( - monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - dataset = _make_dataset( - doc_form=vector_service_module.IndexStructureType.PARENT_CHILD_INDEX, - indexing_technique=IndexTechniqueType.ECONOMY, - ) - segment = _make_segment() - _persist_parent_child_rows(sqlite_session, segment=segment) - - with pytest.raises(ValueError, match="not high quality"): - VectorService.create_segments_vector( - None, - [segment], - dataset, - vector_service_module.IndexStructureType.PARENT_CHILD_INDEX, - session=sqlite_session, - ) - - -def test_update_segment_vector_high_quality_uses_vector( - monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - dataset = _make_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY) - segment = _make_segment() - - vector_instance = MagicMock() - vector_cls = MagicMock(return_value=vector_instance) - monkeypatch.setattr(vector_service_module, "Vector", vector_cls) - - VectorService.update_segment_vector(["k"], segment, dataset, session=sqlite_session) - - vector_cls.assert_called_once_with(dataset=dataset, session=sqlite_session) - vector_instance.delete_by_ids.assert_called_once_with([segment.index_node_id]) - vector_instance.add_texts.assert_called_once() - add_args, add_kwargs = vector_instance.add_texts.call_args - assert len(add_args[0]) == 1 - assert add_kwargs["duplicate_check"] is True - - -def test_update_segment_vector_economy_uses_keyword_with_keywords_list( - monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - dataset = _make_dataset(indexing_technique=IndexTechniqueType.ECONOMY) - segment = _make_segment() - - keyword_instance = MagicMock() - monkeypatch.setattr(vector_service_module, "Keyword", MagicMock(return_value=keyword_instance)) - - VectorService.update_segment_vector(["a", "b"], segment, dataset, session=sqlite_session) - - keyword_instance.delete_by_ids.assert_called_once_with([segment.index_node_id], sqlite_session) - keyword_instance.add_texts.assert_called_once() - args, kwargs = keyword_instance.add_texts.call_args - assert len(args[0]) == 1 - assert args[1] is sqlite_session - assert kwargs["keywords_list"] == [["a", "b"]] - - -def test_update_segment_vector_economy_uses_keyword_without_keywords_list( - monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - dataset = _make_dataset(indexing_technique=IndexTechniqueType.ECONOMY) - segment = _make_segment() - - keyword_instance = MagicMock() - monkeypatch.setattr(vector_service_module, "Keyword", MagicMock(return_value=keyword_instance)) - - VectorService.update_segment_vector(None, segment, dataset, session=sqlite_session) - keyword_instance.add_texts.assert_called_once() - args, kwargs = keyword_instance.add_texts.call_args - assert len(args[0]) == 1 - assert args[1] is sqlite_session - assert "keywords_list" not in kwargs - - -def test_generate_child_chunks_regenerate_cleans_then_saves_children( - monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - dataset = _make_dataset(doc_form=IndexStructureType.PARAGRAPH_INDEX, tenant_id="tenant-1", dataset_id="dataset-1") - segment = _make_segment(segment_id="seg-1") - - dataset_document = DatasetDocument( - id=segment.document_id, - tenant_id=segment.tenant_id, - dataset_id=segment.dataset_id, - position=1, - data_source_type=DataSourceType.UPLOAD_FILE, - batch="batch-1", - name="Document", - created_from=DocumentCreatedFrom.API, - created_by="user-1", - doc_language="en", - ) - processing_rule = DatasetProcessRule( - dataset_id=segment.dataset_id, - mode=ProcessRuleMode.HIERARCHICAL, - rules="{}", - created_by="user-1", - ) - - child1 = _ChildDocStub(page_content="c1", metadata={"doc_id": "c1-id", "doc_hash": "c1-h"}) - child2 = _ChildDocStub(page_content="c2", metadata={"doc_id": "c2-id", "doc_hash": "c2-h"}) - transformed = [_ParentDocStub(children=[child1, child2])] - - index_processor = MagicMock() - index_processor.transform.return_value = transformed - factory_instance = MagicMock() - factory_instance.init_index_processor.return_value = index_processor - monkeypatch.setattr(vector_service_module, "IndexProcessorFactory", MagicMock(return_value=factory_instance)) - - VectorService.generate_child_chunks( - segment=segment, - dataset_document=dataset_document, - dataset=dataset, - embedding_model_instance=MagicMock(), - processing_rule=processing_rule, - regenerate=True, - session=sqlite_session, - ) - - index_processor.clean.assert_called_once() - _, transform_kwargs = index_processor.transform.call_args - assert transform_kwargs["process_rule"]["rules"]["parent_mode"] == vector_service_module.ParentMode.FULL_DOC - index_processor.load.assert_called_once() - stored = sqlite_session.query(ChildChunk).order_by(ChildChunk.position).all() - assert [chunk.content for chunk in stored] == ["c1", "c2"] - - -def test_generate_child_chunks_flushes_even_when_no_children( - monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - dataset = _make_dataset(doc_form=IndexStructureType.PARAGRAPH_INDEX) - segment = _make_segment() - dataset_document = DatasetDocument( - id=segment.document_id, - tenant_id=segment.tenant_id, - dataset_id=segment.dataset_id, - position=1, - data_source_type=DataSourceType.UPLOAD_FILE, - batch="batch-1", - name="Document", - created_from=DocumentCreatedFrom.API, - created_by="user-1", - doc_language="en", - ) - processing_rule = DatasetProcessRule( - dataset_id=segment.dataset_id, - mode=ProcessRuleMode.HIERARCHICAL, - rules="{}", - created_by="user-1", - ) - - index_processor = MagicMock() - index_processor.transform.return_value = [_ParentDocStub(children=[])] - factory_instance = MagicMock() - factory_instance.init_index_processor.return_value = index_processor - monkeypatch.setattr(vector_service_module, "IndexProcessorFactory", MagicMock(return_value=factory_instance)) - - VectorService.generate_child_chunks( - segment=segment, - dataset_document=dataset_document, - dataset=dataset, - embedding_model_instance=MagicMock(), - processing_rule=processing_rule, - regenerate=False, - session=sqlite_session, - ) - - index_processor.load.assert_not_called() - assert sqlite_session.query(ChildChunk).count() == 0 - - -def test_create_child_chunk_vector_high_quality_adds_texts( - monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - dataset = _make_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY) - child_chunk = _make_child_chunk(index_node_id="id", index_node_hash="h") - - vector_instance = MagicMock() - vector_cls = MagicMock(return_value=vector_instance) - monkeypatch.setattr(vector_service_module, "Vector", vector_cls) - - VectorService.create_child_chunk_vector(child_chunk, dataset, session=sqlite_session) - vector_cls.assert_called_once_with(dataset=dataset, session=sqlite_session) - vector_instance.add_texts.assert_called_once() - - -def test_create_child_chunk_vector_economy_noop(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None: - dataset = _make_dataset(indexing_technique=IndexTechniqueType.ECONOMY) - vector_cls = MagicMock() - monkeypatch.setattr(vector_service_module, "Vector", vector_cls) - - child_chunk = _make_child_chunk(index_node_id="id", index_node_hash="h") - - VectorService.create_child_chunk_vector(child_chunk, dataset, session=sqlite_session) - vector_cls.assert_not_called() - - -def test_update_child_chunk_vector_high_quality_updates_vector( - monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - dataset = _make_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY) - - new_chunk = _make_child_chunk( - content="n", index_node_id="nid", index_node_hash="nh", document_id="d", dataset_id="ds" - ) - upd_chunk = _make_child_chunk( - content="u", index_node_id="uid", index_node_hash="uh", document_id="d", dataset_id="ds" - ) - del_chunk = _make_child_chunk(index_node_id="did") - - vector_instance = MagicMock() - vector_cls = MagicMock(return_value=vector_instance) - monkeypatch.setattr(vector_service_module, "Vector", vector_cls) - - VectorService.update_child_chunk_vector([new_chunk], [upd_chunk], [del_chunk], dataset, session=sqlite_session) - - vector_cls.assert_called_once_with(dataset=dataset, session=sqlite_session) - vector_instance.delete_by_ids.assert_called_once_with(["uid", "did"]) - vector_instance.add_texts.assert_called_once() - docs = vector_instance.add_texts.call_args.args[0] - assert len(docs) == 2 - - -def test_update_child_chunk_vector_economy_noop(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None: - dataset = _make_dataset(indexing_technique=IndexTechniqueType.ECONOMY) - vector_cls = MagicMock() - monkeypatch.setattr(vector_service_module, "Vector", vector_cls) - VectorService.update_child_chunk_vector([], [], [], dataset, session=sqlite_session) - vector_cls.assert_not_called() - - -def test_delete_child_chunk_vector_deletes_by_id(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None: - dataset = _make_dataset() - child_chunk = _make_child_chunk(index_node_id="cid") - - vector_instance = MagicMock() - vector_cls = MagicMock(return_value=vector_instance) - monkeypatch.setattr(vector_service_module, "Vector", vector_cls) - - VectorService.delete_child_chunk_vector(child_chunk, dataset, session=sqlite_session) - vector_cls.assert_called_once_with(dataset=dataset, session=sqlite_session) - vector_instance.delete_by_ids.assert_called_once_with(["cid"]) - - -# --------------------------------------------------------------------------- -# update_multimodel_vector (missing coverage in previous suites) -# --------------------------------------------------------------------------- - - -def test_update_multimodel_vector_returns_when_not_high_quality( - monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - dataset = _make_dataset(indexing_technique=IndexTechniqueType.ECONOMY, is_multimodal=True) - segment = _make_segment(tenant_id="t") - - vector_cls = MagicMock() - monkeypatch.setattr(vector_service_module, "Vector", vector_cls) - - VectorService.update_multimodel_vector( - session=sqlite_session, segment=segment, attachment_ids=["a"], dataset=dataset - ) - vector_cls.assert_not_called() - assert not sqlite_session.in_transaction() - - -def test_update_multimodel_vector_returns_when_no_actual_change( - monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - dataset = _make_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY, is_multimodal=True) - segment = _make_segment(tenant_id="t", session=sqlite_session, attachments=[{"id": "a"}, {"id": "b"}]) - - vector_cls = MagicMock() - monkeypatch.setattr(vector_service_module, "Vector", vector_cls) - - VectorService.update_multimodel_vector( - session=sqlite_session, segment=segment, attachment_ids=["b", "a"], dataset=dataset - ) - vector_cls.assert_not_called() - assert sqlite_session.in_transaction() - - -def test_update_multimodel_vector_deletes_bindings_and_commits_on_empty_new_ids( - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, -) -> None: - dataset = _make_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY, is_multimodal=True) - segment = _make_segment( - tenant_id="tenant-1", - session=sqlite_session, - attachments=[{"id": "old-1"}, {"id": "old-2"}], - ) - - vector_instance = MagicMock(name="vector_instance") - vector_cls = MagicMock(return_value=vector_instance) - - monkeypatch.setattr(vector_service_module, "Vector", vector_cls) - - VectorService.update_multimodel_vector(segment=segment, attachment_ids=[], dataset=dataset, session=sqlite_session) - - vector_cls.assert_called_once_with(dataset=dataset, session=sqlite_session) - vector_instance.delete_by_ids.assert_called_once_with(["old-1", "old-2"]) - assert sqlite_session.query(SegmentAttachmentBinding).count() == 0 - vector_instance.add_texts.assert_not_called() - - -def test_update_multimodel_vector_flushes_when_no_upload_files_found( - monkeypatch: pytest.MonkeyPatch, sqlite_session: Session -) -> None: - dataset = _make_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY, is_multimodal=True) - segment = _make_segment(tenant_id="tenant-1", session=sqlite_session, attachments=[{"id": "old-1"}]) - - vector_instance = MagicMock() - monkeypatch.setattr(vector_service_module, "Vector", MagicMock(return_value=vector_instance)) - VectorService.update_multimodel_vector( - session=sqlite_session, segment=segment, attachment_ids=["new-1"], dataset=dataset - ) - - assert sqlite_session.query(SegmentAttachmentBinding).count() == 0 - vector_instance.add_texts.assert_not_called() - - -def test_update_multimodel_vector_adds_bindings_and_vectors_and_skips_missing_upload_files( - monkeypatch: pytest.MonkeyPatch, - caplog: pytest.LogCaptureFixture, - sqlite_session: Session, -) -> None: - dataset = _make_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY, is_multimodal=True) - segment = _make_segment( - segment_id="seg-1", - tenant_id="tenant-1", - session=sqlite_session, - attachments=[{"id": "old-1"}], - ) - - vector_instance = MagicMock() - monkeypatch.setattr(vector_service_module, "Vector", MagicMock(return_value=vector_instance)) - sqlite_session.add(_upload_file()) - sqlite_session.flush() - - with caplog.at_level(logging.WARNING, logger="services.vector_service"): - VectorService.update_multimodel_vector( - session=sqlite_session, segment=segment, attachment_ids=["file-1", "missing"], dataset=dataset - ) - assert any(r.levelno >= logging.WARNING for r in caplog.records) - bindings = sqlite_session.query(SegmentAttachmentBinding).all() - assert len(bindings) == 1 - assert bindings[0].attachment_id == "file-1" - - vector_instance.create_multimodal.assert_called_once() - documents = vector_instance.create_multimodal.call_args.args[0] - assert len(documents) == 1 - assert documents[0].page_content == "img.png" - assert documents[0].metadata["doc_id"] == "file-1" - - -def test_update_multimodel_vector_updates_bindings_without_multimodal_vector_ops( - monkeypatch: pytest.MonkeyPatch, - sqlite_session: Session, -) -> None: - dataset = _make_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY, is_multimodal=False) - segment = _make_segment(tenant_id="tenant-1", session=sqlite_session, attachments=[{"id": "old-1"}]) - - vector_instance = MagicMock() - monkeypatch.setattr(vector_service_module, "Vector", MagicMock(return_value=vector_instance)) - sqlite_session.add(_upload_file()) - sqlite_session.flush() - - VectorService.update_multimodel_vector( - session=sqlite_session, segment=segment, attachment_ids=["file-1"], dataset=dataset - ) - - vector_instance.delete_by_ids.assert_not_called() - vector_instance.add_texts.assert_not_called() - binding = sqlite_session.query(SegmentAttachmentBinding).one() - assert binding.attachment_id == "file-1" - - -def test_update_multimodel_vector_rolls_back_and_reraises_on_error( - monkeypatch: pytest.MonkeyPatch, - caplog: pytest.LogCaptureFixture, - sqlite_session: Session, -) -> None: - dataset = _make_dataset(indexing_technique=IndexTechniqueType.HIGH_QUALITY, is_multimodal=True) - segment = _make_segment( - segment_id="seg-1", - tenant_id="tenant-1", - session=sqlite_session, - attachments=[{"id": "old-1"}], - ) - - vector_instance = MagicMock() - monkeypatch.setattr(vector_service_module, "Vector", MagicMock(return_value=vector_instance)) - sqlite_session.add(_upload_file()) - sqlite_session.flush() - rollback_events: list[str] = [] - event.listen(sqlite_session, "after_rollback", lambda _session: rollback_events.append("rollback")) - monkeypatch.setattr(sqlite_session, "flush", MagicMock(side_effect=RuntimeError("boom"))) - - with caplog.at_level(logging.ERROR, logger="services.vector_service"): - with sqlite_session.no_autoflush: - with pytest.raises(RuntimeError, match="boom"): - VectorService.update_multimodel_vector( - session=sqlite_session, segment=segment, attachment_ids=["file-1"], dataset=dataset - ) - - assert any(r.levelno >= logging.ERROR for r in caplog.records) - assert rollback_events == ["rollback"] diff --git a/api/tests/unit_tests/services/test_website_service.py b/api/tests/unit_tests/services/test_website_service.py index 571e8e99903..cde5d32a8ab 100644 --- a/api/tests/unit_tests/services/test_website_service.py +++ b/api/tests/unit_tests/services/test_website_service.py @@ -1,4 +1,4 @@ -"""Unit tests for services.website_service. +"""Unit tests for services.data_source.website_service. Focuses on provider dispatching, argument validation, and provider-specific branches without making any real network/storage/redis calls. @@ -10,12 +10,13 @@ import json from dataclasses import dataclass from datetime import UTC, datetime from typing import Any -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock, create_autospec, patch import pytest -import services.website_service as website_service_module -from services.website_service import ( +import services.data_source.website_service as website_service_module +from services.data_source.credential_gateway import DatasourceProviderCredentialStore +from services.data_source.website_service import ( CrawlOptions, WebsiteCrawlApiRequest, WebsiteCrawlStatusApiRequest, @@ -111,7 +112,11 @@ def test_get_credentials_and_config_selects_plugin_id_and_key_firecrawl(monkeypa service_instance.get_datasource_credentials.return_value = {"firecrawl_api_key": "k", "base_url": "b"} monkeypatch.setattr(website_service_module, "DatasourceProviderService", MagicMock(return_value=service_instance)) - api_key, config = WebsiteService._get_credentials_and_config("tenant-1", "firecrawl") + api_key, config = WebsiteService( + providers=website_service_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ) + )._get_credentials_and_config("tenant-1", "firecrawl") assert api_key == "k" assert config["base_url"] == "b" @@ -136,7 +141,11 @@ def test_get_credentials_and_config_selects_plugin_id_and_key_api_key( service_instance.get_datasource_credentials.return_value = {"api_key": "enc-key", "base_url": "b"} monkeypatch.setattr(website_service_module, "DatasourceProviderService", MagicMock(return_value=service_instance)) - api_key, config = WebsiteService._get_credentials_and_config("tenant-1", provider) + api_key, config = WebsiteService( + providers=website_service_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ) + )._get_credentials_and_config("tenant-1", provider) assert api_key == "enc-key" assert config["base_url"] == "b" @@ -149,7 +158,11 @@ def test_get_credentials_and_config_selects_plugin_id_and_key_api_key( def test_get_credentials_and_config_rejects_invalid_provider() -> None: with pytest.raises(ValueError, match="Invalid provider"): - WebsiteService._get_credentials_and_config("tenant-1", "unknown") + WebsiteService( + providers=website_service_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ) + )._get_credentials_and_config("tenant-1", "unknown") def test_get_credentials_and_config_hits_unreachable_guard_branch(monkeypatch: pytest.MonkeyPatch) -> None: @@ -174,7 +187,11 @@ def test_get_credentials_and_config_hits_unreachable_guard_branch(monkeypatch: p monkeypatch.setattr(website_service_module, "DatasourceProviderService", MagicMock(return_value=service_instance)) with pytest.raises(ValueError, match="Invalid provider"): - WebsiteService._get_credentials_and_config("tenant-1", FlakyProvider()) # type: ignore[arg-type] + WebsiteService( + providers=website_service_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ) + )._get_credentials_and_config("tenant-1", FlakyProvider()) # type: ignore[arg-type] def test_get_decrypted_api_key_requires_api_key(monkeypatch: pytest.MonkeyPatch) -> None: @@ -204,7 +221,11 @@ def test_crawl_url_dispatches_by_provider(monkeypatch: pytest.MonkeyPatch) -> No firecrawl_mock = MagicMock(return_value={"status": "active", "job_id": "j1"}) monkeypatch.setattr(WebsiteService, "_crawl_with_firecrawl", firecrawl_mock) - result = WebsiteService.crawl_url(api_request) + result = WebsiteService( + providers=website_service_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ) + ).crawl_url(api_request) assert result == {"status": "active", "job_id": "j1"} firecrawl_mock.assert_called_once() @@ -225,7 +246,11 @@ def test_crawl_url_dispatches_other_providers(monkeypatch: pytest.MonkeyPatch, p impl_mock = MagicMock(return_value={"status": "active"}) monkeypatch.setattr(WebsiteService, method_name, impl_mock) - assert WebsiteService.crawl_url(api_request) == {"status": "active"} + assert WebsiteService( + providers=website_service_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ) + ).crawl_url(api_request) == {"status": "active"} impl_mock.assert_called_once() @@ -234,7 +259,11 @@ def test_crawl_url_rejects_invalid_provider(monkeypatch: pytest.MonkeyPatch) -> monkeypatch.setattr(WebsiteService, "_get_credentials_and_config", MagicMock(return_value=("k", {}))) with pytest.raises(ValueError, match="Invalid provider"): - WebsiteService.crawl_url(api_request) + WebsiteService( + providers=website_service_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ) + ).crawl_url(api_request) def test_crawl_with_firecrawl_builds_params_single_page_and_sets_redis(monkeypatch: pytest.MonkeyPatch) -> None: @@ -406,25 +435,41 @@ def test_get_crawl_status_dispatches(monkeypatch: pytest.MonkeyPatch) -> None: firecrawl_status = MagicMock(return_value={"status": "active"}) monkeypatch.setattr(WebsiteService, "_get_firecrawl_status", firecrawl_status) - result = WebsiteService.get_crawl_status("job-1", "firecrawl") + result = WebsiteService( + providers=website_service_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ) + ).get_crawl_status("job-1", "firecrawl") assert result == {"status": "active"} firecrawl_status.assert_called_once_with("job-1", "k", {"base_url": "b"}) watercrawl_status = MagicMock(return_value={"status": "active", "job_id": "w"}) monkeypatch.setattr(WebsiteService, "_get_watercrawl_status", watercrawl_status) - assert WebsiteService.get_crawl_status("job-2", "watercrawl") == {"status": "active", "job_id": "w"} + assert WebsiteService( + providers=website_service_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ) + ).get_crawl_status("job-2", "watercrawl") == {"status": "active", "job_id": "w"} watercrawl_status.assert_called_once_with("job-2", "k", {"base_url": "b"}) jinareader_status = MagicMock(return_value={"status": "active", "job_id": "j"}) monkeypatch.setattr(WebsiteService, "_get_jinareader_status", jinareader_status) - assert WebsiteService.get_crawl_status("job-3", "jinareader") == {"status": "active", "job_id": "j"} + assert WebsiteService( + providers=website_service_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ) + ).get_crawl_status("job-3", "jinareader") == {"status": "active", "job_id": "j"} jinareader_status.assert_called_once_with("job-3", "k") def test_get_crawl_status_typed_rejects_invalid_provider(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(WebsiteService, "_get_credentials_and_config", MagicMock(return_value=("k", {}))) with pytest.raises(ValueError, match="Invalid provider"): - WebsiteService.get_crawl_status_typed(WebsiteCrawlStatusApiRequest(provider="bad", job_id="j")) + WebsiteService( + providers=website_service_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ) + ).get_crawl_status_typed(WebsiteCrawlStatusApiRequest(provider="bad", job_id="j")) def test_get_firecrawl_status_adds_time_consuming_when_completed_and_cached(monkeypatch: pytest.MonkeyPatch) -> None: @@ -532,7 +577,11 @@ def test_get_jinareader_status_completed_formats_processed_items(monkeypatch: py def test_get_crawl_url_data_dispatches_invalid_provider() -> None: with pytest.raises(ValueError, match="Invalid provider"): - WebsiteService.get_crawl_url_data("job-1", "bad", "https://example.com", "tenant-1") + WebsiteService( + providers=website_service_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ) + ).get_crawl_url_data("job-1", "bad", "https://example.com", "tenant-1") def test_get_crawl_url_data_hits_invalid_provider_branch_when_credentials_stubbed( @@ -540,7 +589,11 @@ def test_get_crawl_url_data_hits_invalid_provider_branch_when_credentials_stubbe ) -> None: monkeypatch.setattr(WebsiteService, "_get_credentials_and_config", MagicMock(return_value=("k", {}))) with pytest.raises(ValueError, match="Invalid provider"): - WebsiteService.get_crawl_url_data("job-1", object(), "u", "tenant-1") # type: ignore[arg-type] + WebsiteService( + providers=website_service_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ) + ).get_crawl_url_data("job-1", object(), "u", "tenant-1") # type: ignore[arg-type] @pytest.mark.parametrize( @@ -556,7 +609,11 @@ def test_get_crawl_url_data_dispatches(monkeypatch: pytest.MonkeyPatch, provider impl_mock = MagicMock(return_value={"ok": True}) monkeypatch.setattr(WebsiteService, method_name, impl_mock) - result = WebsiteService.get_crawl_url_data("job-1", provider, "u", "tenant-1") + result = WebsiteService( + providers=website_service_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ) + ).get_crawl_url_data("job-1", provider, "u", "tenant-1") assert result == {"ok": True} impl_mock.assert_called_once() @@ -676,16 +733,28 @@ def test_get_scrape_url_data_dispatches_and_rejects_invalid_provider(monkeypatch scrape_mock = MagicMock(return_value={"data": "x"}) monkeypatch.setattr(WebsiteService, "_scrape_with_firecrawl", scrape_mock) - assert WebsiteService.get_scrape_url_data("firecrawl", "u", "tenant-1", True) == {"data": "x"} + assert WebsiteService( + providers=website_service_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ) + ).get_scrape_url_data("firecrawl", "u", "tenant-1", True) == {"data": "x"} scrape_mock.assert_called_once() watercrawl_mock = MagicMock(return_value={"data": "y"}) monkeypatch.setattr(WebsiteService, "_scrape_with_watercrawl", watercrawl_mock) - assert WebsiteService.get_scrape_url_data("watercrawl", "u", "tenant-1", False) == {"data": "y"} + assert WebsiteService( + providers=website_service_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ) + ).get_scrape_url_data("watercrawl", "u", "tenant-1", False) == {"data": "y"} watercrawl_mock.assert_called_once() with pytest.raises(ValueError, match="Invalid provider"): - WebsiteService.get_scrape_url_data("jinareader", "u", "tenant-1", True) + WebsiteService( + providers=website_service_module.DatasourceProviderService( + credentials=create_autospec(DatasourceProviderCredentialStore, instance=True) + ) + ).get_scrape_url_data("jinareader", "u", "tenant-1", True) def test_scrape_with_firecrawl_calls_app(monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/api/tests/unit_tests/services/tools/test_builtin_tools_manage_service.py b/api/tests/unit_tests/services/tools/test_builtin_tools_manage_service.py index bcaabd7b815..4fef00d32a4 100644 --- a/api/tests/unit_tests/services/tools/test_builtin_tools_manage_service.py +++ b/api/tests/unit_tests/services/tools/test_builtin_tools_manage_service.py @@ -6,10 +6,11 @@ from unittest.mock import MagicMock import pytest from sqlalchemy import select -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, sessionmaker from core.plugin.entities.plugin_daemon import CredentialType from models.tools import BuiltinToolProvider, ToolOAuthSystemClient, ToolOAuthTenantClient +from repositories.credentials.query_repository import CredentialQueryRepository from services.tools import builtin_tools_manage_service as service_module from services.tools.builtin_tools_manage_service import BuiltinToolManageService @@ -391,7 +392,11 @@ class TestGetBuiltinToolProviderCredentialInfo: ) result = BuiltinToolManageService.get_builtin_tool_provider_credential_info( - "tenant-1", "google", session=repository_session + "tenant-1", + "google", + credential_query=CredentialQueryRepository( + session_factory=sessionmaker(bind=repository_session.get_bind(), expire_on_commit=False) + ), ) assert result.credentials == [] @@ -404,7 +409,11 @@ class TestGetBuiltinToolProviderCredentials: _persist_provider(repository_session, credential_id="other", tenant_id="other-tenant") result = BuiltinToolManageService.get_builtin_tool_provider_credentials( - "tenant-1", "google", session=repository_session + "tenant-1", + "google", + credential_query=CredentialQueryRepository( + session_factory=sessionmaker(bind=repository_session.get_bind(), expire_on_commit=False) + ), ) assert result == [] @@ -426,23 +435,23 @@ class TestGetBuiltinToolProviderCredentials: "create_tool_encrypter", MagicMock(return_value=(encrypter, MagicMock())), ) - credential_entity = MagicMock() - convert = MagicMock(return_value=credential_entity) - monkeypatch.setattr( - service_module.ToolTransformService, - "convert_builtin_provider_to_credential_entity", - convert, - ) result = BuiltinToolManageService.get_builtin_tool_provider_credentials( - "tenant-1", "google", session=repository_session + "tenant-1", + "google", + credential_query=CredentialQueryRepository( + session_factory=sessionmaker(bind=repository_session.get_bind(), expire_on_commit=False) + ), ) - assert result == [credential_entity] - converted_provider = convert.call_args.kwargs["provider"] - assert isinstance(converted_provider, BuiltinToolProvider) - assert converted_provider.tenant_id == "tenant-1" - assert converted_provider.is_default is True + assert len(result) == 1 + assert result[0].id == "cred-1" + assert result[0].credentials == {"key": "***"} + assert result[0].is_default is True + repository_session.expire_all() + stored = repository_session.get(BuiltinToolProvider, "cred-1") + assert stored is not None + assert stored.is_default is False class TestGetBuiltinProvider: diff --git a/api/tests/unit_tests/tasks/test_batch_clean_document_task.py b/api/tests/unit_tests/tasks/test_batch_clean_document_task.py index 8f80c5a7ec2..588148d1c60 100644 --- a/api/tests/unit_tests/tasks/test_batch_clean_document_task.py +++ b/api/tests/unit_tests/tasks/test_batch_clean_document_task.py @@ -53,7 +53,7 @@ def test_successful_vector_cleanup_schedules_billing_refresh(cleanup_rows: tuple with ( patch("tasks.batch_clean_document_task.get_image_upload_file_ids", return_value=[]), - patch("tasks.batch_clean_document_task.IndexProcessorFactory") as processor_factory, + patch("tasks.batch_clean_document_task.clean_document_indexes", return_value=tenant_id) as index_cleanup, patch("tasks.batch_clean_document_task.schedule_billing_vector_space_refresh") as schedule_refresh, ): batch_clean_document_task( @@ -63,7 +63,7 @@ def test_successful_vector_cleanup_schedules_billing_refresh(cleanup_rows: tuple file_ids=[], ) - processor_factory.return_value.init_index_processor.return_value.clean.assert_called_once() + index_cleanup.assert_called_once() schedule_refresh.assert_called_once_with(tenant_id) @@ -72,12 +72,10 @@ def test_failed_vector_cleanup_does_not_schedule_billing_refresh(cleanup_rows: t with ( patch("tasks.batch_clean_document_task.get_image_upload_file_ids", return_value=[]), - patch("tasks.batch_clean_document_task.IndexProcessorFactory") as processor_factory, + patch("tasks.batch_clean_document_task.clean_document_indexes") as index_cleanup, patch("tasks.batch_clean_document_task.schedule_billing_vector_space_refresh") as schedule_refresh, ): - processor_factory.return_value.init_index_processor.return_value.clean.side_effect = RuntimeError( - "vector cleanup failed" - ) + index_cleanup.side_effect = RuntimeError("vector cleanup failed") batch_clean_document_task( document_ids=[document_id], dataset_id=dataset_id, @@ -119,7 +117,7 @@ def test_cleans_segment_attachment_bindings_and_files(cleanup_rows: tuple[str, s with ( patch("tasks.batch_clean_document_task.get_image_upload_file_ids", return_value=[]), - patch("tasks.batch_clean_document_task.IndexProcessorFactory"), + patch("tasks.batch_clean_document_task.clean_document_indexes"), patch("tasks.batch_clean_document_task.schedule_billing_vector_space_refresh"), patch("tasks.batch_clean_document_task.storage.delete") as storage_delete, ): diff --git a/api/tests/unit_tests/tasks/test_batch_create_segment_to_index_task.py b/api/tests/unit_tests/tasks/test_batch_create_segment_to_index_task.py index c3ea579e365..5df433759d7 100644 --- a/api/tests/unit_tests/tasks/test_batch_create_segment_to_index_task.py +++ b/api/tests/unit_tests/tasks/test_batch_create_segment_to_index_task.py @@ -3,6 +3,7 @@ from pathlib import Path from unittest.mock import patch from uuid import uuid4 +import pytest from sqlalchemy import event, select from sqlalchemy.engine import Engine from sqlalchemy.orm import Session @@ -12,12 +13,15 @@ from extensions.storage.storage_type import StorageType from models.dataset import Dataset, Document, DocumentSegment from models.enums import CreatorUserRole, DataSourceType, DocumentCreatedFrom, IndexingStatus, SegmentStatus from models.model import UploadFile +from services.knowledge.resource_scope import DatasetRef, DocumentRef from tasks.batch_create_segment_to_index_task import batch_create_segment_to_index_task -def test_batch_create_segment_to_index_task_queries_max_position_once_per_batch( +@pytest.mark.parametrize("foreign_owner", [None, "dataset", "document_tenant", "document_dataset", "upload"]) +def test_batch_create_segment_to_index_task_validates_scope_and_indexes_committed_batch( sqlite_session: Session, sqlite_engine: Engine, + foreign_owner: str | None, ) -> None: tenant_id = str(uuid4()) user_id = str(uuid4()) @@ -71,6 +75,14 @@ def test_batch_create_segment_to_index_task_queries_max_position_once_per_batch( index_node_hash="existing-hash", status=SegmentStatus.COMPLETED, ) + if foreign_owner == "dataset": + dataset.tenant_id = str(uuid4()) + elif foreign_owner == "document_tenant": + document.tenant_id = str(uuid4()) + elif foreign_owner == "document_dataset": + document.dataset_id = str(uuid4()) + elif foreign_owner == "upload": + upload_file.tenant_id = str(uuid4()) sqlite_session.add_all([dataset, document, upload_file, existing_segment]) sqlite_session.commit() @@ -94,8 +106,8 @@ def test_batch_create_segment_to_index_task_queries_max_position_once_per_batch( event.listen(sqlite_engine, "before_cursor_execute", count_max_position_query) try: with ( - patch("tasks.batch_create_segment_to_index_task.storage.download", side_effect=mock_download), - patch("tasks.batch_create_segment_to_index_task.VectorService.create_segments_vector"), + patch("tasks.batch_create_segment_to_index_task.storage.download", side_effect=mock_download) as download, + patch("extensions.ext_application_services.application_services") as services, ): batch_create_segment_to_index_task( job_id=str(uuid4()), @@ -108,6 +120,13 @@ def test_batch_create_segment_to_index_task_queries_max_position_once_per_batch( finally: event.remove(sqlite_engine, "before_cursor_execute", count_max_position_query) + if foreign_owner is not None: + download.assert_not_called() + services.assert_not_called() + assert not max_position_queries + assert sqlite_session.scalars(select(DocumentSegment.id)).all() == [existing_segment.id] + return + positions = sqlite_session.scalars( select(DocumentSegment.position) .where(DocumentSegment.document_id == document.id) @@ -115,3 +134,7 @@ def test_batch_create_segment_to_index_task_queries_max_position_once_per_batch( ).all() assert positions == [7, 8, 9, 10] assert len(max_position_queries) == 1 + indexing = services.return_value.knowledge.segments.mutations.index_segments + indexing.assert_called_once() + assert indexing.call_args.args == (DocumentRef(DatasetRef(tenant_id, dataset.id), document.id),) + assert len(indexing.call_args.kwargs["segment_ids"]) == 3 diff --git a/api/tests/unit_tests/tasks/test_clean_document_task.py b/api/tests/unit_tests/tasks/test_clean_document_task.py index 259dfd6e6a1..a9d7b646c4f 100644 --- a/api/tests/unit_tests/tasks/test_clean_document_task.py +++ b/api/tests/unit_tests/tasks/test_clean_document_task.py @@ -6,7 +6,7 @@ starts from the production incident shape: the caller has already deleted the """ import uuid -from unittest.mock import MagicMock, patch +from unittest.mock import patch import pytest from sqlalchemy import select @@ -71,20 +71,10 @@ def mock_storage(): @pytest.fixture -def mock_index_processor_factory(): - """Mock the vector/index boundary so cleanup behavior is deterministic.""" - with patch("tasks.clean_document_task.IndexProcessorFactory", autospec=True) as factory_cls: - processor = MagicMock() - processor.clean.return_value = None - factory_instance = MagicMock() - factory_instance.init_index_processor.return_value = processor - factory_cls.return_value = factory_instance - - yield { - "factory_cls": factory_cls, - "factory_instance": factory_instance, - "processor": processor, - } +def mock_index_cleanup(tenant_id: str): + """Keep task resilience tests independent of the external index backend.""" + with patch("tasks.clean_document_task.clean_document_indexes", return_value=tenant_id) as cleanup: + yield cleanup def _document(*, document_id: str, dataset_id: str, tenant_id: str, created_by: str) -> Document: @@ -209,7 +199,7 @@ class TestVectorCleanupResilience: sqlite_session: Session, bind_task_sessions: None, mock_storage, - mock_index_processor_factory, + mock_index_cleanup, ) -> None: """A transient billing failure leaves only the unrelated document's rows.""" other_document_id, survivor_segment_id = _persist_deleted_document_state( @@ -219,7 +209,7 @@ class TestVectorCleanupResilience: tenant_id=tenant_id, target_segment_ids=["seg-1", "seg-2"], ) - mock_index_processor_factory["processor"].clean.side_effect = ValueError( + mock_index_cleanup.side_effect = ValueError( "Unable to retrieve billing information. Please try again later or contact support." ) @@ -232,7 +222,7 @@ class TestVectorCleanupResilience: file_id=None, ) - mock_index_processor_factory["processor"].clean.assert_called_once() + mock_index_cleanup.assert_called_once() _assert_relational_cleanup( sqlite_session, document_id=document_id, @@ -249,7 +239,7 @@ class TestVectorCleanupResilience: sqlite_session: Session, bind_task_sessions: None, mock_storage, - mock_index_processor_factory, + mock_index_cleanup, ) -> None: """The happy path calls the index boundary and completes scoped cleanup.""" other_document_id, survivor_segment_id = _persist_deleted_document_state( @@ -268,15 +258,15 @@ class TestVectorCleanupResilience: file_id=None, ) - mock_index_processor_factory["processor"].clean.assert_called_once() - _, kwargs = mock_index_processor_factory["processor"].clean.call_args - cleanup_session = kwargs.pop("session") - assert isinstance(cleanup_session, Session) - assert cleanup_session.get_bind() is sqlite_session.get_bind() + mock_index_cleanup.assert_called_once() + _, kwargs = mock_index_cleanup.call_args + new_session = kwargs.pop("new_session") + with new_session() as cleanup_session: + assert cleanup_session.get_bind() is sqlite_session.get_bind() assert kwargs == { - "with_keywords": True, - "delete_child_chunks": True, - "delete_summaries": True, + "dataset_id": dataset_id, + "document_ids": [document_id], + "doc_form": "paragraph", } _assert_relational_cleanup( sqlite_session, @@ -286,7 +276,7 @@ class TestVectorCleanupResilience: ) schedule_refresh.assert_called_once_with(tenant_id) - def test_no_segments_skips_vector_cleanup( + def test_no_index_state_does_not_refresh_billing( self, document_id: str, dataset_id: str, @@ -294,9 +284,10 @@ class TestVectorCleanupResilience: sqlite_session: Session, bind_task_sessions: None, mock_storage, - mock_index_processor_factory, + mock_index_cleanup, ) -> None: - """A target document without segments skips the vector/index boundary.""" + """Cleanup checks for orphan summaries even when no segments remain.""" + mock_index_cleanup.return_value = None other_document_id, survivor_segment_id = _persist_deleted_document_state( sqlite_session, document_id=document_id, @@ -313,7 +304,7 @@ class TestVectorCleanupResilience: file_id=None, ) - mock_index_processor_factory["factory_cls"].assert_not_called() + mock_index_cleanup.assert_called_once() _assert_relational_cleanup( sqlite_session, document_id=document_id, diff --git a/api/tests/unit_tests/tasks/test_dataset_indexing_task.py b/api/tests/unit_tests/tasks/test_dataset_indexing_task.py index 7ef6aaac1bf..3b9b5103947 100644 --- a/api/tests/unit_tests/tasks/test_dataset_indexing_task.py +++ b/api/tests/unit_tests/tasks/test_dataset_indexing_task.py @@ -8,18 +8,19 @@ ORM rows so each phase observes only committed database state. import uuid from contextlib import nullcontext from types import SimpleNamespace -from unittest.mock import MagicMock, Mock, patch +from unittest.mock import ANY, MagicMock, Mock, patch import pytest from sqlalchemy.orm import Session -from core.indexing_runner import DocumentIsPausedError from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType from enums import CloudPlan, DeploymentEdition from extensions.ext_redis import redis_client from models.dataset import Dataset, Document from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus from services.document_indexing_proxy.document_indexing_task_proxy import DocumentIndexingTaskProxy +from services.knowledge.indexing.errors import DocumentIsPausedError +from services.knowledge.resource_scope import DocumentRef from tasks.document_indexing_task import ( _document_indexing, _document_indexing_with_tenant_queue, @@ -27,9 +28,50 @@ from tasks.document_indexing_task import ( normal_document_indexing_task, priority_document_indexing_task, ) +from tasks.recover_document_indexing_task import recover_document_indexing_task from tests.unit_tests.config_override import apply_config_overrides, config_overrides_context +@pytest.mark.parametrize( + ("status", "entry"), + [ + ("waiting", "run"), + ("parsing", "run"), + ("cleaning", "run"), + ("splitting", "run_in_splitting_status"), + ("indexing", "run_in_indexing_status"), + ("completed", None), + ], +) +def test_recovery_dispatches_detached_owner_reference( + sqlite_session, tenant_id, dataset_id, document_ids, status, entry +): + _persist_indexing_rows(sqlite_session, tenant_id=tenant_id, dataset_id=dataset_id, document_ids=document_ids[:1]) + row = sqlite_session.get(Document, document_ids[0]) + row.indexing_status = status + sqlite_session.commit() + with patch("tasks.recover_document_indexing_task.build_document_indexing_service") as factory: + recover_document_indexing_task.run(dataset_id, document_ids[0]) + service = factory.return_value + if entry is None: + assert service.mock_calls == [] + else: + method = { + "run": service.run, + "run_in_splitting_status": service.run_in_splitting_status, + "run_in_indexing_status": service.run_in_indexing_status, + }[entry] + method.assert_called_once() + (argument,) = method.call_args.args + ref = argument[0] if entry == "run" else argument + assert isinstance(ref, DocumentRef) + assert (ref.dataset.tenant_id, ref.dataset.dataset_id, ref.document_id) == ( + tenant_id, + dataset_id, + document_ids[0], + ) + + @pytest.fixture def tenant_id() -> str: return str(uuid.uuid4()) @@ -61,7 +103,7 @@ def mock_redis() -> MagicMock: def indexing_runner(monkeypatch: pytest.MonkeyPatch) -> MagicMock: runner = MagicMock() runner_class = MagicMock(return_value=runner) - monkeypatch.setattr("tasks.document_indexing_task.IndexingRunner", runner_class) + monkeypatch.setattr("tasks.document_indexing_task.build_document_indexing_service", runner_class) runner._constructor_mock = runner_class return runner @@ -213,19 +255,22 @@ class TestDocumentIndexing: ) _patch_features(monkeypatch, _features()) - def assert_committed_parsing(documents: list[Document], session: Session) -> None: - assert all(document.indexing_status == IndexingStatus.PARSING for document in documents) - assert all(document.processing_started_at is not None for document in documents) - assert all(session.get(Document, document.id) is document for document in documents) + def assert_committed_parsing(refs: list[DocumentRef]) -> None: + with Session(sqlite_session.get_bind()) as observer: + documents = [observer.get(Document, ref.document_id) for ref in refs] + assert all(document.indexing_status == IndexingStatus.PARSING for document in documents) + assert all(document.processing_started_at is not None for document in documents) indexing_runner.run.side_effect = assert_committed_parsing document_indexing_task.run(dataset_id, document_ids) persisted = _persisted_documents(sqlite_session, document_ids) assert [document.indexing_status for document in persisted] == [IndexingStatus.PARSING] * 3 - indexing_runner._constructor_mock.assert_called_once_with(enforce_vector_space_admission=True) + indexing_runner._constructor_mock.assert_called_once_with( + session_factory=ANY, enforce_vector_space_admission=True + ) indexing_runner.run.assert_called_once() - assert isinstance(indexing_runner.run.call_args.args[1], Session) + assert len(indexing_runner.run.call_args.args) == 1 def test_only_existing_documents_are_processed( self, @@ -248,7 +293,7 @@ class TestDocumentIndexing: _document_indexing(dataset_id, document_ids) processed = indexing_runner.run.call_args.args[0] - assert {document.id for document in processed} == set(existing_ids) + assert {ref.document_id for ref in processed} == set(existing_ids) assert sqlite_session.get(Document, document_ids[1]) is None def test_empty_batch_still_reaches_runner( @@ -270,14 +315,14 @@ class TestDocumentIndexing: _document_indexing(dataset_id, []) assert indexing_runner.run.call_args.args[0] == [] - assert isinstance(indexing_runner.run.call_args.args[1], Session) + assert len(indexing_runner.run.call_args.args) == 1 def test_missing_dataset_returns_before_feature_lookup( self, dataset_id: str, document_ids: list[str], monkeypatch: pytest.MonkeyPatch ) -> None: get_features = _patch_features(monkeypatch, _features()) runner_class = MagicMock() - monkeypatch.setattr("tasks.document_indexing_task.IndexingRunner", runner_class) + monkeypatch.setattr("tasks.document_indexing_task.build_document_indexing_service", runner_class) _document_indexing(dataset_id, document_ids) @@ -387,10 +432,12 @@ class TestSummaryDispatch: ) _patch_features(monkeypatch, _features()) - def finish_documents(documents: list[Document], _session: Session) -> None: - documents[0].indexing_status = IndexingStatus.COMPLETED - documents[1].indexing_status = IndexingStatus.COMPLETED - documents[2].indexing_status = IndexingStatus.INDEXING + def finish_documents(refs: list[DocumentRef]) -> None: + with Session(sqlite_session.get_bind()) as writer, writer.begin(): + for ref, status in zip( + refs, [IndexingStatus.COMPLETED, IndexingStatus.COMPLETED, IndexingStatus.INDEXING], strict=True + ): + writer.get(Document, ref.document_id).indexing_status = status indexing_runner.run.side_effect = finish_documents summary_delay = MagicMock() @@ -418,9 +465,12 @@ class TestSummaryDispatch: need_summary=[True], ) _patch_features(monkeypatch, _features()) - indexing_runner.run.side_effect = lambda documents, _session: setattr( - documents[0], "indexing_status", IndexingStatus.COMPLETED - ) + + def complete_documents(refs: list[DocumentRef]) -> None: + with Session(sqlite_session.get_bind()) as writer, writer.begin(): + writer.get(Document, refs[0].document_id).indexing_status = IndexingStatus.COMPLETED + + indexing_runner.run.side_effect = complete_documents summary_delay = MagicMock(side_effect=RuntimeError("queue unavailable")) monkeypatch.setattr("tasks.document_indexing_task.generate_summary_index_task.delay", summary_delay) @@ -449,9 +499,12 @@ class TestSummaryDispatch: need_summary=[True], ) _patch_features(monkeypatch, _features()) - indexing_runner.run.side_effect = lambda documents, _session: setattr( - documents[0], "indexing_status", IndexingStatus.COMPLETED - ) + + def complete_documents(refs: list[DocumentRef]) -> None: + with Session(sqlite_session.get_bind()) as writer, writer.begin(): + writer.get(Document, refs[0].document_id).indexing_status = IndexingStatus.COMPLETED + + indexing_runner.run.side_effect = complete_documents summary_delay = MagicMock() monkeypatch.setattr("tasks.document_indexing_task.generate_summary_index_task.delay", summary_delay) @@ -478,10 +531,11 @@ class TestSummaryDispatch: ) _patch_features(monkeypatch, _features()) - def remove_dataset(_documents: list[Document], session: Session) -> None: - dataset = session.get(Dataset, dataset_id) - assert dataset is not None - session.delete(dataset) + def remove_dataset(_refs: list[DocumentRef]) -> None: + with Session(sqlite_session.get_bind()) as writer, writer.begin(): + dataset = writer.get(Dataset, dataset_id) + assert dataset is not None + writer.delete(dataset) indexing_runner.run.side_effect = remove_dataset summary_delay = MagicMock() diff --git a/api/tests/unit_tests/tasks/test_document_indexing_sync_task.py b/api/tests/unit_tests/tasks/test_document_indexing_sync_task.py index 4fe36b11b81..be8e8b3c022 100644 --- a/api/tests/unit_tests/tasks/test_document_indexing_sync_task.py +++ b/api/tests/unit_tests/tasks/test_document_indexing_sync_task.py @@ -6,7 +6,7 @@ Notion extraction, index cleanup, and indexing remain mocked I/O boundaries. import json import uuid -from unittest.mock import MagicMock, patch +from unittest.mock import ANY, MagicMock, patch import pytest from sqlalchemy.orm import Session, sessionmaker @@ -126,10 +126,10 @@ def _bind_sqlite_session_factory(monkeypatch: pytest.MonkeyPatch, sqlite_session @pytest.fixture def mock_datasource_provider_service(): """Mock datasource credential provider.""" - with patch("tasks.document_indexing_sync_task.DatasourceProviderService", autospec=True) as mock_service_class: + with patch("tasks.document_indexing_sync_task.build_data_source_credentials", autospec=True) as mock_service_class: mock_service = MagicMock() mock_service.get_datasource_credentials.return_value = {"integration_secret": "test_token"} - mock_service_class.return_value = mock_service + mock_service_class.return_value.providers = mock_service yield mock_service @@ -165,6 +165,7 @@ class TestDocumentIndexingSyncTaskCollaboratorParams: # Assert mock_notion_extractor["class"].assert_called_once_with( + notion_token_loader=ANY, notion_workspace_id=notion_workspace_id, notion_obj_id=notion_page_id, notion_page_type="page", @@ -246,15 +247,15 @@ class TestDataSourceInfoSerialization: ): """data_source_info must be serialized with json.dumps before DB write.""" with ( - patch("tasks.document_indexing_sync_task.DatasourceProviderService") as mock_service_class, + patch("tasks.document_indexing_sync_task.build_data_source_credentials") as mock_service_class, patch("tasks.document_indexing_sync_task.NotionExtractor") as mock_extractor_class, patch("tasks.document_indexing_sync_task.IndexProcessorFactory") as mock_ipf, - patch("tasks.document_indexing_sync_task.IndexingRunner") as mock_runner_class, + patch("tasks.document_indexing_sync_task.build_document_indexing_service") as mock_runner_class, ): # External collaborators mock_service = MagicMock() mock_service.get_datasource_credentials.return_value = {"integration_secret": "token"} - mock_service_class.return_value = mock_service + mock_service_class.return_value.providers = mock_service mock_extractor = MagicMock() # Return a *different* timestamp so the task enters the sync/update branch diff --git a/api/tests/unit_tests/tasks/test_document_indexing_update_task.py b/api/tests/unit_tests/tasks/test_document_indexing_update_task.py index 3756a02a194..c38b3440b64 100644 --- a/api/tests/unit_tests/tasks/test_document_indexing_update_task.py +++ b/api/tests/unit_tests/tasks/test_document_indexing_update_task.py @@ -5,6 +5,7 @@ from __future__ import annotations import uuid from collections.abc import Callable from datetime import UTC, datetime +from functools import partial from unittest.mock import MagicMock import pytest @@ -12,12 +13,12 @@ from sqlalchemy import select from sqlalchemy.orm import Session import tasks.document_indexing_update_task as task_module -from core.indexing_runner import DocumentIsPausedError from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType from extensions.storage.storage_type import StorageType from models.dataset import Dataset, Document, DocumentSegment, SegmentAttachmentBinding from models.enums import CreatorUserRole, DataSourceType, DocumentCreatedFrom, IndexingStatus from models.model import UploadFile +from services.knowledge.indexing.errors import DocumentIsPausedError from tasks.document_indexing_update_task import document_indexing_update_task @@ -35,7 +36,7 @@ def task_harness( ) runner = MagicMock() processor = MagicMock() - monkeypatch.setattr(task_module, "IndexingRunner", MagicMock(return_value=runner)) + monkeypatch.setattr(task_module, "build_document_indexing_service", MagicMock(return_value=runner)) monkeypatch.setattr( task_module, "IndexProcessorFactory", @@ -100,9 +101,10 @@ def _persist_rows( return dataset, document -def _complete_indexing(documents: list[Document], _session: Session) -> None: - for document in documents: - document.indexing_status = IndexingStatus.COMPLETED +def _complete_indexing(refs, *, bind) -> None: + with Session(bind) as writer, writer.begin(): + for ref in refs: + writer.get(Document, ref.document_id).indexing_status = IndexingStatus.COMPLETED def _persist_attachment( @@ -145,7 +147,7 @@ def test_queues_summary_when_all_persisted_conditions_match( ) -> None: runner, _processor = task_harness dataset, document = _persist_rows(sqlite_session, summary_index_setting={"enable": True}) - runner.run.side_effect = _complete_indexing + runner.run.side_effect = partial(_complete_indexing, bind=sqlite_session.get_bind()) delay = MagicMock() monkeypatch.setattr(task_module.generate_summary_index_task, "delay", delay) @@ -186,7 +188,7 @@ def test_skips_summary_when_persisted_eligibility_does_not_match( for key, value in document_changes.items(): setattr(document, key, value) sqlite_session.commit() - runner.run.side_effect = _complete_indexing + runner.run.side_effect = partial(_complete_indexing, bind=sqlite_session.get_bind()) delay = MagicMock() monkeypatch.setattr(task_module.generate_summary_index_task, "delay", delay) @@ -276,7 +278,7 @@ def test_queue_failure_is_swallowed_after_successful_indexing( ) -> None: runner, _processor = task_harness dataset, document = _persist_rows(sqlite_session, summary_index_setting={"enable": True}) - runner.run.side_effect = _complete_indexing + runner.run.side_effect = partial(_complete_indexing, bind=sqlite_session.get_bind()) monkeypatch.setattr( task_module.generate_summary_index_task, "delay", @@ -300,7 +302,7 @@ def test_cleans_and_deletes_persisted_segments_with_real_session( summary_index_setting={"enable": True}, with_segment=True, ) - runner.run.side_effect = _complete_indexing + runner.run.side_effect = partial(_complete_indexing, bind=sqlite_session.get_bind()) delay = MagicMock() monkeypatch.setattr(task_module.generate_summary_index_task, "delay", delay) diff --git a/api/tests/unit_tests/tasks/test_enable_segment_index_tasks.py b/api/tests/unit_tests/tasks/test_enable_segment_index_tasks.py index fb64f5b5ba4..260e8a15ba2 100644 --- a/api/tests/unit_tests/tasks/test_enable_segment_index_tasks.py +++ b/api/tests/unit_tests/tasks/test_enable_segment_index_tasks.py @@ -94,7 +94,7 @@ def test_enable_segment_commits_index_rows_after_loading( _record_transaction_events(sqlite_session_factory, phase_events), patch("tasks.enable_segment_to_index_task.IndexProcessorFactory") as processor_factory, patch( - "services.summary_index_service.SummaryIndexService.enable_summaries_for_segments", + "services.knowledge.summaries.adapters.SummaryIndexAdapter.enable_summaries_for_segments", enable_summaries, ), patch("tasks.enable_segment_to_index_task.redis_client.delete"), @@ -126,7 +126,9 @@ def test_enable_segment_rolls_back_before_error_compensation( with ( _record_transaction_events(sqlite_session_factory, phase_events), patch("tasks.enable_segment_to_index_task.IndexProcessorFactory") as processor_factory, - patch("services.summary_index_service.SummaryIndexService.enable_summaries_for_segments") as enable_summaries, + patch( + "services.knowledge.summaries.adapters.SummaryIndexAdapter.enable_summaries_for_segments" + ) as enable_summaries, patch("tasks.enable_segment_to_index_task.redis_client.delete"), ): processor_factory.return_value.init_index_processor.return_value = index_processor @@ -157,7 +159,7 @@ def test_enable_segments_commits_index_rows_after_loading( _record_transaction_events(sqlite_session_factory, phase_events), patch("tasks.enable_segments_to_index_task.IndexProcessorFactory") as processor_factory, patch( - "services.summary_index_service.SummaryIndexService.enable_summaries_for_segments", + "services.knowledge.summaries.adapters.SummaryIndexAdapter.enable_summaries_for_segments", enable_summaries, ), patch("tasks.enable_segments_to_index_task.redis_client.delete"), @@ -189,7 +191,9 @@ def test_enable_segments_rolls_back_before_error_compensation( with ( _record_transaction_events(sqlite_session_factory, phase_events), patch("tasks.enable_segments_to_index_task.IndexProcessorFactory") as processor_factory, - patch("services.summary_index_service.SummaryIndexService.enable_summaries_for_segments") as enable_summaries, + patch( + "services.knowledge.summaries.adapters.SummaryIndexAdapter.enable_summaries_for_segments" + ) as enable_summaries, patch("tasks.enable_segments_to_index_task.redis_client.delete"), ): processor_factory.return_value.init_index_processor.return_value = index_processor diff --git a/api/tests/unit_tests/tasks/test_retry_document_indexing_task.py b/api/tests/unit_tests/tasks/test_retry_document_indexing_task.py index 47d02381da9..2adf8cd37e5 100644 --- a/api/tests/unit_tests/tasks/test_retry_document_indexing_task.py +++ b/api/tests/unit_tests/tasks/test_retry_document_indexing_task.py @@ -1,8 +1,10 @@ -from unittest.mock import MagicMock, patch +from unittest.mock import ANY, MagicMock, patch from uuid import uuid4 +import pytest from sqlalchemy.orm import Session +from core.app.apps.pipeline.pipeline_generator import PipelineGenerator from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType from enums import DeploymentEdition from models import Account, Tenant, TenantAccountJoin @@ -14,7 +16,8 @@ from tests.unit_tests.config_override import config_overrides_context @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) -def test_retry_enforces_vector_space_admission(sqlite_session: Session) -> None: +@pytest.mark.parametrize("runtime_mode", [DatasetRuntimeMode.GENERAL, DatasetRuntimeMode.RAG_PIPELINE]) +def test_retry_uses_dataset_runtime(sqlite_session: Session, runtime_mode: DatasetRuntimeMode) -> None: tenant = Tenant(name="Retry tenant") user = Account(name="Retry user", email=f"retry-{uuid4()}@example.com") membership = TenantAccountJoin( @@ -31,7 +34,7 @@ def test_retry_enforces_vector_space_admission(sqlite_session: Session) -> None: data_source_type=DataSourceType.UPLOAD_FILE, indexing_technique=IndexTechniqueType.ECONOMY, chunk_structure=IndexStructureType.PARAGRAPH_INDEX, - runtime_mode=DatasetRuntimeMode.GENERAL, + runtime_mode=runtime_mode, ) document = Document( id=str(uuid4()), @@ -56,12 +59,22 @@ def test_retry_enforces_vector_space_admission(sqlite_session: Session) -> None: with ( patch("tasks.retry_document_indexing_task.FeatureService.get_features", return_value=features), patch("tasks.retry_document_indexing_task.IndexProcessorFactory"), - patch("tasks.retry_document_indexing_task.IndexingRunner") as indexing_runner, + patch("tasks.retry_document_indexing_task.build_document_indexing_service") as indexing_runner, + patch("tasks.retry_document_indexing_task.RagPipelineService") as pipeline_service, patch("tasks.retry_document_indexing_task.redis_client"), ): retry_document_indexing_task.run(dataset.id, [document.id], user.id) - indexing_runner.assert_called_once_with(enforce_vector_space_admission=True) - run_documents, run_session = indexing_runner.return_value.run.call_args.args - assert [item.id for item in run_documents] == [document.id] - assert isinstance(run_session, Session) + if runtime_mode == DatasetRuntimeMode.GENERAL: + pipeline_service.assert_not_called() + indexing_runner.assert_called_once_with(session_factory=ANY, enforce_vector_space_admission=True) + (run_documents,) = indexing_runner.return_value.run.call_args.args + assert [item.document_id for item in run_documents] == [document.id] + assert run_documents[0].dataset.tenant_id == dataset.tenant_id + else: + indexing_runner.assert_not_called() + retry = pipeline_service.return_value.retry_error_document + retry.assert_called_once() + retry_dataset, retry_document, retry_user = retry.call_args.args + assert (retry_dataset.id, retry_document.id, retry_user.id) == (dataset.id, document.id, user.id) + assert isinstance(retry.call_args.kwargs["generator"], PipelineGenerator) diff --git a/api/tests/unit_tests/tasks/test_segment_index_cleanup_tasks.py b/api/tests/unit_tests/tasks/test_segment_index_cleanup_tasks.py index 6a256656ad3..910cf8abf34 100644 --- a/api/tests/unit_tests/tasks/test_segment_index_cleanup_tasks.py +++ b/api/tests/unit_tests/tasks/test_segment_index_cleanup_tasks.py @@ -94,7 +94,7 @@ def test_disable_segment_commits_index_cleanup( _record_transaction_events(sqlite_session_factory, phase_events), patch("tasks.disable_segment_from_index_task.IndexProcessorFactory") as processor_factory, patch( - "services.summary_index_service.SummaryIndexService.disable_summaries_for_segments", + "services.knowledge.summaries.adapters.SummaryIndexAdapter.disable_summaries_for_segments", disable_summaries, ), patch("tasks.disable_segment_from_index_task.redis_client.delete"), @@ -123,7 +123,7 @@ def test_disable_segments_commits_index_cleanup( _record_transaction_events(sqlite_session_factory, phase_events), patch("tasks.disable_segments_from_index_task.IndexProcessorFactory") as processor_factory, patch( - "services.summary_index_service.SummaryIndexService.disable_summaries_for_segments", + "services.knowledge.summaries.adapters.SummaryIndexAdapter.disable_summaries_for_segments", disable_summaries, ), patch("tasks.disable_segments_from_index_task.redis_client.delete"), diff --git a/api/tests/unit_tests/tasks/test_sync_website_document_indexing_task.py b/api/tests/unit_tests/tasks/test_sync_website_document_indexing_task.py index 594c4295305..0cfa3d62742 100644 --- a/api/tests/unit_tests/tasks/test_sync_website_document_indexing_task.py +++ b/api/tests/unit_tests/tasks/test_sync_website_document_indexing_task.py @@ -87,7 +87,7 @@ def test_cleanup_is_owner_scoped_and_skips_empty_vector_ids(sqlite_session: Sess with ( patch("tasks.sync_website_document_indexing_task.FeatureService.get_features", return_value=features), patch("tasks.sync_website_document_indexing_task.IndexProcessorFactory") as processor_factory, - patch("tasks.sync_website_document_indexing_task.IndexingRunner") as indexing_runner, + patch("tasks.sync_website_document_indexing_task.build_document_indexing_service") as indexing_runner, patch("tasks.sync_website_document_indexing_task.redis_client"), ): sync_website_document_indexing_task(dataset.id, document.id) diff --git a/packages/contracts/generated/api/console/data-source/orpc.gen.ts b/packages/contracts/generated/api/console/data-source/orpc.gen.ts index 692089d1845..69711d6a5b7 100644 --- a/packages/contracts/generated/api/console/data-source/orpc.gen.ts +++ b/packages/contracts/generated/api/console/data-source/orpc.gen.ts @@ -3,25 +3,11 @@ import { oc } from '@orpc/contract' import * as z from 'zod' import { - zGetDataSourceIntegratesByBindingIdByActionPath, - zGetDataSourceIntegratesByBindingIdByActionResponse, zGetDataSourceIntegratesResponse, zPatchDataSourceIntegratesByBindingIdByActionPath, zPatchDataSourceIntegratesByBindingIdByActionResponse, - zPatchDataSourceIntegratesResponse, } from './zod.gen.ts' -export const get = oc - .route({ - inputStructure: 'detailed', - method: 'GET', - operationId: 'getDataSourceIntegratesByBindingIdByAction', - path: '/data-source/integrates/{binding_id}/{action}', - tags: ['console'], - }) - .input(z.object({ params: zGetDataSourceIntegratesByBindingIdByActionPath })) - .output(zGetDataSourceIntegratesByBindingIdByActionResponse) - export const patch = oc .route({ inputStructure: 'detailed', @@ -34,7 +20,6 @@ export const patch = oc .output(zPatchDataSourceIntegratesByBindingIdByActionResponse) export const byAction = { - get, patch, } @@ -42,7 +27,7 @@ export const byBindingId = { byAction, } -export const get2 = oc +export const get = oc .route({ inputStructure: 'detailed', method: 'GET', @@ -52,19 +37,8 @@ export const get2 = oc }) .output(zGetDataSourceIntegratesResponse) -export const patch2 = oc - .route({ - inputStructure: 'detailed', - method: 'PATCH', - operationId: 'patchDataSourceIntegrates', - path: '/data-source/integrates', - tags: ['console'], - }) - .output(zPatchDataSourceIntegratesResponse) - export const integrates = { - get: get2, - patch: patch2, + get, byBindingId, } diff --git a/packages/contracts/generated/api/console/data-source/types.gen.ts b/packages/contracts/generated/api/console/data-source/types.gen.ts index 1c1051369c4..a19681a65e1 100644 --- a/packages/contracts/generated/api/console/data-source/types.gen.ts +++ b/packages/contracts/generated/api/console/data-source/types.gen.ts @@ -4,6 +4,48 @@ export type ClientOptions = { baseUrl: `${string}://${string}/console/api` | (string & {}) } +export type DataSourceIntegrateListResponse = { + data: Array +} + +export type SimpleResultResponse = { + result: string +} + +export type DataSourceIntegrateResponse = { + created_at: number | null + disabled: boolean | null + id: string | null + is_bound: boolean + link: string + provider: string + source_info: DataSourceIntegrateWorkspaceResponse | null +} + +export type DataSourceIntegrateWorkspaceResponse = { + pages: Array + total: number + workspace_icon: string | null + workspace_id: string | null + workspace_name: string | null +} + +export type DataSourceIntegratePageResponse = { + page_icon: DataSourceIntegrateIconResponse | null + page_id: string + page_name: string + parent_id: string + type: NotionPageType +} + +export type DataSourceIntegrateIconResponse = { + emoji?: string | null + type?: string | null + url?: string | null +} + +export type NotionPageType = 'database' | 'page' + export type GetDataSourceIntegratesData = { body?: never path?: never @@ -12,49 +54,12 @@ export type GetDataSourceIntegratesData = { } export type GetDataSourceIntegratesResponses = { - 200: { - [key: string]: unknown - } + 200: DataSourceIntegrateListResponse } export type GetDataSourceIntegratesResponse = GetDataSourceIntegratesResponses[keyof GetDataSourceIntegratesResponses] -export type PatchDataSourceIntegratesData = { - body?: never - path?: never - query?: never - url: '/data-source/integrates' -} - -export type PatchDataSourceIntegratesResponses = { - 200: { - [key: string]: unknown - } -} - -export type PatchDataSourceIntegratesResponse = - PatchDataSourceIntegratesResponses[keyof PatchDataSourceIntegratesResponses] - -export type GetDataSourceIntegratesByBindingIdByActionData = { - body?: never - path: { - action: string - binding_id: string - } - query?: never - url: '/data-source/integrates/{binding_id}/{action}' -} - -export type GetDataSourceIntegratesByBindingIdByActionResponses = { - 200: { - [key: string]: unknown - } -} - -export type GetDataSourceIntegratesByBindingIdByActionResponse = - GetDataSourceIntegratesByBindingIdByActionResponses[keyof GetDataSourceIntegratesByBindingIdByActionResponses] - export type PatchDataSourceIntegratesByBindingIdByActionData = { body?: never path: { @@ -66,9 +71,7 @@ export type PatchDataSourceIntegratesByBindingIdByActionData = { } export type PatchDataSourceIntegratesByBindingIdByActionResponses = { - 200: { - [key: string]: unknown - } + 200: SimpleResultResponse } export type PatchDataSourceIntegratesByBindingIdByActionResponse = diff --git a/packages/contracts/generated/api/console/data-source/zod.gen.ts b/packages/contracts/generated/api/console/data-source/zod.gen.ts index 8236a06d2e4..3e4b58997ce 100644 --- a/packages/contracts/generated/api/console/data-source/zod.gen.ts +++ b/packages/contracts/generated/api/console/data-source/zod.gen.ts @@ -3,24 +3,72 @@ import * as z from 'zod' /** - * Success + * SimpleResultResponse */ -export const zGetDataSourceIntegratesResponse = z.record(z.string(), z.unknown()) +export const zSimpleResultResponse = z.object({ + result: z.string(), +}) /** - * Success + * DataSourceIntegrateIconResponse */ -export const zPatchDataSourceIntegratesResponse = z.record(z.string(), z.unknown()) +export const zDataSourceIntegrateIconResponse = z.object({ + emoji: z.string().nullish(), + type: z.string().nullish(), + url: z.string().nullish(), +}) -export const zGetDataSourceIntegratesByBindingIdByActionPath = z.object({ - action: z.string(), - binding_id: z.uuid(), +/** + * NotionPageType + */ +export const zNotionPageType = z.enum(['database', 'page']) + +/** + * DataSourceIntegratePageResponse + */ +export const zDataSourceIntegratePageResponse = z.object({ + page_icon: zDataSourceIntegrateIconResponse.nullable(), + page_id: z.string(), + page_name: z.string(), + parent_id: z.string(), + type: zNotionPageType, +}) + +/** + * DataSourceIntegrateWorkspaceResponse + */ +export const zDataSourceIntegrateWorkspaceResponse = z.object({ + pages: z.array(zDataSourceIntegratePageResponse), + total: z.int(), + workspace_icon: z.string().nullable(), + workspace_id: z.string().nullable(), + workspace_name: z.string().nullable(), +}) + +/** + * DataSourceIntegrateResponse + */ +export const zDataSourceIntegrateResponse = z.object({ + created_at: z.int().nullable(), + disabled: z.boolean().nullable(), + id: z.string().nullable(), + is_bound: z.boolean(), + link: z.string(), + provider: z.string(), + source_info: zDataSourceIntegrateWorkspaceResponse.nullable(), +}) + +/** + * DataSourceIntegrateListResponse + */ +export const zDataSourceIntegrateListResponse = z.object({ + data: z.array(zDataSourceIntegrateResponse), }) /** * Success */ -export const zGetDataSourceIntegratesByBindingIdByActionResponse = z.record(z.string(), z.unknown()) +export const zGetDataSourceIntegratesResponse = zDataSourceIntegrateListResponse export const zPatchDataSourceIntegratesByBindingIdByActionPath = z.object({ action: z.string(), @@ -30,7 +78,4 @@ export const zPatchDataSourceIntegratesByBindingIdByActionPath = z.object({ /** * Success */ -export const zPatchDataSourceIntegratesByBindingIdByActionResponse = z.record( - z.string(), - z.unknown(), -) +export const zPatchDataSourceIntegratesByBindingIdByActionResponse = zSimpleResultResponse diff --git a/packages/contracts/generated/api/console/datasets/orpc.gen.ts b/packages/contracts/generated/api/console/datasets/orpc.gen.ts index 2cbafe6d030..8ed04583d05 100644 --- a/packages/contracts/generated/api/console/datasets/orpc.gen.ts +++ b/packages/contracts/generated/api/console/datasets/orpc.gen.ts @@ -47,8 +47,6 @@ import { zGetDatasetsByDatasetIdDocumentsByDocumentIdPipelineExecutionLogResponse, zGetDatasetsByDatasetIdDocumentsByDocumentIdQuery, zGetDatasetsByDatasetIdDocumentsByDocumentIdResponse, - zGetDatasetsByDatasetIdDocumentsByDocumentIdSegmentsBatchImportPath, - zGetDatasetsByDatasetIdDocumentsByDocumentIdSegmentsBatchImportResponse, zGetDatasetsByDatasetIdDocumentsByDocumentIdSegmentsBySegmentIdChildChunksPath, zGetDatasetsByDatasetIdDocumentsByDocumentIdSegmentsBySegmentIdChildChunksQuery, zGetDatasetsByDatasetIdDocumentsByDocumentIdSegmentsBySegmentIdChildChunksResponse, @@ -127,9 +125,6 @@ import { zPatchDatasetsExternalKnowledgeApiByExternalKnowledgeApiIdResponse, zPostDatasetsApiKeysBody, zPostDatasetsApiKeysResponse, - zPostDatasetsBatchImportStatusByJobIdBody, - zPostDatasetsBatchImportStatusByJobIdPath, - zPostDatasetsBatchImportStatusByJobIdResponse, zPostDatasetsBody, zPostDatasetsByDatasetIdApiKeysByStatusPath, zPostDatasetsByDatasetIdApiKeysByStatusResponse, @@ -269,25 +264,8 @@ export const get3 = oc .input(z.object({ params: zGetDatasetsBatchImportStatusByJobIdPath })) .output(zGetDatasetsBatchImportStatusByJobIdResponse) -export const post2 = oc - .route({ - inputStructure: 'detailed', - method: 'POST', - operationId: 'postDatasetsBatchImportStatusByJobId', - path: '/datasets/batch_import_status/{job_id}', - tags: ['console'], - }) - .input( - z.object({ - body: zPostDatasetsBatchImportStatusByJobIdBody, - params: zPostDatasetsBatchImportStatusByJobIdPath, - }), - ) - .output(zPostDatasetsBatchImportStatusByJobIdResponse) - export const byJobId = { get: get3, - post: post2, } export const batchImportStatus = { @@ -297,7 +275,7 @@ export const batchImportStatus = { /** * Create external knowledge dataset */ -export const post3 = oc +export const post2 = oc .route({ description: 'Create external knowledge dataset', inputStructure: 'detailed', @@ -311,7 +289,7 @@ export const post3 = oc .output(zPostDatasetsExternalResponse) export const external = { - post: post3, + post: post2, } /** @@ -405,7 +383,7 @@ export const get6 = oc /** * Create external knowledge API template */ -export const post4 = oc +export const post3 = oc .route({ description: 'Create external knowledge API template', inputStructure: 'detailed', @@ -420,14 +398,14 @@ export const post4 = oc export const externalKnowledgeApi = { get: get6, - post: post4, + post: post3, byExternalKnowledgeApiId, } /** * Estimate dataset indexing cost */ -export const post5 = oc +export const post4 = oc .route({ description: 'Estimate dataset indexing cost', inputStructure: 'detailed', @@ -440,13 +418,13 @@ export const post5 = oc .output(zPostDatasetsIndexingEstimateResponse) export const indexingEstimate = { - post: post5, + post: post4, } /** * Initialize dataset with documents */ -export const post6 = oc +export const post5 = oc .route({ description: 'Initialize dataset with documents', inputStructure: 'detailed', @@ -459,7 +437,7 @@ export const post6 = oc .output(zPostDatasetsInitResponse) export const init = { - post: post6, + post: post5, } export const get7 = oc @@ -480,7 +458,7 @@ export const metadata = { builtIn, } -export const post7 = oc +export const post6 = oc .route({ inputStructure: 'detailed', method: 'POST', @@ -492,7 +470,7 @@ export const post7 = oc .output(zPostDatasetsNotionIndexingEstimateResponse) export const notionIndexingEstimate = { - post: post7, + post: post6, } /** @@ -552,7 +530,7 @@ export const retrievalSetting = { byVectorType, } -export const post8 = oc +export const post7 = oc .route({ inputStructure: 'detailed', method: 'POST', @@ -564,7 +542,7 @@ export const post8 = oc .output(zPostDatasetsByDatasetIdApiKeysByStatusResponse) export const byStatus = { - post: post8, + post: post7, } export const apiKeys2 = { @@ -634,7 +612,7 @@ export const batch = { * * Download selected dataset documents as a single ZIP archive (upload-file only) */ -export const post9 = oc +export const post8 = oc .route({ description: 'Download selected dataset documents as a single ZIP archive (upload-file only)', inputStructure: 'detailed', @@ -653,7 +631,7 @@ export const post9 = oc .output(zPostDatasetsByDatasetIdDocumentsDownloadZipResponse) export const downloadZip = { - post: post9, + post: post8, } /** @@ -664,7 +642,7 @@ export const downloadZip = { * (indexing_technique must be 'high_quality' and summary_index_setting.enable must be true), * then asynchronously generates summary indexes for the provided documents. */ -export const post10 = oc +export const post9 = oc .route({ description: "Generate summary index for documents\nThis endpoint checks if the dataset configuration supports summary generation\n(indexing_technique must be 'high_quality' and summary_index_setting.enable must be true),\nthen asynchronously generates summary indexes for the provided documents.", @@ -684,10 +662,10 @@ export const post10 = oc .output(zPostDatasetsByDatasetIdDocumentsGenerateSummaryResponse) export const generateSummary = { - post: post10, + post: post9, } -export const post11 = oc +export const post10 = oc .route({ inputStructure: 'detailed', method: 'POST', @@ -705,7 +683,7 @@ export const post11 = oc .output(zPostDatasetsByDatasetIdDocumentsMetadataResponse) export const metadata2 = { - post: post11, + post: post10, } export const patch2 = oc @@ -911,7 +889,7 @@ export const processing = { byAction: byAction2, } -export const post12 = oc +export const post11 = oc .route({ inputStructure: 'detailed', method: 'POST', @@ -928,7 +906,7 @@ export const post12 = oc .output(zPostDatasetsByDatasetIdDocumentsByDocumentIdRenameResponse) export const rename = { - post: post12, + post: post11, } export const patch6 = oc @@ -951,7 +929,7 @@ export const byAction3 = { patch: patch6, } -export const post13 = oc +export const post12 = oc .route({ inputStructure: 'detailed', method: 'POST', @@ -968,22 +946,11 @@ export const post13 = oc .output(zPostDatasetsByDatasetIdDocumentsByDocumentIdSegmentResponse) export const segment = { - post: post13, + post: post12, byAction: byAction3, } -export const get19 = oc - .route({ - inputStructure: 'detailed', - method: 'GET', - operationId: 'getDatasetsByDatasetIdDocumentsByDocumentIdSegmentsBatchImport', - path: '/datasets/{dataset_id}/documents/{document_id}/segments/batch_import', - tags: ['console'], - }) - .input(z.object({ params: zGetDatasetsByDatasetIdDocumentsByDocumentIdSegmentsBatchImportPath })) - .output(zGetDatasetsByDatasetIdDocumentsByDocumentIdSegmentsBatchImportResponse) - -export const post14 = oc +export const post13 = oc .route({ inputStructure: 'detailed', method: 'POST', @@ -1000,8 +967,7 @@ export const post14 = oc .output(zPostDatasetsByDatasetIdDocumentsByDocumentIdSegmentsBatchImportResponse) export const batchImport = { - get: get19, - post: post14, + post: post13, } export const delete3 = oc @@ -1049,7 +1015,7 @@ export const byChildChunkId = { patch: patch7, } -export const get20 = oc +export const get19 = oc .route({ inputStructure: 'detailed', method: 'GET', @@ -1082,7 +1048,7 @@ export const patch8 = oc ) .output(zPatchDatasetsByDatasetIdDocumentsByDocumentIdSegmentsBySegmentIdChildChunksResponse) -export const post15 = oc +export const post14 = oc .route({ inputStructure: 'detailed', method: 'POST', @@ -1099,9 +1065,9 @@ export const post15 = oc .output(zPostDatasetsByDatasetIdDocumentsByDocumentIdSegmentsBySegmentIdChildChunksResponse) export const childChunks = { - get: get20, + get: get19, patch: patch8, - post: post15, + post: post14, byChildChunkId, } @@ -1158,7 +1124,7 @@ export const delete5 = oc ) .output(zDeleteDatasetsByDatasetIdDocumentsByDocumentIdSegmentsResponse) -export const get21 = oc +export const get20 = oc .route({ inputStructure: 'detailed', method: 'GET', @@ -1176,7 +1142,7 @@ export const get21 = oc export const segments = { delete: delete5, - get: get21, + get: get20, batchImport, bySegmentId, } @@ -1195,7 +1161,7 @@ export const segments = { * - timeout: Number of summaries that timed out * - summaries: List of summary records with status and content preview */ -export const get22 = oc +export const get21 = oc .route({ description: 'Get summary index generation status for a document\nReturns:\n- total_segments: Total number of segments in the document\n- summary_status: Dictionary with status counts\n - completed: Number of summaries completed\n - generating: Number of summaries being generated\n - error: Number of summaries with errors\n - not_started: Number of segments without summary records\n - timeout: Number of summaries that timed out\n- summaries: List of summary records with status and content preview', @@ -1210,13 +1176,13 @@ export const get22 = oc .output(zGetDatasetsByDatasetIdDocumentsByDocumentIdSummaryStatusResponse) export const summaryStatus = { - get: get22, + get: get21, } /** * sync website document */ -export const get23 = oc +export const get22 = oc .route({ inputStructure: 'detailed', method: 'GET', @@ -1229,7 +1195,7 @@ export const get23 = oc .output(zGetDatasetsByDatasetIdDocumentsByDocumentIdWebsiteSyncResponse) export const websiteSync = { - get: get23, + get: get22, } export const delete6 = oc @@ -1247,7 +1213,7 @@ export const delete6 = oc /** * Get document details */ -export const get24 = oc +export const get23 = oc .route({ description: 'Get document details', inputStructure: 'detailed', @@ -1266,7 +1232,7 @@ export const get24 = oc export const byDocumentId = { delete: delete6, - get: get24, + get: get23, download, indexingEstimate: indexingEstimate3, indexingStatus: indexingStatus2, @@ -1296,7 +1262,7 @@ export const delete7 = oc /** * Get documents in a dataset */ -export const get25 = oc +export const get24 = oc .route({ description: 'Get documents in a dataset', inputStructure: 'detailed', @@ -1313,7 +1279,7 @@ export const get25 = oc ) .output(zGetDatasetsByDatasetIdDocumentsResponse) -export const post16 = oc +export const post15 = oc .route({ inputStructure: 'detailed', method: 'POST', @@ -1331,8 +1297,8 @@ export const post16 = oc export const documents = { delete: delete7, - get: get25, - post: post16, + get: get24, + post: post15, downloadZip, generateSummary, metadata: metadata2, @@ -1343,7 +1309,7 @@ export const documents = { /** * Get dataset error documents */ -export const get26 = oc +export const get25 = oc .route({ description: 'Get dataset error documents', inputStructure: 'detailed', @@ -1356,13 +1322,13 @@ export const get26 = oc .output(zGetDatasetsByDatasetIdErrorDocsResponse) export const errorDocs = { - get: get26, + get: get25, } /** * Test external knowledge retrieval for dataset */ -export const post17 = oc +export const post16 = oc .route({ description: 'Test external knowledge retrieval for dataset', inputStructure: 'detailed', @@ -1380,13 +1346,13 @@ export const post17 = oc .output(zPostDatasetsByDatasetIdExternalHitTestingResponse) export const externalHitTesting = { - post: post17, + post: post16, } /** * Test dataset knowledge retrieval */ -export const post18 = oc +export const post17 = oc .route({ description: 'Test dataset knowledge retrieval', inputStructure: 'detailed', @@ -1404,13 +1370,13 @@ export const post18 = oc .output(zPostDatasetsByDatasetIdHitTestingResponse) export const hitTesting = { - post: post18, + post: post17, } /** * Get dataset indexing status */ -export const get27 = oc +export const get26 = oc .route({ description: 'Get dataset indexing status', inputStructure: 'detailed', @@ -1423,10 +1389,10 @@ export const get27 = oc .output(zGetDatasetsByDatasetIdIndexingStatusResponse) export const indexingStatus3 = { - get: get27, + get: get26, } -export const post19 = oc +export const post18 = oc .route({ inputStructure: 'detailed', method: 'POST', @@ -1439,7 +1405,7 @@ export const post19 = oc .output(zPostDatasetsByDatasetIdMetadataBuiltInByActionResponse) export const byAction4 = { - post: post19, + post: post18, } export const builtIn2 = { @@ -1479,7 +1445,7 @@ export const byMetadataId = { patch: patch10, } -export const get28 = oc +export const get27 = oc .route({ inputStructure: 'detailed', method: 'GET', @@ -1490,7 +1456,7 @@ export const get28 = oc .input(z.object({ params: zGetDatasetsByDatasetIdMetadataPath })) .output(zGetDatasetsByDatasetIdMetadataResponse) -export const post20 = oc +export const post19 = oc .route({ inputStructure: 'detailed', method: 'POST', @@ -1508,13 +1474,13 @@ export const post20 = oc .output(zPostDatasetsByDatasetIdMetadataResponse) export const metadata4 = { - get: get28, - post: post20, + get: get27, + post: post19, builtIn: builtIn2, byMetadataId, } -export const get29 = oc +export const get28 = oc .route({ inputStructure: 'detailed', method: 'GET', @@ -1526,7 +1492,7 @@ export const get29 = oc .output(zGetDatasetsByDatasetIdNotionSyncResponse) export const sync2 = { - get: get29, + get: get28, } export const notion2 = { @@ -1536,7 +1502,7 @@ export const notion2 = { /** * Get dataset permission user list */ -export const get30 = oc +export const get29 = oc .route({ description: 'Get dataset permission user list', inputStructure: 'detailed', @@ -1549,13 +1515,13 @@ export const get30 = oc .output(zGetDatasetsByDatasetIdPermissionPartUsersResponse) export const permissionPartUsers = { - get: get30, + get: get29, } /** * Get dataset query history */ -export const get31 = oc +export const get30 = oc .route({ description: 'Get dataset query history', inputStructure: 'detailed', @@ -1568,13 +1534,13 @@ export const get31 = oc .output(zGetDatasetsByDatasetIdQueriesResponse) export const queries = { - get: get31, + get: get30, } /** * Get applications related to dataset */ -export const get32 = oc +export const get31 = oc .route({ description: 'Get applications related to dataset', inputStructure: 'detailed', @@ -1587,13 +1553,13 @@ export const get32 = oc .output(zGetDatasetsByDatasetIdRelatedAppsResponse) export const relatedApps = { - get: get32, + get: get31, } /** * retry document */ -export const post21 = oc +export const post20 = oc .route({ inputStructure: 'detailed', method: 'POST', @@ -1612,13 +1578,13 @@ export const post21 = oc .output(zPostDatasetsByDatasetIdRetryResponse) export const retry = { - post: post21, + post: post20, } /** * Check if dataset is in use */ -export const get33 = oc +export const get32 = oc .route({ description: 'Check if dataset is in use', inputStructure: 'detailed', @@ -1631,7 +1597,7 @@ export const get33 = oc .output(zGetDatasetsByDatasetIdUseCheckResponse) export const useCheck2 = { - get: get33, + get: get32, } export const delete9 = oc @@ -1649,7 +1615,7 @@ export const delete9 = oc /** * Get dataset details */ -export const get34 = oc +export const get33 = oc .route({ description: 'Get dataset details', inputStructure: 'detailed', @@ -1678,7 +1644,7 @@ export const patch11 = oc export const byDatasetId = { delete: delete9, - get: get34, + get: get33, patch: patch11, apiKeys: apiKeys2, autoDisableLogs, @@ -1720,7 +1686,7 @@ export const byApiKeyId2 = { /** * Get all API keys for a dataset */ -export const get35 = oc +export const get34 = oc .route({ description: 'Get all API keys for a dataset', inputStructure: 'detailed', @@ -1735,7 +1701,7 @@ export const get35 = oc /** * Create a new API key for a dataset */ -export const post22 = oc +export const post21 = oc .route({ description: 'Create a new API key for a dataset', inputStructure: 'detailed', @@ -1749,8 +1715,8 @@ export const post22 = oc .output(zPostDatasetsByResourceIdApiKeysResponse) export const apiKeys3 = { - get: get35, - post: post22, + get: get34, + post: post21, byApiKeyId: byApiKeyId2, } @@ -1761,7 +1727,7 @@ export const byResourceId = { /** * Get list of datasets */ -export const get36 = oc +export const get35 = oc .route({ description: 'Get list of datasets', inputStructure: 'detailed', @@ -1776,7 +1742,7 @@ export const get36 = oc /** * Create a new dataset */ -export const post23 = oc +export const post22 = oc .route({ description: 'Create a new dataset', inputStructure: 'detailed', @@ -1790,8 +1756,8 @@ export const post23 = oc .output(zPostDatasetsResponse) export const datasets = { - get: get36, - post: post23, + get: get35, + post: post22, apiBaseInfo, apiKeys, batchImportStatus, diff --git a/packages/contracts/generated/api/console/datasets/types.gen.ts b/packages/contracts/generated/api/console/datasets/types.gen.ts index fda9408dcd2..6c84febc239 100644 --- a/packages/contracts/generated/api/console/datasets/types.gen.ts +++ b/packages/contracts/generated/api/console/datasets/types.gen.ts @@ -22,7 +22,7 @@ export type DatasetCreatePayload = { provider?: string } -export type DatasetDetailResponse = { +export type DatasetDetailWithPartialMembersResponse = { app_count: number author_name: string | null built_in_field_enabled: boolean @@ -47,6 +47,7 @@ export type DatasetDetailResponse = { is_published: boolean maintainer?: string | null name: string + partial_member_list?: Array | null permission: string permission_keys?: Array pipeline_id: string | null @@ -88,10 +89,6 @@ export type SegmentBatchImportStatusResponse = { job_status: string } -export type BatchImportPayload = { - upload_file_id: string -} - export type ExternalDatasetCreatePayload = { description?: string | null external_knowledge_api_id: string @@ -102,6 +99,46 @@ export type ExternalDatasetCreatePayload = { name: string } +export type DatasetDetailResponse = { + app_count: number + author_name: string | null + built_in_field_enabled: boolean + chunk_structure: string | null + created_at: number + created_by: string + data_source_type: string | null + description: string | null + doc_form: string | null + doc_metadata: Array + document_count: number + embedding_available?: boolean | null + embedding_model: string | null + embedding_model_provider: string | null + enable_api: boolean + external_knowledge_info?: DatasetExternalKnowledgeInfoResponse + external_retrieval_model: DatasetExternalRetrievalModelResponse | null + icon_info?: DatasetIconInfoResponse + id: string + indexing_technique: string | null + is_multimodal: boolean + is_published: boolean + maintainer?: string | null + name: string + permission: string + permission_keys?: Array + pipeline_id: string | null + provider: string + retrieval_model_dict: DatasetRetrievalModelResponse + runtime_mode: string | null + summary_index_setting?: DatasetSummaryIndexSettingResponse + tags: Array + total_available_documents: number + total_documents: number + updated_at: number + updated_by: string | null + word_count: number +} + export type ExternalKnowledgeApiListResponse = { data: Array has_more: boolean @@ -191,9 +228,7 @@ export type DatasetMetadataBuiltInFieldsResponse = { export type NotionEstimatePayload = { doc_form?: string doc_language?: string - notion_info_list: Array<{ - [key: string]: unknown - }> + notion_info_list: Array process_rule: { [key: string]: unknown } @@ -217,47 +252,6 @@ export type RetrievalSettingResponse = { retrieval_method: Array } -export type DatasetDetailWithPartialMembersResponse = { - app_count: number - author_name: string | null - built_in_field_enabled: boolean - chunk_structure: string | null - created_at: number - created_by: string - data_source_type: string | null - description: string | null - doc_form: string | null - doc_metadata: Array - document_count: number - embedding_available?: boolean | null - embedding_model: string | null - embedding_model_provider: string | null - enable_api: boolean - external_knowledge_info?: DatasetExternalKnowledgeInfoResponse - external_retrieval_model: DatasetExternalRetrievalModelResponse | null - icon_info?: DatasetIconInfoResponse - id: string - indexing_technique: string | null - is_multimodal: boolean - is_published: boolean - maintainer?: string | null - name: string - partial_member_list?: Array | null - permission: string - permission_keys?: Array - pipeline_id: string | null - provider: string - retrieval_model_dict: DatasetRetrievalModelResponse - runtime_mode: string | null - summary_index_setting?: DatasetSummaryIndexSettingResponse - tags: Array - total_available_documents: number - total_documents: number - updated_at: number - updated_by: string | null - word_count: number -} - export type DatasetUpdatePayload = { description?: string | null embedding_model?: string | null @@ -443,6 +437,10 @@ export type ConsoleSegmentListResponse = { total_pages: number } +export type BatchImportPayload = { + upload_file_id: string +} + export type SegmentUpdatePayload = { answer?: string | null attachment_ids?: Array | null @@ -713,6 +711,12 @@ export type DatasetMetadataBuiltInFieldResponse = { type: string } +export type NotionEstimateWorkspacePayload = { + credential_id: string + pages: Array + workspace_id: string +} + export type ProcessRuleMode = 'automatic' | 'custom' | 'hierarchical' export type Rule = { @@ -929,9 +933,14 @@ export type WeightModel = { weight_type?: 'customized' | 'keyword_first' | 'semantic_first' | null } +export type NotionEstimatePagePayload = { + page_id: string + type: NotionPageType +} + export type PreProcessingRule = { enabled: boolean - id: 'remove_extra_spaces' | 'remove_stopwords' | 'remove_urls_emails' + id: PreProcessingRuleKey } export type Segmentation = { @@ -1064,6 +1073,10 @@ export type WeightVectorSetting = { vector_weight: number } +export type NotionPageType = 'database' | 'page' + +export type PreProcessingRuleKey = 'remove_extra_spaces' | 'remove_stopwords' | 'remove_urls_emails' + export type HitTestingDocument = { data_source_type: string doc_metadata: unknown | null @@ -1126,7 +1139,7 @@ export type PostDatasetsErrors = { } export type PostDatasetsResponses = { - 201: DatasetDetailResponse + 201: DatasetDetailWithPartialMembersResponse } export type PostDatasetsResponse = PostDatasetsResponses[keyof PostDatasetsResponses] @@ -1209,22 +1222,6 @@ export type GetDatasetsBatchImportStatusByJobIdResponses = { export type GetDatasetsBatchImportStatusByJobIdResponse = GetDatasetsBatchImportStatusByJobIdResponses[keyof GetDatasetsBatchImportStatusByJobIdResponses] -export type PostDatasetsBatchImportStatusByJobIdData = { - body: BatchImportPayload - path: { - job_id: string - } - query?: never - url: '/datasets/batch_import_status/{job_id}' -} - -export type PostDatasetsBatchImportStatusByJobIdResponses = { - 200: SegmentBatchImportStatusResponse -} - -export type PostDatasetsBatchImportStatusByJobIdResponse = - PostDatasetsBatchImportStatusByJobIdResponses[keyof PostDatasetsBatchImportStatusByJobIdResponses] - export type PostDatasetsExternalData = { body: ExternalDatasetCreatePayload path?: never @@ -2028,23 +2025,6 @@ export type GetDatasetsByDatasetIdDocumentsByDocumentIdSegmentsResponses = { export type GetDatasetsByDatasetIdDocumentsByDocumentIdSegmentsResponse = GetDatasetsByDatasetIdDocumentsByDocumentIdSegmentsResponses[keyof GetDatasetsByDatasetIdDocumentsByDocumentIdSegmentsResponses] -export type GetDatasetsByDatasetIdDocumentsByDocumentIdSegmentsBatchImportData = { - body?: never - path: { - dataset_id: string - document_id: string - } - query?: never - url: '/datasets/{dataset_id}/documents/{document_id}/segments/batch_import' -} - -export type GetDatasetsByDatasetIdDocumentsByDocumentIdSegmentsBatchImportResponses = { - 200: SegmentBatchImportStatusResponse -} - -export type GetDatasetsByDatasetIdDocumentsByDocumentIdSegmentsBatchImportResponse = - GetDatasetsByDatasetIdDocumentsByDocumentIdSegmentsBatchImportResponses[keyof GetDatasetsByDatasetIdDocumentsByDocumentIdSegmentsBatchImportResponses] - export type PostDatasetsByDatasetIdDocumentsByDocumentIdSegmentsBatchImportData = { body: BatchImportPayload path: { diff --git a/packages/contracts/generated/api/console/datasets/zod.gen.ts b/packages/contracts/generated/api/console/datasets/zod.gen.ts index ee911690420..0d7a89c3f85 100644 --- a/packages/contracts/generated/api/console/datasets/zod.gen.ts +++ b/packages/contracts/generated/api/console/datasets/zod.gen.ts @@ -43,13 +43,6 @@ export const zSegmentBatchImportStatusResponse = z.object({ job_status: z.string(), }) -/** - * BatchImportPayload - */ -export const zBatchImportPayload = z.object({ - upload_file_id: z.string(), -}) - /** * ExternalDatasetCreatePayload * @@ -96,16 +89,6 @@ export const zIndexingEstimatePayload = z.object({ process_rule: z.record(z.string(), z.unknown()), }) -/** - * NotionEstimatePayload - */ -export const zNotionEstimatePayload = z.object({ - doc_form: z.string().optional().default('text_model'), - doc_language: z.string().optional().default('English'), - notion_info_list: z.array(z.record(z.string(), z.unknown())), - process_rule: z.record(z.string(), z.unknown()), -}) - /** * RetrievalSettingResponse */ @@ -209,6 +192,13 @@ export const zSegmentCreatePayload = z.object({ keywords: z.array(z.string()).nullish(), }) +/** + * BatchImportPayload + */ +export const zBatchImportPayload = z.object({ + upload_file_id: z.string(), +}) + /** * SegmentUpdatePayload */ @@ -828,14 +818,6 @@ export const zRetrievalMethod = z.enum([ 'semantic_search', ]) -/** - * PreProcessingRule - */ -export const zPreProcessingRule = z.object({ - enabled: z.boolean(), - id: z.enum(['remove_extra_spaces', 'remove_stopwords', 'remove_urls_emails']), -}) - /** * Segmentation */ @@ -845,33 +827,6 @@ export const zSegmentation = z.object({ separator: z.string().optional().default('\n'), }) -/** - * Rule - */ -export const zRule = z.object({ - parent_mode: z.enum(['full-doc', 'paragraph']).nullish(), - pre_processing_rules: z.array(zPreProcessingRule).nullish(), - segmentation: zSegmentation.nullish(), - subchunk_segmentation: zSegmentation.nullish(), -}) - -/** - * ProcessRuleResponse - */ -export const zProcessRuleResponse = z.object({ - limits: z.record(z.string(), z.unknown()), - mode: zProcessRuleMode, - rules: zRule.nullish(), -}) - -/** - * ProcessRule - */ -export const zProcessRule = z.object({ - mode: zProcessRuleMode, - rules: zRule.nullish(), -}) - /** * MetadataDetail */ @@ -1024,49 +979,6 @@ export const zDatasetRetrievalModelResponse = z.object({ weights: zDatasetWeightedScoreResponse.nullish(), }) -/** - * DatasetDetailResponse - */ -export const zDatasetDetailResponse = z.object({ - app_count: z.int(), - author_name: z.string().nullable(), - built_in_field_enabled: z.boolean(), - chunk_structure: z.string().nullable(), - created_at: z.int(), - created_by: z.string(), - data_source_type: z.string().nullable(), - description: z.string().nullable(), - doc_form: z.string().nullable(), - doc_metadata: z.array(zDatasetDocMetadataResponse), - document_count: z.int(), - embedding_available: z.boolean().nullish(), - embedding_model: z.string().nullable(), - embedding_model_provider: z.string().nullable(), - enable_api: z.boolean(), - external_knowledge_info: zDatasetExternalKnowledgeInfoResponse.optional(), - external_retrieval_model: zDatasetExternalRetrievalModelResponse.nullable(), - icon_info: zDatasetIconInfoResponse.optional(), - id: z.string(), - indexing_technique: z.string().nullable(), - is_multimodal: z.boolean(), - is_published: z.boolean(), - maintainer: z.string().nullish(), - name: z.string(), - permission: z.string(), - permission_keys: z.array(z.string()).optional(), - pipeline_id: z.string().nullable(), - provider: z.string(), - retrieval_model_dict: zDatasetRetrievalModelResponse, - runtime_mode: z.string().nullable(), - summary_index_setting: zDatasetSummaryIndexSettingResponse.optional(), - tags: z.array(zDatasetTagResponse), - total_available_documents: z.int(), - total_documents: z.int(), - updated_at: z.int(), - updated_by: z.string().nullable(), - word_count: z.int(), -}) - /** * DatasetDetailWithPartialMembersResponse */ @@ -1111,6 +1023,49 @@ export const zDatasetDetailWithPartialMembersResponse = z.object({ word_count: z.int(), }) +/** + * DatasetDetailResponse + */ +export const zDatasetDetailResponse = z.object({ + app_count: z.int(), + author_name: z.string().nullable(), + built_in_field_enabled: z.boolean(), + chunk_structure: z.string().nullable(), + created_at: z.int(), + created_by: z.string(), + data_source_type: z.string().nullable(), + description: z.string().nullable(), + doc_form: z.string().nullable(), + doc_metadata: z.array(zDatasetDocMetadataResponse), + document_count: z.int(), + embedding_available: z.boolean().nullish(), + embedding_model: z.string().nullable(), + embedding_model_provider: z.string().nullable(), + enable_api: z.boolean(), + external_knowledge_info: zDatasetExternalKnowledgeInfoResponse.optional(), + external_retrieval_model: zDatasetExternalRetrievalModelResponse.nullable(), + icon_info: zDatasetIconInfoResponse.optional(), + id: z.string(), + indexing_technique: z.string().nullable(), + is_multimodal: z.boolean(), + is_published: z.boolean(), + maintainer: z.string().nullish(), + name: z.string(), + permission: z.string(), + permission_keys: z.array(z.string()).optional(), + pipeline_id: z.string().nullable(), + provider: z.string(), + retrieval_model_dict: zDatasetRetrievalModelResponse, + runtime_mode: z.string().nullable(), + summary_index_setting: zDatasetSummaryIndexSettingResponse.optional(), + tags: z.array(zDatasetTagResponse), + total_available_documents: z.int(), + total_documents: z.int(), + updated_at: z.int(), + updated_by: z.string().nullable(), + word_count: z.int(), +}) + /** * DatasetListItemResponse */ @@ -1279,6 +1234,79 @@ export const zHitTestingPayload = z.object({ retrieval_model: zRetrievalModel.nullish(), }) +/** + * NotionPageType + */ +export const zNotionPageType = z.enum(['database', 'page']) + +/** + * NotionEstimatePagePayload + */ +export const zNotionEstimatePagePayload = z.object({ + page_id: z.string().min(1), + type: zNotionPageType, +}) + +/** + * NotionEstimateWorkspacePayload + */ +export const zNotionEstimateWorkspacePayload = z.object({ + credential_id: z.string().min(1), + pages: z.array(zNotionEstimatePagePayload).min(1), + workspace_id: z.string().min(1), +}) + +/** + * NotionEstimatePayload + */ +export const zNotionEstimatePayload = z.object({ + doc_form: z.string().optional().default('text_model'), + doc_language: z.string().optional().default('English'), + notion_info_list: z.array(zNotionEstimateWorkspacePayload).min(1), + process_rule: z.record(z.string(), z.unknown()), +}) + +export const zPreProcessingRuleKey = z.enum([ + 'remove_extra_spaces', + 'remove_stopwords', + 'remove_urls_emails', +]) + +/** + * PreProcessingRule + */ +export const zPreProcessingRule = z.object({ + enabled: z.boolean(), + id: zPreProcessingRuleKey, +}) + +/** + * Rule + */ +export const zRule = z.object({ + parent_mode: z.enum(['full-doc', 'paragraph']).nullish(), + pre_processing_rules: z.array(zPreProcessingRule).nullish(), + segmentation: zSegmentation.nullish(), + subchunk_segmentation: zSegmentation.nullish(), +}) + +/** + * ProcessRuleResponse + */ +export const zProcessRuleResponse = z.object({ + limits: z.record(z.string(), z.unknown()), + mode: zProcessRuleMode, + rules: zRule.nullish(), +}) + +/** + * ProcessRule + */ +export const zProcessRule = z.object({ + mode: zProcessRuleMode, + rules: zRule.nullish(), +}) + /** * HitTestingDocument */ @@ -1477,7 +1505,7 @@ export const zPostDatasetsBody = zDatasetCreatePayload /** * Dataset created successfully */ -export const zPostDatasetsResponse = zDatasetDetailResponse +export const zPostDatasetsResponse = zDatasetDetailWithPartialMembersResponse /** * API base info retrieved successfully @@ -1514,17 +1542,6 @@ export const zGetDatasetsBatchImportStatusByJobIdPath = z.object({ */ export const zGetDatasetsBatchImportStatusByJobIdResponse = zSegmentBatchImportStatusResponse -export const zPostDatasetsBatchImportStatusByJobIdBody = zBatchImportPayload - -export const zPostDatasetsBatchImportStatusByJobIdPath = z.object({ - job_id: z.uuid(), -}) - -/** - * Batch import started - */ -export const zPostDatasetsBatchImportStatusByJobIdResponse = zSegmentBatchImportStatusResponse - export const zPostDatasetsExternalBody = zExternalDatasetCreatePayload /** @@ -1990,17 +2007,6 @@ export const zGetDatasetsByDatasetIdDocumentsByDocumentIdSegmentsQuery = z.objec export const zGetDatasetsByDatasetIdDocumentsByDocumentIdSegmentsResponse = zConsoleSegmentListResponse -export const zGetDatasetsByDatasetIdDocumentsByDocumentIdSegmentsBatchImportPath = z.object({ - dataset_id: z.uuid(), - document_id: z.uuid(), -}) - -/** - * Batch import status - */ -export const zGetDatasetsByDatasetIdDocumentsByDocumentIdSegmentsBatchImportResponse = - zSegmentBatchImportStatusResponse - export const zPostDatasetsByDatasetIdDocumentsByDocumentIdSegmentsBatchImportBody = zBatchImportPayload diff --git a/packages/contracts/generated/api/console/notion/types.gen.ts b/packages/contracts/generated/api/console/notion/types.gen.ts index 43545a9721d..92a1ebe6011 100644 --- a/packages/contracts/generated/api/console/notion/types.gen.ts +++ b/packages/contracts/generated/api/console/notion/types.gen.ts @@ -25,7 +25,7 @@ export type NotionIntegratePageResponse = { page_id: string page_name: string parent_id: string | null - type: string + type: NotionPageType } export type DataSourceIntegrateIconResponse = { @@ -34,6 +34,8 @@ export type DataSourceIntegrateIconResponse = { url?: string | null } +export type NotionPageType = 'database' | 'page' + export type GetNotionPagesByPageIdByPageTypePreviewData = { body?: never path: { diff --git a/packages/contracts/generated/api/console/notion/zod.gen.ts b/packages/contracts/generated/api/console/notion/zod.gen.ts index 7f228414f19..f144e14845e 100644 --- a/packages/contracts/generated/api/console/notion/zod.gen.ts +++ b/packages/contracts/generated/api/console/notion/zod.gen.ts @@ -18,6 +18,11 @@ export const zDataSourceIntegrateIconResponse = z.object({ url: z.string().nullish(), }) +/** + * NotionPageType + */ +export const zNotionPageType = z.enum(['database', 'page']) + /** * NotionIntegratePageResponse */ @@ -27,7 +32,7 @@ export const zNotionIntegratePageResponse = z.object({ page_id: z.string(), page_name: z.string(), parent_id: z.string().nullable(), - type: z.string(), + type: zNotionPageType, }) /** diff --git a/packages/contracts/generated/api/service/types.gen.ts b/packages/contracts/generated/api/service/types.gen.ts index b51d4c34d93..31940741fcd 100644 --- a/packages/contracts/generated/api/service/types.gen.ts +++ b/packages/contracts/generated/api/service/types.gen.ts @@ -1777,9 +1777,11 @@ export type PipelineUploadFileResponse = { export type PreProcessingRule = { enabled: boolean - id: 'remove_extra_spaces' | 'remove_stopwords' | 'remove_urls_emails' + id: PreProcessingRuleKey } +export type PreProcessingRuleKey = 'remove_extra_spaces' | 'remove_stopwords' | 'remove_urls_emails' + export type ProcessRule = { mode: ProcessRuleMode rules?: Rule | null diff --git a/packages/contracts/generated/api/service/zod.gen.ts b/packages/contracts/generated/api/service/zod.gen.ts index 9b2bcb161b9..24cfe744072 100644 --- a/packages/contracts/generated/api/service/zod.gen.ts +++ b/packages/contracts/generated/api/service/zod.gen.ts @@ -1664,12 +1664,18 @@ export const zPipelineUploadFileResponse = z.object({ size: z.int(), }) +export const zPreProcessingRuleKey = z.enum([ + 'remove_extra_spaces', + 'remove_stopwords', + 'remove_urls_emails', +]) + /** * PreProcessingRule */ export const zPreProcessingRule = z.object({ enabled: z.boolean(), - id: z.enum(['remove_extra_spaces', 'remove_stopwords', 'remove_urls_emails']), + id: zPreProcessingRuleKey, }) /** diff --git a/web/app/components/base/notion-page-selector/__tests__/base.spec.tsx b/web/app/components/base/notion-page-selector/__tests__/base.spec.tsx index 4995e87fc40..8975198bbc8 100644 --- a/web/app/components/base/notion-page-selector/__tests__/base.spec.tsx +++ b/web/app/components/base/notion-page-selector/__tests__/base.spec.tsx @@ -145,6 +145,7 @@ describe('NotionPageSelector Base', () => { expect(screen.getByTestId('notion-page-selector-base')).toBeInTheDocument() expect(screen.getByText('Root 1')).toBeInTheDocument() + expect(screen.getByText('External 1')).toBeInTheDocument() const checkbox = screen.getByRole('checkbox', { name: 'Root 1' }) await user.click(checkbox) diff --git a/web/app/components/base/notion-page-selector/base.tsx b/web/app/components/base/notion-page-selector/base.tsx index 6fccd5e6b13..c0762bb0cef 100644 --- a/web/app/components/base/notion-page-selector/base.tsx +++ b/web/app/components/base/notion-page-selector/base.tsx @@ -100,6 +100,7 @@ const NotionPageSelector = ({ isFetching: isFetchingNotionPages, isError: isFetchingNotionPagesError, } = usePreImportNotionPages({ datasetId, credentialId: currentCredentialId }) + const notionPages = notionsPages?.notion_info?.flatMap((workspace) => workspace.pages) || [] const pagesMapAndSelectedPagesId: [DataSourceNotionPageMap, Set, Set] = useMemo(() => { @@ -206,7 +207,7 @@ const NotionPageSelector = ({ value={selectedPagesId} disabledValue={pagesMapAndSelectedPagesId[2]} searchValue={searchValue} - list={notionsPages!.notion_info?.[0]!.pages || []} + list={notionPages} pagesMap={pagesMapAndSelectedPagesId[0]} onSelect={handleSelectPages} canPreview={canPreview}