Mocking VectorStore class

[sunnstix]

For my testing, I want to avoid creating an instance of postgres to test my code. Do you have any suggestions on how to create a stateful MockVectorStore class that replicates the behavior of the class? Primarily I want to mock the aadd_documents and asimilarity_search methods.

[dosubot[bot]]

To create a stateful MockVectorStore class that replicates the behavior of the VectorStore class, including the aadd_documents and asimilarity_search methods, you can use the following approach:

class MockVectorStore:
    def __init__(self):
        self.documents = []
        self.index = {}

    async def aadd_documents(self, docs):
        ids = []
        for doc in docs:
            doc_id = f"doc_{len(self.documents)}"
            self.documents.append(doc)
            self.index[doc_id] = doc
            ids.append(doc_id)
        return ids

    async def asimilarity_search(self, query, k=None, filter=None):
        # For simplicity, return all documents as similar
        results = [
            {
                "page_content": doc,
                "metadata": {"some": "metadata", "value": {"body": doc}}
            }
            for doc in self.documents
        ]
        if filter:
            # Apply filtering logic based on metadata
            results = [result for result in results if all(
                result["metadata"].get(key) == value for key, value in filter.items()
            )]
        return results[:k] if k else results

# Example usage
mock_store = MockVectorStore()
await mock_store.aadd_documents(["This is a new test", "This is a second test"])
results = await mock_store.asimilarity_search("Who was inspired by Ada Lovelace?", k=2)
print(results)

This MockVectorStore class maintains a list of documents and an index to simulate adding documents and performing a similarity search. The aadd_documents method assigns a unique ID to each document, and the asimilarity_search method returns all documents with mock metadata. The method also supports filtering based on metadata and limiting the number of results with the k parameter ^[1][2].

To continue talking to Dosu, mention @dosu.

---

^{Help Dosu learn! Give it feedback: Great Response | Irrelevant Answer | Incorrect Sources | Too Verbose | Hallucination | Bug Report | Other}

[sunnstix]

@dosu Also mock the init method

[sunnstix]

@dosu Also mock the init method

[dosubot[bot]]

To mock the __init__ method of the VectorStore class in your MockVectorStore, you can initialize any necessary attributes that the original class would set up. Here's how you can modify the MockVectorStore to include a mocked __init__ method:

class MockVectorStore:
    def __init__(self, embedding=None, connection_string=None, table=None, schema=None, logger=None, drop=False):
        # Initialize attributes as needed
        self.documents = []
        self.index = {}
        self.embedding = embedding
        self.connection_string = connection_string
        self.table = table
        self.schema = schema
        self.logger = logger
        self.drop = drop

    async def aadd_documents(self, docs):
        ids = []
        for doc in docs:
            doc_id = f"doc_{len(self.documents)}"
            self.documents.append(doc)
            self.index[doc_id] = doc
            ids.append(doc_id)
        return ids

    async def asimilarity_search(self, query, k=None, filter=None):
        # For simplicity, return all documents as similar
        results = [
            {
                "page_content": doc,
                "metadata": {"some": "metadata", "value": {"body": doc}}
            }
            for doc in self.documents
        ]
        if filter:
            # Apply filtering logic based on metadata
            results = [result for result in results if all(
                result["metadata"].get(key) == value for key, value in filter.items()
            )]
        return results[:k] if k else results

# Example usage
mock_store = MockVectorStore()
await mock_store.aadd_documents(["This is a new test", "This is a second test"])
results = await mock_store.asimilarity_search("Who was inspired by Ada Lovelace?", k=2)
print(results)

This mock class now includes an __init__ method that initializes attributes similar to those in the VectorStore class, allowing you to simulate its behavior without needing a real database connection. Adjust the parameters and attributes as necessary to fit your specific testing needs.

To continue talking to Dosu, mention @dosu.

---

^{Help Dosu learn! Give it feedback: Great Response | Irrelevant Answer | Incorrect Sources | Too Verbose | Hallucination | Bug Report | Other}

[sunnstix]

@dosu also mock the adelete method

[dosubot[bot]]

To mock the adelete method in your MockVectorStore class, you can implement it to simulate the deletion of documents by their IDs. Here's how you can extend the MockVectorStore class to include a mocked adelete method:

class MockVectorStore:
    def __init__(self, embedding=None, connection_string=None, table=None, schema=None, logger=None, drop=False):
        # Initialize attributes as needed
        self.documents = []
        self.index = {}
        self.embedding = embedding
        self.connection_string = connection_string
        self.table = table
        self.schema = schema
        self.logger = logger
        self.drop = drop

    async def aadd_documents(self, docs):
        ids = []
        for doc in docs:
            doc_id = f"doc_{len(self.documents)}"
            self.documents.append(doc)
            self.index[doc_id] = doc
            ids.append(doc_id)
        return ids

    async def asimilarity_search(self, query, k=None, filter=None):
        # For simplicity, return all documents as similar
        results = [
            {
                "page_content": doc,
                "metadata": {"some": "metadata", "value": {"body": doc}}
            }
            for doc in self.documents
        ]
        if filter:
            # Apply filtering logic based on metadata
            results = [result for result in results if all(
                result["metadata"].get(key) == value for key, value in filter.items()
            )]
        return results[:k] if k else results

    async def adelete(self, ids=None, **kwargs):
        if ids is None:
            raise ValueError("No ids provided to delete.")
        for doc_id in ids:
            if doc_id in self.index:
                self.documents.remove(self.index[doc_id])
                del self.index[doc_id]
        return True

# Example usage
mock_store = MockVectorStore()
await mock_store.aadd_documents(["This is a new test", "This is a second test"])
await mock_store.adelete(["doc_0"])
results = await mock_store.asimilarity_search("Who was inspired by Ada Lovelace?", k=2)
print(results)

This adelete method simulates the deletion of documents by their IDs from the documents list and index dictionary. It raises a ValueError if no IDs are provided, similar to the behavior of the actual adelete method in the VectorStore class ^[1].

To continue talking to Dosu, mention @dosu.

---

^{Help Dosu learn! Give it feedback: Great Response | Irrelevant Answer | Incorrect Sources | Too Verbose | Hallucination | Bug Report | Other}