vectorstores.py — langchain Source File

Source Code

"""Test suite to test `VectorStore` integrations."""

from abc import abstractmethod

import pytest
from langchain_core.documents import Document
from langchain_core.embeddings import DeterministicFakeEmbedding, Embeddings
from langchain_core.vectorstores import VectorStore

from langchain_tests.base import BaseStandardTests

# Arbitrarily chosen. Using a small embedding size
# so tests are faster and easier to debug.
EMBEDDING_SIZE = 6


def _sort_by_id(documents: list[Document]) -> list[Document]:
    return sorted(documents, key=lambda doc: doc.id or "")


class VectorStoreIntegrationTests(BaseStandardTests):
    """Base class for vector store integration tests.

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

    The fixture should use the `get_embeddings` method to get a pre-defined
    embeddings model that should be used for this test suite.

    Here is a template:

    ```python
    from typing import Generator

    import pytest
    from langchain_core.vectorstores import VectorStore
    from langchain_parrot_link.vectorstores import ParrotVectorStore
    from langchain_tests.integration_tests.vectorstores import VectorStoreIntegrationTests


    class TestParrotVectorStore(VectorStoreIntegrationTests):
        @pytest.fixture()
        def vectorstore(self) -> Generator[VectorStore, None, None]:  # type: ignore
            \"\"\"Get an empty vectorstore.\"\"\"
            store = ParrotVectorStore(self.get_embeddings())
            # note: store should be EMPTY at this point
            # if you need to delete data, you may do so here
            try:
                yield store
            finally:
                # cleanup operations, or deleting data
                pass
    ```

    In the fixture, before the `yield` we instantiate an empty vector store. In the
    `finally` block, we call whatever logic is necessary to bring the vector store
    to a clean state.

    ```python
    from typing import Generator
// ... (783 more lines)