| from typing import Optional | from typing import Optional | ||||
| from flask import Flask, current_app | from flask import Flask, current_app | ||||
| from sqlalchemy import and_, or_ | |||||
| from sqlalchemy.orm import load_only | from sqlalchemy.orm import load_only | ||||
| from sqlalchemy.sql.expression import false | |||||
| from configs import dify_config | from configs import dify_config | ||||
| from core.rag.data_post_processor.data_post_processor import DataPostProcessor | from core.rag.data_post_processor.data_post_processor import DataPostProcessor | ||||
| child_chunks = db.session.query(ChildChunk).filter(ChildChunk.index_node_id.in_(child_index_node_ids)).all() | child_chunks = db.session.query(ChildChunk).filter(ChildChunk.index_node_id.in_(child_index_node_ids)).all() | ||||
| child_chunk_map = {chunk.index_node_id: chunk for chunk in child_chunks} | child_chunk_map = {chunk.index_node_id: chunk for chunk in child_chunks} | ||||
| # Batch query DocumentSegment with unified conditions | |||||
| segment_ids_from_child = [chunk.segment_id for chunk in child_chunks] | |||||
| segment_conditions = [] | |||||
| if index_node_ids: | |||||
| segment_conditions.append(DocumentSegment.index_node_id.in_(index_node_ids)) | |||||
| if segment_ids_from_child: | |||||
| segment_conditions.append(DocumentSegment.id.in_(segment_ids_from_child)) | |||||
| if segment_conditions: | |||||
| filter_expr = or_(*segment_conditions) | |||||
| else: | |||||
| filter_expr = false() | |||||
| segment_map = { | segment_map = { | ||||
| segment.id: segment | segment.id: segment | ||||
| for segment in db.session.query(DocumentSegment) | for segment in db.session.query(DocumentSegment) | ||||
| .filter( | .filter( | ||||
| ( | |||||
| DocumentSegment.index_node_id.in_(index_node_ids) | |||||
| | DocumentSegment.id.in_([chunk.segment_id for chunk in child_chunks]) | |||||
| ), | |||||
| DocumentSegment.enabled == True, | |||||
| DocumentSegment.status == "completed", | |||||
| and_( | |||||
| filter_expr, | |||||
| DocumentSegment.enabled == True, | |||||
| DocumentSegment.status == "completed", | |||||
| ) | |||||
| ) | ) | ||||
| .options( | .options( | ||||
| load_only( | load_only( |