opensource/dify
2
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-07-22 03:08:33 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor: pass db.session explicitly in DatasetIndexToolCallbackHandler (#38082)

Browse Source
This commit is contained in:
Checo 2026-07-01 14:16:13 +08:00 committed by GitHub
parent edf21269b3
commit 3390be3978
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 45 additions and 47 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
api/core/callback_handler/index_tool_callback_handler.py
View File

@ -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]):

4
api/core/tools/utils/dataset_retriever/dataset_multi_retriever_tool.py
View File

@ -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

4
api/core/tools/utils/dataset_retriever/dataset_retriever_tool.py
View File

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

60
api/tests/unit_tests/core/callback_handler/test_index_tool_callback_handler.py
View File

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

4
api/tests/unit_tests/core/tools/utils/test_misc_utils_extra.py
View File

@ -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):