fix: segments query missing idx hit (#1072)

2191312b · takatost · GitHub · fcc6b41a · 2191312b
Unverified Commit 2191312b authored Aug 31, 2023 by takatost Committed by GitHub Aug 31, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

dataset_retriever_tool.py api/core/tool/dataset_retriever_tool.py +2 -1

No files found.
--- a/api/core/tool/dataset_retriever_tool.py
+++ b/api/core/tool/dataset_retriever_tool.py
@@ -98,7 +98,8 @@ class DatasetRetrieverTool(BaseTool):
            hit_callback.on_tool_end(documents)
            document_context_list = []
            index_node_ids = [document.metadata['doc_id'] for document in documents]
-            segments = DocumentSegment.query.filter(DocumentSegment.completed_at.isnot(None),
+            segments = DocumentSegment.query.filter(DocumentSegment.dataset_id == self.dataset_id,
+                                                    DocumentSegment.completed_at.isnot(None),
                                                    DocumentSegment.status == 'completed',
                                                    DocumentSegment.enabled == True,
                                                    DocumentSegment.index_node_id.in_(index_node_ids)