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