Back to Repositories

Testing TreeIndex Retriever Implementation in llama_index

A comprehensive test suite for validating TreeIndex retriever functionality in LlamaIndex, focusing on query operations and document retrieval mechanisms.

The tests verify core retrieval capabilities and summarization features of the TreeIndex implementation.

Test Coverage Overview

The test suite provides coverage for TreeIndex retriever operations with specific focus on:

Basic query functionality validation
Summarization query testing
Document retrieval verification
Tree building configuration options

Implementation Analysis

The testing approach employs fixture-based setup with mocked LLM predictors and token splitters.

Key patterns include:

Structured kwargs handling for index and query configurations
Multiple retriever modes testing
Document list manipulation
Node count verification

Technical Details

Testing infrastructure includes:

Python typing for strict type checking
Pytest fixtures for dependency injection
Mocked components for isolated testing
TreeIndex core functionality testing
Document schema implementation

Best Practices Demonstrated

The test suite exemplifies robust testing practices through:

Clear test function separation and naming
Comprehensive fixture usage
Explicit assertion checking
Configuration isolation between tests
Modular test structure

run-llama/llama_index

llama-index-core/tests/indices/tree/test_retrievers.py

            
from typing import Dict, List

from llama_index.core.indices.tree.base import TreeIndex
from llama_index.core.schema import Document


def test_query(
    documents: List[Document],
    patch_llm_predictor,
    patch_token_text_splitter,
    struct_kwargs: Dict,
) -> None:
    """Test query."""
    index_kwargs, query_kwargs = struct_kwargs
    tree = TreeIndex.from_documents(documents, **index_kwargs)

    # test default query
    query_str = "What is?"
    retriever = tree.as_retriever()
    nodes = retriever.retrieve(query_str)
    assert len(nodes) == 1


def test_summarize_query(
    documents: List[Document],
    patch_llm_predictor,
    patch_token_text_splitter,
    struct_kwargs: Dict,
) -> None:
    """Test summarize query."""
    # create tree index without building tree
    index_kwargs, orig_query_kwargs = struct_kwargs
    index_kwargs = index_kwargs.copy()
    index_kwargs.update({"build_tree": False})
    tree = TreeIndex.from_documents(documents, **index_kwargs)

    # test retrieve all leaf
    query_str = "What is?"
    retriever = tree.as_retriever(retriever_mode="all_leaf")
    nodes = retriever.retrieve(query_str)
    assert len(nodes) == 4