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

This commit is contained in:
takatost 2023-08-02 21:31:54 +08:00 committed by GitHub
parent 0d791839e6
commit 654985177f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -109,7 +109,11 @@ class DatasetRetrieverTool(BaseTool):
).all() ).all()
if segments: 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: if segment.answer:
document_context_list.append(f'question:{segment.content} \nanswer:{segment.answer}') document_context_list.append(f'question:{segment.content} \nanswer:{segment.answer}')
else: else: