Skip to content

Commit 17ad425

Browse files
committed
Handle missing config keys during checkpoint search
When queried for checkpoints of a root graph, the `config` object the checkpointer receives will not have a `checkpoint_id` or `checkpoint_ns` key. Currently, we reinterpret their absence as presence but with either an empty string or the string "None", neither of which will find anything when queried. One because we aren't saving the string "None," and the other because this index doesn't support querying for empty strings on these fields (would need the INDEXEMPTY option on fields, which RedisVL doesn't currently support and only recent versions of RediSearch support). This commit doesn't entirely solve the problem, however. Now, we are searching for any checkpoint for the thread, not checkpoints for the thread for the root graph. We need to store a sentinel value like "__empty__" and then query for it when the keys are not present in config.
1 parent fed21c6 commit 17ad425

File tree

4 files changed

+129
-19
lines changed

4 files changed

+129
-19
lines changed

langgraph/checkpoint/redis/__init__.py

Lines changed: 5 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -258,19 +258,14 @@ def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
258258
Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found.
259259
"""
260260
thread_id = config["configurable"]["thread_id"]
261-
checkpoint_id = str(get_checkpoint_id(config))
261+
checkpoint_id = get_checkpoint_id(config)
262262
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
263263

264+
checkpoint_filter_expression = Tag("thread_id") == thread_id
264265
if checkpoint_id:
265-
checkpoint_filter_expression = (
266-
(Tag("thread_id") == thread_id)
267-
& (Tag("checkpoint_ns") == checkpoint_ns)
268-
& (Tag("checkpoint_id") == checkpoint_id)
269-
)
270-
else:
271-
checkpoint_filter_expression = (Tag("thread_id") == thread_id) & (
272-
Tag("checkpoint_ns") == checkpoint_ns
273-
)
266+
checkpoint_filter_expression &= Tag("checkpoint_id") == str(checkpoint_id)
267+
if checkpoint_ns:
268+
checkpoint_filter_expression &= Tag("checkpoint_ns") == checkpoint_ns
274269

275270
# Construct the query
276271
checkpoints_query = FilterQuery(

langgraph/checkpoint/redis/aio.py

Lines changed: 4 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -132,16 +132,11 @@ async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
132132
checkpoint_id = get_checkpoint_id(config)
133133
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
134134

135+
checkpoint_filter_expression = Tag("thread_id") == thread_id
135136
if checkpoint_id:
136-
checkpoint_filter_expression = (
137-
(Tag("thread_id") == thread_id)
138-
& (Tag("checkpoint_ns") == checkpoint_ns)
139-
& (Tag("checkpoint_id") == checkpoint_id)
140-
)
141-
else:
142-
checkpoint_filter_expression = (Tag("thread_id") == thread_id) & (
143-
Tag("checkpoint_ns") == checkpoint_ns
144-
)
137+
checkpoint_filter_expression &= Tag("checkpoint_id") == str(checkpoint_id)
138+
if checkpoint_ns:
139+
checkpoint_filter_expression &= Tag("checkpoint_ns") == checkpoint_ns
145140

146141
# Construct the query
147142
checkpoints_query = FilterQuery(

tests/test_async.py

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -592,3 +592,65 @@ async def test_async_redis_checkpointer(
592592
checkpoints = [c async for c in checkpointer.alist(config)]
593593
assert len(checkpoints) > 0
594594
assert checkpoints[-1].checkpoint["id"] == latest["id"]
595+
596+
597+
@pytest.mark.requires_api_keys
598+
@pytest.mark.asyncio
599+
async def test_root_graph_checkpoint(
600+
redis_url: str, tools: List[BaseTool], model: ChatOpenAI
601+
) -> None:
602+
"""
603+
A regression test for a bug where queries for checkpoints from the
604+
root graph were failing to find valid checkpoints. When called from
605+
a root graph, the `checkpoint_id` and `checkpoint_ns` keys are not
606+
in the config object.
607+
"""
608+
609+
async with AsyncRedisSaver.from_conn_string(redis_url) as checkpointer:
610+
await checkpointer.asetup()
611+
# Create agent with checkpointer
612+
graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)
613+
614+
# Test initial query
615+
config: RunnableConfig = {
616+
"configurable": {
617+
"thread_id": "test1",
618+
"checkpoint_ns": "",
619+
"checkpoint_id": "",
620+
}
621+
}
622+
res = await graph.ainvoke(
623+
{"messages": [("human", "what's the weather in sf")]}, config
624+
)
625+
626+
assert res is not None
627+
628+
# Test checkpoint retrieval
629+
latest = await checkpointer.aget(config)
630+
631+
assert latest is not None
632+
assert all(
633+
k in latest
634+
for k in [
635+
"v",
636+
"ts",
637+
"id",
638+
"channel_values",
639+
"channel_versions",
640+
"versions_seen",
641+
]
642+
)
643+
assert "messages" in latest["channel_values"]
644+
assert (
645+
len(latest["channel_values"]["messages"]) == 4
646+
) # Initial + LLM + Tool + Final
647+
648+
# Test checkpoint tuple
649+
tuple_result = await checkpointer.aget_tuple(config)
650+
assert tuple_result is not None
651+
assert tuple_result.checkpoint == latest
652+
653+
# Test listing checkpoints
654+
checkpoints = [c async for c in checkpointer.alist(config)]
655+
assert len(checkpoints) > 0
656+
assert checkpoints[-1].checkpoint["id"] == latest["id"]

tests/test_sync.py

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -486,3 +486,61 @@ def test_sync_redis_checkpointer(
486486
checkpoints = list(checkpointer.list(config))
487487
assert len(checkpoints) > 0
488488
assert checkpoints[-1].checkpoint["id"] == latest["id"]
489+
490+
491+
@pytest.mark.requires_api_keys
492+
def test_root_graph_checkpoint(
493+
tools: list[BaseTool], model: ChatOpenAI, redis_url: str
494+
) -> None:
495+
"""
496+
A regression test for a bug where queries for checkpoints from the
497+
root graph were failing to find valid checkpoints. When called from
498+
a root graph, the `checkpoint_id` and `checkpoint_ns` keys are not
499+
in the config object.
500+
"""
501+
with RedisSaver.from_conn_string(redis_url) as checkpointer:
502+
checkpointer.setup()
503+
# Create agent with checkpointer
504+
graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)
505+
506+
# Test initial query
507+
config: RunnableConfig = {
508+
"configurable": {
509+
"thread_id": "test1",
510+
}
511+
}
512+
res = graph.invoke(
513+
{"messages": [("human", "what's the weather in sf")]}, config
514+
)
515+
516+
assert res is not None
517+
518+
# Test checkpoint retrieval
519+
latest = checkpointer.get(config)
520+
521+
assert latest is not None
522+
assert all(
523+
k in latest
524+
for k in [
525+
"v",
526+
"ts",
527+
"id",
528+
"channel_values",
529+
"channel_versions",
530+
"versions_seen",
531+
]
532+
)
533+
assert "messages" in latest["channel_values"]
534+
assert (
535+
len(latest["channel_values"]["messages"]) == 4
536+
) # Initial + LLM + Tool + Final
537+
538+
# Test checkpoint tuple
539+
tuple_result = checkpointer.get_tuple(config)
540+
assert tuple_result is not None
541+
assert tuple_result.checkpoint == latest
542+
543+
# Test listing checkpoints
544+
checkpoints = list(checkpointer.list(config))
545+
assert len(checkpoints) > 0
546+
assert checkpoints[-1].checkpoint["id"] == latest["id"]

0 commit comments

Comments
 (0)