test_combine_documents.py — langchain Source File

Source Code

"""Test functionality related to combining documents."""

import re
from typing import Any

import pytest
from langchain_core.documents import Document
from langchain_core.prompts import PromptTemplate, aformat_document, format_document

from langchain_classic.chains.combine_documents.reduce import (
    collapse_docs,
    split_list_of_docs,
)
from langchain_classic.chains.qa_with_sources import load_qa_with_sources_chain
from tests.unit_tests.llms.fake_llm import FakeLLM


def _fake_docs_len_func(docs: list[Document]) -> int:
    return len(_fake_combine_docs_func(docs))


def _fake_combine_docs_func(docs: list[Document], **_: Any) -> str:
    return "".join([d.page_content for d in docs])


def test_multiple_input_keys() -> None:
    chain = load_qa_with_sources_chain(FakeLLM(), chain_type="stuff")
    assert chain.input_keys == ["input_documents", "question"]


def test__split_list_long_single_doc() -> None:
    """Test splitting of a long single doc."""
    docs = [Document(page_content="foo" * 100)]
    with pytest.raises(
        ValueError, match="A single document was longer than the context length"
    ):
        split_list_of_docs(docs, _fake_docs_len_func, 100)


def test__split_list_single_doc() -> None:
    """Test splitting works with just a single doc."""
    docs = [Document(page_content="foo")]
    doc_list = split_list_of_docs(docs, _fake_docs_len_func, 100)
    assert doc_list == [docs]


def test__split_list_double_doc() -> None:
    """Test splitting works with just two docs."""
    docs = [Document(page_content="foo"), Document(page_content="bar")]
    doc_list = split_list_of_docs(docs, _fake_docs_len_func, 100)
    assert doc_list == [docs]


def test__split_list_works_correctly() -> None:
    """Test splitting works correctly."""
    docs = [
        Document(page_content="foo"),
        Document(page_content="bar"),
        Document(page_content="baz"),
        Document(page_content="foo" * 2),
// ... (100 more lines)