Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
100 changes: 67 additions & 33 deletions langgraph/checkpoint/redis/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,8 +31,11 @@
from langgraph.checkpoint.redis.base import (
CHECKPOINT_PREFIX,
CHECKPOINT_WRITE_PREFIX,
INDEXED_CHECKPOINT_FILTER_KEYS,
REDIS_KEY_SEPARATOR,
BaseRedisSaver,
_checkpoint_id_filter_matches,
_metadata_filter_matches,
)
from langgraph.checkpoint.redis.key_registry import SyncCheckpointKeyRegistry
from langgraph.checkpoint.redis.message_exporter import (
Expand Down Expand Up @@ -250,9 +253,12 @@ def list(

# Search for checkpoints with any namespace, including an empty
# string, while `checkpoint_id` has to have a value.
if checkpoint_ns := config["configurable"].get("checkpoint_ns"):
if "checkpoint_ns" in config["configurable"]:
filter_expression.append(
Tag("checkpoint_ns") == to_storage_safe_str(checkpoint_ns)
Tag("checkpoint_ns")
== to_storage_safe_str(
config["configurable"].get("checkpoint_ns", "")
)
)
if checkpoint_id := get_checkpoint_id(config):
filter_expression.append(
Expand All @@ -269,20 +275,20 @@ def list(
filter_expression.append(Tag("thread_id") == to_storage_safe_id(v))
elif k == "run_id":
filter_expression.append(Tag("run_id") == to_storage_safe_id(v))
else:
raise ValueError(f"Unsupported filter key: {k}")

if before:
before_checkpoint_id = get_checkpoint_id(before)
if before_checkpoint_id:
try:
before_ulid = ULID.from_str(before_checkpoint_id)
before_ts = before_ulid.timestamp
# Use numeric range query: checkpoint_ts < before_ts
filter_expression.append(Num("checkpoint_ts") < before_ts)
except Exception:
# If not a valid ULID, ignore the before filter
pass
requires_post_filter = bool(
filter and any(k not in INDEXED_CHECKPOINT_FILTER_KEYS for k in filter)
)
before_checkpoint_id = get_checkpoint_id(before) if before else None

if before_checkpoint_id and not requires_post_filter:
try:
before_ulid = ULID.from_str(before_checkpoint_id)
before_ts = before_ulid.timestamp
# Use numeric range query: checkpoint_ts < before_ts
filter_expression.append(Num("checkpoint_ts") < before_ts)
except Exception:
requires_post_filter = True

# Combine all filter expressions
combined_filter = filter_expression[0] if filter_expression else "*"
Expand All @@ -291,30 +297,48 @@ def list(

# Construct the Redis query
# Sort by checkpoint_id in descending order to get most recent checkpoints first
query = FilterQuery(
filter_expression=combined_filter,
return_fields=[
"thread_id",
"checkpoint_ns",
"checkpoint_id",
"parent_checkpoint_id",
"$.checkpoint",
"$.metadata",
"has_writes", # Include has_writes to optimize pending_writes loading
],
num_results=limit or 10000,
sort_by=("checkpoint_id", "DESC"),
)
return_fields = [
"thread_id",
"checkpoint_ns",
"checkpoint_id",
"parent_checkpoint_id",
"$.checkpoint",
"$.metadata",
"has_writes", # Include has_writes to optimize pending_writes loading
]

# Execute the query
results = self.checkpoints_index.search(query)
if requires_post_filter:
result_docs = []
offset = 0
page_size = 10000
while True:
query = FilterQuery(
filter_expression=combined_filter,
return_fields=return_fields,
num_results=page_size,
sort_by=("checkpoint_id", "DESC"),
)
query.paging(offset, page_size)
results = self.checkpoints_index.search(query)
result_docs.extend(results.docs)
if len(results.docs) < page_size:
break
offset += page_size
else:
query = FilterQuery(
filter_expression=combined_filter,
return_fields=return_fields,
num_results=limit or 10000,
sort_by=("checkpoint_id", "DESC"),
)
result_docs = self.checkpoints_index.search(query).docs

# Pre-process all docs to collect batch query requirements
all_docs_data = []
pending_sends_batch_keys = []
pending_writes_batch_keys = []

for doc in results.docs:
for doc in result_docs:
# Extract all attributes once
doc_dict = doc.__dict__ if hasattr(doc, "__dict__") else {}

Expand Down Expand Up @@ -430,6 +454,12 @@ def list(
else:
metadata = cast(CheckpointMetadata, metadata_dict)

if requires_post_filter and (
not _checkpoint_id_filter_matches(checkpoint_id, before_checkpoint_id)
or not _metadata_filter_matches(metadata, filter)
):
continue

# Pre-create the config structure more efficiently
config_param: RunnableConfig = {
"configurable": {
Expand Down Expand Up @@ -474,6 +504,10 @@ def list(
parent_config=parent_config,
pending_writes=pending_writes,
)
if limit is not None:
limit -= 1
if limit <= 0:
break

def put(
self,
Expand Down Expand Up @@ -1692,7 +1726,7 @@ def prune(

# Validate input
if not thread_ids:
raise ValueError("``thread_ids`` must be a non-empty sequence")
return
if keep_last < 0:
raise ValueError(f"``keep_last`` must be >= 0, got {keep_last}")
if max_results < 1:
Expand Down
100 changes: 67 additions & 33 deletions langgraph/checkpoint/redis/aio.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,8 +44,11 @@
from langgraph.checkpoint.redis.base import (
CHECKPOINT_PREFIX,
CHECKPOINT_WRITE_PREFIX,
INDEXED_CHECKPOINT_FILTER_KEYS,
REDIS_KEY_SEPARATOR,
BaseRedisSaver,
_checkpoint_id_filter_matches,
_metadata_filter_matches,
)
from langgraph.checkpoint.redis.key_registry import (
AsyncCheckpointKeyRegistry as AsyncKeyRegistry,
Expand Down Expand Up @@ -638,9 +641,12 @@ async def alist(

# Search for checkpoints with any namespace, including an empty
# string, while `checkpoint_id` has to have a value.
if checkpoint_ns := config["configurable"].get("checkpoint_ns"):
if "checkpoint_ns" in config["configurable"]:
filter_expression.append(
Tag("checkpoint_ns") == to_storage_safe_str(checkpoint_ns)
Tag("checkpoint_ns")
== to_storage_safe_str(
config["configurable"].get("checkpoint_ns", "")
)
)
if checkpoint_id := get_checkpoint_id(config):
filter_expression.append(
Expand All @@ -657,20 +663,20 @@ async def alist(
filter_expression.append(Tag("thread_id") == to_storage_safe_id(v))
elif k == "run_id":
filter_expression.append(Tag("run_id") == to_storage_safe_id(v))
else:
raise ValueError(f"Unsupported filter key: {k}")

if before:
before_checkpoint_id = get_checkpoint_id(before)
if before_checkpoint_id:
try:
before_ulid = ULID.from_str(before_checkpoint_id)
before_ts = before_ulid.timestamp
# Use numeric range query: checkpoint_ts < before_ts
filter_expression.append(Num("checkpoint_ts") < before_ts)
except Exception:
# If not a valid ULID, ignore the before filter
pass
requires_post_filter = bool(
filter and any(k not in INDEXED_CHECKPOINT_FILTER_KEYS for k in filter)
)
before_checkpoint_id = get_checkpoint_id(before) if before else None

if before_checkpoint_id and not requires_post_filter:
try:
before_ulid = ULID.from_str(before_checkpoint_id)
before_ts = before_ulid.timestamp
# Use numeric range query: checkpoint_ts < before_ts
filter_expression.append(Num("checkpoint_ts") < before_ts)
except Exception:
requires_post_filter = True

# Combine all filter expressions
combined_filter = filter_expression[0] if filter_expression else "*"
Expand All @@ -679,30 +685,48 @@ async def alist(

# Construct the Redis query
# Sort by checkpoint_id in descending order to get most recent checkpoints first
query = FilterQuery(
filter_expression=combined_filter,
return_fields=[
"thread_id",
"checkpoint_ns",
"checkpoint_id",
"parent_checkpoint_id",
"$.checkpoint",
"$.metadata",
"has_writes", # Include has_writes to optimize pending_writes loading
],
num_results=limit or 10000,
sort_by=("checkpoint_id", "DESC"),
)
return_fields = [
"thread_id",
"checkpoint_ns",
"checkpoint_id",
"parent_checkpoint_id",
"$.checkpoint",
"$.metadata",
"has_writes", # Include has_writes to optimize pending_writes loading
]

# Execute the query asynchronously
results = await self.checkpoints_index.search(query)
if requires_post_filter:
result_docs = []
offset = 0
page_size = 10000
while True:
query = FilterQuery(
filter_expression=combined_filter,
return_fields=return_fields,
num_results=page_size,
sort_by=("checkpoint_id", "DESC"),
)
query.paging(offset, page_size)
results = await self.checkpoints_index.search(query)
result_docs.extend(results.docs)
if len(results.docs) < page_size:
break
offset += page_size
else:
query = FilterQuery(
filter_expression=combined_filter,
return_fields=return_fields,
num_results=limit or 10000,
sort_by=("checkpoint_id", "DESC"),
)
result_docs = (await self.checkpoints_index.search(query)).docs

# Pre-process all docs to collect batch query requirements
all_docs_data = []
pending_sends_batch_keys = []
pending_writes_batch_keys = []

for doc in results.docs:
for doc in result_docs:
# Extract all attributes once
doc_dict = doc.__dict__ if hasattr(doc, "__dict__") else {}

Expand Down Expand Up @@ -822,6 +846,12 @@ async def alist(
else:
metadata = cast(CheckpointMetadata, metadata_dict)

if requires_post_filter and (
not _checkpoint_id_filter_matches(checkpoint_id, before_checkpoint_id)
or not _metadata_filter_matches(metadata, filter)
):
continue

# Pre-create the config structure more efficiently
config_param: RunnableConfig = {
"configurable": {
Expand Down Expand Up @@ -866,6 +896,10 @@ async def alist(
parent_config=parent_config,
pending_writes=pending_writes,
)
if limit is not None:
limit -= 1
if limit <= 0:
break

async def aput(
self,
Expand Down Expand Up @@ -2086,7 +2120,7 @@ async def aprune(

# Validate inputs
if not thread_ids:
raise ValueError("``thread_ids`` must be a non-empty sequence")
return
if keep_last < 0:
raise ValueError(f"``keep_last`` must be >= 0, got {keep_last}")
if max_results < 1:
Expand Down
28 changes: 28 additions & 0 deletions langgraph/checkpoint/redis/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,34 @@
REDIS_KEY_SEPARATOR = ":"
CHECKPOINT_PREFIX = "checkpoint"
CHECKPOINT_WRITE_PREFIX = "checkpoint_write"
INDEXED_CHECKPOINT_FILTER_KEYS = {"source", "step", "thread_id", "run_id"}


def _metadata_filter_matches(
metadata: dict[str, Any], filters: Optional[dict[str, Any]]
) -> bool:
"""Return whether checkpoint metadata matches user-supplied filters."""
if not filters:
return True
for key, expected in filters.items():
if key in {"thread_id", "run_id"}:
continue
actual = metadata.get(key)
if isinstance(expected, list):
if actual not in expected:
return False
elif actual != expected:
return False
return True


def _checkpoint_id_filter_matches(
checkpoint_id: str, before_checkpoint_id: Optional[str]
) -> bool:
"""Apply list(before=...) using the same descending checkpoint_id order."""
if not before_checkpoint_id:
return True
return checkpoint_id < before_checkpoint_id


class BaseRedisSaver(BaseCheckpointSaver[str], Generic[RedisClientType, IndexType]):
Expand Down
Loading
Loading