Merge remote-tracking branch 'origin/main' into feat/RAAE-1287/checkpoint-custom-prefix

Author: vishal-bala

langgraph/store/redis/__init__.py

@@ -257,7 +257,9 @@ def _batch_get_ops(
             for idx, key in items:
                 if key in key_to_row:
                     data, doc_id = key_to_row[key]
-                    results[idx] = _row_to_item(namespace, data)
+                    results[idx] = _row_to_item(
+                        namespace, data, deserialize_fn=self._deserialize_value
+                    )
 
                     # Find the corresponding operation by looking it up in the operation list
                     # This is needed because idx is the index in the overall operation list
@@ -587,6 +589,7 @@ def _batch_search_ops(
                                 _decode_ns(store_doc["prefix"]),
                                 store_doc,
                                 score=score,
+                                deserialize_fn=self._deserialize_value,
                             )
                         )
 
@@ -653,7 +656,13 @@ def _batch_search_ops(
                         )
                         refresh_keys.append(vector_key)
 
-                    items.append(_row_to_search_item(_decode_ns(data["prefix"]), data))
+                    items.append(
+                        _row_to_search_item(
+                            _decode_ns(data["prefix"]),
+                            data,
+                            deserialize_fn=self._deserialize_value,
+                        )
+                    )
 
                 # Refresh TTL if requested
                 if op.refresh_ttl and refresh_keys and self.ttl_config:

langgraph/store/redis/aio.py

@@ -472,7 +472,9 @@ async def _batch_get_ops(
             for idx, key in items:
                 if key in key_to_row:
                     data, doc_id = key_to_row[key]
-                    results[idx] = _row_to_item(namespace, data)
+                    results[idx] = _row_to_item(
+                        namespace, data, deserialize_fn=self._deserialize_value
+                    )
 
                     # Find the corresponding operation by looking it up in the operation list
                     # This is needed because idx is the index in the overall operation list
@@ -578,7 +580,7 @@ async def _aprepare_batch_PUT_queries(
                 doc = RedisDocument(
                     prefix=_namespace_to_text(op.namespace),
                     key=op.key,
-                    value=op.value,
+                    value=self._serialize_value(op.value),
                     created_at=now,
                     updated_at=now,
                     ttl_minutes=ttl_minutes,
@@ -872,6 +874,7 @@ async def _batch_search_ops(
                                 _decode_ns(store_doc["prefix"]),
                                 store_doc,
                                 score=score,
+                                deserialize_fn=self._deserialize_value,
                             )
                         )
 
@@ -939,7 +942,13 @@ async def _batch_search_ops(
                         )
                         refresh_keys.append(vector_key)
 
-                    items.append(_row_to_search_item(_decode_ns(data["prefix"]), data))
+                    items.append(
+                        _row_to_search_item(
+                            _decode_ns(data["prefix"]),
+                            data,
+                            deserialize_fn=self._deserialize_value,
+                        )
+                    )
 
                 # Refresh TTL if requested
                 if op.refresh_ttl and refresh_keys and self.ttl_config:

langgraph/store/redis/base.py

@@ -3,12 +3,14 @@
 from __future__ import annotations
 
 import copy
+import json
 import logging
 import threading
 from collections import defaultdict
 from datetime import datetime, timedelta, timezone
 from typing import (
     Any,
+    Callable,
     Dict,
     Generic,
     Iterable,
@@ -40,6 +42,8 @@
 from redisvl.query.filter import Tag, Text
 from redisvl.utils.token_escaper import TokenEscaper
 
+from langgraph.checkpoint.redis.jsonplus_redis import JsonPlusRedisSerializer
+
 from .token_unescaper import TokenUnescaper
 from .types import IndexType, RedisClientType
 
@@ -124,6 +128,9 @@ class BaseRedisStore(Generic[RedisClientType, IndexType]):
     supports_ttl: bool = True
     ttl_config: Optional[TTLConfig] = None
 
+    # Serializer for handling complex objects like LangChain messages
+    _serde: JsonPlusRedisSerializer
+
     def _apply_ttl_to_keys(
         self,
         main_key: str,
@@ -223,6 +230,8 @@ def __init__(
         self._redis = conn
         # Store cluster_mode; None means auto-detect in RedisStore or AsyncRedisStore
         self.cluster_mode = cluster_mode
+        # Initialize the serializer for handling complex objects like LangChain messages
+        self._serde = JsonPlusRedisSerializer()
 
         # Store custom prefixes
         self.store_prefix = store_prefix
@@ -357,6 +366,109 @@ async def aset_client_info(self) -> None:
                 # Silently fail if even echo doesn't work
                 pass
 
+    def _serialize_value(self, value: Any) -> Any:
+        """Serialize a value for storage in Redis.
+
+        This method handles complex objects like LangChain messages by
+        serializing them to a JSON-compatible format.
+
+        The method is smart about serialization:
+        - If the value is a simple JSON-serializable dict/list, it's stored as-is
+        - If the value contains complex objects (HumanMessage, etc.), it uses
+          the serde wrapper format with __serde_type__ and __serde_data__ keys
+
+        Note: Values containing LangChain messages will be wrapped in a serde format,
+        which means filters on nested fields won't work for such values.
+
+        Args:
+            value: The value to serialize (can contain HumanMessage, AIMessage, etc.)
+
+        Returns:
+            A JSON-serializable representation of the value
+        """
+        if value is None:
+            return None
+
+        # First, try standard JSON serialization to check if it's needed
+        try:
+            json.dumps(value)
+            # Value is already JSON-serializable, return as-is for backward
+            # compatibility and to preserve filter functionality
+            return value
+        except TypeError:
+            # Value contains non-JSON-serializable objects, use serde wrapper
+            pass
+
+        # Use the serializer to handle complex objects
+        type_str, data_bytes = self._serde.dumps_typed(value)
+        # Store the serialized data with type info for proper deserialization
+        # Handle different type formats explicitly for clarity
+        if type_str == "json":
+            data_encoded = data_bytes.decode("utf-8")
+        else:
+            # bytes, bytearray, msgpack, and other types are hex-encoded
+            data_encoded = data_bytes.hex()
+
+        return {
+            "__serde_type__": type_str,
+            "__serde_data__": data_encoded,
+        }
+
+    def _deserialize_value(self, value: Any) -> Any:
+        """Deserialize a value from Redis storage.
+
+        This method handles both new serialized format and legacy plain values
+        for backward compatibility.
+
+        Args:
+            value: The value from Redis (may be serialized or plain)
+
+        Returns:
+            The deserialized value with proper Python objects (HumanMessage, etc.)
+        """
+        if value is None:
+            return None
+
+        # Check if this is a serialized value (new format)
+        # Use exact key check to prevent collisions with user data
+        if isinstance(value, dict) and set(value.keys()) == {
+            "__serde_type__",
+            "__serde_data__",
+        }:
+            type_str = value["__serde_type__"]
+            data_str = value["__serde_data__"]
+
+            try:
+                # Convert back to bytes based on type
+                if type_str == "json":
+                    data_bytes = data_str.encode("utf-8")
+                else:
+                    # bytes, bytearray, msgpack types are hex-encoded
+                    data_bytes = bytes.fromhex(data_str)
+
+                return self._serde.loads_typed((type_str, data_bytes))
+            except (ValueError, TypeError) as e:
+                # Handle hex decoding errors or deserialization failures
+                logger.error(
+                    "Failed to deserialize value from Redis: type=%r, error=%s",
+                    type_str,
+                    e,
+                )
+                # Return None to indicate deserialization failure
+                return None
+            except Exception as e:
+                # Handle any other unexpected errors during deserialization
+                logger.error(
+                    "Unexpected error deserializing value from Redis: type=%r, error=%s",
+                    type_str,
+                    e,
+                )
+                return None
+
+        # Legacy format: value is stored as-is (plain JSON-serializable data)
+        # Return as-is for backward compatibility
+        return value
+
     def _get_batch_GET_ops_queries(
         self,
         get_ops: Sequence[tuple[int, GetOp]],
@@ -433,7 +545,7 @@ def _prepare_batch_PUT_queries(
                 doc = RedisDocument(
                     prefix=_namespace_to_text(op.namespace),
                     key=op.key,
-                    value=op.value,
+                    value=self._serialize_value(op.value),
                     created_at=now,
                     updated_at=now,
                     ttl_minutes=ttl_minutes,
@@ -568,10 +680,27 @@ def _decode_ns(ns: str) -> tuple[str, ...]:
     return tuple(_token_unescaper.unescape(ns).split("."))
 
 
-def _row_to_item(namespace: tuple[str, ...], row: dict[str, Any]) -> Item:
-    """Convert a row from Redis to an Item."""
+def _row_to_item(
+    namespace: tuple[str, ...],
+    row: dict[str, Any],
+    deserialize_fn: Optional[Callable[[Any], Any]] = None,
+) -> Item:
+    """Convert a row from Redis to an Item.
+
+    Args:
+        namespace: The namespace tuple for this item
+        row: The raw row data from Redis
+        deserialize_fn: Optional function to deserialize the value (handles
+            LangChain messages, etc.)
+
+    Returns:
+        An Item with properly deserialized value
+    """
+    value = row["value"]
+    if deserialize_fn is not None:
+        value = deserialize_fn(value)
     return Item(
-        value=row["value"],
+        value=value,
         key=row["key"],
         namespace=namespace,
         created_at=datetime.fromtimestamp(row["created_at"] / 1_000_000, timezone.utc),
@@ -583,10 +712,25 @@ def _row_to_search_item(
     namespace: tuple[str, ...],
     row: dict[str, Any],
     score: Optional[float] = None,
+    deserialize_fn: Optional[Callable[[Any], Any]] = None,
 ) -> SearchItem:
-    """Convert a row from Redis to a SearchItem."""
+    """Convert a row from Redis to a SearchItem.
+
+    Args:
+        namespace: The namespace tuple for this item
+        row: The raw row data from Redis
+        score: Optional similarity score from vector search
+        deserialize_fn: Optional function to deserialize the value (handles
+            LangChain messages, etc.)
+
+    Returns:
+        A SearchItem with properly deserialized value
+    """
+    value = row["value"]
+    if deserialize_fn is not None:
+        value = deserialize_fn(value)
     return SearchItem(
-        value=row["value"],
+        value=value,
         key=row["key"],
         namespace=namespace,
         created_at=datetime.fromtimestamp(row["created_at"] / 1_000_000, timezone.utc),

poetry.lock

@@ -20,7 +20,7 @@ packages = [{ include = "langgraph" }]
 python = ">=3.10,<3.14"
 langgraph-checkpoint = ">=3.0.0,<4.0.0"
 redisvl = ">=0.11.0,<1.0.0"
-redis = ">=5.2.1,<7.0.0"
+redis = ">=5.2.1"
 orjson = "^3.9.0"
 tomli = { version = "^2.0.1", python = "<3.11" }