Mocking vector store for downstream app unit tests

[tasansal]

Hi, we are building an application with langchain and in the process of writing unit tests. What is the best way to mock the vector store? We are currently using PGVector, however, this probably applies to all other vector stores. Do you have any recommendations on how to do this?

The modules I want to test are:

• Connect langchain's PGVector vectorstore
• Use it as retriever to get pre-defined results

[dosubot[bot]]

🤖

To mock the vector store for unit tests in LangChain, you can use a fake embedding class. This class can generate fake embeddings for the documents and queries, which can then be used to initialize the PGVector object in your test functions. Here's an example of how you can do this:

class FakeEmbeddingsWithAdaDimension(FakeEmbeddings):
    """Fake embeddings functionality for testing."""

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """Return simple embeddings."""
        return [
            [float(1.0)] * (ADA_TOKEN_COUNT - 1) + [float(i)] for i in range(len(texts))
        ]

    def embed_query(self, text: str) -> List[float]:
        """Return simple embeddings."""
        return [float(1.0)] * (ADA_TOKEN_COUNT - 1) + [float(0.0)]

def test_pgvector() -> None:
    """Test end to end construction and search."""
    texts = ["foo", "bar", "baz"]
    docsearch = PGVector.from_texts(
        texts=texts,
        collection_name="test_collection",
        embedding=FakeEmbeddingsWithAdaDimension(),
        connection_string=CONNECTION_STRING,
        pre_delete_collection=True,
    )
    output = docsearch.similarity_search("foo", k=1)
    assert output == [Document(page_content="foo")]

In the test_pgvector function, the PGVector object is initialized with the FakeEmbeddingsWithAdaDimension class as the embedding parameter. This allows the PGVector object to use the fake embeddings for its operations, effectively mocking the actual embeddings that would be used in a real scenario.

You can find more examples of how to use the FakeEmbeddingsWithAdaDimension class for more complex scenarios in the LangChain test files, such as when the PGVector object is initialized with metadata or when the PGVector object is used as a retriever. These examples can be used as a guide for how to mock the PGVector object in different testing scenarios.

You can find these examples in the following files:

• libs/langchain/tests/integration_tests/vectorstores/test_pgvector.py
• tests/integration_tests/vectorstores/test_pgvector.py

I hope this helps! If you have any other questions, feel free to ask.

Sources

• libs/langchain/tests/integration_tests/vectorstores/test_pgvector.py
• tests/integration_tests/vectorstores/test_pgvector.py
• libs/langchain/langchain/vectorstores/pgvector.py

---

This response is meant to be useful and save you time. It is not meant to be a precise solution, but rather a starting point for your own research.

Help me be more useful! Please leave a 👍 if this is helpful and 👎 if it is irrelevant.

If you want to continue the conversation, start your reply with @dosu-bot.

[tasansal]

@dosu-bot I am more concerned about constructing the PGVector instance without a real database running.

[dosubot[bot]]

🤖

To construct a PGVector instance for testing purposes without a real database running, you can use a fake embedding class like FakeEmbeddingsWithAdaDimension. This class overrides the embed_documents and embed_query methods to return simple embeddings. You can then use the PGVector.from_texts or PGVector.from_embeddings methods to create a PGVector instance with these fake embeddings. Here's an example:

from langchain.vectorstores.pgvector import PGVector
from tests.integration_tests.vectorstores.fake_embeddings import FakeEmbeddingsWithAdaDimension

class FakeEmbeddingsWithAdaDimension(FakeEmbeddings):
    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        return [[float(1.0)] * (ADA_TOKEN_COUNT - 1) + [float(i)] for i in range(len(texts))]

    def embed_query(self, text: str) -> List[float]:
        return [float(1.0)] * (ADA_TOKEN_COUNT - 1) + [float(0.0)]

# Create the PGVector instance
pg_vector = PGVector.from_texts(
    texts=["foo", "bar", "baz"],
    collection_name="test_collection",
    embedding=FakeEmbeddingsWithAdaDimension(),
    connection_string="postgresql+psycopg2://user:password@localhost:5432/mydatabase",
    pre_delete_collection=True,
)

In this example, FakeEmbeddingsWithAdaDimension is a fake embedding class that returns simple embeddings, and pg_vector is a PGVector instance created with these fake embeddings. This way, you don't need a real database to be running for testing.

In the context shared, you can also see how to use the PGVector.from_embeddings method to create a PGVector instance with pre-computed embeddings. This can be useful if you want to test the behavior of the PGVector class with specific embeddings. Here's an example:

texts = ["foo", "bar", "baz"]
text_embeddings = FakeEmbeddingsWithAdaDimension().embed_documents(texts)
text_embedding_pairs = list(zip(texts, text_embeddings))
docsearch = PGVector.from_embeddings(
    text_embeddings=text_embedding_pairs,
    collection_name="test_collection",
    embedding=FakeEmbeddingsWithAdaDimension(),
    connection_string="postgresql+psycopg2://user:password@localhost:5432/mydatabase",
    pre_delete_collection=True,
)

In this example, text_embeddings are pre-computed embeddings for the texts, and docsearch is a PGVector instance created with these embeddings.

You can find more information about this in the LangChain codebase, specifically in the libs/langchain/langchain/vectorstores/pgvector.py and tests/integration_tests/vectorstores/test_pgvector.py files.

Sources

• libs/langchain/langchain/vectorstores/pgvector.py
• libs/langchain/tests/integration_tests/vectorstores/test_pgvector.py
• tests/integration_tests/vectorstores/test_pgvector.py

---

This response is meant to be useful and save you time. It is not meant to be a precise solution, but rather a starting point for your own research.

Help me be more useful! Please leave a 👍 if this is helpful and 👎 if it is irrelevant.

If you want to continue the conversation, start your reply with @dosu-bot.

[sunnstix]

@dosu-bot your response illustrates how to mock the embeddings endpoint. I am interested in mocking the PGVector class as a whole, not in mocking the embeddings. Suggest a solution which implements a MockPGVector class which has a mock implementation for PGVector's methods