Back to Repositories

Testing Token Counter Implementation in llama_index

This test suite validates the TokenCountingHandler functionality in llama_index, focusing on accurate token counting for both LLM and embedding operations. It ensures proper event handling and token count tracking across different callback events.

Test Coverage Overview

The test suite provides comprehensive coverage of the TokenCountingHandler class, focusing on event handling and token counting accuracy.

Tests event start functionality for both LLM and embedding operations
Validates token counting for different event types
Verifies proper accumulation of token counts
Covers edge cases in event handling

Implementation Analysis

The testing approach employs pytest fixtures and isolated test cases to validate token counting behavior. The implementation uses mock payloads and event IDs to simulate real-world usage scenarios.

Separate test functions for event start and end operations
Verification of token count accumulation
Event type differentiation testing

Technical Details

Testing Framework: pytest
Key Components: TokenCountingHandler, CBEventType
Test Data: Predefined payload and event ID constants
Assertion Patterns: Length checks and count verification

Best Practices Demonstrated

The test suite exemplifies several testing best practices in Python unit testing.

Clear test function naming and documentation
Isolated test cases with specific assertions
Consistent test data usage
Comprehensive event type coverage

run-llama/llama_index

llama-index-core/tests/callbacks/test_token_counter.py

            
"""Embeddings."""

from llama_index.core.callbacks.schema import CBEventType
from llama_index.core.callbacks.token_counting import TokenCountingHandler

TEST_PAYLOAD = {"chunks": ["one"], "formatted_prompt": "two", "response": "three"}
TEST_ID = "my id"


def test_on_event_start() -> None:
    """Test event start."""
    handler = TokenCountingHandler()

    event_id = handler.on_event_start(
        CBEventType.LLM, payload=TEST_PAYLOAD, event_id=TEST_ID
    )

    assert event_id == TEST_ID

    event_id = handler.on_event_start(
        CBEventType.EMBEDDING, payload=TEST_PAYLOAD, event_id=TEST_ID
    )

    assert event_id == TEST_ID
    assert len(handler.llm_token_counts) == 0
    assert len(handler.embedding_token_counts) == 0


def test_on_event_end() -> None:
    """Test event end."""
    handler = TokenCountingHandler()

    handler.on_event_end(CBEventType.LLM, payload=TEST_PAYLOAD, event_id=TEST_ID)

    assert len(handler.llm_token_counts) == 1
    assert len(handler.embedding_token_counts) == 0

    handler.on_event_end(CBEventType.EMBEDDING, payload=TEST_PAYLOAD, event_id=TEST_ID)

    assert len(handler.llm_token_counts) == 1
    assert len(handler.embedding_token_counts) == 1

    assert handler.embedding_token_counts[0].total_token_count == 1
    assert handler.llm_token_counts[0].total_token_count == 2

    # test actual counts
    # LLM should be two (prompt plus response)
    # Embedding should be one (single token chunk)
    assert handler.total_llm_token_count == 2
    assert handler.total_embedding_token_count == 1