mirror of
https://github.com/langgenius/dify.git
synced 2026-09-24 23:22:26 +08:00
fix(api): isolate side-effect session writes in multimodal and RAG handlers (#38210)
Co-authored-by: FFXN <31929997+FFXN@users.noreply.github.com>
This commit is contained in:
@@ -5,6 +5,8 @@ from collections.abc import Generator, Mapping, Sequence
|
||||
from mimetypes import guess_extension
|
||||
from typing import TYPE_CHECKING, Any, Union
|
||||
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
|
||||
from core.app.app_config.entities import ExternalDataVariableEntity, PromptTemplateEntity
|
||||
from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom
|
||||
from core.app.apps.exc import GenerateTaskStoppedError
|
||||
@@ -423,7 +425,9 @@ class AppRunner:
|
||||
_logger.exception("Failed to save image file")
|
||||
return
|
||||
|
||||
# Create MessageFile record
|
||||
# Create MessageFile record.
|
||||
# Use an independent session so this side-effect write does not
|
||||
# commit or close the caller's request-scoped session.
|
||||
message_file = MessageFile(
|
||||
message_id=message_id,
|
||||
type=FileType.IMAGE,
|
||||
@@ -437,9 +441,8 @@ class AppRunner:
|
||||
created_by=user_id,
|
||||
)
|
||||
|
||||
db.session.add(message_file)
|
||||
db.session.commit()
|
||||
db.session.refresh(message_file)
|
||||
with sessionmaker(bind=db.engine, expire_on_commit=False).begin() as session:
|
||||
session.add(message_file)
|
||||
|
||||
# Publish QueueMessageFileEvent
|
||||
queue_manager.publish(
|
||||
|
||||
@@ -2,7 +2,7 @@ import logging
|
||||
from collections.abc import Sequence
|
||||
|
||||
from sqlalchemy import select, update
|
||||
from sqlalchemy.orm import scoped_session
|
||||
from sqlalchemy.orm import Session, scoped_session, sessionmaker
|
||||
|
||||
from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom
|
||||
from core.app.entities.app_invoke_entities import InvokeFrom
|
||||
@@ -10,6 +10,7 @@ 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
|
||||
@@ -46,47 +47,52 @@ class DatasetIndexToolCallbackHandler:
|
||||
created_by=self._user_id,
|
||||
)
|
||||
|
||||
session.add(dataset_query)
|
||||
session.commit()
|
||||
# Use an independent session so this audit-log side effect does
|
||||
# not commit or close the caller's request-scoped session.
|
||||
with sessionmaker(bind=db.engine, expire_on_commit=False).begin() as independent_session:
|
||||
independent_session.add(dataset_query)
|
||||
|
||||
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 = session.scalar(dataset_document_stmt)
|
||||
if not dataset_document:
|
||||
_logger.warning(
|
||||
"Expected DatasetDocument record to exist, but none was found, document_id=%s",
|
||||
document_id,
|
||||
)
|
||||
continue
|
||||
if dataset_document.doc_form == IndexStructureType.PARENT_CHILD_INDEX:
|
||||
child_chunk_stmt = select(ChildChunk).where(
|
||||
ChildChunk.index_node_id == document.metadata["doc_id"],
|
||||
ChildChunk.dataset_id == dataset_document.dataset_id,
|
||||
ChildChunk.document_id == dataset_document.id,
|
||||
)
|
||||
child_chunk = session.scalar(child_chunk_stmt)
|
||||
if child_chunk:
|
||||
session.execute(
|
||||
update(DocumentSegment)
|
||||
.where(DocumentSegment.id == child_chunk.segment_id)
|
||||
.values(hit_count=DocumentSegment.hit_count + 1)
|
||||
# Use an independent session so hit-count updates do not
|
||||
# interfere with the caller's request-scoped session.
|
||||
with Session(db.engine, expire_on_commit=False) as independent_session:
|
||||
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 = independent_session.scalar(dataset_document_stmt)
|
||||
if not dataset_document:
|
||||
_logger.warning(
|
||||
"Expected DatasetDocument record to exist, but none was found, document_id=%s",
|
||||
document_id,
|
||||
)
|
||||
else:
|
||||
conditions = [DocumentSegment.index_node_id == document.metadata["doc_id"]]
|
||||
continue
|
||||
if dataset_document.doc_form == IndexStructureType.PARENT_CHILD_INDEX:
|
||||
child_chunk_stmt = select(ChildChunk).where(
|
||||
ChildChunk.index_node_id == document.metadata["doc_id"],
|
||||
ChildChunk.dataset_id == dataset_document.dataset_id,
|
||||
ChildChunk.document_id == dataset_document.id,
|
||||
)
|
||||
child_chunk = independent_session.scalar(child_chunk_stmt)
|
||||
if child_chunk:
|
||||
independent_session.execute(
|
||||
update(DocumentSegment)
|
||||
.where(DocumentSegment.id == child_chunk.segment_id)
|
||||
.values(hit_count=DocumentSegment.hit_count + 1)
|
||||
)
|
||||
else:
|
||||
conditions = [DocumentSegment.index_node_id == document.metadata["doc_id"]]
|
||||
|
||||
if "dataset_id" in document.metadata:
|
||||
conditions.append(DocumentSegment.dataset_id == document.metadata["dataset_id"])
|
||||
if "dataset_id" in document.metadata:
|
||||
conditions.append(DocumentSegment.dataset_id == document.metadata["dataset_id"])
|
||||
|
||||
# add hit count to document segment
|
||||
session.execute(
|
||||
update(DocumentSegment).where(*conditions).values(hit_count=DocumentSegment.hit_count + 1)
|
||||
)
|
||||
# add hit count to document segment
|
||||
independent_session.execute(
|
||||
update(DocumentSegment).where(*conditions).values(hit_count=DocumentSegment.hit_count + 1)
|
||||
)
|
||||
|
||||
session.commit()
|
||||
independent_session.commit()
|
||||
|
||||
# TODO(-LAN-): Improve type check
|
||||
def return_retriever_resource_info(self, resource: Sequence[RetrievalSourceMetadata]):
|
||||
|
||||
Reference in New Issue
Block a user