opensource
/
dify
mirror of https://github.com/langgenius/dify.git
2
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix multimodal embedding retrival test

This commit is contained in:
jyong 2025-12-26 17:04:46 +08:00
parent 894a3c03a2
commit 901cc64ac9
1 changed files with 69 additions and 76 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

@ -381,10 +381,9 @@ class RetrievalService:
records = [] records = []
include_segment_ids = set() include_segment_ids = set()
segment_child_map = {} segment_child_map = {}
segment_file_map = {}
valid_dataset_documents = {} valid_dataset_documents = {}
image_doc_ids = [] image_doc_ids: list[Any] = []
child_index_node_ids = [] child_index_node_ids = []
index_node_ids = [] index_node_ids = []
doc_to_document_map = {} doc_to_document_map = {}
@ -421,7 +420,7 @@ class RetrievalService:
index_node_segments: list[DocumentSegment] = [] index_node_segments: list[DocumentSegment] = []
segments: list[DocumentSegment] = [] segments: list[DocumentSegment] = []
attachment_map = {} attachment_map = {}
child_chunk_map = {} child_chunk_map: dict[Any, Any] = {}
doc_segment_map = {} doc_segment_map = {}
with session_factory.create_session() as session: with session_factory.create_session() as session:
@ -429,16 +428,27 @@ class RetrievalService:
for attachment in attachments: for attachment in attachments:
segment_ids.append(attachment["segment_id"]) segment_ids.append(attachment["segment_id"])
attachment_map[attachment["segment_id"]] = attachment if attachment["segment_id"] in attachment_map:
doc_segment_map[attachment["segment_id"]] = attachment["attachment_id"] attachment_map[attachment["segment_id"]].append(attachment["attachment_info"])
else:
attachment_map[attachment["segment_id"]] = [attachment["attachment_info"]]
if attachment["attachment_id"] in doc_segment_map:
doc_segment_map[attachment["segment_id"]].append(attachment["attachment_id"])
else:
doc_segment_map[attachment["segment_id"]] = [attachment["attachment_id"]]
child_chunk_stmt = select(ChildChunk).where(ChildChunk.index_node_id.in_(child_index_node_ids)) child_chunk_stmt = select(ChildChunk).where(ChildChunk.index_node_id.in_(child_index_node_ids))
child_index_nodes = session.execute(child_chunk_stmt).scalars().all() child_index_nodes = session.execute(child_chunk_stmt).scalars().all()
for i in child_index_nodes: for i in child_index_nodes:
segment_ids.append(i.segment_id) segment_ids.append(i.segment_id)
child_chunk_map[i.segment_id] = i if i.segment_id in child_chunk_map:
doc_segment_map[i.segment_id] = i.index_node_id child_chunk_map[i.segment_id].append(i)
else:
child_chunk_map[i.segment_id] = [i]
if i.segment_id in doc_segment_map:
doc_segment_map[i.segment_id].append(i.index_node_id)
else:
doc_segment_map[i.segment_id] = [i.index_node_id]
if index_node_ids: if index_node_ids:
document_segment_stmt = select(DocumentSegment).where( document_segment_stmt = select(DocumentSegment).where(
@ -448,7 +458,7 @@ class RetrievalService:
) )
index_node_segments = session.execute(document_segment_stmt).scalars().all() # type: ignore index_node_segments = session.execute(document_segment_stmt).scalars().all() # type: ignore
for index_node_segment in index_node_segments: for index_node_segment in index_node_segments:
doc_segment_map[index_node_segment.id] = index_node_segment.index_node_id doc_segment_map[index_node_segment.id] = [index_node_segment.index_node_id]
if segment_ids: if segment_ids:
document_segment_stmt = select(DocumentSegment).where( document_segment_stmt = select(DocumentSegment).where(
DocumentSegment.enabled == True, DocumentSegment.enabled == True,
@ -461,85 +471,65 @@ class RetrievalService:
segments.extend(index_node_segments) segments.extend(index_node_segments)
for segment in segments: for segment in segments:
doc_id = doc_segment_map.get(segment.id) child_chunks: list[ChildChunk] = child_chunk_map.get(segment.id, [])
child_chunk = child_chunk_map.get(segment.id) attachment_infos: list[dict[str, Any]] = attachment_map.get(segment.id, [])
attachment_info = attachment_map.get(segment.id) ds_dataset_document: DatasetDocument | None = valid_dataset_documents.get(
segment.document_id
)
if doc_id: if ds_dataset_document and ds_dataset_document.doc_form == IndexStructureType.PARENT_CHILD_INDEX:
document = doc_to_document_map[doc_id] if segment.id not in include_segment_ids:
ds_dataset_document: DatasetDocument | None = valid_dataset_documents.get( include_segment_ids.add(segment.id)
document.metadata.get("document_id") if child_chunks or attachment_infos:
) child_chunk_details = []
max_score = 0.0
if ds_dataset_document and ds_dataset_document.doc_form == IndexStructureType.PARENT_CHILD_INDEX: for child_chunk in child_chunks:
if segment.id not in include_segment_ids: document = doc_to_document_map[child_chunk.index_node_id]
include_segment_ids.add(segment.id)
if child_chunk:
child_chunk_detail = { child_chunk_detail = {
"id": child_chunk.id, "id": child_chunk.id,
"content": child_chunk.content, "content": child_chunk.content,
"position": child_chunk.position, "position": child_chunk.position,
"score": document.metadata.get("score", 0.0) if document else 0.0, "score": document.metadata.get("score", 0.0) if document else 0.0,
} }
map_detail = { child_chunk_details.append(child_chunk_detail)
"max_score": document.metadata.get("score", 0.0) if document else 0.0, max_score = max(max_score, document.metadata.get("score", 0.0) if document else 0.0)
"child_chunks": [child_chunk_detail], for attachment_info in attachment_infos:
} file_document = doc_to_document_map[attachment_info["id"]]
segment_child_map[segment.id] = map_detail max_score = max(max_score, file_document.metadata.get("score", 0.0) if file_document else 0.0)
record = {
"segment": segment, map_detail = {
"max_score": max_score,
"child_chunks": child_chunk_details,
} }
if attachment_info: segment_child_map[segment.id] = map_detail
segment_file_map[segment.id] = [attachment_info] record = {
records.append(record) "segment": segment,
else: }
if child_chunk: records.append(record)
child_chunk_detail = { else:
"id": child_chunk.id, if segment.id not in include_segment_ids:
"content": child_chunk.content, include_segment_ids.add(segment.id)
"position": child_chunk.position, max_score = 0.0
"score": document.metadata.get("score", 0.0), document = doc_to_document_map.get(segment.index_node_id)
} if document:
if segment.id in segment_child_map: max_score = max(max_score, document.metadata.get("score", 0.0))
segment_child_map[segment.id]["child_chunks"].append(child_chunk_detail) # type: ignore for attachment_info in attachment_infos:
segment_child_map[segment.id]["max_score"] = max( file_document = doc_to_document_map.get(attachment_info["id"])
segment_child_map[segment.id]["max_score"], if file_document:
document.metadata.get("score", 0.0) if document else 0.0, max_score = max(max_score, file_document.metadata.get("score", 0.0))
) record = {
else: "segment": segment,
segment_child_map[segment.id] = { "score": max_score,
"max_score": document.metadata.get("score", 0.0) if document else 0.0, }
"child_chunks": [child_chunk_detail], records.append(record)
}
if attachment_info:
if segment.id in segment_file_map:
segment_file_map[segment.id].append(attachment_info)
else:
segment_file_map[segment.id] = [attachment_info]
else:
if segment.id not in include_segment_ids:
include_segment_ids.add(segment.id)
record = {
"segment": segment,
"score": document.metadata.get("score", 0.0), # type: ignore
}
if attachment_info:
segment_file_map[segment.id] = [attachment_info]
records.append(record)
else:
if attachment_info:
attachment_infos = segment_file_map.get(segment.id, [])
if attachment_info not in attachment_infos:
attachment_infos.append(attachment_info)
segment_file_map[segment.id] = attachment_infos
# Add child chunks information to records # Add child chunks information to records
for record in records: for record in records:
if record["segment"].id in segment_child_map: if record["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[record["segment"].id].get("child_chunks") # type: ignore
record["score"] = segment_child_map[record["segment"].id]["max_score"] # type: ignore record["score"] = segment_child_map[record["segment"].id]["max_score"] # type: ignore
if record["segment"].id in segment_file_map: if record["segment"].id in attachment_map:
record["files"] = segment_file_map[record["segment"].id] # type: ignore[assignment] record["files"] = attachment_map[record["segment"].id] # type: ignore[assignment]
result = [] result = []
for record in records: for record in records:
@ -551,6 +541,9 @@ class RetrievalService:
if not isinstance(child_chunks, list): if not isinstance(child_chunks, list):
child_chunks = None child_chunks = None
if child_chunks:
child_chunks = sorted(child_chunks, key=lambda x: x.get, reverse=True)
# Extract files, ensuring it's a list or None # Extract files, ensuring it's a list or None
files = record.get("files") files = record.get("files")
if not isinstance(files, list): if not isinstance(files, list):
@ -570,7 +563,7 @@ class RetrievalService:
) )
result.append(retrieval_segment) result.append(retrieval_segment)
return result return sorted(result, key=lambda x: x.score, reverse=True)
except Exception as e: except Exception as e:
db.session.rollback() db.session.rollback()
raise e raise e