Преглед изворни кода

fix segment query index not effect (#17704)

tags/1.3.0
Jyong пре 6 месеци
родитељ
комит
abfcd9d3b6
No account linked to committer's email address
1 измењених фајлова са 21 додато и 7 уклоњено
  1. 21
    7
      api/core/rag/datasource/retrieval_service.py

+ 21
- 7
api/core/rag/datasource/retrieval_service.py Прегледај датотеку

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(

Loading…
Откажи
Сачувај