Fix for llama_index.packs.raptor tree_traversal retrieval (#17406)

run-llama · Jan 1, 2025 · 3ce83b1 · 3ce83b1
1 parent ba743fc
commit 3ce83b1
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 4 deletions.
diff --git a/llama-index-packs/llama-index-packs-raptor/llama_index/packs/raptor/base.py b/llama-index-packs/llama-index-packs-raptor/llama_index/packs/raptor/base.py
@@ -239,7 +239,8 @@ async def tree_traversal_retrieval(self, query_str: str) -> Response:
         """Query the index as a tree, traversing the tree from the top down."""
         # get top k nodes for each level, starting with the top
         parent_ids = None
-        nodes = []
+        selected_node_ids = set()
+        selected_nodes = []
         level = self.tree_depth - 1
         while level >= 0:
             # retrieve nodes at the current level
@@ -251,6 +252,11 @@ async def tree_traversal_retrieval(self, query_str: str) -> Response:
                     ),
                 ).aretrieve(query_str)
 
+                for node in nodes:
+                    if node.id_ not in selected_node_ids:
+                        selected_nodes.append(node)
+                        selected_node_ids.add(node.id_)
+
                 parent_ids = [node.id_ for node in nodes]
                 if self._verbose:
                     print(f"Retrieved parent IDs from level {level}: {parent_ids!s}")
@@ -269,14 +275,18 @@ async def tree_traversal_retrieval(self, query_str: str) -> Response:
                 )
 
                 nodes = [node for nested in nested_nodes for node in nested]
+                for node in nodes:
+                    if node.id_ not in selected_node_ids:
+                        selected_nodes.append(node)
+                        selected_node_ids.add(node.id_)
 
                 if self._verbose:
                     print(f"Retrieved {len(nodes)} from parents at level {level}.")
 
                 level -= 1
                 parent_ids = None
 
-        return nodes
+        return selected_nodes
 
     def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
         """Retrieve nodes given query and mode."""

diff --git a/llama-index-packs/llama-index-packs-raptor/pyproject.toml b/llama-index-packs/llama-index-packs-raptor/pyproject.toml
@@ -31,7 +31,7 @@ license = "MIT"
 name = "llama-index-packs-raptor"
 packages = [{include = "llama_index/"}]
 readme = "README.md"
-version = "0.3.0"
+version = "0.3.1"
 
 [tool.poetry.dependencies]
 python = ">=3.9,<4.0"

diff --git a/llama-index-packs/llama-index-packs-raptor/tests/test_packs_raptor.py b/llama-index-packs/llama-index-packs-raptor/tests/test_packs_raptor.py
@@ -27,4 +27,4 @@ def test_raptor() -> None:
     assert len(nodes) == 2
 
     nodes = retriever.retrieve("text", mode="tree_traversal")
-    assert len(nodes) == 2
+    assert len(nodes) == 5