[Bug]: ValueError: Query ID Not Found on Re-execution with ChromaDB Embeddings

Bug Description

I am encountering a ValueError when trying to re-execute code that uses previously generated ChromaDB embeddings. The code works perfectly on the first run, generating responses from the LLM based on user input. However, when I stop the execution and run the code again to reuse those embeddings, I get the following error:

ValueError: Query id d28e3de4-1c4f-420a-9184-97bf8556b11b not found in either retriever_dict or query_engine_dict.

Actual Behavior On re-execution, the code raises a ValueError stating that the Query ID is not found in either retriever_dict or query_engine_dict.

Expected Behavior The code should run successfully using previous embeddings without raising any ValueError

Version

llama-index==0.10.12

Steps to Reproduce

class DenseXRetrievalPack(BaseLlamaPack):
def init(
self,
documents: List[Document],
proposition_llm: Optional[LLM] = None,
query_llm: Optional[LLM] = None,
embed_model: Optional[BaseEmbedding] = None,
text_splitter: TextSplitter = SentenceSplitter(),
vector_store: Optional[ElasticsearchStore] = None,
similarity_top_k: int = 4,
) -> None:
"""Init params."""
self._proposition_llm = llm

    embed_model = embed_model

    nodes = text_splitter.get_nodes_from_documents(documents)
    sub_nodes = self._gen_propositions(nodes)
    all_nodes = nodes + sub_nodes
    all_nodes_dict = {n.node_id: n for n in all_nodes}

    service_context = ServiceContext.from_defaults(
        llm=query_llm ,
        embed_model=embed_model,
        num_output=self._proposition_llm.metadata.num_output,
    )

    if os.path.exists('./chroma_db'):
        chroma_client = chromadb.PersistentClient(path="./chroma_db")
        chroma_collection = chroma_client.get_or_create_collection("quickstart")
        vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
        storage_context = StorageContext.from_defaults(vector_store=vector_store)
        self.vector_index = VectorStoreIndex.from_vector_store(vector_store,service_context=service_context)
    else:
       chroma_client = chromadb.PersistentClient(path="./chroma_db")
       chroma_collection = chroma_client.get_or_create_collection("quickstart")
       vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
       storage_context = StorageContext.from_defaults(vector_store=vector_store)
       self.vector_index = VectorStoreIndex(
            all_nodes, service_context=service_context, show_progress=True,storage_context=storage_context,store_nodes_override=True
            )
    self.retriever = RecursiveRetriever(
        "vector",
        retriever_dict={
            "vector": self.vector_index.as_retriever(
                similarity_top_k=similarity_top_k
            )
        },
        node_dict=all_nodes_dict,
    )

    self.query_engine = RetrieverQueryEngine.from_args(
        self.retriever, service_context=service_context
    )

async def _aget_proposition(self, node: TextNode) -> List[TextNode]:
    """Get proposition."""
    inital_output = await self._proposition_llm.apredict(
        PROPOSITIONS_PROMPT, node_text=node.text
    )
    outputs = inital_output.split("\n")

    all_propositions = []

    for output in outputs:
        if not output.strip():
            continue
        if not output.strip().endswith("]"):
            if not output.strip().endswith('"') and not output.strip().endswith(
                ","
            ):
                output = output + '"'
            output = output + " ]"
        if not output.strip().startswith("["):
            if not output.strip().startswith('"'):
                output = '"' + output
            output = "[ " + output

        try:
            propositions = json.loads(output)
        except Exception:
            # fallback to yaml
            try:
                propositions = yaml.safe_load(output)
            except Exception:
                # fallback to next output
                continue

        if not isinstance(propositions, list):
            continue

        all_propositions.extend(propositions)

    assert isinstance(all_propositions, list)
    nodes = [TextNode(text=prop) for prop in all_propositions if prop]

    return [IndexNode.from_text_node(n, node.node_id) for n in nodes]

def _gen_propositions(self, nodes: List[TextNode]) -> List[TextNode]:
    """Get propositions."""
    sub_nodes = asyncio.run(
        run_jobs(
            [self._aget_proposition(node) for node in nodes],
            show_progress=True,
            workers=8,
        )
    )
    print(sub_nodes)

    # Flatten list
    return [node for sub_node in sub_nodes for node in sub_node]

def get_modules(self) -> Dict[str, Any]:
    """Get modules."""
    return {
        "query_engine": self.query_engine,
        "retriever": self.retriever,
    }

def run(self, query_str: str, **kwargs: Any) -> RESPONSE_TYPE:
    """Run the pipeline."""
    return self.query_engine.query(query_str)

Relevant Logs/Tracbacks

No response

To resolve the ValueError stating that the Query ID is not found in either retriever_dict or query_engine_dict when re-executing code that uses previously generated ChromaDB embeddings, ensure that the retriever and query_engine are correctly initialized and accessible. The RecursiveRetriever class initializes these components, so you need to make sure that the retriever_dict and query_engine_dict contain the necessary entries.

Here is an example of how you can initialize and use the RecursiveRetriever class to avoid the ValueError:

from llama_index.core import Document, ServiceContext, VectorStoreIndex
from llama_index.core.retrievers import RecursiveRetriever
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.llms.openai import OpenAI
from llama_index.core.node_parser.text import SentenceSplitter
from llama_index.core.schema import TextNode, IndexNode

# Initialize the documents and other components
documents = [Document(text="your_document_text")]  # Replace with your actual documents
text_splitter = SentenceSplitter()
nodes = text_splitter.get_nodes_from_documents(documents)

# Initialize the embedding model and service context
embed_model = OpenAIEmbedding(embed_batch_size=128)
service_context = ServiceContext.from_defaults(
    llm=OpenAI(),
    embed_model=embed_model,
    num_output=750,
)

# Create the VectorStoreIndex
vector_index = VectorStoreIndex(
    nodes, service_context=service_context, show_progress=True
)

# Initialize the retriever and query engine dictionaries
retriever_dict = {
    "vector": vector_index.as_retriever(similarity_top_k=4)
}
all_nodes_dict = {n.node_id: n for n in nodes}

# Initialize the RecursiveRetriever
recursive_retriever = RecursiveRetriever(
    "vector",
    retriever_dict=retriever_dict,
    node_dict=all_nodes_dict,
)

# Initialize the RetrieverQueryEngine
query_engine = RetrieverQueryEngine.from_args(
    recursive_retriever,
    service_context=service_context,
    streaming=False,
)

# Create a query bundle
query_str = "your_query_string"

# Use the query engine to run a query
result = query_engine.query(query_str)

print(result)

This code ensures that the retriever and query_engine are correctly initialized and used, which should prevent the ValueError related to missing Query IDs. Make sure that the root_id is present in the retriever_dict and that there are no overlapping keys between retriever_dict and query_engine_dict ^[1]^[2]^[3].

To continue talking to Dosu, mention @dosu.

run-llama / llama_index