refactor(api): inject retrieval dependencies into APIs and annotation tasks

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")),

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,

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
         ),

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

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={}),
             )

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: ...

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

api/tasks/annotation/__init__.py

@@ -0,0 +1 @@
+"""Annotation indexing and import tasks."""

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,

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,

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(