Back to Repositories

Testing Dataset Generation Implementation in LlamaIndex

This test suite focuses on validating the dataset generation functionality in LlamaIndex, specifically testing the DatasetGenerator class’s ability to create question-answer pairs from text nodes.

Test Coverage Overview

The test suite provides comprehensive coverage of the dataset generation process in LlamaIndex.

Tests creation of question-answer pairs from text nodes
Validates prompt template processing
Verifies correct generation of questions and answers
Ensures proper handling of multiple text nodes

Implementation Analysis

The testing approach uses a mock LLM predictor to simulate question and answer generation.

Key implementation patterns include:

Custom prompt template configuration
Text node processing verification
Question-answer pair validation
Mock LLM response handling

Technical Details

Testing infrastructure includes:

patch_llm_predictor fixture for LLM simulation
DatasetGenerator class implementation
PromptTemplate configuration
TextNode data structures
Question-answer pair validation logic

Best Practices Demonstrated

The test exhibits strong testing practices through isolation of components and thorough validation.

Clear test case organization
Comprehensive assertion checks
Proper mock object usage
Isolated component testing
Explicit expected outcome verification

run-llama/llama_index

llama-index-core/tests/evaluation/test_dataset_generation.py

            
"""Test dataset generation."""

from llama_index.core.evaluation.dataset_generation import DatasetGenerator
from llama_index.core.prompts.base import PromptTemplate
from llama_index.core.prompts.prompt_type import PromptType
from llama_index.core.schema import TextNode


def test_dataset_generation(patch_llm_predictor) -> None:
    """Test dataset generation."""
    test_nodes = [TextNode(text="hello_world"), TextNode(text="foo_bar")]

    question_gen_prompt = PromptTemplate(
        """\
Context information is below.
---------------------
{context_str}
---------------------
Given the context information and not prior knowledge.
generate only questions based on the below query.
{query_str}
""",
        prompt_type=PromptType.QUESTION_ANSWER,
    )

    dataset_generator = DatasetGenerator(
        test_nodes,
        text_question_template=question_gen_prompt,
        question_gen_query="gen_question",
    )
    eval_dataset = dataset_generator.generate_dataset_from_nodes()
    qr_pairs = eval_dataset.qr_pairs
    assert len(qr_pairs) == 2
    # the mock LLM concatenates query with context with ":"
    # the first call is to generate the question
    assert qr_pairs[0][0] == "gen_question:hello_world"
    # the second call is to generate the answer
    assert qr_pairs[0][1] == "gen_question:hello_world:hello_world"
    assert qr_pairs[1][0] == "gen_question:foo_bar"
    assert qr_pairs[1][1] == "gen_question:foo_bar:foo_bar"