From 3390be39784c408b45e39dbf0263f5caae244c9c Mon Sep 17 00:00:00 2001 From: Checo <104541981+sergioperezcheco@users.noreply.github.com> Date: Wed, 1 Jul 2026 14:16:13 +0800 Subject: [PATCH] refactor: pass db.session explicitly in DatasetIndexToolCallbackHandler (#38082) --- .../index_tool_callback_handler.py | 20 +++---- .../dataset_multi_retriever_tool.py | 4 +- .../dataset_retriever_tool.py | 4 +- .../test_index_tool_callback_handler.py | 60 +++++++++---------- .../core/tools/utils/test_misc_utils_extra.py | 4 +- 5 files changed, 45 insertions(+), 47 deletions(-) diff --git a/api/core/callback_handler/index_tool_callback_handler.py b/api/core/callback_handler/index_tool_callback_handler.py index 205e0042901..5494769082e 100644 --- a/api/core/callback_handler/index_tool_callback_handler.py +++ b/api/core/callback_handler/index_tool_callback_handler.py @@ -2,6 +2,7 @@ import logging from collections.abc import Sequence from sqlalchemy import select, update +from sqlalchemy.orm import scoped_session from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom from core.app.entities.app_invoke_entities import InvokeFrom @@ -9,7 +10,6 @@ from core.app.entities.queue_entities import QueueRetrieverResourcesEvent from core.rag.entities import RetrievalSourceMetadata from core.rag.index_processor.constant.index_type import IndexStructureType 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 CreatorUserRole, DatasetQuerySource @@ -29,7 +29,7 @@ class DatasetIndexToolCallbackHandler: self._user_id = user_id self._invoke_from = invoke_from - def on_query(self, query: str, dataset_id: str): + def on_query(self, query: str, dataset_id: str, session: scoped_session): """ Handle query. """ @@ -46,16 +46,16 @@ class DatasetIndexToolCallbackHandler: created_by=self._user_id, ) - db.session.add(dataset_query) - db.session.commit() + session.add(dataset_query) + session.commit() - def on_tool_end(self, documents: list[Document]): + def on_tool_end(self, documents: list[Document], session: scoped_session): """Handle tool end.""" for document in documents: if document.metadata is not None: document_id = document.metadata["document_id"] dataset_document_stmt = select(DatasetDocument).where(DatasetDocument.id == document_id) - dataset_document = db.session.scalar(dataset_document_stmt) + dataset_document = session.scalar(dataset_document_stmt) if not dataset_document: _logger.warning( "Expected DatasetDocument record to exist, but none was found, document_id=%s", @@ -68,9 +68,9 @@ class DatasetIndexToolCallbackHandler: ChildChunk.dataset_id == dataset_document.dataset_id, ChildChunk.document_id == dataset_document.id, ) - child_chunk = db.session.scalar(child_chunk_stmt) + child_chunk = session.scalar(child_chunk_stmt) if child_chunk: - db.session.execute( + session.execute( update(DocumentSegment) .where(DocumentSegment.id == child_chunk.segment_id) .values(hit_count=DocumentSegment.hit_count + 1) @@ -82,11 +82,11 @@ class DatasetIndexToolCallbackHandler: conditions.append(DocumentSegment.dataset_id == document.metadata["dataset_id"]) # add hit count to document segment - db.session.execute( + session.execute( update(DocumentSegment).where(*conditions).values(hit_count=DocumentSegment.hit_count + 1) ) - db.session.commit() + session.commit() # TODO(-LAN-): Improve type check def return_retriever_resource_info(self, resource: Sequence[RetrievalSourceMetadata]): diff --git a/api/core/tools/utils/dataset_retriever/dataset_multi_retriever_tool.py b/api/core/tools/utils/dataset_retriever/dataset_multi_retriever_tool.py index beb8c5d005a..55a406482a3 100644 --- a/api/core/tools/utils/dataset_retriever/dataset_multi_retriever_tool.py +++ b/api/core/tools/utils/dataset_retriever/dataset_multi_retriever_tool.py @@ -79,7 +79,7 @@ class DatasetMultiRetrieverTool(DatasetRetrieverBaseTool): all_documents = rerank_runner.run(query, all_documents, self.score_threshold, self.top_k) for hit_callback in self.hit_callbacks: - hit_callback.on_tool_end(all_documents) + hit_callback.on_tool_end(all_documents, db.session) document_score_list = {} for item in all_documents: @@ -166,7 +166,7 @@ class DatasetMultiRetrieverTool(DatasetRetrieverBaseTool): return [] for hit_callback in hit_callbacks: - hit_callback.on_query(query, dataset.id) + hit_callback.on_query(query, dataset.id, db.session) # get retrieval model , if the model is not setting , using default retrieval_model = dataset.retrieval_model or default_retrieval_model diff --git a/api/core/tools/utils/dataset_retriever/dataset_retriever_tool.py b/api/core/tools/utils/dataset_retriever/dataset_retriever_tool.py index 85a6e57b4ce..d9f99884d0c 100644 --- a/api/core/tools/utils/dataset_retriever/dataset_retriever_tool.py +++ b/api/core/tools/utils/dataset_retriever/dataset_retriever_tool.py @@ -64,7 +64,7 @@ class DatasetRetrieverTool(DatasetRetrieverBaseTool): if not dataset: return "" for hit_callback in self.hit_callbacks: - hit_callback.on_query(query, dataset.id) + hit_callback.on_query(query, dataset.id, db.session) dataset_retrieval = DatasetRetrieval() metadata_filter_document_ids, metadata_condition = dataset_retrieval.get_metadata_filter_condition( [dataset.id], @@ -159,7 +159,7 @@ class DatasetRetrieverTool(DatasetRetrieverBaseTool): else: documents = [] for hit_callback in self.hit_callbacks: - hit_callback.on_tool_end(documents) + hit_callback.on_tool_end(documents, db.session) document_score_list = {} if dataset.indexing_technique != IndexTechniqueType.ECONOMY: for item in documents: 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 f23669c3c7b..4912badfc55 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 @@ -14,9 +14,6 @@ def mock_queue_manager(mocker: MockerFixture): @pytest.fixture def handler(mock_queue_manager, mocker: MockerFixture): - mocker.patch( - "core.callback_handler.index_tool_callback_handler.db", - ) return DatasetIndexToolCallbackHandler( queue_manager=mock_queue_manager, app_id="app-1", @@ -37,7 +34,7 @@ class TestOnQuery: ) def test_on_query_success_roles(self, mocker: MockerFixture, mock_queue_manager, invoke_from, expected_role): # Arrange - mock_db = mocker.patch("core.callback_handler.index_tool_callback_handler.db") + mock_session = mocker.Mock() handler = DatasetIndexToolCallbackHandler( queue_manager=mock_queue_manager, @@ -50,16 +47,16 @@ class TestOnQuery: handler._invoke_from = invoke_from # Act - handler.on_query("test query", "dataset-1") + handler.on_query("test query", "dataset-1", mock_session) # Assert - mock_db.session.add.assert_called_once() - dataset_query = mock_db.session.add.call_args.args[0] + mock_session.add.assert_called_once() + dataset_query = mock_session.add.call_args.args[0] assert dataset_query.created_by_role == expected_role - mock_db.session.commit.assert_called_once() + mock_session.commit.assert_called_once() def test_on_query_none_values(self, mocker: MockerFixture, mock_queue_manager): - mock_db = mocker.patch("core.callback_handler.index_tool_callback_handler.db") + mock_session = mocker.Mock() handler = DatasetIndexToolCallbackHandler( queue_manager=mock_queue_manager, @@ -69,40 +66,40 @@ class TestOnQuery: invoke_from=None, ) - handler.on_query(None, None) + handler.on_query(None, None, mock_session) - mock_db.session.add.assert_called_once() - mock_db.session.commit.assert_called_once() + mock_session.add.assert_called_once() + mock_session.commit.assert_called_once() class TestOnToolEnd: def test_on_tool_end_no_metadata(self, handler: DatasetIndexToolCallbackHandler, mocker: MockerFixture): - mock_db = mocker.patch("core.callback_handler.index_tool_callback_handler.db") + mock_session = mocker.Mock() document = mocker.Mock() document.metadata = None - handler.on_tool_end([document]) + handler.on_tool_end([document], mock_session) - mock_db.session.commit.assert_not_called() + mock_session.commit.assert_not_called() def test_on_tool_end_dataset_document_not_found( self, handler: DatasetIndexToolCallbackHandler, mocker: MockerFixture ): - mock_db = mocker.patch("core.callback_handler.index_tool_callback_handler.db") - mock_db.session.scalar.return_value = None + mock_session = mocker.Mock() + mock_session.scalar.return_value = None document = mocker.Mock() document.metadata = {"document_id": "doc-1", "doc_id": "node-1"} - handler.on_tool_end([document]) + handler.on_tool_end([document], mock_session) - mock_db.session.scalar.assert_called_once() + mock_session.scalar.assert_called_once() def test_on_tool_end_parent_child_index_with_child( self, handler: DatasetIndexToolCallbackHandler, mocker: MockerFixture ): - mock_db = mocker.patch("core.callback_handler.index_tool_callback_handler.db") + mock_session = mocker.Mock() mock_dataset_doc = mocker.Mock() from core.callback_handler.index_tool_callback_handler import IndexStructureType @@ -114,23 +111,23 @@ class TestOnToolEnd: mock_child_chunk = mocker.Mock() mock_child_chunk.segment_id = "segment-1" - mock_db.session.scalar.side_effect = [mock_dataset_doc, mock_child_chunk] + mock_session.scalar.side_effect = [mock_dataset_doc, mock_child_chunk] document = mocker.Mock() document.metadata = {"document_id": "doc-1", "doc_id": "node-1"} - handler.on_tool_end([document]) + handler.on_tool_end([document], mock_session) - mock_db.session.execute.assert_called_once() - mock_db.session.commit.assert_called_once() + mock_session.execute.assert_called_once() + mock_session.commit.assert_called_once() def test_on_tool_end_non_parent_child_index(self, handler: DatasetIndexToolCallbackHandler, mocker: MockerFixture): - mock_db = mocker.patch("core.callback_handler.index_tool_callback_handler.db") + mock_session = mocker.Mock() mock_dataset_doc = mocker.Mock() mock_dataset_doc.doc_form = "OTHER" - mock_db.session.scalar.return_value = mock_dataset_doc + mock_session.scalar.return_value = mock_dataset_doc document = mocker.Mock() document.metadata = { @@ -139,13 +136,14 @@ class TestOnToolEnd: "dataset_id": "dataset-1", } - handler.on_tool_end([document]) + handler.on_tool_end([document], mock_session) - mock_db.session.execute.assert_called_once() - mock_db.session.commit.assert_called_once() + mock_session.execute.assert_called_once() + mock_session.commit.assert_called_once() - def test_on_tool_end_empty_documents(self, handler: DatasetIndexToolCallbackHandler): - handler.on_tool_end([]) + def test_on_tool_end_empty_documents(self, handler: DatasetIndexToolCallbackHandler, mocker: MockerFixture): + mock_session = mocker.Mock() + handler.on_tool_end([], mock_session) class TestReturnRetrieverResourceInfo: diff --git a/api/tests/unit_tests/core/tools/utils/test_misc_utils_extra.py b/api/tests/unit_tests/core/tools/utils/test_misc_utils_extra.py index a93624123e2..cb3be81ab61 100644 --- a/api/tests/unit_tests/core/tools/utils/test_misc_utils_extra.py +++ b/api/tests/unit_tests/core/tools/utils/test_misc_utils_extra.py @@ -48,10 +48,10 @@ class _TestHitCallback(DatasetIndexToolCallbackHandler): self.documents: list[RagDocument] | None = None self.resources = None - def on_query(self, query: str, dataset_id: str): + def on_query(self, query: str, dataset_id: str, session=None): self.queries.append((query, dataset_id)) - def on_tool_end(self, documents: list[RagDocument]): + def on_tool_end(self, documents: list[RagDocument], session=None): self.documents = documents def return_retriever_resource_info(self, resource):