From c14239b1ba1ad29b77e67a792489c8717074b0b0 Mon Sep 17 00:00:00 2001
From: Jorge <jzaldivar@google.com>
Date: Wed, 13 Mar 2024 12:31:41 +0100
Subject: [PATCH] Add stream updating

---
 .../langchain_google_vertexai/vectorstores/vectorstores.py  | 6 +++++-
 libs/vertexai/tests/integration_tests/test_vectorstores.py  | 2 +-
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/libs/vertexai/langchain_google_vertexai/vectorstores/vectorstores.py b/libs/vertexai/langchain_google_vertexai/vectorstores/vectorstores.py
index 77fa08f7..582010a5 100644
--- a/libs/vertexai/langchain_google_vertexai/vectorstores/vectorstores.py
+++ b/libs/vertexai/langchain_google_vertexai/vectorstores/vectorstores.py
@@ -269,6 +269,7 @@ def from_components(  # Implemented in order to keep the current API
         endpoint_id: str,
         credentials_path: Optional[str] = None,
         embedding: Optional[Embeddings] = None,
+        stream_update: bool = False,
         **kwargs: Any,
     ) -> "VectorSearchVectorStore":
         """Takes the object creation out of the constructor.
@@ -284,6 +285,8 @@ def from_components(  # Implemented in order to keep the current API
             the local file system.
             embedding: The :class:`Embeddings` that will be used for
             embedding the texts.
+            stream_update: Whether to update with streaming or batching. VectorSearch
+                index must be compatible with stream/batch updates.
             kwargs: Additional keyword arguments to pass to
                 VertexAIVectorSearch.__init__().
         Returns:
@@ -300,7 +303,8 @@ def from_components(  # Implemented in order to keep the current API
         return cls(
             document_storage=GCSDocumentStorage(bucket=bucket),
             searcher=VectorSearchSearcher(
-                endpoint=endpoint, index=index, staging_bucket=bucket
+                endpoint=endpoint, index=index, staging_bucket=bucket,
+                stream_update=stream_update
             ),
             embbedings=embedding,
         )
diff --git a/libs/vertexai/tests/integration_tests/test_vectorstores.py b/libs/vertexai/tests/integration_tests/test_vectorstores.py
index d3e769ca..eada386b 100644
--- a/libs/vertexai/tests/integration_tests/test_vectorstores.py
+++ b/libs/vertexai/tests/integration_tests/test_vectorstores.py
@@ -170,7 +170,7 @@ def test_vector_store_filtering(vector_store: VectorSearchVectorStore):
     )
 
     assert len(documents) > 0
-
+    assert all(document.metadata["color"] == "blue" for document in documents)
     assert all(document.metadata["price"] < 20.0 for document in documents)