diff --git a/api/controllers/console/datasets/hit_testing_base.py b/api/controllers/console/datasets/hit_testing_base.py index ef8a34bccd89a3..eeb727c3bf442e 100644 --- a/api/controllers/console/datasets/hit_testing_base.py +++ b/api/controllers/console/datasets/hit_testing_base.py @@ -3,9 +3,9 @@ from pydantic import BaseModel, Field from sqlalchemy.orm import Session -from werkzeug.exceptions import Forbidden, InternalServerError, NotFound import services.errors.base +from controllers.common.errors import AccessDeniedError, InternalServerError, NotFoundError from controllers.console.app.error import ( CompletionRequestError, ProviderModelCurrentlyNotSupportError, @@ -19,6 +19,7 @@ ProviderTokenNotInitError, QuotaExceededError, ) +from extensions.ext_application_services import application_services from graphon.model_runtime.errors.invoke import InvokeError from libs.login import resolve_account_fallback from models.account import Account @@ -93,12 +94,12 @@ def get_and_validate_dataset( current_user, _ = resolve_account_fallback(current_user, current_tenant_id) dataset = DatasetService.get_dataset(dataset_id, session) if dataset is None: - raise NotFound("Dataset not found.") + raise NotFoundError("Dataset not found.") try: DatasetService.check_dataset_permission(dataset, current_user, session) except services.errors.base.NoPermissionError as e: - raise Forbidden(str(e)) from e + raise AccessDeniedError(str(e)) from e return dataset @@ -123,6 +124,7 @@ def perform_hit_testing( try: current_user, _ = resolve_account_fallback(current_user, current_tenant_id) response = HitTestingService.retrieve( + retrieval=application_services().workflow_runtime.retrieval, session=session, dataset=dataset, query=cast(str, args.get("query")), diff --git a/api/controllers/inner_api/knowledge/retrieval.py b/api/controllers/inner_api/knowledge/retrieval.py index 9d83007aa3d727..dc48b3dee3497d 100644 --- a/api/controllers/inner_api/knowledge/retrieval.py +++ b/api/controllers/inner_api/knowledge/retrieval.py @@ -9,17 +9,15 @@ from flask_restx import Resource from pydantic import ValidationError -from sqlalchemy.orm import Session from controllers.common.schema import register_response_schema_models, register_schema_models -from controllers.console.app.wraps import with_session from controllers.inner_api import inner_api_ns from controllers.inner_api.wraps import plugin_inner_api_only from core.workflow.nodes.knowledge_retrieval import exc as retrieval_exc +from extensions.ext_application_services import application_services from libs.exception import BaseHTTPException from services.entities.knowledge_retrieval_inner import InnerKnowledgeRetrieveRequest, InnerKnowledgeRetrieveResponse from services.errors.knowledge_retrieval import ExternalKnowledgeRetrievalError, InnerKnowledgeRetrievalServiceError -from services.knowledge_retrieval_inner_service import InnerKnowledgeRetrievalService class InnerKnowledgeRetrievalHttpError(BaseHTTPException): @@ -71,8 +69,7 @@ class InnerKnowledgeRetrieveApi(Resource): 500: "Unexpected knowledge retrieval failure", } ) - @with_session - def post(self, session: Session) -> dict[str, object]: + def post(self) -> dict[str, object]: """Validate the payload, run retrieval, and return workflow-style sources.""" try: payload = InnerKnowledgeRetrieveRequest.model_validate(inner_api_ns.payload or {}) @@ -84,7 +81,7 @@ def post(self, session: Session) -> dict[str, object]: ) from exc try: - response = InnerKnowledgeRetrievalService().retrieve(payload, session=session) + response = application_services().knowledge.inner_retrieval.retrieve(payload) except InnerKnowledgeRetrievalServiceError as exc: raise InnerKnowledgeRetrievalHttpError( error_code=exc.error_code, diff --git a/api/extensions/application_services/knowledge.py b/api/extensions/application_services/knowledge.py index e632ea6c82b56a..3e0f98412269b0 100644 --- a/api/extensions/application_services/knowledge.py +++ b/api/extensions/application_services/knowledge.py @@ -8,12 +8,14 @@ from configs import dify_config from core.app.apps.pipeline.pipeline_generator import PipelineGenerator from core.rag.extractor.entity.datasource_type import DatasourceType +from extensions.application_services.retrieval import build_dataset_retrieval from libs.helper import generate_text_hash from repositories.app.console_repository import ConsoleAppRepository 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.retrieval_repository import KnowledgeRetrievalRepository from repositories.knowledge.segment_repository import SQLAlchemySegmentRepository from repositories.knowledge.upload_file_repository import SQLAlchemyKnowledgeUploadRepository from services.api_token_service import ApiTokenCache, get_effective_token_last_used_at @@ -52,6 +54,7 @@ from services.knowledge.segments.application import DatasetSegmentApplicationService from services.knowledge.segments.indexing import SegmentIndexingGateway from services.knowledge_fs.feature_service import get_knowledge_fs_features +from services.knowledge_retrieval_inner_service import InnerKnowledgeRetrievalService from services.rbac.members import MemberService from services.tag_application_service import TagTargetQuery from services.workflow_run_agg import WorkflowRunAgg @@ -64,6 +67,7 @@ @dataclass(frozen=True, slots=True) class KnowledgeServices: + inner_retrieval: InnerKnowledgeRetrievalService metadata: MetadataService datasets: DatasetApplicationService external: ExternalKnowledgeApplicationService @@ -133,6 +137,9 @@ def build_knowledge_services( members=rbac_members, ) return KnowledgeServices( + inner_retrieval=InnerKnowledgeRetrievalService( + scopes=KnowledgeRetrievalRepository(database_client), retrieval=build_dataset_retrieval(database_client) + ), metadata=MetadataService( store=SQLAlchemyMetadataRepository(session_factory=database_client), dataset_access=dataset_access ), diff --git a/api/extensions/ext_application_services.py b/api/extensions/ext_application_services.py index 906e4cb840dec0..60ff42b12edff4 100644 --- a/api/extensions/ext_application_services.py +++ b/api/extensions/ext_application_services.py @@ -40,7 +40,7 @@ from extensions.application_services.rbac import RBACServices, build_rbac_services from extensions.application_services.resource_access_token import build_resource_access_token_service from extensions.application_services.trial_app import TrialAppServices, build_trial_app_services -from extensions.application_services.workflow import build_workflow_suggestions +from extensions.application_services.workflow import build_workflow_execution_dependencies, build_workflow_suggestions from extensions.application_services.workflow_variables import ( build_console_workflow_variables, build_workflow_variable_service, @@ -123,6 +123,7 @@ from services.api_based_extension_application_service import APIBasedExtensionApplicationService from services.app.advanced_prompt_template_service import AdvancedPromptTemplateService from services.app.api_key_service import AppApiKeyService +from services.app.generation.runtime import AppGenerationRuntime from services.app.mcp_server_service import AppMCPServerService from services.app.workflow_comment_mention_gateway import CeleryWorkflowCommentMentionNotifier from services.app.workflow_comment_service import WorkflowCommentService @@ -302,6 +303,7 @@ class AppScopedEndUserServices: @dataclass(frozen=True, slots=True) class ApplicationServices: + workflow_runtime: AppGenerationRuntime workflow_variables: WorkflowVariableService console_workflow_variables: ConsoleWorkflowVariableService workflow_suggestions: WorkflowInstructionSuggestions @@ -614,6 +616,7 @@ def build_application_services( ) return ApplicationServices( + workflow_runtime=build_workflow_execution_dependencies(database_client), workflow_variables=workflow_variables, console_workflow_variables=build_console_workflow_variables( database_client=database_client, variables=workflow_variables diff --git a/api/services/hit_testing_service.py b/api/services/hit_testing_service.py index 628fec5534fe51..a66babff4dec91 100644 --- a/api/services/hit_testing_service.py +++ b/api/services/hit_testing_service.py @@ -12,13 +12,13 @@ from core.rag.index_processor.constant.index_type import IndexTechniqueType from core.rag.index_processor.constant.query_type import QueryType from core.rag.models.document import Document -from core.rag.retrieval.dataset_retrieval import DatasetRetrieval from core.rag.retrieval.retrieval_methods import RetrievalMethod from graphon.model_runtime.entities import LLMMode from models import Account from models.dataset import Dataset, DatasetQuery from models.dataset import Document as DatasetDocument from models.enums import CreatorUserRole, DatasetQuerySource +from services.knowledge.retrieval.ports import DatasetRetrievalFactory logger = logging.getLogger(__name__) @@ -113,6 +113,7 @@ def retrieve( limit: int = 10, *, session: Session, + retrieval: DatasetRetrievalFactory, ): start = time.perf_counter() @@ -124,20 +125,19 @@ def retrieve( document_ids_filter = None metadata_filtering_conditions_raw = resolved_retrieval_model.get("metadata_filtering_conditions", {}) if metadata_filtering_conditions_raw and query: - dataset_retrieval = DatasetRetrieval() + dataset_retrieval = retrieval() from core.rag.entities import MetadataFilteringCondition metadata_filtering_conditions = MetadataFilteringCondition.model_validate(metadata_filtering_conditions_raw) metadata_filter_document_ids, metadata_condition = dataset_retrieval.get_metadata_filter_condition( - session=session, dataset_ids=[dataset.id], query=query, metadata_filtering_mode="manual", metadata_filtering_conditions=metadata_filtering_conditions, inputs={}, - tenant_id="", + tenant_id=dataset.tenant_id, user_id="", metadata_model_config=ModelConfig(provider="", name="", mode=LLMMode.CHAT, completion_params={}), ) diff --git a/api/services/knowledge/retrieval/ports.py b/api/services/knowledge/retrieval/ports.py index 90121b48aa8791..b7ebeeb2b4f984 100644 --- a/api/services/knowledge/retrieval/ports.py +++ b/api/services/knowledge/retrieval/ports.py @@ -78,7 +78,7 @@ def __call__( class DatasetRetriever(Protocol): - """Retrieval operations consumed by the workflow execution layer.""" + """Retrieval operations consumed by workflow execution and hit testing.""" @property def llm_usage(self) -> LLMUsage: ... diff --git a/api/services/knowledge_retrieval_inner_service.py b/api/services/knowledge_retrieval_inner_service.py index c10c413d4db643..7a235eda3513fe 100644 --- a/api/services/knowledge_retrieval_inner_service.py +++ b/api/services/knowledge_retrieval_inner_service.py @@ -1,7 +1,7 @@ """Service wrapper for the inner knowledge retrieval API. This service keeps the internal HTTP contract small while reusing the workflow -retrieval stack in ``core.rag.retrieval.dataset_retrieval.DatasetRetrieval``. +retrieval stack in ``services.knowledge.retrieval.dataset_retrieval.DatasetRetrieval``. The only authorization enforced here is tenant ownership of the caller app and requested datasets. @@ -12,36 +12,28 @@ of a separate validation error. """ -from sqlalchemy import select -from sqlalchemy.orm import Session - from core.rag.entities.metadata_entities import Condition, MetadataFilteringCondition -from core.rag.retrieval.dataset_retrieval import DatasetRetrieval from core.workflow.nodes.knowledge_retrieval.retrieval import KnowledgeRetrievalRequest from graphon.model_runtime.utils.encoders import jsonable_encoder from graphon.nodes.llm.entities import ModelConfig -from models.dataset import Dataset -from models.model import App from services.entities.knowledge_retrieval_inner import ( InnerKnowledgeRetrieveRequest, InnerKnowledgeRetrieveResponse, InnerKnowledgeRetrieveUsage, ) -from services.errors.knowledge_retrieval import ( - InnerKnowledgeRetrieveAppNotFoundError, - InnerKnowledgeRetrieveAppTenantMismatchError, - InnerKnowledgeRetrieveDatasetNotFoundError, - InnerKnowledgeRetrieveDatasetTenantMismatchError, -) +from services.knowledge.retrieval.ports import DatasetRetrievalFactory, RetrievalScopeQueries class InnerKnowledgeRetrievalService: """Validate inner caller scope and delegate to workflow dataset retrieval.""" + def __init__(self, *, scopes: RetrievalScopeQueries, retrieval: DatasetRetrievalFactory): + self._scopes = scopes + self._retrieval = retrieval + def retrieve( self, request: InnerKnowledgeRetrieveRequest, - session: Session, ) -> InnerKnowledgeRetrieveResponse: """Run tenant-scoped retrieval for a trusted internal caller. @@ -60,39 +52,17 @@ def retrieve( InnerKnowledgeRetrieveDatasetTenantMismatchError: At least one requested dataset is outside the caller tenant. """ - self._validate_caller_app(tenant_id=request.caller.tenant_id, app_id=request.caller.app_id, session=session) - self._validate_datasets(tenant_id=request.caller.tenant_id, dataset_ids=request.dataset_ids, session=session) + self._scopes.validate_scope( + tenant_id=request.caller.tenant_id, app_id=request.caller.app_id, dataset_ids=request.dataset_ids + ) - rag = DatasetRetrieval() - results = rag.knowledge_retrieval(session=session, request=self._to_rag_request(request)) + rag = self._retrieval() + results = rag.knowledge_retrieval(request=self._to_rag_request(request)) return InnerKnowledgeRetrieveResponse( results=results, usage=InnerKnowledgeRetrieveUsage.model_validate(jsonable_encoder(rag.llm_usage)), ) - def _validate_caller_app(self, *, tenant_id: str, app_id: str, session: Session) -> None: - app = session.scalar(select(App).where(App.id == app_id).limit(1)) - if app is None: - raise InnerKnowledgeRetrieveAppNotFoundError(f"App '{app_id}' not found") - if app.tenant_id != tenant_id: - raise InnerKnowledgeRetrieveAppTenantMismatchError( - f"App '{app_id}' does not belong to tenant '{tenant_id}'" - ) - - def _validate_datasets(self, *, tenant_id: str, dataset_ids: list[str], session: Session) -> None: - datasets = session.scalars(select(Dataset).where(Dataset.id.in_(dataset_ids))).all() - - found_ids = {dataset.id for dataset in datasets} - missing_ids = sorted(set(dataset_ids) - found_ids) - if missing_ids: - raise InnerKnowledgeRetrieveDatasetNotFoundError(f"Datasets not found: {', '.join(missing_ids)}") - - mismatched_ids = sorted(dataset.id for dataset in datasets if dataset.tenant_id != tenant_id) - if mismatched_ids: - raise InnerKnowledgeRetrieveDatasetTenantMismatchError( - f"Datasets do not belong to tenant '{tenant_id}': {', '.join(mismatched_ids)}" - ) - def _to_rag_request(self, request: InnerKnowledgeRetrieveRequest) -> KnowledgeRetrievalRequest: metadata_model_config = request.metadata_filtering.metadata_model_config metadata_conditions = request.metadata_filtering.conditions diff --git a/api/tasks/annotation/__init__.py b/api/tasks/annotation/__init__.py new file mode 100644 index 00000000000000..39cc12bf6cb876 --- /dev/null +++ b/api/tasks/annotation/__init__.py @@ -0,0 +1 @@ +"""Annotation indexing and import tasks.""" diff --git a/api/tasks/annotation/add_annotation_to_index_task.py b/api/tasks/annotation/add_annotation_to_index_task.py index 509b9a0355c234..ddf9e760dd078d 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.index_processor.constant.index_type import IndexTechniqueType from core.rag.models.document import Document from models.dataset import Dataset -from services.knowledge.dataset_service import DatasetCollectionBindingService +from repositories.knowledge.collection_binding_repository import DatasetCollectionBindingRepository logger = logging.getLogger(__name__) @@ -33,8 +33,10 @@ def add_annotation_to_index_task( try: with session_factory.create_session() as session: - dataset_collection_binding = DatasetCollectionBindingService.get_dataset_collection_binding_by_id_and_type( - collection_binding_id, session, "annotation" + dataset_collection_binding = ( + DatasetCollectionBindingRepository.get_dataset_collection_binding_by_id_and_type( + collection_binding_id, session, "annotation" + ) ) dataset = Dataset( id=app_id, diff --git a/api/tasks/annotation/batch_import_annotations_task.py b/api/tasks/annotation/batch_import_annotations_task.py index 7ec50de79a8f8a..2c9e3653a1bf28 100644 --- a/api/tasks/annotation/batch_import_annotations_task.py +++ b/api/tasks/annotation/batch_import_annotations_task.py @@ -4,7 +4,6 @@ import click from celery import shared_task from sqlalchemy import select -from werkzeug.exceptions import NotFound from core.db.session_factory import session_factory from core.rag.datasource.vdb.vector_factory import Vector @@ -14,8 +13,9 @@ from models.dataset import Dataset from models.model import App, AppAnnotationSetting, MessageAnnotation from repositories.annotation_import_job_repository import annotation_import_job_owner_key +from repositories.knowledge.collection_binding_repository import DatasetCollectionBindingRepository +from services.annotation.errors import AnnotationResourceNotFoundError from services.annotation_import_service import AnnotationImportRecord -from services.knowledge.dataset_service import DatasetCollectionBindingService logger = logging.getLogger(__name__) @@ -66,12 +66,12 @@ def batch_import_annotations_task( if app_annotation_setting: dataset_collection_binding = ( - DatasetCollectionBindingService.get_dataset_collection_binding_by_id_and_type( + DatasetCollectionBindingRepository.get_dataset_collection_binding_by_id_and_type( app_annotation_setting.collection_binding_id, session, "annotation" ) ) if not dataset_collection_binding: - raise NotFound("App annotation setting not found") + raise AnnotationResourceNotFoundError("App annotation setting not found") dataset = Dataset( id=app_id, tenant_id=tenant_id, diff --git a/api/tasks/annotation/delete_annotation_index_task.py b/api/tasks/annotation/delete_annotation_index_task.py index f86d06e153fc7a..ada2ead3a043db 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.rag.datasource.vdb.vector_factory import Vector from core.rag.index_processor.constant.index_type import IndexTechniqueType from models.dataset import Dataset -from services.knowledge.dataset_service import DatasetCollectionBindingService +from repositories.knowledge.collection_binding_repository import DatasetCollectionBindingRepository logger = logging.getLogger(__name__) @@ -22,8 +22,10 @@ def delete_annotation_index_task(annotation_id: str, app_id: str, tenant_id: str start_at = time.perf_counter() try: with session_factory.create_session() as session: - dataset_collection_binding = DatasetCollectionBindingService.get_dataset_collection_binding_by_id_and_type( - collection_binding_id, session, "annotation" + dataset_collection_binding = ( + DatasetCollectionBindingRepository.get_dataset_collection_binding_by_id_and_type( + collection_binding_id, session, "annotation" + ) ) dataset = Dataset( diff --git a/api/tasks/annotation/enable_annotation_reply_task.py b/api/tasks/annotation/enable_annotation_reply_task.py index 56a1a49cf8d40b..c0edb87b3e60aa 100644 --- a/api/tasks/annotation/enable_annotation_reply_task.py +++ b/api/tasks/annotation/enable_annotation_reply_task.py @@ -19,7 +19,7 @@ RedisAnnotationReplyJobRepository, annotation_reply_job_owner_key, ) -from services.knowledge.dataset_service import DatasetCollectionBindingService +from repositories.knowledge.collection_binding_repository import DatasetCollectionBindingRepository logger = logging.getLogger(__name__) @@ -57,7 +57,7 @@ def enable_annotation_reply_task( try: documents = [] - dataset_collection_binding = DatasetCollectionBindingService.get_dataset_collection_binding( + dataset_collection_binding = DatasetCollectionBindingRepository.get_dataset_collection_binding( embedding_provider_name, embedding_model_name, session, CollectionBindingType.ANNOTATION ) annotation_setting = session.scalar( @@ -66,7 +66,7 @@ def enable_annotation_reply_task( if annotation_setting: if dataset_collection_binding.id != annotation_setting.collection_binding_id: old_dataset_collection_binding = ( - DatasetCollectionBindingService.get_dataset_collection_binding_by_id_and_type( + DatasetCollectionBindingRepository.get_dataset_collection_binding_by_id_and_type( annotation_setting.collection_binding_id, session, CollectionBindingType.ANNOTATION ) ) diff --git a/api/tasks/annotation/update_annotation_to_index_task.py b/api/tasks/annotation/update_annotation_to_index_task.py index 611934c5cb1dfc..e53f657704c6cd 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.index_processor.constant.index_type import IndexTechniqueType from core.rag.models.document import Document from models.dataset import Dataset -from services.knowledge.dataset_service import DatasetCollectionBindingService +from repositories.knowledge.collection_binding_repository import DatasetCollectionBindingRepository logger = logging.getLogger(__name__) @@ -33,8 +33,10 @@ def update_annotation_to_index_task( try: with session_factory.create_session() as session: - dataset_collection_binding = DatasetCollectionBindingService.get_dataset_collection_binding_by_id_and_type( - collection_binding_id, session, "annotation" + dataset_collection_binding = ( + DatasetCollectionBindingRepository.get_dataset_collection_binding_by_id_and_type( + collection_binding_id, session, "annotation" + ) ) dataset = Dataset( diff --git a/api/tests/e2e_tests/containers/core/rag/retrieval/test_dataset_retrieval_integration.py b/api/tests/e2e_tests/containers/core/rag/retrieval/test_dataset_retrieval_integration.py index 3976ea9fec4112..e218d1260d2d2f 100644 --- a/api/tests/e2e_tests/containers/core/rag/retrieval/test_dataset_retrieval_integration.py +++ b/api/tests/e2e_tests/containers/core/rag/retrieval/test_dataset_retrieval_integration.py @@ -3,17 +3,24 @@ import pytest from faker import Faker -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, sessionmaker from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType -from core.rag.retrieval.dataset_retrieval import DatasetRetrieval from core.workflow.nodes.knowledge_retrieval.retrieval import KnowledgeRetrievalRequest +from extensions.application_services.retrieval import build_dataset_retrieval from models.dataset import Dataset, Document from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus +from services.knowledge.retrieval.dataset_retrieval import DatasetRetrieval from tests.e2e_tests.containers.helpers import accounts as account_fixtures from tests.e2e_tests.containers.helpers import generate_valid_password +def _build_dataset_retrieval(session: Session) -> DatasetRetrieval: + retrieval = build_dataset_retrieval(sessionmaker(bind=session.get_bind()))() + assert isinstance(retrieval, DatasetRetrieval) + return retrieval + + class TestGetAvailableDatasetsIntegration: def test_returns_datasets_with_available_documents( self, db_session_with_containers: Session, mock_external_service_dependencies @@ -69,8 +76,8 @@ def test_returns_datasets_with_available_documents( db_session_with_containers.commit() # Act - dataset_retrieval = DatasetRetrieval() - result = dataset_retrieval._get_available_datasets(tenant.id, [dataset.id]) + dataset_retrieval = _build_dataset_retrieval(db_session_with_containers) + result = dataset_retrieval._records.available_datasets(tenant.id, [dataset.id]) # Assert assert len(result) == 1 @@ -126,8 +133,8 @@ def test_filters_out_datasets_with_only_archived_documents( db_session_with_containers.commit() # Act - dataset_retrieval = DatasetRetrieval() - result = dataset_retrieval._get_available_datasets(tenant.id, [dataset.id]) + dataset_retrieval = _build_dataset_retrieval(db_session_with_containers) + result = dataset_retrieval._records.available_datasets(tenant.id, [dataset.id]) # Assert assert len(result) == 0 @@ -180,8 +187,8 @@ def test_filters_out_datasets_with_only_disabled_documents( db_session_with_containers.commit() # Act - dataset_retrieval = DatasetRetrieval() - result = dataset_retrieval._get_available_datasets(tenant.id, [dataset.id]) + dataset_retrieval = _build_dataset_retrieval(db_session_with_containers) + result = dataset_retrieval._records.available_datasets(tenant.id, [dataset.id]) # Assert assert len(result) == 0 @@ -234,8 +241,8 @@ def test_filters_out_datasets_with_non_completed_documents( db_session_with_containers.commit() # Act - dataset_retrieval = DatasetRetrieval() - result = dataset_retrieval._get_available_datasets(tenant.id, [dataset.id]) + dataset_retrieval = _build_dataset_retrieval(db_session_with_containers) + result = dataset_retrieval._records.available_datasets(tenant.id, [dataset.id]) # Assert assert len(result) == 0 @@ -278,8 +285,8 @@ def test_includes_external_datasets_without_documents( db_session_with_containers.commit() # Act - dataset_retrieval = DatasetRetrieval() - result = dataset_retrieval._get_available_datasets(tenant.id, [dataset.id]) + dataset_retrieval = _build_dataset_retrieval(db_session_with_containers) + result = dataset_retrieval._records.available_datasets(tenant.id, [dataset.id]) # Assert assert len(result) == 1 @@ -355,8 +362,8 @@ def test_filters_by_tenant_id(self, db_session_with_containers: Session, mock_ex db_session_with_containers.commit() # Act - request from tenant1, should only get tenant1's dataset - dataset_retrieval = DatasetRetrieval() - result = dataset_retrieval._get_available_datasets(tenant1.id, [dataset1.id, dataset2.id]) + dataset_retrieval = _build_dataset_retrieval(db_session_with_containers) + result = dataset_retrieval._records.available_datasets(tenant1.id, [dataset1.id, dataset2.id]) # Assert assert len(result) == 1 @@ -382,8 +389,8 @@ def test_returns_empty_list_when_no_datasets_found( # Don't create any datasets # Act - dataset_retrieval = DatasetRetrieval() - result = dataset_retrieval._get_available_datasets(tenant.id, [str(uuid.uuid4())]) + dataset_retrieval = _build_dataset_retrieval(db_session_with_containers) + result = dataset_retrieval._records.available_datasets(tenant.id, [str(uuid.uuid4())]) # Assert assert result == [] @@ -439,9 +446,9 @@ def test_returns_only_requested_dataset_ids( db_session_with_containers.commit() # Act - request only dataset 0 and 2, not dataset 1 - dataset_retrieval = DatasetRetrieval() + dataset_retrieval = _build_dataset_retrieval(db_session_with_containers) requested_ids = [datasets[0].id, datasets[2].id] - result = dataset_retrieval._get_available_datasets(tenant.id, requested_ids) + result = dataset_retrieval._records.available_datasets(tenant.id, requested_ids) # Assert assert len(result) == 2 @@ -506,14 +513,14 @@ def test_knowledge_retrieval_with_available_datasets( top_k=5, ) - dataset_retrieval = DatasetRetrieval() + dataset_retrieval = _build_dataset_retrieval(db_session_with_containers) # Mock rate limit check and retrieval with patch.object(dataset_retrieval, "_check_knowledge_rate_limit"): with patch.object(dataset_retrieval, "get_metadata_filter_condition", return_value=(None, None)): with patch.object(dataset_retrieval, "multiple_retrieve", return_value=[]): # Act - result = dataset_retrieval.knowledge_retrieval(db_session_with_containers, request) + result = dataset_retrieval.knowledge_retrieval(request) # Assert assert isinstance(result, list) @@ -557,12 +564,12 @@ def test_knowledge_retrieval_no_available_datasets( top_k=5, ) - dataset_retrieval = DatasetRetrieval() + dataset_retrieval = _build_dataset_retrieval(db_session_with_containers) # Mock rate limit check with patch.object(dataset_retrieval, "_check_knowledge_rate_limit"): # Act - result = dataset_retrieval.knowledge_retrieval(db_session_with_containers, request) + result = dataset_retrieval.knowledge_retrieval(request) # Assert assert result == [] @@ -605,7 +612,7 @@ def test_knowledge_retrieval_rate_limit_exceeded( top_k=5, ) - dataset_retrieval = DatasetRetrieval() + dataset_retrieval = _build_dataset_retrieval(db_session_with_containers) # Mock rate limit check to raise exception with patch.object( @@ -615,7 +622,7 @@ def test_knowledge_retrieval_rate_limit_exceeded( ): # Act & Assert with pytest.raises(Exception, match="Rate limit exceeded"): - dataset_retrieval.knowledge_retrieval(db_session_with_containers, request) + dataset_retrieval.knowledge_retrieval(request) @pytest.fixture diff --git a/api/tests/e2e_tests/containers/services/test_hit_testing_service.py b/api/tests/e2e_tests/containers/services/test_hit_testing_service.py index 0d989d0f3b07ae..1ad2c353d884a4 100644 --- a/api/tests/e2e_tests/containers/services/test_hit_testing_service.py +++ b/api/tests/e2e_tests/containers/services/test_hit_testing_service.py @@ -8,11 +8,12 @@ import pytest from pydantic import BaseModel, ConfigDict, TypeAdapter from sqlalchemy import func, select -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, sessionmaker from core.rag.embedding.retrieval import RetrievalSegments from core.rag.models.document import Document from core.rag.retrieval.retrieval_methods import RetrievalMethod +from extensions.application_services.retrieval import build_dataset_retrieval from models.dataset import Dataset, DatasetQuery, DocumentSegment from models.dataset import Document as DatasetDocument from models.enums import DataSourceType, DocumentCreatedFrom, DocumentDocType, SegmentStatus @@ -277,12 +278,13 @@ def test_retrieve_should_use_default_model_when_none_provided( response = _RetrieveResponse.model_validate( HitTestingService.retrieve( + session=db_session_with_containers, + retrieval=build_dataset_retrieval(sessionmaker(bind=db_session_with_containers.get_bind())), dataset=dataset, query="test query", account=account, retrieval_model=None, external_retrieval_model=external_retrieval_model, - session=db_session_with_containers, ) ) @@ -295,7 +297,7 @@ def test_retrieve_should_use_default_model_when_none_provided( assert after_count == before_count + 1 @patch("core.rag.datasource.retrieval_service.RetrievalService.retrieve") - @patch("core.rag.retrieval.dataset_retrieval.DatasetRetrieval.get_metadata_filter_condition") + @patch("services.knowledge.retrieval.dataset_retrieval.DatasetRetrieval.get_metadata_filter_condition") def test_retrieve_should_handle_metadata_filtering( self, mock_get_meta: MagicMock, mock_retrieve: MagicMock, db_session_with_containers: Session ) -> None: @@ -320,12 +322,13 @@ def test_retrieve_should_handle_metadata_filtering( mock_retrieve.return_value = retrieved_documents HitTestingService.retrieve( + session=db_session_with_containers, + retrieval=build_dataset_retrieval(sessionmaker(bind=db_session_with_containers.get_bind())), dataset=dataset, query="test query", account=account, retrieval_model=retrieval_model, external_retrieval_model=external_retrieval_model, - session=db_session_with_containers, ) mock_get_meta.assert_called_once() @@ -333,7 +336,7 @@ def test_retrieve_should_handle_metadata_filtering( assert mock_retrieve.call_args.kwargs["document_ids_filter"] == ["doc_id1"] @patch("core.rag.datasource.retrieval_service.RetrievalService.retrieve") - @patch("core.rag.retrieval.dataset_retrieval.DatasetRetrieval.get_metadata_filter_condition") + @patch("services.knowledge.retrieval.dataset_retrieval.DatasetRetrieval.get_metadata_filter_condition") def test_retrieve_should_return_empty_if_metadata_filtering_fails( self, mock_get_meta: MagicMock, mock_retrieve: MagicMock, db_session_with_containers: Session ) -> None: @@ -357,12 +360,13 @@ def test_retrieve_should_return_empty_if_metadata_filtering_fails( response = _RetrieveResponse.model_validate( HitTestingService.retrieve( + session=db_session_with_containers, + retrieval=build_dataset_retrieval(sessionmaker(bind=db_session_with_containers.get_bind())), dataset=dataset, query="test query", account=account, retrieval_model=retrieval_model, external_retrieval_model=external_retrieval_model, - session=db_session_with_containers, ) ) @@ -389,13 +393,14 @@ def test_retrieve_should_handle_attachments( mock_retrieve.return_value = retrieved_documents HitTestingService.retrieve( + session=db_session_with_containers, + retrieval=build_dataset_retrieval(sessionmaker(bind=db_session_with_containers.get_bind())), dataset=dataset, query="test query", account=account, retrieval_model=retrieval_model, external_retrieval_model=external_retrieval_model, attachment_ids=attachment_ids, - session=db_session_with_containers, ) mock_retrieve.assert_called_once_with( @@ -449,12 +454,13 @@ def test_retrieve_should_handle_reranking_and_threshold( mock_retrieve.return_value = retrieved_documents HitTestingService.retrieve( + session=db_session_with_containers, + retrieval=build_dataset_retrieval(sessionmaker(bind=db_session_with_containers.get_bind())), dataset=dataset, query="test query", account=account, retrieval_model=retrieval_model, external_retrieval_model=external_retrieval_model, - session=db_session_with_containers, ) mock_retrieve.assert_called_once() diff --git a/api/tests/helpers/unit/services/hit_service.py b/api/tests/helpers/unit/services/hit_service.py index c576d6fa646256..adb62a3a6d6449 100644 --- a/api/tests/helpers/unit/services/hit_service.py +++ b/api/tests/helpers/unit/services/hit_service.py @@ -12,10 +12,11 @@ import pytest from sqlalchemy import select -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, sessionmaker from core.rag.models.document import Document from core.rag.retrieval.retrieval_methods import RetrievalMethod +from extensions.application_services.retrieval import build_dataset_retrieval from models import Account, Tenant from models.dataset import Dataset, DatasetQuery from services.hit_testing_service import HitTestingService @@ -166,7 +167,13 @@ def test_retrieve_success_with_default_retrieval_model(self, sqlite_session: Ses # Act result = HitTestingService.retrieve( - dataset, query, account, retrieval_model, external_retrieval_model, session=sqlite_session + dataset, + query, + account, + retrieval_model, + external_retrieval_model, + session=sqlite_session, + retrieval=build_dataset_retrieval(sessionmaker(bind=sqlite_session.get_bind())), ) # Assert @@ -217,7 +224,13 @@ def test_retrieve_success_with_custom_retrieval_model(self, sqlite_session: Sess # Act result = HitTestingService.retrieve( - dataset, query, account, retrieval_model, external_retrieval_model, session=sqlite_session + dataset, + query, + account, + retrieval_model, + external_retrieval_model, + session=sqlite_session, + retrieval=build_dataset_retrieval(sessionmaker(bind=sqlite_session.get_bind())), ) # Assert @@ -265,17 +278,21 @@ def test_retrieve_with_metadata_filtering(self, sqlite_session: Session): patch( "services.hit_testing_service.RetrievalService.format_retrieval_documents", autospec=True ) as mock_format, - patch("services.hit_testing_service.DatasetRetrieval", autospec=True) as mock_dataset_retrieval_class, patch("services.hit_testing_service.time.perf_counter", autospec=True) as mock_perf_counter, ): mock_perf_counter.side_effect = [0.0, 0.1] - mock_dataset_retrieval_class.return_value = mock_dataset_retrieval mock_retrieve.return_value = documents mock_format.return_value = mock_records # Act result = HitTestingService.retrieve( - dataset, query, account, retrieval_model, external_retrieval_model, session=sqlite_session + dataset, + query, + account, + retrieval_model, + external_retrieval_model, + session=sqlite_session, + retrieval=lambda: mock_dataset_retrieval, ) # Assert @@ -310,17 +327,21 @@ def test_retrieve_with_metadata_filtering_no_documents(self, sqlite_session: Ses mock_dataset_retrieval.get_metadata_filter_condition.return_value = ({}, True) with ( - patch("services.hit_testing_service.DatasetRetrieval", autospec=True) as mock_dataset_retrieval_class, patch( "services.hit_testing_service.RetrievalService.format_retrieval_documents", autospec=True ) as mock_format, ): - mock_dataset_retrieval_class.return_value = mock_dataset_retrieval mock_format.return_value = [] # Act result = HitTestingService.retrieve( - dataset, query, account, retrieval_model, external_retrieval_model, session=sqlite_session + dataset, + query, + account, + retrieval_model, + external_retrieval_model, + session=sqlite_session, + retrieval=lambda: mock_dataset_retrieval, ) # Assert @@ -363,7 +384,13 @@ def test_retrieve_with_dataset_retrieval_model(self, sqlite_session: Session): # Act result = HitTestingService.retrieve( - dataset, query, account, retrieval_model, external_retrieval_model, session=sqlite_session + dataset, + query, + account, + retrieval_model, + external_retrieval_model, + session=sqlite_session, + retrieval=build_dataset_retrieval(sessionmaker(bind=sqlite_session.get_bind())), ) # Assert @@ -450,7 +477,12 @@ def test_external_retrieve_non_external_provider(self, sqlite_session: Session): # Act result = HitTestingService.external_retrieve( - dataset, query, account, external_retrieval_model, metadata_filtering_conditions, session=sqlite_session + dataset, + query, + account, + external_retrieval_model, + metadata_filtering_conditions, + session=sqlite_session, ) # Assert 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 b3d9a6a1ec46b8..dfdb6f4514ca32 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 @@ -1,9 +1,9 @@ from unittest.mock import Mock, patch import pytest -from werkzeug.exceptions import Forbidden, InternalServerError, NotFound import services.errors.base +from controllers.common.errors import AccessDeniedError, InternalServerError, NotFoundError from controllers.console.app.error import ( CompletionRequestError, ProviderModelCurrentlyNotSupportError, @@ -26,6 +26,8 @@ from services.hit_testing_service import HitTestingService from services.knowledge.dataset_service import DatasetService +pytestmark = pytest.mark.usefixtures("workflow_application") + @pytest.fixture def account(): @@ -104,7 +106,7 @@ def test_dataset_not_found(self, account): "get_dataset", return_value=None, ): - with pytest.raises(NotFound, match="Dataset not found"): + with pytest.raises(NotFoundError, match="Dataset not found"): DatasetsHitTestingBase.get_and_validate_dataset(Mock(), "dataset-1", account, "tenant-1") def test_permission_denied(self, dataset, account): @@ -120,7 +122,7 @@ def test_permission_denied(self, dataset, account): side_effect=services.errors.base.NoPermissionError("no access"), ), ): - with pytest.raises(Forbidden, match="no access"): + with pytest.raises(AccessDeniedError, match="no access"): DatasetsHitTestingBase.get_and_validate_dataset(Mock(), "dataset-1", account, "tenant-1") diff --git a/api/tests/unit_tests/controllers/inner_api/knowledge/test_retrieval.py b/api/tests/unit_tests/controllers/inner_api/knowledge/test_retrieval.py index 947dd3fe0d092f..8c380faa148354 100644 --- a/api/tests/unit_tests/controllers/inner_api/knowledge/test_retrieval.py +++ b/api/tests/unit_tests/controllers/inner_api/knowledge/test_retrieval.py @@ -2,7 +2,8 @@ from collections.abc import Iterator from contextlib import contextmanager -from unittest.mock import patch +from types import SimpleNamespace +from unittest.mock import MagicMock, create_autospec import pytest from flask import Flask @@ -16,6 +17,7 @@ InnerKnowledgeRetrieveAppNotFoundError, InnerKnowledgeRetrieveDatasetTenantMismatchError, ) +from services.knowledge_retrieval_inner_service import InnerKnowledgeRetrievalService from tests.helpers.unit.config_override import config_overrides_context @@ -27,6 +29,17 @@ def inner_api_app() -> Flask: return app +@pytest.fixture +def mock_retrieve(monkeypatch: pytest.MonkeyPatch) -> MagicMock: + inner_retrieval = create_autospec(InnerKnowledgeRetrievalService, instance=True, spec_set=True) + services = SimpleNamespace(knowledge=SimpleNamespace(inner_retrieval=inner_retrieval)) + monkeypatch.setattr( + "controllers.inner_api.knowledge.retrieval.application_services", + lambda: services, + ) + return inner_retrieval.retrieve + + def _headers(api_key: str | None = "inner-key") -> dict[str, str]: headers = {"Content-Type": "application/json"} if api_key is not None: @@ -96,8 +109,7 @@ def test_post_returns_400_for_invalid_body(self, inner_api_app: Flask): assert response.status_code == 400 assert response.get_json()["code"] == "invalid_request" - @patch("controllers.inner_api.knowledge.retrieval.InnerKnowledgeRetrievalService.retrieve") - def test_post_returns_404_for_service_not_found_error(self, mock_retrieve, inner_api_app: Flask): + def test_post_returns_404_for_service_not_found_error(self, mock_retrieve: MagicMock, inner_api_app: Flask) -> None: mock_retrieve.side_effect = InnerKnowledgeRetrieveAppNotFoundError("app missing") with _plugin_inner_auth(): @@ -110,8 +122,7 @@ def test_post_returns_404_for_service_not_found_error(self, mock_retrieve, inner assert response.status_code == 404 assert response.get_json()["code"] == "app_not_found" - @patch("controllers.inner_api.knowledge.retrieval.InnerKnowledgeRetrievalService.retrieve") - def test_post_returns_403_for_service_forbidden_error(self, mock_retrieve, inner_api_app: Flask): + def test_post_returns_403_for_service_forbidden_error(self, mock_retrieve: MagicMock, inner_api_app: Flask) -> None: mock_retrieve.side_effect = InnerKnowledgeRetrieveDatasetTenantMismatchError("wrong tenant") with _plugin_inner_auth(): @@ -124,8 +135,9 @@ def test_post_returns_403_for_service_forbidden_error(self, mock_retrieve, inner assert response.status_code == 403 assert response.get_json()["code"] == "dataset_tenant_mismatch" - @patch("controllers.inner_api.knowledge.retrieval.InnerKnowledgeRetrievalService.retrieve") - def test_post_returns_422_for_retrieval_config_value_error(self, mock_retrieve, inner_api_app: Flask): + def test_post_returns_422_for_retrieval_config_value_error( + self, mock_retrieve: MagicMock, inner_api_app: Flask + ) -> None: mock_retrieve.side_effect = ValueError("invalid reranking config") with _plugin_inner_auth(): @@ -138,8 +150,7 @@ def test_post_returns_422_for_retrieval_config_value_error(self, mock_retrieve, assert response.status_code == 422 assert response.get_json()["code"] == "retrieval_config_invalid" - @patch("controllers.inner_api.knowledge.retrieval.InnerKnowledgeRetrievalService.retrieve") - def test_post_returns_429_for_rate_limit_error(self, mock_retrieve, inner_api_app: Flask): + def test_post_returns_429_for_rate_limit_error(self, mock_retrieve: MagicMock, inner_api_app: Flask) -> None: mock_retrieve.side_effect = RateLimitExceededError("knowledge rate limited") with _plugin_inner_auth(): @@ -180,8 +191,9 @@ def test_post_returns_400_for_automatic_metadata_without_model_config(self, inne assert response.status_code == 400 assert response.get_json()["code"] == "invalid_request" - @patch("controllers.inner_api.knowledge.retrieval.InnerKnowledgeRetrievalService.retrieve") - def test_post_returns_502_for_external_knowledge_failure(self, mock_retrieve, inner_api_app: Flask): + def test_post_returns_502_for_external_knowledge_failure( + self, mock_retrieve: MagicMock, inner_api_app: Flask + ) -> None: mock_retrieve.side_effect = ExternalKnowledgeRetrievalError("upstream failed") with _plugin_inner_auth(): @@ -194,8 +206,7 @@ def test_post_returns_502_for_external_knowledge_failure(self, mock_retrieve, in assert response.status_code == 502 assert response.get_json()["code"] == "external_knowledge_failed" - @patch("controllers.inner_api.knowledge.retrieval.InnerKnowledgeRetrievalService.retrieve") - def test_post_returns_service_response(self, mock_retrieve, inner_api_app: Flask): + def test_post_returns_service_response(self, mock_retrieve: MagicMock, inner_api_app: Flask) -> None: mock_retrieve.return_value = InnerKnowledgeRetrieveResponse( results=[ Source( 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 5c853e53ddc335..777663373c9df0 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 @@ -15,18 +15,29 @@ """ import uuid +from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest from flask import Flask, g -from werkzeug.exceptions import Forbidden, NotFound import services.errors.base +from controllers.common.errors import AccessDeniedError, NotFoundError +from controllers.console.datasets import hit_testing_base as hit_testing_base_module from controllers.service_api.dataset.hit_testing import HitTestingApi, HitTestingPayload from models.account import Account, Tenant, TenantAccountRole from models.dataset import Dataset from services.knowledge.entities.knowledge_entities import RetrievalModel + +@pytest.fixture +def workflow_retrieval(monkeypatch: pytest.MonkeyPatch) -> MagicMock: + retrieval = MagicMock(name="dataset_retrieval_factory") + services = SimpleNamespace(workflow_runtime=SimpleNamespace(retrieval=retrieval)) + monkeypatch.setattr(hit_testing_base_module, "application_services", lambda: services) + return retrieval + + # --------------------------------------------------------------------------- # HitTestingPayload Model Tests # --------------------------------------------------------------------------- @@ -150,6 +161,7 @@ def test_post_success( mock_hit_svc: MagicMock, mock_ns: MagicMock, app: Flask, + workflow_retrieval: MagicMock, ) -> None: """Test successful hit testing request.""" dataset_id = str(uuid.uuid4()) @@ -175,6 +187,7 @@ def test_post_success( assert response["query"] == {"content": "test query"} mock_hit_svc.retrieve.assert_called_once() + assert mock_hit_svc.retrieve.call_args.kwargs["retrieval"] is workflow_retrieval @patch("controllers.service_api.dataset.hit_testing.service_api_ns") @patch("controllers.console.datasets.hit_testing_base.HitTestingService") @@ -185,6 +198,7 @@ def test_post_with_retrieval_model( mock_hit_svc: MagicMock, mock_ns: MagicMock, app: Flask, + workflow_retrieval: MagicMock, ) -> None: """Test hit testing with custom retrieval model.""" dataset_id = str(uuid.uuid4()) @@ -221,6 +235,7 @@ def test_post_with_retrieval_model( assert response["query"] == {"content": "complex query"} call_kwargs = mock_hit_svc.retrieve.call_args + assert call_kwargs.kwargs["retrieval"] is workflow_retrieval # retrieval_model is serialized via model_dump, verify key fields passed_retrieval_model = call_kwargs.kwargs.get("retrieval_model") assert passed_retrieval_model is not None @@ -236,6 +251,7 @@ def test_post_preserves_retrieval_model_metadata_filtering_conditions( mock_hit_svc: MagicMock, mock_ns: MagicMock, app: Flask, + workflow_retrieval: MagicMock, ) -> None: """Service API retrieval payload should not drop metadata filters.""" dataset_id = str(uuid.uuid4()) @@ -277,6 +293,7 @@ def test_post_preserves_retrieval_model_metadata_filtering_conditions( HitTestingApi.post.__wrapped__(api, tenant_id, dataset_id) passed_retrieval_model = mock_hit_svc.retrieve.call_args.kwargs.get("retrieval_model") + assert mock_hit_svc.retrieve.call_args.kwargs["retrieval"] is workflow_retrieval assert passed_retrieval_model is not None assert passed_retrieval_model["metadata_filtering_conditions"] == metadata_filtering_conditions @@ -289,6 +306,7 @@ def test_post_prepares_nullable_list_fields( mock_hit_svc: MagicMock, mock_ns: MagicMock, app: Flask, + workflow_retrieval: MagicMock, ) -> None: """Test service API prepares nullable list fields from retrieval records.""" dataset_id = str(uuid.uuid4()) @@ -323,6 +341,7 @@ def test_post_prepares_nullable_list_fields( assert record["score"] == 0.9 assert record["tsne_position"] is None assert record["summary"] is None + assert mock_hit_svc.retrieve.call_args.kwargs["retrieval"] is workflow_retrieval @patch("controllers.service_api.dataset.hit_testing.service_api_ns") @patch("controllers.console.datasets.hit_testing_base.DatasetService") @@ -344,7 +363,7 @@ def test_post_dataset_not_found( # TODO: the service APIs are NOT migrated yet, so we have to do the very dirty hack g._login_user = account api = HitTestingApi() - with pytest.raises(NotFound): + with pytest.raises(NotFoundError): HitTestingApi.post.__wrapped__(api, tenant_id, dataset_id) @patch("controllers.service_api.dataset.hit_testing.service_api_ns") @@ -370,5 +389,5 @@ def test_post_no_dataset_permission( # TODO: the service APIs are NOT migrated yet, so we have to do the very dirty hack g._login_user = account api = HitTestingApi() - with pytest.raises(Forbidden): + with pytest.raises(AccessDeniedError): HitTestingApi.post.__wrapped__(api, tenant_id, dataset_id) diff --git a/api/tests/unit_tests/core/callback_handler/test_index_tool_callback_handler.py b/api/tests/unit_tests/core/callback_handler/test_index_tool_callback_handler.py index 76b0af1fdfa9d8..0bef9b69770c0c 100644 --- a/api/tests/unit_tests/core/callback_handler/test_index_tool_callback_handler.py +++ b/api/tests/unit_tests/core/callback_handler/test_index_tool_callback_handler.py @@ -11,12 +11,17 @@ import core.callback_handler.index_tool_callback_handler as callback_module from core.app.entities.app_invoke_entities import InvokeFrom -from core.callback_handler.index_tool_callback_handler import DatasetIndexToolCallbackHandler +from core.callback_handler.index_tool_callback_handler import ( + DatasetIndexToolCallbackHandler as DatasetQueryCallbackHandler, +) from core.rag.index_processor.constant.index_type import IndexStructureType from core.rag.models.document import Document from models.dataset import ChildChunk, DatasetQuery, DocumentSegment from models.dataset import Document as DatasetDocument from models.enums import CreatorUserRole, DatasetQuerySource, DataSourceType, DocumentCreatedFrom +from services.knowledge.retrieval.adapters.resource_events import ( + DatasetIndexToolCallbackHandler as RetrievalResourceCallbackHandler, +) @dataclass(eq=False, repr=False) @@ -41,7 +46,7 @@ def mock_queue_manager(mocker: MockerFixture): @pytest.fixture def handler(mock_queue_manager, sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr(callback_module, "db", _DatabaseBinding(sqlite_engine)) - return DatasetIndexToolCallbackHandler( + return DatasetQueryCallbackHandler( queue_manager=mock_queue_manager, app_id=str(uuid4()), message_id=str(uuid4()), @@ -131,7 +136,7 @@ def test_on_query_success_roles( expected_role: str, ) -> None: monkeypatch.setattr(callback_module, "db", _DatabaseBinding(sqlite_engine)) - handler = DatasetIndexToolCallbackHandler( + handler = DatasetQueryCallbackHandler( queue_manager=mock_queue_manager, app_id=str(uuid4()), message_id=str(uuid4()), @@ -154,7 +159,7 @@ def test_on_query_none_values_roll_back_independent_transaction( mock_queue_manager, ) -> None: monkeypatch.setattr(callback_module, "db", _DatabaseBinding(sqlite_engine)) - handler = DatasetIndexToolCallbackHandler( + handler = DatasetQueryCallbackHandler( queue_manager=mock_queue_manager, app_id=None, message_id=None, @@ -171,7 +176,7 @@ def test_on_query_none_values_roll_back_independent_transaction( class TestOnToolEnd: def test_on_tool_end_no_metadata( self, - handler: DatasetIndexToolCallbackHandler, + handler: DatasetQueryCallbackHandler, caller_session_boundary: _CallerSessionBoundary, ) -> None: document = Document.model_construct(page_content="content", metadata=None, provider="dify") @@ -180,7 +185,7 @@ def test_on_tool_end_no_metadata( def test_on_tool_end_dataset_document_not_found( self, - handler: DatasetIndexToolCallbackHandler, + handler: DatasetQueryCallbackHandler, sqlite_session: Session, caller_session_boundary: _CallerSessionBoundary, ) -> None: @@ -195,7 +200,7 @@ def test_on_tool_end_dataset_document_not_found( def test_on_tool_end_parent_child_index_with_child( self, - handler: DatasetIndexToolCallbackHandler, + handler: DatasetQueryCallbackHandler, sqlite_session: Session, caller_session_boundary: _CallerSessionBoundary, ) -> None: @@ -230,7 +235,7 @@ def test_on_tool_end_parent_child_index_with_child( def test_on_tool_end_non_parent_child_index( self, - handler: DatasetIndexToolCallbackHandler, + handler: DatasetQueryCallbackHandler, sqlite_session: Session, caller_session_boundary: _CallerSessionBoundary, ) -> None: @@ -256,15 +261,16 @@ def test_on_tool_end_non_parent_child_index( def test_on_tool_end_empty_documents( self, - handler: DatasetIndexToolCallbackHandler, + handler: DatasetQueryCallbackHandler, caller_session_boundary: _CallerSessionBoundary, ) -> None: handler.on_tool_end([], caller_session_boundary.session) class TestReturnRetrieverResourceInfo: - def test_publish_called(self, handler: DatasetIndexToolCallbackHandler, mock_queue_manager, mocker: MockerFixture): - mocker.patch("core.callback_handler.index_tool_callback_handler.QueueRetrieverResourcesEvent") + def test_publish_called(self, mock_queue_manager, mocker: MockerFixture): + handler = RetrievalResourceCallbackHandler(mock_queue_manager) + mocker.patch("services.knowledge.retrieval.adapters.resource_events.QueueRetrieverResourcesEvent") resources = [mocker.Mock()] handler.return_retriever_resource_info(resources) diff --git a/api/tests/unit_tests/core/workflow/nodes/knowledge_retrieval/test_retrieval.py b/api/tests/unit_tests/core/workflow/nodes/knowledge_retrieval/test_retrieval.py index f9ad0c98e3a8e8..79fa020ae4b0f1 100644 --- a/api/tests/unit_tests/core/workflow/nodes/knowledge_retrieval/test_retrieval.py +++ b/api/tests/unit_tests/core/workflow/nodes/knowledge_retrieval/test_retrieval.py @@ -8,8 +8,8 @@ from sqlalchemy.orm import Session from core.rag.rerank.rerank_type import RerankMode -from core.rag.retrieval.dataset_retrieval import DatasetRetrieval from core.workflow.nodes.knowledge_retrieval.retrieval import KnowledgeRetrievalRequest +from extensions.application_services.retrieval import build_dataset_retrieval from models.dataset import Dataset from tests.unit_tests.extensions.otel.conftest import memory_span_exporter as memory_span_exporter # noqa: PLC0414 from tests.unit_tests.extensions.otel.conftest import mock_account_user as mock_account_user # noqa: PLC0414 @@ -28,9 +28,7 @@ def _otel_enabled(config_overrides: Callable[..., None]) -> None: def test_knowledge_retrieval_creates_a_child_otel_span( - memory_span_exporter, - tracer_provider_with_memory_exporter, - sqlite_session: Session, + memory_span_exporter, tracer_provider_with_memory_exporter, sqlite_session: Session, sqlite_session_factory ) -> None: """The retrieval entry point must be visible beneath its workflow node span.""" request = KnowledgeRetrievalRequest( @@ -42,19 +40,19 @@ def test_knowledge_retrieval_creates_a_child_otel_span( retrieval_mode="multiple", query="test query", ) - retrieval = DatasetRetrieval() + retrieval = build_dataset_retrieval(sqlite_session_factory)() with ( patch.object(retrieval, "_check_knowledge_rate_limit"), - patch.object(retrieval, "_get_available_datasets", return_value=[]), + patch.object(retrieval._records, "available_datasets", return_value=[]), get_tracer(__name__).start_as_current_span("knowledge-retrieval-node") as node_span, ): - assert retrieval.knowledge_retrieval(sqlite_session, request) == [] + assert retrieval.knowledge_retrieval(request) == [] retrieval_span = next( span for span in memory_span_exporter.get_finished_spans() - if span.name == "core.rag.retrieval.dataset_retrieval.DatasetRetrieval.knowledge_retrieval" + if span.name == "services.knowledge.retrieval.dataset_retrieval.DatasetRetrieval.knowledge_retrieval" ) node_span_context = node_span.get_span_context() assert retrieval_span.context.trace_id == node_span_context.trace_id @@ -63,11 +61,10 @@ def test_knowledge_retrieval_creates_a_child_otel_span( def test_multiple_retrieve_preserves_otel_context_in_dataset_thread( - app, - tracer_provider_with_memory_exporter, + app, tracer_provider_with_memory_exporter, sqlite_session_factory ) -> None: """Per-dataset retrieval spans must remain in the workflow node trace.""" - retrieval = DatasetRetrieval() + retrieval = build_dataset_retrieval(sqlite_session_factory)() dataset = Dataset( id=str(uuid4()), indexing_technique="high_quality", @@ -102,11 +99,9 @@ def record_active_trace(**_kwargs: object) -> None: def test_retriever_thread_exception_sets_error_span_and_is_collected( - app, - memory_span_exporter, - tracer_provider_with_memory_exporter, + app, memory_span_exporter, tracer_provider_with_memory_exporter, sqlite_session_factory ) -> None: - retrieval = DatasetRetrieval() + retrieval = build_dataset_retrieval(sqlite_session_factory)() cancel_event = threading.Event() thread_exceptions: list[Exception] = [] expected_error = RuntimeError("retrieval failed") @@ -115,7 +110,7 @@ def test_retriever_thread_exception_sets_error_span_and_is_collected( patch.object(retrieval, "_retriever", side_effect=expected_error), ): retrieval._run_retriever_thread_safely( - flask_app=app, + tenant_id="tenant-1", dataset_id=str(uuid4()), query="test query", top_k=4, @@ -138,11 +133,9 @@ def test_retriever_thread_exception_sets_error_span_and_is_collected( def test_retriever_thread_exception_emits_skip_event_when_requested( - app, - memory_span_exporter, - tracer_provider_with_memory_exporter, + app, memory_span_exporter, tracer_provider_with_memory_exporter, sqlite_session_factory ) -> None: - retrieval = DatasetRetrieval() + retrieval = build_dataset_retrieval(sqlite_session_factory)() cancel_event = threading.Event() thread_exceptions: list[Exception] = [] expected_error = RuntimeError("retrieval failed") @@ -153,7 +146,7 @@ def test_retriever_thread_exception_emits_skip_event_when_requested( get_tracer(__name__).start_as_current_span("dataset-retrieval-parent") as parent_span, ): retrieval._run_retriever_thread_safely( - flask_app=app, + tenant_id="tenant-1", dataset_id=dataset_id, query="test query", top_k=4, diff --git a/api/tests/unit_tests/services/knowledge/retrieval/test_attachments.py b/api/tests/unit_tests/services/knowledge/retrieval/test_attachments.py index d149c1d266d7c7..84dd5d075bc244 100644 --- a/api/tests/unit_tests/services/knowledge/retrieval/test_attachments.py +++ b/api/tests/unit_tests/services/knowledge/retrieval/test_attachments.py @@ -6,21 +6,23 @@ import pytest from sqlalchemy import Engine, event -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, sessionmaker 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.index_processor.constant.index_type import IndexStructureType 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.application_services.retrieval import build_dataset_retrieval 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 +from repositories.knowledge import retrieval_repository as retrieval_repository_module +from repositories.knowledge.dataset_read_repository import authorize_retrieved_segment +from services.knowledge.retrieval import dataset_retrieval as retrieval_module +from services.knowledge.retrieval.dataset_retrieval import DatasetRetrieval @pytest.fixture @@ -91,15 +93,16 @@ def graph(sqlite_session: Session) -> tuple[Dataset, Document, DocumentSegment, @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, + sqlite_session_factory: sessionmaker[Session], ) -> None: dataset, document, segment, file_ids = graph - retrieval = DatasetRetrieval() + retrieval = build_dataset_retrieval(sqlite_session_factory)() + assert isinstance(retrieval, DatasetRetrieval) monkeypatch.setattr(retrieval, "_check_knowledge_rate_limit", MagicMock()) - monkeypatch.setattr(retrieval, "_get_available_datasets", MagicMock(return_value=[dataset])) + monkeypatch.setattr(retrieval._records, "available_datasets", MagicMock(return_value=[dataset])) monkeypatch.setattr(retrieval, "get_metadata_filter_condition", MagicMock(return_value=(None, None))) monkeypatch.setattr( retrieval, @@ -139,7 +142,6 @@ def test_published_retrieval_grants_bound_images_before_signing( 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, @@ -159,7 +161,6 @@ def test_published_retrieval_grants_bound_images_before_signing( 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, @@ -225,14 +226,16 @@ def test_workflow_retrieval_releases_owned_connections_and_preserves_caller_tran monkeypatch: pytest.MonkeyPatch, with_image: bool, signing_fails: bool, + sqlite_session_factory: sessionmaker[Session], ) -> None: dataset, document, segment, _ = graph if not with_image: segment.content = "Plain text" sqlite_session.commit() - retrieval = DatasetRetrieval() + retrieval = build_dataset_retrieval(sqlite_session_factory)() + assert isinstance(retrieval, DatasetRetrieval) monkeypatch.setattr(retrieval, "_check_knowledge_rate_limit", MagicMock()) - monkeypatch.setattr(retrieval, "_get_available_datasets", MagicMock(return_value=[dataset])) + monkeypatch.setattr(retrieval._records, "available_datasets", MagicMock(return_value=[dataset])) monkeypatch.setattr( retrieval, "multiple_retrieve", @@ -248,7 +251,7 @@ def test_workflow_retrieval_releases_owned_connections_and_preserves_caller_tran ) if signing_fails: monkeypatch.setattr( - retrieval_module, "sign_segment_content", MagicMock(side_effect=RuntimeError("signing failed")) + retrieval_repository_module, "sign_segment_content", MagicMock(side_effect=RuntimeError("signing failed")) ) request = KnowledgeRetrievalRequest( tenant_id=dataset.tenant_id, @@ -287,9 +290,9 @@ def track_checkin(*_args: object) -> None: 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) + retrieval.knowledge_retrieval(request) else: - sources = retrieval.knowledge_retrieval(caller_session, request) + sources = retrieval.knowledge_retrieval(request) assert len(sources) == 1 assert sources[0].metadata.segment_id == segment.id assert caller_session.get_transaction() is caller_txn diff --git a/api/tests/unit_tests/services/test_hit_testing_service.py b/api/tests/unit_tests/services/test_hit_testing_service.py index f8f9cdea897a27..598ba1a4e63d44 100644 --- a/api/tests/unit_tests/services/test_hit_testing_service.py +++ b/api/tests/unit_tests/services/test_hit_testing_service.py @@ -10,6 +10,7 @@ from core.rag.index_processor.constant.index_type import IndexTechniqueType from core.rag.models.document import Document from core.rag.retrieval.retrieval_methods import RetrievalMethod +from extensions.application_services.workflow import WorkflowExecutionDependencies from models import Account from models.dataset import Dataset, DatasetQuery from services.hit_testing_service import HitTestingService @@ -68,7 +69,12 @@ def model_dump(self) -> dict[str, str | float]: class TestHitTestingServiceRetrieveTransaction: """Regression tests for https://github.com/langgenius/dify/issues/38998.""" - def test_retrieve_should_not_close_caller_managed_transaction(self, sqlite_session: Session) -> None: + def test_retrieve_should_not_close_caller_managed_transaction( + self, + sqlite_session: Session, + *, + workflow_runtime: WorkflowExecutionDependencies, + ) -> None: """retrieve() must not commit the caller-managed session mid-flow. When the caller manages the session inside a ``session.begin()`` block, an @@ -94,7 +100,15 @@ def test_retrieve_should_not_close_caller_managed_transaction(self, sqlite_sessi mock_format.return_value = [record] with sqlite_session.begin(): - result = HitTestingService.retrieve(dataset, query, account, None, {}, session=sqlite_session) + result = HitTestingService.retrieve( + dataset, + query, + account, + None, + {}, + session=sqlite_session, + retrieval=workflow_runtime.retrieval, + ) assert result["query"]["content"] == query query_log = sqlite_session.scalar(select(DatasetQuery)) @@ -128,6 +142,8 @@ def test_retrieve_should_use_keyword_search_only_for_economy_datasets( indexing_technique: IndexTechniqueType, retrieval_model: dict[str, object] | None, expected_method: RetrievalMethod, + *, + workflow_runtime: WorkflowExecutionDependencies, ) -> None: dataset = Dataset( id="dataset-123", @@ -145,6 +161,14 @@ def test_retrieve_should_use_keyword_search_only_for_economy_datasets( ): no_documents: list[Document] = [] mock_retrieve.return_value = no_documents - HitTestingService.retrieve(dataset, query, account, retrieval_model, {}, session=sqlite_session) + HitTestingService.retrieve( + dataset, + query, + account, + retrieval_model, + {}, + session=sqlite_session, + retrieval=workflow_runtime.retrieval, + ) assert mock_retrieve.call_args.kwargs["retrieval_method"] == expected_method diff --git a/api/tests/unit_tests/services/test_knowledge_retrieval_inner_service.py b/api/tests/unit_tests/services/test_knowledge_retrieval_inner_service.py index 5e6937505385bf..0876a8f5bb229a 100644 --- a/api/tests/unit_tests/services/test_knowledge_retrieval_inner_service.py +++ b/api/tests/unit_tests/services/test_knowledge_retrieval_inner_service.py @@ -1,13 +1,15 @@ """Unit tests for the inner knowledge retrieval service.""" -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock import pytest -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, sessionmaker from core.workflow.nodes.knowledge_retrieval.retrieval import Source, SourceMetadata +from extensions.application_services.retrieval import build_dataset_retrieval from models.dataset import Dataset from models.model import App, AppMode +from repositories.knowledge.retrieval_repository import KnowledgeRetrievalRepository from services.entities.knowledge_retrieval_inner import InnerKnowledgeRetrieveRequest from services.errors.knowledge_retrieval import ( InnerKnowledgeRetrieveAppNotFoundError, @@ -113,10 +115,8 @@ def _build_source() -> Source: class TestInnerKnowledgeRetrievalService: @pytest.mark.parametrize("sqlite_session", [(App, Dataset)], indirect=True) - @patch("services.knowledge_retrieval_inner_service.DatasetRetrieval") def test_retrieve_maps_multiple_request_and_skips_enable_api_check( self, - mock_rag_cls, sqlite_session: Session, ): request = _build_request() @@ -145,9 +145,10 @@ def test_retrieve_maps_multiple_request_and_skips_enable_api_check( "time_to_first_token": None, "time_to_generate": None, } - mock_rag_cls.return_value = rag - response = InnerKnowledgeRetrievalService().retrieve(request, sqlite_session) + response = InnerKnowledgeRetrievalService( + scopes=KnowledgeRetrievalRepository(sessionmaker(bind=sqlite_session.get_bind())), retrieval=lambda: rag + ).retrieve(request) rag_request = rag.knowledge_retrieval.call_args.kwargs["request"] assert rag_request.tenant_id == TENANT_ID @@ -171,12 +172,11 @@ def test_retrieve_maps_multiple_request_and_skips_enable_api_check( assert rag_request.attachment_ids == ["attachment-1"] assert response.results[0].title == "FAQ.md" assert response.usage.currency == "USD" - assert rag.knowledge_retrieval.call_args.kwargs["session"] is sqlite_session - assert sqlite_session.in_transaction() + assert "session" not in rag.knowledge_retrieval.call_args.kwargs + assert not sqlite_session.in_transaction() @pytest.mark.parametrize("sqlite_session", [(App, Dataset)], indirect=True) - @patch("services.knowledge_retrieval_inner_service.DatasetRetrieval") - def test_retrieve_maps_single_request(self, mock_rag_cls, sqlite_session: Session): + def test_retrieve_maps_single_request(self, sqlite_session: Session): request = _build_request( dataset_ids=[DATASET_1_ID], retrieval={ @@ -217,9 +217,10 @@ def test_retrieve_maps_single_request(self, mock_rag_cls, sqlite_session: Sessio "currency": "USD", "latency": 1, } - mock_rag_cls.return_value = rag - InnerKnowledgeRetrievalService().retrieve(request, sqlite_session) + InnerKnowledgeRetrievalService( + scopes=KnowledgeRetrievalRepository(sessionmaker(bind=sqlite_session.get_bind())), retrieval=lambda: rag + ).retrieve(request) rag_request = rag.knowledge_retrieval.call_args.kwargs["request"] assert rag_request.retrieval_mode == "single" @@ -230,29 +231,38 @@ def test_retrieve_maps_single_request(self, mock_rag_cls, sqlite_session: Sessio assert rag_request.metadata_filtering_mode == "automatic" assert rag_request.metadata_model_config is not None assert rag_request.metadata_model_config.provider == "openai" - assert sqlite_session.in_transaction() + assert not sqlite_session.in_transaction() @pytest.mark.parametrize("sqlite_session", [(App, Dataset)], indirect=True) def test_retrieve_raises_when_app_missing(self, sqlite_session: Session): with pytest.raises(InnerKnowledgeRetrieveAppNotFoundError): - InnerKnowledgeRetrievalService().retrieve(_build_request(), sqlite_session) - assert sqlite_session.in_transaction() + InnerKnowledgeRetrievalService( + scopes=KnowledgeRetrievalRepository(sessionmaker(bind=sqlite_session.get_bind())), + retrieval=build_dataset_retrieval(sessionmaker(bind=sqlite_session.get_bind())), + ).retrieve(_build_request()) + assert not sqlite_session.in_transaction() @pytest.mark.parametrize("sqlite_session", [(App, Dataset)], indirect=True) def test_retrieve_raises_when_app_belongs_to_other_tenant(self, sqlite_session: Session): _persist_state(sqlite_session, _app(tenant_id=OTHER_TENANT_ID)) with pytest.raises(InnerKnowledgeRetrieveAppTenantMismatchError): - InnerKnowledgeRetrievalService().retrieve(_build_request(), sqlite_session) - assert sqlite_session.in_transaction() + InnerKnowledgeRetrievalService( + scopes=KnowledgeRetrievalRepository(sessionmaker(bind=sqlite_session.get_bind())), + retrieval=build_dataset_retrieval(sessionmaker(bind=sqlite_session.get_bind())), + ).retrieve(_build_request()) + assert not sqlite_session.in_transaction() @pytest.mark.parametrize("sqlite_session", [(App, Dataset)], indirect=True) def test_retrieve_raises_when_dataset_missing(self, sqlite_session: Session): _persist_state(sqlite_session, _app(), _dataset(dataset_id=DATASET_1_ID)) with pytest.raises(InnerKnowledgeRetrieveDatasetNotFoundError): - InnerKnowledgeRetrievalService().retrieve(_build_request(), sqlite_session) - assert sqlite_session.in_transaction() + InnerKnowledgeRetrievalService( + scopes=KnowledgeRetrievalRepository(sessionmaker(bind=sqlite_session.get_bind())), + retrieval=build_dataset_retrieval(sessionmaker(bind=sqlite_session.get_bind())), + ).retrieve(_build_request()) + assert not sqlite_session.in_transaction() @pytest.mark.parametrize("sqlite_session", [(App, Dataset)], indirect=True) def test_retrieve_raises_when_dataset_belongs_to_other_tenant(self, sqlite_session: Session): @@ -264,5 +274,8 @@ def test_retrieve_raises_when_dataset_belongs_to_other_tenant(self, sqlite_sessi ) with pytest.raises(InnerKnowledgeRetrieveDatasetTenantMismatchError): - InnerKnowledgeRetrievalService().retrieve(_build_request(), sqlite_session) - assert sqlite_session.in_transaction() + InnerKnowledgeRetrievalService( + scopes=KnowledgeRetrievalRepository(sessionmaker(bind=sqlite_session.get_bind())), + retrieval=build_dataset_retrieval(sessionmaker(bind=sqlite_session.get_bind())), + ).retrieve(_build_request()) + assert not sqlite_session.in_transaction()