WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-08 02:37:52 +02:00
Code Issues Packages Projects Releases Wiki Activity

Checkpoint inputs before starting the first step for easier error recovery

Browse Source 
- this enables easier retrying, for any error just do .invoke(None, config) no matter which step the error happened on
This commit is contained in:
Nuno Campos
2024-05-06 11:53:20 -07:00
parent f160f82912
commit 611ecdb1cd
5 changed files with 175 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
langgraph/checkpoint/sqlite.py
+10 -2
View File
@@ -1,5 +1,6 @@
import pickle
import sqlite3
import threading
from contextlib import AbstractContextManager, contextmanager
from types import TracebackType
from typing import Any, Iterator, Optional
@@ -94,6 +95,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
super().__init__(serde=serde)
self.conn = conn
self.is_setup = False
self.lock = threading.Lock()
@classmethod
def from_conn_string(cls, conn_string: str) -> "SqliteSaver":
@@ -115,7 +117,13 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
memory = SqliteSaver.from_conn_string("checkpoints.sqlite")
"""
return SqliteSaver(conn=sqlite3.connect(conn_string))
return SqliteSaver(
conn=sqlite3.connect(
conn_string,
# https://ricardoanderegg.com/posts/python-sqlite-thread-safety/
check_same_thread=False,
)
)
def __enter__(self) -> Self:
return self
@@ -348,7 +356,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
saved_config
) # Output: {"configurable": {"thread_id": "1", "thread_ts": 2024-05-04T06:32:42.235444+00:00"}}
"""
with self.cursor() as cur:
with self.lock, self.cursor() as cur:
cur.execute(
"INSERT OR REPLACE INTO checkpoints (thread_id, thread_ts, parent_ts, checkpoint, metadata) VALUES (?, ?, ?, ?, ?)",
(
langgraph/pregel/__init__.py
+100 -32
View File
@@ -342,7 +342,7 @@ class Pregel(
read_channels(channels, self.stream_channels_asis),
tuple(name for name, _ in next_tasks),
config,
saved.metadata,
saved.metadata if saved else None,
)
async def aget_state(self, config: RunnableConfig) -> StateSnapshot:
@@ -361,7 +361,7 @@ class Pregel(
read_channels(channels, self.stream_channels_asis),
tuple(name for name, _ in next_tasks),
config,
saved.metadata,
saved.metadata if saved else None,
)
def get_state_history(
@@ -623,6 +623,7 @@ class Pregel(
run_id=config.get("run_id"),
)
try:
bg: list[concurrent.futures.Future] = []
if config["recursion_limit"] < 1:
raise ValueError("recursion_limit must be at least 1")
if self.checkpointer and not config.get("configurable"):
@@ -656,6 +657,7 @@ class Pregel(
else None
)
checkpoint = saved.checkpoint if saved else empty_checkpoint()
start = saved.metadata.get("step", -2) + 1 if saved else -1
# create channels from checkpoint
with ChannelsManager(
self.channels, checkpoint
@@ -668,6 +670,27 @@ class Pregel(
)
# apply input writes
_apply_writes(checkpoint, channels, input_writes)
# save input checkpoint
if self.checkpointer is not None:
checkpoint = create_checkpoint(checkpoint, channels)
bg.append(
executor.submit(
self.checkpointer.put,
checkpoint_config,
copy_checkpoint(checkpoint),
{"source": "input", "step": start},
)
)
checkpoint_config = {
"configurable": {
"thread_id": checkpoint_config["configurable"][
"thread_id"
],
"thread_ts": checkpoint["ts"],
}
}
# increment start to 0
start += 1
else:
# if received no input, take that as signal to proceed
# past previous interrupt, if any
@@ -681,7 +704,6 @@ class Pregel(
# channel updates from step N are only visible in step N+1
# channels are guaranteed to be immutable for the duration of the step,
# with channel updates applied only at the transition between steps
start = saved.metadata.get("step", -1) + 1 if saved else 0
stop = start + config["recursion_limit"] + 1
for step in range(start, stop):
next_checkpoint, next_tasks = _prepare_next_tasks(
@@ -786,21 +808,29 @@ class Pregel(
# save end of step checkpoint
if self.checkpointer is not None:
checkpoint = create_checkpoint(checkpoint, channels)
checkpoint_config = self.checkpointer.put(
checkpoint_config,
checkpoint,
{"source": "loop", "step": step},
)
if stream_mode == "debug":
yield map_debug_checkpoint(
step,
bg.append(
executor.submit(
self.checkpointer.put,
checkpoint_config,
channels,
self.stream_channels_asis,
copy_checkpoint(checkpoint),
{"source": "loop", "step": step},
)
elif stream_mode == "debug":
)
checkpoint_config = {
"configurable": {
"thread_id": checkpoint_config["configurable"][
"thread_id"
],
"thread_ts": checkpoint["ts"],
}
}
# yield debug checkpoint
if stream_mode == "debug":
yield map_debug_checkpoint(
step, None, channels, self.stream_channels_asis
step,
checkpoint_config if self.checkpointer else None,
channels,
self.stream_channels_asis,
)
# after execution, check if we should interrupt
@@ -824,6 +854,12 @@ class Pregel(
task.cancel()
except NameError:
pass
# wait for all background tasks to finish
done, _ = concurrent.futures.wait(
bg, return_when=concurrent.futures.ALL_COMPLETED
)
for task in done:
task.result()
async def astream(
self,
@@ -855,7 +891,7 @@ class Pregel(
None,
)
try:
tasks: list[asyncio.Task] = []
bg: list[asyncio.Task] = []
if config["recursion_limit"] < 1:
raise ValueError("recursion_limit must be at least 1")
if self.checkpointer and not config.get("configurable"):
@@ -889,6 +925,7 @@ class Pregel(
else None
)
checkpoint = saved.checkpoint if saved else empty_checkpoint()
start = saved.metadata.get("step", -2) + 1 if saved else -1
# create channels from checkpoint
async with AsyncChannelsManager(self.channels, checkpoint) as channels:
# map inputs to channel updates
@@ -899,6 +936,28 @@ class Pregel(
)
# apply input writes
_apply_writes(checkpoint, channels, input_writes)
# save input checkpoint
if self.checkpointer is not None:
checkpoint = create_checkpoint(checkpoint, channels)
bg.append(
asyncio.create_task(
self.checkpointer.aput(
checkpoint_config,
copy_checkpoint(checkpoint),
{"source": "input", "step": start},
)
)
)
checkpoint_config = {
"configurable": {
"thread_id": checkpoint_config["configurable"][
"thread_id"
],
"thread_ts": checkpoint["ts"],
}
}
# increment start to 0
start += 1
else:
# if received no input, take that as signal to proceed
# past previous interrupt, if any
@@ -1027,21 +1086,30 @@ class Pregel(
# save end of step checkpoint
if self.checkpointer is not None:
checkpoint = create_checkpoint(checkpoint, channels)
checkpoint_config = await self.checkpointer.aput(
checkpoint_config,
checkpoint,
{"source": "loop", "step": step},
)
if stream_mode == "debug":
yield map_debug_checkpoint(
step,
checkpoint_config,
channels,
self.stream_channels_asis,
bg.append(
asyncio.create_task(
self.checkpointer.aput(
checkpoint_config,
checkpoint,
{"source": "loop", "step": step},
)
)
elif stream_mode == "debug":
)
checkpoint_config = {
"configurable": {
"thread_id": checkpoint_config["configurable"][
"thread_id"
],
"thread_ts": checkpoint["ts"],
}
}
# yield debug checkpoint
if stream_mode == "debug":
yield map_debug_checkpoint(
step, None, channels, self.stream_channels_asis
step,
checkpoint_config if self.checkpointer else None,
channels,
self.stream_channels_asis,
)
# after execution, check if we should interrupt
@@ -1063,11 +1131,11 @@ class Pregel(
try:
for task in futures:
task.cancel()
tasks.append(task)
bg.append(task)
except NameError:
pass
# wait for all tasks to finish
await asyncio.gather(*tasks, return_exceptions=True)
# wait for all background tasks to finish
await asyncio.gather(*bg)
def invoke(
self,
langgraph/pregel/types.py
+3 -1
View File
@@ -3,6 +3,8 @@ from typing import Any, Literal, NamedTuple, Optional, Union
from langchain_core.runnables import Runnable, RunnableConfig
from langgraph.checkpoint.base import CheckpointMetadata
class PregelTaskDescription(NamedTuple):
name: str
@@ -25,7 +27,7 @@ class StateSnapshot(NamedTuple):
"""Nodes to execute in the next step, if any"""
config: RunnableConfig
"""Config used to fetch this snapshot"""
metadata: dict[str, Any]
metadata: CheckpointMetadata
"""Metadata associated with this snapshot"""
parent_config: Optional[RunnableConfig] = None
"""Config used to fetch the parent snapshot, if any"""
tests/test_pregel.py
+34 -13
View File
@@ -2,6 +2,7 @@ import json
import operator
import time
import warnings
from collections import Counter
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager
from typing import Annotated, Any, Generator, Literal, Optional, TypedDict, Union
@@ -672,7 +673,7 @@ def test_invoke_checkpoint(mocker: MockerFixture) -> None:
def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"])
adder = mocker.Mock(side_effect=lambda x: x["total"] + x["input"])
def raise_if_above_10(input: int) -> int:
if input > 10:
@@ -681,7 +682,7 @@ def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None:
one = (
Channel.subscribe_to(["input"]).join(["total"])
| add_one
| adder
| Channel.write_to("output", "total")
| raise_if_above_10
)
@@ -701,10 +702,11 @@ def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None:
thread_1 = {"configurable": {"thread_id": "1"}}
# total starts out as 0, so output is 0+2=2
assert app.invoke(2, thread_1) == 2
assert app.invoke(2, thread_1, debug=1) == 2
state = app.get_state(thread_1)
assert state is not None
assert state.values.get("total") == 2
assert state.next == ()
assert state.config["configurable"]["thread_ts"] == memory.get(thread_1)["ts"]
# total is now 2, so output is 2+3=5
assert app.invoke(3, thread_1) == 5
@@ -719,14 +721,22 @@ def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None:
state = app.get_state(thread_1)
assert state is not None
assert state.values.get("total") == 7
assert state.next == ("one",)
"""we checkpoint inputs and it failed on "one", so the next node is one"""
# we can recover from error by sending new inputs
assert app.invoke(2, thread_1) == 9
state = app.get_state(thread_1)
assert state is not None
assert state.values.get("total") == 16, "total is now 7+9=16"
assert state.next == ()
thread_2 = {"configurable": {"thread_id": "2"}}
# on a new thread, total starts out as 0, so output is 0+5=5
assert app.invoke(5, thread_2) == 5
assert app.invoke(5, thread_2, debug=True) == 5
state = app.get_state({"configurable": {"thread_id": "1"}})
assert state is not None
assert state.values.get("total") == 7
assert state.next == ()
assert state.values.get("total") == 16
assert state.next == (), "checkpoint of other thread not touched"
state = app.get_state(thread_2)
assert state is not None
assert state.values.get("total") == 5
@@ -735,8 +745,12 @@ def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None:
assert len(list(app.get_state_history(thread_1, limit=1))) == 1
# list all checkpoints for thread 1
thread_1_history = [c for c in app.get_state_history(thread_1)]
# there are 2: one for each successful ainvoke()
assert len(thread_1_history) == 2
# there are 7 checkpoints
assert len(thread_1_history) == 7
assert Counter(c.metadata["source"] for c in thread_1_history) == {
"input": 4,
"loop": 3,
}
# sorted descending
assert (
thread_1_history[0].config["configurable"]["thread_ts"]
@@ -748,10 +762,10 @@ def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None:
)
assert len(cursored) == 1
assert cursored[0].config == thread_1_history[1].config
# the second checkpoint
assert thread_1_history[0].values["total"] == 7
# the first checkpoint
assert thread_1_history[1].values["total"] == 2
# the last checkpoint
assert thread_1_history[0].values["total"] == 16
# the first "loop" checkpoint
assert thread_1_history[-2].values["total"] == 2
# can get each checkpoint using aget with config
assert (
memory.get(thread_1_history[0].config)["ts"]
@@ -769,7 +783,14 @@ def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None:
> thread_1_history[0].config["configurable"]["thread_ts"]
)
# 1 more checkpoint in history
assert len(list(app.get_state_history(thread_1))) == 3
assert len(list(app.get_state_history(thread_1))) == 8
assert Counter(
c.metadata["source"] for c in app.get_state_history(thread_1)
) == {
"update": 1,
"input": 4,
"loop": 3,
}
# the latest checkpoint is the updated one
assert app.get_state(thread_1) == app.get_state(thread_1_next_config)
tests/test_pregel_async.py
+28 -8
View File
@@ -1,6 +1,7 @@
import asyncio
import json
import operator
from collections import Counter
from contextlib import asynccontextmanager, contextmanager
from typing import (
Annotated,
@@ -712,13 +713,21 @@ async def test_invoke_checkpoint_aiosqlite(mocker: MockerFixture) -> None:
state = await app.aget_state(thread_1)
assert state is not None
assert state.values.get("total") == 7
assert state.next == ("one",)
"""we checkpoint inputs and it failed on "one", so the next node is one"""
# we can recover from error by sending new inputs
assert await app.ainvoke(2, thread_1) == 9
state = await app.aget_state(thread_1)
assert state is not None
assert state.values.get("total") == 16, "total is now 7+9=16"
assert state.next == ()
thread_2 = {"configurable": {"thread_id": "2"}}
# on a new thread, total starts out as 0, so output is 0+5=5
assert await app.ainvoke(5, thread_2) == 5
state = await app.aget_state({"configurable": {"thread_id": "1"}})
assert state is not None
assert state.values.get("total") == 7
assert state.values.get("total") == 16
assert state.next == ()
state = await app.aget_state(thread_2)
assert state is not None
@@ -728,8 +737,12 @@ async def test_invoke_checkpoint_aiosqlite(mocker: MockerFixture) -> None:
assert len([c async for c in app.aget_state_history(thread_1, limit=1)]) == 1
# list all checkpoints for thread 1
thread_1_history = [c async for c in app.aget_state_history(thread_1)]
# there are 2: one for each successful ainvoke()
assert len(thread_1_history) == 2
# there are 7 checkpoints
assert len(thread_1_history) == 7
assert Counter(c.metadata["source"] for c in thread_1_history) == {
"input": 4,
"loop": 3,
}
# sorted descending
assert (
thread_1_history[0].config["configurable"]["thread_ts"]
@@ -744,10 +757,10 @@ async def test_invoke_checkpoint_aiosqlite(mocker: MockerFixture) -> None:
]
assert len(cursored) == 1
assert cursored[0].config == thread_1_history[1].config
# the second checkpoint
assert thread_1_history[0].values["total"] == 7
# the first checkpoint
assert thread_1_history[1].values["total"] == 2
# the last checkpoint
assert thread_1_history[0].values["total"] == 16
# the first "loop" checkpoint
assert thread_1_history[-2].values["total"] == 2
# can get each checkpoint using aget with config
assert (await memory.aget(thread_1_history[0].config))[
"ts"
@@ -763,7 +776,14 @@ async def test_invoke_checkpoint_aiosqlite(mocker: MockerFixture) -> None:
> thread_1_history[0].config["configurable"]["thread_ts"]
)
# 1 more checkpoint in history
assert len([h async for h in app.aget_state_history(thread_1)]) == 3
assert len([c async for c in app.aget_state_history(thread_1)]) == 8
assert Counter(
[c.metadata["source"] async for c in app.aget_state_history(thread_1)]
) == {
"update": 1,
"input": 4,
"loop": 3,
}
# the latest checkpoint is the updated one
assert await app.aget_state(thread_1) == await app.aget_state(
thread_1_next_config