Issue with Tree Traversal Mode in RAPTOR using Azure AI Search

[ZeinabSh12]

I'm using RAPTOR with llama-index and Azure AI Search as the vector database. The tree_traversal retrieval mode fails to return nodes at level 0, resulting in no final nodes being retrieved. This happens across varied test sets.

Setup:

Model: gpt-4
Embedding Model: text-embedding-ada-002
Vector Database: Azure AI Search
Code Snippet: `raptor_pack = RaptorPack(
documents,
embed_model=OpenAIEmbedding(model="text-embedding-ada-002"),
llm=OpenAI(model="gpt-4", temperature=0.1),
vector_store=vector_store, # Azure AI Search
similarity_top_k=2,
mode="tree_traversal",
transformations=[SentenceSplitter(chunk_size=800, chunk_overlap=200)],
verbose=True
)

nodes = raptor_pack.run("What baselines is raptor compared against?", mode="tree_traversal")
print(len(nodes))
print(nodes[0].text if nodes else "No nodes retrieved")
` Issue:

Nodes are inserted correctly at each level.
Retrieval fails at level 0 with no parent IDs retrieved. Logs: Starting retrieval at level 2 Retrieved parent IDs from level 2: ['id1', 'id2'] Starting retrieval at level 1 Retrieved parent IDs from level 1: ['id3', 'id4'] Starting retrieval at level 0 Retrieved parent IDs from level 0: [] No parent IDs retrieved at level 0 Final nodes retrieved: [] Could you please provide any insights or suggestions for resolving this issue?

[dosubot[bot]]

The issue seems to be related to the retrieval process at level 0 in the tree_traversal_retrieval method. Specifically, the retrieval of parent IDs at level 0 is failing, which results in no final nodes being retrieved.

Here are some steps to help resolve this issue:

1. Verify Node Metadata: Ensure that nodes at level 0 have the correct parent_id set during insertion and that the level metadata is correctly assigned.

2. Check Insertion Code: Make sure the parent_id and level metadata are correctly set when inserting nodes. Here is an example:

   # insert the nodes with their embeddings and parent_id
   nodes_with_embeddings = []
   for cluster, summary_doc in zip(nodes_per_cluster, new_nodes):
       for node in cluster:
           node.metadata["parent_id"] = summary_doc.id_
           node.excluded_embed_metadata_keys.append("parent_id")
           node.excluded_llm_metadata_keys.append("parent_id")
           node.embedding = id_to_embedding[node.id_]
           nodes_with_embeddings.append(node)
   
   self.index.insert_nodes(nodes_with_embeddings)
   
   # set the current nodes to the new nodes
   cur_nodes = new_nodes

3. Ensure Correct Level Assignment: Make sure the level metadata is correctly assigned to nodes at each level:

   new_nodes = [
       TextNode(
           text=summary,
           metadata={"level": level},
           excluded_embed_metadata_keys=["level"],
           excluded_llm_metadata_keys=["level"],
       )
       for summary in summaries_per_cluster
   ]

4. Add Logging: Add additional logging to verify that nodes at level 0 are correctly inserted and that their parent_id and level metadata are correctly set.

For more details, you can refer to the example provided in the LlamaIndex documentation for setting up and using RAPTOR with the tree_traversal mode ^[1][2].

To continue talking to Dosu, mention @dosu.