Move sentinel replacement to serializer

Author: abrookins

langgraph/checkpoint/redis/__init__.py

@@ -78,13 +78,19 @@ def list(
         # Construct the filter expression
         filter_expression = []
         if config:
-            filter_expression.append(
-                Tag("thread_id") == config["configurable"]["thread_id"]
-            )
-            if checkpoint_ns := config["configurable"].get("checkpoint_ns"):
+            thread_id = config["configurable"]["thread_id"]
+            checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
+            checkpoint_id = get_checkpoint_id(config)
+            filter_expression.append(Tag("thread_id") == thread_id)
+
+            # Following the Postgres implementation, we only want to filter by
+            # checkpoint_ns if it's set. This is slightly different than the
+            # get_tuple() logic, where we always query for checkpoint_ns.
+            if checkpoint_ns:
                 filter_expression.append(Tag("checkpoint_ns") == checkpoint_ns)
-            if checkpoint_id := get_checkpoint_id(config):
-                filter_expression.append(Tag("checkpoint_id") == checkpoint_id)
+            # We want to find all checkpoints for the thread matching the other
+            # filters, with any checkpoint_id.
+            filter_expression.append(Tag("checkpoint_id") == checkpoint_id)
 
         if filter:
             for k, v in filter.items():
@@ -195,9 +201,9 @@ def put(
         """Store a checkpoint to Redis."""
         configurable = config["configurable"].copy()
         thread_id = configurable.pop("thread_id")
-        checkpoint_ns = configurable.pop("checkpoint_ns")
+        checkpoint_ns = configurable.pop("checkpoint_ns", "")
         checkpoint_id = configurable.pop(
-            "checkpoint_id", configurable.pop("thread_ts", None)
+            "checkpoint_id", configurable.pop("thread_ts", "")
         )
 
         copy = checkpoint.copy()
@@ -211,10 +217,10 @@ def put(
 
         # Store checkpoint data
         checkpoint_data = {
-            "thread_id": thread_id,
-            "checkpoint_ns": checkpoint_ns,
-            "checkpoint_id": checkpoint["id"],
-            "parent_checkpoint_id": checkpoint_id,
+            "thread_id": thread_id or "",
+            "checkpoint_ns": checkpoint_ns or "",
+            "checkpoint_id": checkpoint["id"] or "",
+            "parent_checkpoint_id": checkpoint_id or "",
             "checkpoint": self._dump_checkpoint(copy),
             "metadata": self._dump_metadata(metadata),
         }
@@ -261,11 +267,16 @@ def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
         checkpoint_id = get_checkpoint_id(config)
         checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
 
-        checkpoint_filter_expression = Tag("thread_id") == thread_id
         if checkpoint_id:
-            checkpoint_filter_expression &= Tag("checkpoint_id") == str(checkpoint_id)
-        if checkpoint_ns:
-            checkpoint_filter_expression &= Tag("checkpoint_ns") == checkpoint_ns
+            checkpoint_filter_expression = (
+                (Tag("thread_id") == thread_id)
+                & (Tag("checkpoint_ns") == checkpoint_ns)
+                & (Tag("checkpoint_id") == str(checkpoint_id))
+            )
+        else:
+            checkpoint_filter_expression = (Tag("thread_id") == thread_id) & (
+                Tag("checkpoint_ns") == checkpoint_ns
+            )
 
         # Construct the query
         checkpoints_query = FilterQuery(
@@ -289,20 +300,25 @@ def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
 
         doc = results.docs[0]
 
+        doc_thread_id = doc["thread_id"]
+        doc_checkpoint_ns = doc["checkpoint_ns"]
+        doc_checkpoint_id = doc["checkpoint_id"]
+        doc_parent_checkpoint_id = doc["parent_checkpoint_id"]
+
         # Fetch channel_values
         channel_values = self.get_channel_values(
-            thread_id=doc["thread_id"],
-            checkpoint_ns=doc["checkpoint_ns"],
-            checkpoint_id=doc["checkpoint_id"],
+            thread_id=doc_thread_id,
+            checkpoint_ns=doc_checkpoint_ns,
+            checkpoint_id=doc_checkpoint_id,
         )
 
         # Fetch pending_sends from parent checkpoint
         pending_sends = []
-        if doc["parent_checkpoint_id"]:
+        if doc_parent_checkpoint_id:
             pending_sends = self._load_pending_sends(
-                thread_id=thread_id,
-                checkpoint_ns=checkpoint_ns,
-                parent_checkpoint_id=doc["parent_checkpoint_id"],
+                thread_id=doc_thread_id,
+                checkpoint_ns=doc_checkpoint_ns,
+                parent_checkpoint_id=doc_parent_checkpoint_id,
             )
 
         # Fetch and parse metadata
@@ -324,7 +340,7 @@ def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
             "configurable": {
                 "thread_id": thread_id,
                 "checkpoint_ns": checkpoint_ns,
-                "checkpoint_id": doc["checkpoint_id"],
+                "checkpoint_id": doc_checkpoint_id,
             }
         }
 
@@ -335,7 +351,7 @@ def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
         )
 
         pending_writes = self._load_pending_writes(
-            thread_id, checkpoint_ns, checkpoint_id
+            thread_id, checkpoint_ns, doc_checkpoint_id
         )
 
         return CheckpointTuple(

langgraph/checkpoint/redis/jsonplus_redis.py

@@ -1,17 +1,46 @@
 import base64
+import logging
 from typing import Any, Union
 
 from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
 
 
+logger = logging.getLogger(__name__)
+
+
+# RediSearch versions below 2.10 don't support indexing and querying
+# empty strings, so we use a sentinel value to represent empty strings.
+EMPTY_STRING_SENTINEL = "__empty__"
+
+
 class JsonPlusRedisSerializer(JsonPlusSerializer):
     """Redis-optimized serializer that stores strings directly."""
 
+    SENTINEL_FIELDS = [
+        "thread_id",
+        "checkpoint_id",
+        "checkpoint_ns",
+        "parent_checkpoint_id",
+    ]
+
     def dumps_typed(self, obj: Any) -> tuple[str, str]:  # type: ignore[override]
         if isinstance(obj, (bytes, bytearray)):
             return "base64", base64.b64encode(obj).decode("utf-8")
         else:
-            return "json", self.dumps(obj).decode("utf-8")
+            for field in self.SENTINEL_FIELDS:
+                try:
+                    if field in obj and not obj[field]:
+                        obj[field] = EMPTY_STRING_SENTINEL
+                except (KeyError, AttributeError):
+                    try:
+                        if hasattr(obj, field) and not getattr(obj, field, None):
+                            setattr(obj, field, EMPTY_STRING_SENTINEL)
+                    except Exception as e:
+                        logger.debug(
+                            f"Error setting {field} from empty string to sentinel: {e}"
+                        )
+            results = self.dumps(obj).decode("utf-8")
+            return "json", results
 
     def loads_typed(self, data: tuple[str, Union[str, bytes]]) -> Any:
         type_, data_ = data
@@ -22,4 +51,21 @@ def loads_typed(self, data: tuple[str, Union[str, bytes]]) -> Any:
             return decoded
         elif type_ == "json":
             data_bytes = data_ if isinstance(data_, bytes) else data_.encode()
-            return self.loads(data_bytes)
+            results = self.loads(data_bytes)
+            for field in self.SENTINEL_FIELDS:
+                try:
+                    if field in results and results[field] == EMPTY_STRING_SENTINEL:
+                        results[field] = ""
+                except (KeyError, AttributeError):
+                    try:
+                        if (
+                            hasattr(results, field)
+                            and getattr(results, field) == EMPTY_STRING_SENTINEL
+                        ):
+                            setattr(results, field, "")
+                    except Exception as e:
+                        logger.debug(
+                            f"Error setting {field} from sentinel to empty string: {e}"
+                        )
+                        pass
+            return results