From c7a3a4fc0e6451b31167e8855c94e38177c4ff42 Mon Sep 17 00:00:00 2001 From: -LAN- Date: Tue, 6 Jan 2026 19:00:21 +0800 Subject: [PATCH] Updated dataset query attribution to use the `UserFrom` enum in the callback and hit-testing paths, matching the retrieval flow and removing literal role strings. Adjusted `api/core/callback_handler/index_tool_callback_handler.py` and `api/services/hit_testing_service.py` where `DatasetQuery` records are created. fix(api): use UserFrom for dataset query created_by_role Tests not run (not requested). --- api/core/callback_handler/index_tool_callback_handler.py | 3 ++- api/services/hit_testing_service.py | 5 +++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/api/core/callback_handler/index_tool_callback_handler.py b/api/core/callback_handler/index_tool_callback_handler.py index d0279349ca..1387967e9b 100644 --- a/api/core/callback_handler/index_tool_callback_handler.py +++ b/api/core/callback_handler/index_tool_callback_handler.py @@ -12,6 +12,7 @@ from core.rag.models.document import Document from extensions.ext_database import db from models.dataset import ChildChunk, DatasetQuery, DocumentSegment from models.dataset import Document as DatasetDocument +from models.enums import UserFrom _logger = logging.getLogger(__name__) @@ -38,7 +39,7 @@ class DatasetIndexToolCallbackHandler: source="app", source_app_id=self._app_id, created_by_role=( - "account" if self._invoke_from in {InvokeFrom.EXPLORE, InvokeFrom.DEBUGGER} else "end_user" + UserFrom.ACCOUNT if self._invoke_from in {InvokeFrom.EXPLORE, InvokeFrom.DEBUGGER} else UserFrom.END_USER ), created_by=self._user_id, ) diff --git a/api/services/hit_testing_service.py b/api/services/hit_testing_service.py index 8cbf3a25c3..13224d9c66 100644 --- a/api/services/hit_testing_service.py +++ b/api/services/hit_testing_service.py @@ -13,6 +13,7 @@ from core.rag.retrieval.retrieval_methods import RetrievalMethod from extensions.ext_database import db from models import Account from models.dataset import Dataset, DatasetQuery +from models.enums import UserFrom logger = logging.getLogger(__name__) @@ -98,7 +99,7 @@ class HitTestingService: content=json.dumps(dataset_queries), source="hit_testing", source_app_id=None, - created_by_role="account", + created_by_role=UserFrom.ACCOUNT, created_by=account.id, ) db.session.add(dataset_query) @@ -138,7 +139,7 @@ class HitTestingService: content=query, source="hit_testing", source_app_id=None, - created_by_role="account", + created_by_role=UserFrom.ACCOUNT, created_by=account.id, )