From 03023a968ea4656db4a4f4218a2f395e197fc6ca Mon Sep 17 00:00:00 2001
From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com>
Date: Mon, 25 Nov 2024 16:45:23 -0800
Subject: [PATCH] Add optoin to not index

---
 .../langgraph/store/postgres/base.py               |  2 ++
 libs/checkpoint-postgres/tests/test_async_store.py |  9 +++++++++
 libs/checkpoint-postgres/tests/test_store.py       |  5 +++++
 libs/checkpoint/langgraph/store/base/__init__.py   | 14 ++++++++++++--
 libs/checkpoint/langgraph/store/base/batch.py      |  3 ++-
 5 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/base.py b/libs/checkpoint-postgres/langgraph/store/postgres/base.py
index b4e11b268..6a2c2a613 100644
--- a/libs/checkpoint-postgres/langgraph/store/postgres/base.py
+++ b/libs/checkpoint-postgres/langgraph/store/postgres/base.py
@@ -350,6 +350,8 @@ class BasePostgresStore(Generic[C]):
                 elif text_fields is None:
                     text_fields = ["__root__"]
                 for op in inserts:
+                    if op.index is False:
+                        continue
                     value = op.value
                     ns = _namespace_to_text(op.namespace)
                     k = op.key
diff --git a/libs/checkpoint-postgres/tests/test_async_store.py b/libs/checkpoint-postgres/tests/test_async_store.py
index 805c5425f..e0e9211de 100644
--- a/libs/checkpoint-postgres/tests/test_async_store.py
+++ b/libs/checkpoint-postgres/tests/test_async_store.py
@@ -307,6 +307,15 @@ async def test_vector_update_with_embedding(vector_store: AsyncPostgresStore) ->
         if r.key == "doc1":
             assert r.response_metadata["score"] > after_score
 
+    # Don't index this one
+    await vector_store.aput(
+        ("test",), "doc4", {"text": "new text about dogs"}, index=False
+    )
+    results_new = await vector_store.asearch(
+        ("test",), query="new text about dogs", limit=3
+    )
+    assert not any(r.key == "doc4" for r in results_new)
+
 
 async def test_vector_search_with_filters(vector_store: AsyncPostgresStore) -> None:
     """Test combining vector search with filters."""
diff --git a/libs/checkpoint-postgres/tests/test_store.py b/libs/checkpoint-postgres/tests/test_store.py
index c8de856e9..39f9455a7 100644
--- a/libs/checkpoint-postgres/tests/test_store.py
+++ b/libs/checkpoint-postgres/tests/test_store.py
@@ -458,6 +458,11 @@ def test_vector_update_with_embedding(vector_store: PostgresStore) -> None:
         if r.key == "doc1":
             assert r.response_metadata["score"] > after_score
 
+    # Don't index this one
+    vector_store.put(("test",), "doc4", {"text": "new text about dogs"}, index=False)
+    results_new = vector_store.search(("test",), query="new text about dogs", limit=3)
+    assert not any(r.key == "doc4" for r in results_new)
+
 
 def test_vector_search_with_filters(vector_store: PostgresStore) -> None:
     """Test combining vector search with filters."""
diff --git a/libs/checkpoint/langgraph/store/base/__init__.py b/libs/checkpoint/langgraph/store/base/__init__.py
index cf1f75a41..6356c5d5b 100644
--- a/libs/checkpoint/langgraph/store/base/__init__.py
+++ b/libs/checkpoint/langgraph/store/base/__init__.py
@@ -176,6 +176,12 @@ class PutOp(NamedTuple):
       - Values can be of any serializable type
     - If None, it indicates that the item should be deleted
     """
+    index: Optional[bool] = None
+    """Whether to index the item (if supported by the store).
+    
+    Defaults to True if the store supports indexing. This will embed the document
+    so it can be queried using search.
+    """
 
 
 NameSpacePath = tuple[Union[str, Literal["*"]], ...]
@@ -351,9 +357,11 @@ class BaseStore(ABC):
             namespace: Hierarchical path for the item.
             key: Unique identifier within the namespace.
             value: Dictionary containing the item's data.
+            index: Whether to index the item (if supported by the store).
+                Defaults to True if the store supports indexing.
         """
         _validate_namespace(namespace)
-        self.batch([PutOp(namespace, key, value)])
+        self.batch([PutOp(namespace, key, value, index=index)])
 
     def delete(self, namespace: tuple[str, ...], key: str) -> None:
         """Delete an item.
@@ -468,9 +476,11 @@ class BaseStore(ABC):
             namespace: Hierarchical path for the item.
             key: Unique identifier within the namespace.
             value: Dictionary containing the item's data.
+            index: Whether to index the item (if supported by the store).
+                Defaults to True if the store supports indexing.
         """
         _validate_namespace(namespace)
-        await self.abatch([PutOp(namespace, key, value)])
+        await self.abatch([PutOp(namespace, key, value, index)])
 
     async def adelete(self, namespace: tuple[str, ...], key: str) -> None:
         """Asynchronously delete an item.
diff --git a/libs/checkpoint/langgraph/store/base/batch.py b/libs/checkpoint/langgraph/store/base/batch.py
index 90f837622..c2898b947 100644
--- a/libs/checkpoint/langgraph/store/base/batch.py
+++ b/libs/checkpoint/langgraph/store/base/batch.py
@@ -57,10 +57,11 @@ class AsyncBatchedBaseStore(BaseStore):
         namespace: tuple[str, ...],
         key: str,
         value: dict[str, Any],
+        index: Optional[bool] = None,
     ) -> None:
         _validate_namespace(namespace)
         fut = self._loop.create_future()
-        self._aqueue[fut] = PutOp(namespace, key, value)
+        self._aqueue[fut] = PutOp(namespace, key, value, index)
         return await fut
 
     async def adelete(