fix: session unbound during parent-child retrieval (#29396)

This commit is contained in:
Jyong
2025-12-10 14:08:55 +08:00
committed by GitHub
parent e8720de9ad
commit b49e2646ff
3 changed files with 9 additions and 16 deletions

View File

@@ -371,7 +371,7 @@ class RetrievalService:
include_segment_ids = set() include_segment_ids = set()
segment_child_map = {} segment_child_map = {}
segment_file_map = {} segment_file_map = {}
with Session(db.engine) as session: with Session(bind=db.engine, expire_on_commit=False) as session:
# Process documents # Process documents
for document in documents: for document in documents:
segment_id = None segment_id = None
@@ -395,7 +395,7 @@ class RetrievalService:
session, session,
) )
if attachment_info_dict: if attachment_info_dict:
attachment_info = attachment_info_dict["attchment_info"] attachment_info = attachment_info_dict["attachment_info"]
segment_id = attachment_info_dict["segment_id"] segment_id = attachment_info_dict["segment_id"]
else: else:
child_index_node_id = document.metadata.get("doc_id") child_index_node_id = document.metadata.get("doc_id")
@@ -417,13 +417,6 @@ class RetrievalService:
DocumentSegment.status == "completed", DocumentSegment.status == "completed",
DocumentSegment.id == segment_id, DocumentSegment.id == segment_id,
) )
.options(
load_only(
DocumentSegment.id,
DocumentSegment.content,
DocumentSegment.answer,
)
)
.first() .first()
) )
@@ -475,7 +468,7 @@ class RetrievalService:
session, session,
) )
if attachment_info_dict: if attachment_info_dict:
attachment_info = attachment_info_dict["attchment_info"] attachment_info = attachment_info_dict["attachment_info"]
segment_id = attachment_info_dict["segment_id"] segment_id = attachment_info_dict["segment_id"]
document_segment_stmt = select(DocumentSegment).where( document_segment_stmt = select(DocumentSegment).where(
DocumentSegment.dataset_id == dataset_document.dataset_id, DocumentSegment.dataset_id == dataset_document.dataset_id,
@@ -684,7 +677,7 @@ class RetrievalService:
.first() .first()
) )
if attachment_binding: if attachment_binding:
attchment_info = { attachment_info = {
"id": upload_file.id, "id": upload_file.id,
"name": upload_file.name, "name": upload_file.name,
"extension": "." + upload_file.extension, "extension": "." + upload_file.extension,
@@ -692,5 +685,5 @@ class RetrievalService:
"source_url": sign_upload_file(upload_file.id, upload_file.extension), "source_url": sign_upload_file(upload_file.id, upload_file.extension),
"size": upload_file.size, "size": upload_file.size,
} }
return {"attchment_info": attchment_info, "segment_id": attachment_binding.segment_id} return {"attachment_info": attachment_info, "segment_id": attachment_binding.segment_id}
return None return None

View File

@@ -266,7 +266,7 @@ class DatasetRetrieval:
).all() ).all()
if attachments_with_bindings: if attachments_with_bindings:
for _, upload_file in attachments_with_bindings: for _, upload_file in attachments_with_bindings:
attchment_info = File( attachment_info = File(
id=upload_file.id, id=upload_file.id,
filename=upload_file.name, filename=upload_file.name,
extension="." + upload_file.extension, extension="." + upload_file.extension,
@@ -280,7 +280,7 @@ class DatasetRetrieval:
storage_key=upload_file.key, storage_key=upload_file.key,
url=sign_upload_file(upload_file.id, upload_file.extension), url=sign_upload_file(upload_file.id, upload_file.extension),
) )
context_files.append(attchment_info) context_files.append(attachment_info)
if show_retrieve_source: if show_retrieve_source:
for record in records: for record in records:
segment = record.segment segment = record.segment

View File

@@ -697,7 +697,7 @@ class LLMNode(Node[LLMNodeData]):
).all() ).all()
if attachments_with_bindings: if attachments_with_bindings:
for _, upload_file in attachments_with_bindings: for _, upload_file in attachments_with_bindings:
attchment_info = File( attachment_info = File(
id=upload_file.id, id=upload_file.id,
filename=upload_file.name, filename=upload_file.name,
extension="." + upload_file.extension, extension="." + upload_file.extension,
@@ -711,7 +711,7 @@ class LLMNode(Node[LLMNodeData]):
storage_key=upload_file.key, storage_key=upload_file.key,
url=sign_upload_file(upload_file.id, upload_file.extension), url=sign_upload_file(upload_file.id, upload_file.extension),
) )
context_files.append(attchment_info) context_files.append(attachment_info)
yield RunRetrieverResourceEvent( yield RunRetrieverResourceEvent(
retriever_resources=original_retriever_resource, retriever_resources=original_retriever_resource,
context=context_str.strip(), context=context_str.strip(),