opensource/dify
3
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-09-24 23:22:26 +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 06:16:13 +00:00
committed by GitHub
parent edf21269b3
commit 3390be3978
5 changed files with 45 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/core/callback_handler/index_tool_callback_handler.py
+10 -10
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]):