test_llm_chain_extractor() — langchain Function Reference

Source Code

libs/langchain/tests/unit_tests/retrievers/document_compressors/test_chain_extract.py lines 7–46

def test_llm_chain_extractor() -> None:
    documents = [
        Document(
            page_content=(
                "The sky is blue. Candlepin bowling is popular in New England."
            ),
            metadata={"a": 1},
        ),
        Document(
            page_content=(
                "Mercury is the closest planet to the Sun. "
                "Candlepin bowling balls are smaller."
            ),
            metadata={"b": 2},
        ),
        Document(page_content="The moon is round.", metadata={"c": 3}),
    ]
    llm = FakeListChatModel(
        responses=[
            "Candlepin bowling is popular in New England.",
            "Candlepin bowling balls are smaller.",
            "NO_OUTPUT",
        ],
    )
    doc_compressor = LLMChainExtractor.from_llm(llm)
    output = doc_compressor.compress_documents(
        documents,
        "Tell me about Candlepin bowling.",
    )
    expected = documents = [
        Document(
            page_content="Candlepin bowling is popular in New England.",
            metadata={"a": 1},
        ),
        Document(
            page_content="Candlepin bowling balls are smaller.",
            metadata={"b": 2},
        ),
    ]
    assert output == expected