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:
Pranav Agarwal
2026-07-06 05:17:12 +00:00
committed by GitHub
co-authored by FFXN
parent 93eb6d32b5
commit d9c99daf29
4 changed files with 286 additions and 231 deletions
+7 -4
View File
@@ -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]):