AI/dify
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/langgenius/dify.git synced 2025-07-13 04:01:50 +08:00
Code Issues Packages Projects Releases Wiki Activity

revert batch query (#17707)

Browse Source
This commit is contained in:
Jyong 2025-04-09 20:25:36 +08:00 committed by GitHub
parent 1d5c07dedb
commit 8b3be4224d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 74 additions and 109 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

143
api/core/rag/datasource/retrieval_service.py
View File

@ -1,13 +1,9 @@
import concurrent.futures import concurrent.futures
import logging
import time
from concurrent.futures import ThreadPoolExecutor from concurrent.futures import ThreadPoolExecutor
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
@ -182,7 +178,6 @@ class RetrievalService:
if not dataset: if not dataset:
raise ValueError("dataset not found") raise ValueError("dataset not found")
start = time.time()
vector = Vector(dataset=dataset) vector = Vector(dataset=dataset)
documents = vector.search_by_vector( documents = vector.search_by_vector(
query, query,
@ -192,7 +187,6 @@ class RetrievalService:
filter={"group_id": [dataset.id]}, filter={"group_id": [dataset.id]},
document_ids_filter=document_ids_filter, document_ids_filter=document_ids_filter,
) )
logging.debug(f"embedding_search ends at {time.time() - start:.2f} seconds")
if documents: if documents:
if ( if (
@ -276,8 +270,7 @@ class RetrievalService:
return [] return []
try: try:
start_time = time.time() # Collect document IDs
# Collect document IDs with existence check
document_ids = {doc.metadata.get("document_id") for doc in documents if "document_id" in doc.metadata} document_ids = {doc.metadata.get("document_id") for doc in documents if "document_id" in doc.metadata}
if not document_ids: if not document_ids:
return [] return []
@ -295,51 +288,34 @@ class RetrievalService:
include_segment_ids = set() include_segment_ids = set()
segment_child_map = {} segment_child_map = {}
# Precompute doc_forms to avoid redundant checks # Process documents
doc_forms = {} for document in documents:
for doc in documents: document_id = document.metadata.get("document_id")
document_id = doc.metadata.get("document_id") if document_id not in dataset_documents:
dataset_doc = dataset_documents.get(document_id) continue
if dataset_doc:
doc_forms[document_id] = dataset_doc.doc_form
# Batch collect index node IDs with type safety dataset_document = dataset_documents[document_id]
child_index_node_ids = [] if not dataset_document:
index_node_ids = [] continue
for doc in documents:
document_id = doc.metadata.get("document_id")
if doc_forms.get(document_id) == IndexType.PARENT_CHILD_INDEX:
child_index_node_ids.append(doc.metadata.get("doc_id"))
else:
index_node_ids.append(doc.metadata.get("doc_id"))
# Batch query ChildChunk if dataset_document.doc_form == IndexType.PARENT_CHILD_INDEX:
child_chunks = db.session.query(ChildChunk).filter(ChildChunk.index_node_id.in_(child_index_node_ids)).all() # Handle parent-child documents
child_chunk_map = {chunk.index_node_id: chunk for chunk in child_chunks} child_index_node_id = document.metadata.get("doc_id")
segment_ids_from_child = [chunk.segment_id for chunk in child_chunks] child_chunk = (
segment_conditions = [] db.session.query(ChildChunk).filter(ChildChunk.index_node_id == child_index_node_id).first()
)
if index_node_ids: if not child_chunk:
segment_conditions.append(DocumentSegment.index_node_id.in_(index_node_ids)) continue
if segment_ids_from_child: segment = (
segment_conditions.append(DocumentSegment.id.in_(segment_ids_from_child)) db.session.query(DocumentSegment)
if segment_conditions:
filter_expr = or_(*segment_conditions)
else:
filter_expr = false()
segment_map = {
segment.id: segment
for segment in db.session.query(DocumentSegment)
.filter( .filter(
and_( DocumentSegment.dataset_id == dataset_document.dataset_id,
filter_expr,
DocumentSegment.enabled == True, DocumentSegment.enabled == True,
DocumentSegment.status == "completed", DocumentSegment.status == "completed",
) DocumentSegment.id == child_chunk.segment_id,
) )
.options( .options(
load_only( load_only(
@ -348,37 +324,30 @@ class RetrievalService:
DocumentSegment.answer, DocumentSegment.answer,
) )
) )
.all() .first()
} )
for document in documents:
document_id = document.metadata.get("document_id")
dataset_document = dataset_documents.get(document_id)
if not dataset_document:
continue
doc_form = doc_forms.get(document_id)
if doc_form == IndexType.PARENT_CHILD_INDEX:
# Handle parent-child documents using preloaded data
child_index_node_id = document.metadata.get("doc_id")
if not child_index_node_id:
continue
child_chunk = child_chunk_map.get(child_index_node_id)
if not child_chunk:
continue
segment = segment_map.get(child_chunk.segment_id)
if not segment: if not segment:
continue continue
if segment.id not in include_segment_ids: if segment.id not in include_segment_ids:
include_segment_ids.add(segment.id) include_segment_ids.add(segment.id)
map_detail = {"max_score": document.metadata.get("score", 0.0), "child_chunks": []} child_chunk_detail = {
"id": child_chunk.id,
"content": child_chunk.content,
"position": child_chunk.position,
"score": document.metadata.get("score", 0.0),
}
map_detail = {
"max_score": document.metadata.get("score", 0.0),
"child_chunks": [child_chunk_detail],
}
segment_child_map[segment.id] = map_detail segment_child_map[segment.id] = map_detail
records.append({"segment": segment}) record = {
"segment": segment,
# Append child chunk details }
records.append(record)
else:
child_chunk_detail = { child_chunk_detail = {
"id": child_chunk.id, "id": child_chunk.id,
"content": child_chunk.content, "content": child_chunk.content,
@ -389,44 +358,40 @@ class RetrievalService:
segment_child_map[segment.id]["max_score"] = max( segment_child_map[segment.id]["max_score"] = max(
segment_child_map[segment.id]["max_score"], document.metadata.get("score", 0.0) segment_child_map[segment.id]["max_score"], document.metadata.get("score", 0.0)
) )
else: else:
# Handle normal documents # Handle normal documents
index_node_id = document.metadata.get("doc_id") index_node_id = document.metadata.get("doc_id")
if not index_node_id: if not index_node_id:
continue continue
segment = next( segment = (
( db.session.query(DocumentSegment)
s .filter(
for s in segment_map.values() DocumentSegment.dataset_id == dataset_document.dataset_id,
if s.index_node_id == index_node_id and s.dataset_id == dataset_document.dataset_id DocumentSegment.enabled == True,
), DocumentSegment.status == "completed",
None, DocumentSegment.index_node_id == index_node_id,
)
.first()
) )
if not segment: if not segment:
continue continue
if segment.id not in include_segment_ids:
include_segment_ids.add(segment.id) include_segment_ids.add(segment.id)
records.append( record = {
{
"segment": segment, "segment": segment,
"score": document.metadata.get("score", 0.0), "score": document.metadata.get("score"), # type: ignore
} }
) records.append(record)
# Merge child chunks information # Add child chunks information to records
for record in records: for record in records:
segment_id = record["segment"].id if record["segment"].id in segment_child_map:
if segment_id in segment_child_map: record["child_chunks"] = segment_child_map[record["segment"].id].get("child_chunks") # type: ignore
record["child_chunks"] = segment_child_map[segment_id]["child_chunks"] record["score"] = segment_child_map[record["segment"].id]["max_score"]
record["score"] = segment_child_map[segment_id]["max_score"]
logging.debug(f"Formatting retrieval documents took {time.time() - start_time:.2f} seconds")
return [RetrievalSegments(**record) for record in records] return [RetrievalSegments(**record) for record in records]
except Exception as e: except Exception as e:
# Only rollback if there were write operations
db.session.rollback() db.session.rollback()
raise e raise e