AsyncDocumentIndexTestSuite Class — langchain Architecture

Source Code

libs/standard-tests/langchain_tests/integration_tests/indexer.py lines 209–398

class AsyncDocumentIndexTestSuite(ABC):
    """Test suite for checking the read-write of a document index.

    Implementers should subclass this test suite and provide a fixture
    that returns an empty index for each test.
    """

    @abstractmethod
    @pytest.fixture
    async def index(self) -> AsyncGenerator[DocumentIndex, None]:
        """Get the index."""

    async def test_upsert_documents_has_no_ids(self, index: DocumentIndex) -> None:
        """Verify that there is not parameter called IDs in upsert."""
        signature = inspect.signature(index.upsert)
        assert "ids" not in signature.parameters

    async def test_upsert_no_ids(self, index: DocumentIndex) -> None:
        """Upsert works with documents that do not have IDs.

        At the moment, the ID field in documents is optional.
        """
        documents = [
            Document(page_content="foo", metadata={"id": 1}),
            Document(page_content="bar", metadata={"id": 2}),
        ]
        response = await index.aupsert(documents)
        ids = sorted(response["succeeded"])

        # Ordering is not guaranteed, need to test carefully
        documents = await index.aget(ids)
        sorted_documents = sorted(documents, key=lambda x: x.id or "")

        if sorted_documents[0].page_content == "bar":
            assert sorted_documents[0] == Document(
                page_content="bar", metadata={"id": 2}, id=ids[0]
            )
            assert sorted_documents[1] == Document(
                page_content="foo", metadata={"id": 1}, id=ids[1]
            )
        else:
            assert sorted_documents[0] == Document(
                page_content="foo", metadata={"id": 1}, id=ids[0]
            )
            assert sorted_documents[1] == Document(
                page_content="bar", metadata={"id": 2}, id=ids[1]
            )

    async def test_upsert_some_ids(self, index: DocumentIndex) -> None:
        """Test an upsert where some docs have IDs and some don't."""
        foo_uuid = str(uuid.UUID(int=7))
        documents = [
            Document(id=foo_uuid, page_content="foo", metadata={"id": 1}),
            Document(page_content="bar", metadata={"id": 2}),
        ]
        response = await index.aupsert(documents)
        ids = response["succeeded"]
        other_id = next(iter(set(ids) - {foo_uuid}))
        assert response["failed"] == []
        assert foo_uuid in ids
        # Ordering is not guaranteed, so we use a set.
        documents = await index.aget(ids)
        first_doc = documents[0]
        if first_doc.id == foo_uuid:
            assert documents == [
                Document(page_content="foo", metadata={"id": 1}, id=foo_uuid),
                Document(page_content="bar", metadata={"id": 2}, id=other_id),
            ]
        else:
            assert documents == [
                Document(page_content="bar", metadata={"id": 2}, id=other_id),
                Document(page_content="foo", metadata={"id": 1}, id=foo_uuid),
            ]

    async def test_upsert_overwrites(self, index: DocumentIndex) -> None:
        """Test that upsert overwrites existing content."""
        foo_uuid = str(uuid.UUID(int=7))
        documents = [
            Document(id=foo_uuid, page_content="foo", metadata={"bar": 1}),
        ]
        response = await index.aupsert(documents)