fix: segment resort in dataset retrieve by index_node_id_to_position (#721)

3 years ago · 654985177f
parent 0d791839e6
commit 654985177f
1 changed files with 5 additions and 1 deletions
--- 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: