From 0d19a81bd8b5d3ff1aff303ca05bb2ab79cf8ca3 Mon Sep 17 00:00:00 2001 From: -LAN- Date: Tue, 6 Jan 2026 19:32:30 +0800 Subject: [PATCH] refactor(app): add InvokeFrom.to_creator_user_role and use it in dataset retrieval Tests: not run (not requested). --- api/core/app/entities/app_invoke_entities.py | 6 ++++++ api/core/rag/retrieval/dataset_retrieval.py | 6 +----- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/api/core/app/entities/app_invoke_entities.py b/api/core/app/entities/app_invoke_entities.py index 0cb573cb86..d07152e1ff 100644 --- a/api/core/app/entities/app_invoke_entities.py +++ b/api/core/app/entities/app_invoke_entities.py @@ -9,6 +9,7 @@ from core.app.app_config.entities import EasyUIBasedAppConfig, WorkflowUIBasedAp from core.entities.provider_configuration import ProviderModelBundle from core.file import File, FileUploadConfig from core.model_runtime.entities.model_entities import AIModelEntity +from models.enums import CreatorUserRole if TYPE_CHECKING: from core.ops.ops_trace_manager import TraceQueueManager @@ -79,6 +80,11 @@ class InvokeFrom(StrEnum): return "dev" + def to_creator_user_role(self) -> CreatorUserRole: + if self in {InvokeFrom.EXPLORE, InvokeFrom.DEBUGGER}: + return CreatorUserRole.ACCOUNT + return CreatorUserRole.END_USER + class ModelConfigWithCredentialsEntity(BaseModel): """ diff --git a/api/core/rag/retrieval/dataset_retrieval.py b/api/core/rag/retrieval/dataset_retrieval.py index f40d2d4bd2..99ec6cb5f3 100644 --- a/api/core/rag/retrieval/dataset_retrieval.py +++ b/api/core/rag/retrieval/dataset_retrieval.py @@ -177,11 +177,7 @@ class DatasetRetrieval: ) all_documents = [] - creator_user_role = ( - CreatorUserRole.ACCOUNT - if invoke_from in {InvokeFrom.EXPLORE, InvokeFrom.DEBUGGER} - else CreatorUserRole.END_USER - ) + creator_user_role = invoke_from.to_creator_user_role() if retrieve_config.retrieve_strategy == DatasetRetrieveConfigEntity.RetrieveStrategy.SINGLE: all_documents = self.single_retrieve( app_id,