Back to Repositories

Testing CondensePlusContext Chat Engine Functionality in LlamaIndex

This test suite validates the CondensePlusContextChatEngine functionality in LlamaIndex, focusing on both synchronous and asynchronous chat interactions. The tests verify chat history management, response streaming, and system prompt integration using mock components.

Test Coverage Overview

The test suite provides comprehensive coverage of the CondensePlusContextChatEngine’s core capabilities.

Basic chat functionality with history tracking
Streaming response handling
Asynchronous chat operations
System prompt persistence across interactions
Chat history maintenance and verification

Implementation Analysis

The testing approach utilizes pytest fixtures and mock components to isolate the chat engine functionality.

The implementation leverages MockLLM and MockEmbedding classes for controlled testing environments, with specific patterns for both synchronous and asynchronous operations using pytest.mark.asyncio decorators.

Technical Details

pytest framework with asyncio support
Mock components: MockLLM, MockEmbedding
VectorStoreIndex for document handling
Fixture-based test setup
Async/await pattern implementation
Stream response testing with generators

Best Practices Demonstrated

The test suite exemplifies strong testing practices through systematic verification of both synchronous and asynchronous workflows.

Consistent assertion patterns
Proper fixture usage
Comprehensive edge case coverage
Clear test isolation
Structured async testing patterns

run-llama/llama_index

llama-index-core/tests/chat_engine/test_condense_plus_context.py

            
import pytest

from llama_index.core import MockEmbedding
from llama_index.core.chat_engine.condense_plus_context import (
    CondensePlusContextChatEngine,
)
from llama_index.core.indices import VectorStoreIndex
from llama_index.core.llms.mock import MockLLM
from llama_index.core.schema import Document

SYSTEM_PROMPT = "Talk like a pirate."


@pytest.fixture()
def chat_engine() -> CondensePlusContextChatEngine:
    index = VectorStoreIndex.from_documents(
        [Document.example()], embed_model=MockEmbedding(embed_dim=3)
    )
    retriever = index.as_retriever()
    return CondensePlusContextChatEngine.from_defaults(
        retriever, llm=MockLLM(), system_prompt=SYSTEM_PROMPT
    )


def test_chat(chat_engine: CondensePlusContextChatEngine):
    response = chat_engine.chat("Hello World!")
    assert SYSTEM_PROMPT in str(response)
    assert "Hello World!" in str(response)
    assert len(chat_engine.chat_history) == 2

    response = chat_engine.chat("What is the capital of the moon?")
    assert SYSTEM_PROMPT in str(response)
    assert "Hello World!" in str(response)
    assert "What is the capital of the moon?" in str(response)
    assert len(chat_engine.chat_history) == 4


def test_chat_stream(chat_engine: CondensePlusContextChatEngine):
    response = chat_engine.stream_chat("Hello World!")

    num_iters = 0
    for _ in response.response_gen:
        num_iters += 1

    assert num_iters > 10
    assert SYSTEM_PROMPT in str(response)
    assert "Hello World!" in str(response)
    assert len(chat_engine.chat_history) == 2

    response = chat_engine.stream_chat("What is the capital of the moon?")

    num_iters = 0
    for _ in response.response_gen:
        num_iters += 1

    assert num_iters > 10
    assert SYSTEM_PROMPT in str(response)
    assert "Hello World!" in str(response)
    assert "What is the capital of the moon?" in str(response)
    assert len(chat_engine.chat_history) == 4


@pytest.mark.asyncio()
async def test_achat(chat_engine: CondensePlusContextChatEngine):
    response = await chat_engine.achat("Hello World!")
    assert SYSTEM_PROMPT in str(response)
    assert "Hello World!" in str(response)
    assert len(chat_engine.chat_history) == 2

    response = await chat_engine.achat("What is the capital of the moon?")
    assert SYSTEM_PROMPT in str(response)
    assert "Hello World!" in str(response)
    assert "What is the capital of the moon?" in str(response)
    assert len(chat_engine.chat_history) == 4


@pytest.mark.asyncio()
async def test_chat_astream(chat_engine: CondensePlusContextChatEngine):
    response = await chat_engine.astream_chat("Hello World!")

    num_iters = 0
    async for _ in response.async_response_gen():
        num_iters += 1

    assert num_iters > 10
    assert SYSTEM_PROMPT in str(response)
    assert "Hello World!" in str(response)
    assert len(chat_engine.chat_history) == 2

    response = await chat_engine.astream_chat("What is the capital of the moon?")

    num_iters = 0
    async for _ in response.async_response_gen():
        num_iters += 1

    assert num_iters > 10
    assert SYSTEM_PROMPT in str(response)
    assert "Hello World!" in str(response)
    assert "What is the capital of the moon?" in str(response)
    assert len(chat_engine.chat_history) == 4