From 1c934fff0e206c50ce05be26185839350c8a5ad9 Mon Sep 17 00:00:00 2001
From: Michael Landis <michael@momentohq.com>
Date: Wed, 20 Dec 2023 19:11:43 -0800
Subject: [PATCH] community[patch]: support momento vector index filter
 expressions (#14978)

**Description**

For the Momento Vector Index (MVI) vector store implementation, pass
through `filter_expression` kwarg to the MVI client, if specified. This
change will enable the MVI self query implementation in a future PR.

Also fixes some integration tests.
---
 .../vectorstores/momento_vector_index.py      | 14 +++++++++++--
 .../vectorstores/test_momento_vector_index.py | 21 ++++++++++++++++++-
 2 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/libs/community/langchain_community/vectorstores/momento_vector_index.py b/libs/community/langchain_community/vectorstores/momento_vector_index.py
index b8f5b3e5512..4e2e03dd79f 100644
--- a/libs/community/langchain_community/vectorstores/momento_vector_index.py
+++ b/libs/community/langchain_community/vectorstores/momento_vector_index.py
@@ -306,8 +306,13 @@ class MomentoVectorIndex(VectorStore):
 
         if "top_k" in kwargs:
             k = kwargs["k"]
+        filter_expression = kwargs.get("filter_expression", None)
         response = self._client.search(
-            self.index_name, embedding, top_k=k, metadata_fields=ALL_METADATA
+            self.index_name,
+            embedding,
+            top_k=k,
+            metadata_fields=ALL_METADATA,
+            filter_expression=filter_expression,
         )
 
         if not isinstance(response, Search.Success):
@@ -366,8 +371,13 @@ class MomentoVectorIndex(VectorStore):
         from momento.requests.vector_index import ALL_METADATA
         from momento.responses.vector_index import SearchAndFetchVectors
 
+        filter_expression = kwargs.get("filter_expression", None)
         response = self._client.search_and_fetch_vectors(
-            self.index_name, embedding, top_k=fetch_k, metadata_fields=ALL_METADATA
+            self.index_name,
+            embedding,
+            top_k=fetch_k,
+            metadata_fields=ALL_METADATA,
+            filter_expression=filter_expression,
         )
 
         if isinstance(response, SearchAndFetchVectors.Success):
diff --git a/libs/community/tests/integration_tests/vectorstores/test_momento_vector_index.py b/libs/community/tests/integration_tests/vectorstores/test_momento_vector_index.py
index 0e8e178fd7d..e7d54356f6f 100644
--- a/libs/community/tests/integration_tests/vectorstores/test_momento_vector_index.py
+++ b/libs/community/tests/integration_tests/vectorstores/test_momento_vector_index.py
@@ -1,10 +1,12 @@
+import os
 import time
 import uuid
-from typing import Iterator, List
+from typing import Generator, Iterator, List
 
 import pytest
 from langchain_core.documents import Document
 
+from langchain_community.document_loaders import TextLoader
 from langchain_community.embeddings import OpenAIEmbeddings
 from langchain_community.vectorstores import MomentoVectorIndex
 
@@ -24,6 +26,23 @@ def wait() -> None:
     time.sleep(1)
 
 
+@pytest.fixture(scope="module")
+def embedding_openai() -> OpenAIEmbeddings:
+    if not os.environ.get("OPENAI_API_KEY"):
+        raise ValueError("OPENAI_API_KEY is not set")
+    return OpenAIEmbeddings()
+
+
+@pytest.fixture(scope="function")
+def texts() -> Generator[List[str], None, None]:
+    # Load the documents from a file located in the fixtures directory
+    documents = TextLoader(
+        os.path.join(os.path.dirname(__file__), "fixtures", "sharks.txt")
+    ).load()
+
+    yield [doc.page_content for doc in documents]
+
+
 @pytest.fixture(scope="function")
 def vector_store(
     embedding_openai: OpenAIEmbeddings, random_index_name: str