Question about Tree Traversal in RAPTOR Retrieval

[ZN1010]

In RAPTOR's tree traversal method, we are supposed to retrieve and concatenate top k nodes of each level, so the nodes of higher levels will have summaries as their texts. However, when I ran the llama_index implementation and tried to print the final retrieved nodes, in the tree_traversal_retrieval() function of base.py in RaptorPack, I always get the leaf nodes as the returned output. It seems that all the nodes from higher levels are not incorporated into assembling the relevant context to users' query.

Maybe there is a misalignment between the paper idea of tree traversal and the actual implementation? Or I am missing anything important? Thanks a lot!

[dosubot[bot]]

Hey @ZN1010! I'm here to help you with any issues or questions you have. I can assist with bug fixes, answer technical queries, and guide you on becoming a contributor.

There doesn't appear to be a misalignment between the paper's idea of tree traversal and the actual implementation in llama_index. The tree_traversal_retrieval() function in the RaptorPack implementation is designed to retrieve and concatenate the top k nodes at each level, including higher-level nodes with summaries.

However, the behavior you're experiencing—where only leaf nodes are returned—might be influenced by specific configurations or parameters. Key parameters that affect this behavior include similarity_top_k and tree_depth.

1. similarity_top_k: Determines the number of top nodes to retrieve based on similarity at each level.
2. tree_depth: Defines the depth of the tree to traverse during the retrieval process.

You can adjust these parameters when initializing the RaptorRetriever class or by setting them directly on the instance before calling the tree_traversal_retrieval() function. Here’s an example:

from llama_index.packs.raptor import RaptorRetriever
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.llms.openai import OpenAI
from llama_index.vector_stores.chroma import ChromaVectorStore
import chromadb

# Initialize the vector store
client = chromadb.PersistentClient(path="./raptor_paper_db")
collection = client.get_or_create_collection("raptor")
vector_store = ChromaVectorStore(chroma_collection=collection)

# Initialize the RaptorRetriever
raptor_retriever = RaptorRetriever(
    [],
    embed_model=OpenAIEmbedding(model="text-embedding-3-small"),
    llm=OpenAI(model="gpt-3.5-turbo", temperature=0.1),
    vector_store=vector_store,
    similarity_top_k=10,  # Set the number of top nodes to retrieve
    mode="tree_traversal"  # Set the default mode
)

# Adjust the tree depth if needed
raptor_retriever.tree_depth = 3  # Set the depth of the tree

# Perform tree traversal retrieval
response = await raptor_retriever.tree_traversal_retrieval(query_str="your query here")

Adjusting these parameters should help ensure that higher-level nodes with summaries are incorporated into the retrieval process ^[1][2].

To continue talking to Dosu, mention @dosu.

[ZN1010]

Thanks! I used the same similarity_top_k and tree_depth. But when I print out all the returned nodes of tree_traversal_retrieval (), I still got a bunch of leaf nodes. Check line 271 of base.py in RaptorPack:

nodes = [node for nested in nested_nodes for node in nested]

All the retrieved higher-level nodes are eventually replaced by leaf nodes and it seems that the pipeline is only concatenating the returned leaf nodes of the function. In case I miss anything important, please let me know. @dosu.

[dosubot[bot]]

The issue you're encountering is due to the logic in the tree_traversal_retrieval function, specifically in the while loop where the nodes list is updated. The nodes list is being overwritten in each iteration of the loop, which results in only the leaf nodes being retained by the end of the traversal.

To retain all nodes from each level, you need to accumulate the nodes from each level rather than replacing the nodes list. Here is a modified version of the tree_traversal_retrieval function that accumulates nodes from all levels:

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
    all_nodes = []
    level = self.tree_depth - 1
    while level >= 0:
        # retrieve nodes at the current level
        if parent_ids is None:
            nodes = await self.index.as_retriever(
                similarity_top_k=self.similarity_top_k,
                filters=MetadataFilters(
                    filters=[MetadataFilter(key="level", value=level)]
                ),
            ).aretrieve(query_str)

            parent_ids = [node.id_ for node in nodes]
            if self._verbose:
                print(f"Retrieved parent IDs from level {level}: {parent_ids!s}")
        # retrieve nodes that are children of the nodes at the previous level
        elif parent_ids is not None and len(parent_ids) > 0:
            nested_nodes = await asyncio.gather(
                *[
                    self.index.as_retriever(
                        similarity_top_k=self.similarity_top_k,
                        filters=MetadataFilters(
                            filters=[MetadataFilter(key="parent_id", value=id_)]
                        ),
                    ).aretrieve(query_str)
                    for id_ in parent_ids
                ]
            )

            nodes = [node for nested in nested_nodes for node in nested]

            if self._verbose:
                print(f"Retrieved {len(nodes)} from parents at level {level}.")

            parent_ids = [node.id_ for node in nodes]

        all_nodes.extend(nodes)
        level -= 1

    return all_nodes

In this modified version, the all_nodes list accumulates nodes from each level, ensuring that nodes from all levels are retained and returned ^[1].

To continue talking to Dosu, mention @dosu.

[ZN1010]

The issue you're encountering is due to the logic in the tree_traversal_retrieval function, specifically in the while loop where the nodes list is updated. The nodes list is being overwritten in each iteration of the loop, which results in only the leaf nodes being retained by the end of the traversal.

To retain all nodes from each level, you need to accumulate the nodes from each level rather than replacing the nodes list. Here is a modified version of the tree_traversal_retrieval function that accumulates nodes from all levels:

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
    all_nodes = []
    level = self.tree_depth - 1
    while level >= 0:
        # retrieve nodes at the current level
        if parent_ids is None:
            nodes = await self.index.as_retriever(
                similarity_top_k=self.similarity_top_k,
                filters=MetadataFilters(
                    filters=[MetadataFilter(key="level", value=level)]
                ),
            ).aretrieve(query_str)

            parent_ids = [node.id_ for node in nodes]
            if self._verbose:
                print(f"Retrieved parent IDs from level {level}: {parent_ids!s}")
        # retrieve nodes that are children of the nodes at the previous level
        elif parent_ids is not None and len(parent_ids) > 0:
            nested_nodes = await asyncio.gather(
                *[
                    self.index.as_retriever(
                        similarity_top_k=self.similarity_top_k,
                        filters=MetadataFilters(
                            filters=[MetadataFilter(key="parent_id", value=id_)]
                        ),
                    ).aretrieve(query_str)
                    for id_ in parent_ids
                ]
            )

            nodes = [node for nested in nested_nodes for node in nested]

            if self._verbose:
                print(f"Retrieved {len(nodes)} from parents at level {level}.")

            parent_ids = [node.id_ for node in nodes]

        all_nodes.extend(nodes)
        level -= 1

    return all_nodes

In this modified version, the all_nodes list accumulates nodes from each level, ensuring that nodes from all levels are retained and returned [1].

To continue talking to Dosu, mention @dosu.

Thanks! So does this mean that there is a mistake/misalignment in the llama_index implementation and this function needs to be updated? Based on the paper idea, this function is supposed to concatenate retrieved nodes from all levels instead of just leaf nodes. @dosu.

[dosubot[bot]]

Yes, there is a misalignment in the current llama_index implementation. The function should be updated to align with the paper's idea of concatenating nodes from all levels. The current implementation consolidates chunks recursively in a bottom-up fashion, but it does not concatenate nodes from all levels as described in the paper. Specifically, the build_index_from_nodes method in GPTTreeIndexBuilder consolidates nodes by creating intermediate layers rather than concatenating nodes from all levels. Additionally, the _prepare_node_and_text_chunks and _construct_parent_nodes methods focus on creating and summarizing chunks at each level, but they do not concatenate nodes from all levels either ^[1][2][3][4].

To continue talking to Dosu, mention @dosu.