From 654985177f95ce3792363a1c11a475521521b54c Mon Sep 17 00:00:00 2001 From: takatost Date: Wed, 2 Aug 2023 21:31:54 +0800 Subject: [PATCH] fix: segment resort in dataset retrieve by index_node_id_to_position (#721) --- api/core/tool/dataset_retriever_tool.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/api/core/tool/dataset_retriever_tool.py b/api/core/tool/dataset_retriever_tool.py index 45c5a9f226..35f15bbceb 100644 --- a/api/core/tool/dataset_retriever_tool.py +++ b/api/core/tool/dataset_retriever_tool.py @@ -109,7 +109,11 @@ class DatasetRetrieverTool(BaseTool): ).all() if segments: - for segment in segments: + index_node_id_to_position = {id: position for position, id in enumerate(index_node_ids)} + sorted_segments = sorted(segments, + key=lambda segment: index_node_id_to_position.get(segment.index_node_id, + float('inf'))) + for segment in sorted_segments: if segment.answer: document_context_list.append(f'question:{segment.content} \nanswer:{segment.answer}') else: