tests/unit_tests/chains/test_llm_checker.py

# flake8: noqa E501

"""Test LLMCheckerChain functionality."""

import pytest

from langchain.chains.llm_checker.base import LLMCheckerChain
from langchain.chains.llm_checker.prompt import (
    _CHECK_ASSERTIONS_TEMPLATE,
    _CREATE_DRAFT_ANSWER_TEMPLATE,
    _LIST_ASSERTIONS_TEMPLATE,
    _REVISED_ANSWER_TEMPLATE,
)
from tests.unit_tests.llms.fake_llm import FakeLLM


@pytest.fixture
def fake_llm_checker_chain() -> LLMCheckerChain:
    """Fake LLMCheckerChain for testing."""
    queries = {
        _CREATE_DRAFT_ANSWER_TEMPLATE.format(
            question="Which mammal lays the biggest eggs?"
        ): "I don't know which mammal layers the biggest eggs.",
        _LIST_ASSERTIONS_TEMPLATE.format(
            statement="I don't know which mammal layers the biggest eggs.",
        ): "1) I know that mammals lay eggs.\n2) I know that birds lay eggs.\n3) I know that birds are mammals.",
        _CHECK_ASSERTIONS_TEMPLATE.format(
            assertions="1) I know that mammals lay eggs.\n2) I know that birds lay eggs.\n3) I know that birds are mammals.",
        ): "1) I know that mammals lay eggs. TRUE\n2) I know that birds lay eggs. TRUE\n3) I know that birds are mammals. TRUE",
        _REVISED_ANSWER_TEMPLATE.format(
            checked_assertions="1) I know that mammals lay eggs. TRUE\n2) I know that birds lay eggs. TRUE\n3) I know that birds are mammals. TRUE",
            question="Which mammal lays the biggest eggs?",
        ): "I still don't know.",
    }
    fake_llm = FakeLLM(queries=queries)
    return LLMCheckerChain(llm=fake_llm, input_key="q", output_key="a")


def test_simple_question(fake_llm_checker_chain: LLMCheckerChain) -> None:
    """Test simple question that should not need python."""
    question = "Which mammal lays the biggest eggs?"
    output = fake_llm_checker_chain.run(question)
    assert output == "I still don't know."