| @@ -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) | |||