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

Merge pull request #1550 from langchain-ai/nc/30aug/replay-without-fork

Browse Source 
Remove the need to fork a thread to replay a past state
This commit is contained in:
Nuno Campos
2024-08-30 10:39:57 -07:00
committed by GitHub
parent 99a845128e 6036b4a7d4
commit 36d1570c0f
5 changed files with 24 additions and 110 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/pregel/algo.py
+3 -9
View File
@@ -50,7 +50,7 @@ from langgraph.pregel.io import read_channel, read_channels
from langgraph.pregel.log import logger
from langgraph.pregel.manager import ChannelsManager
from langgraph.pregel.read import PregelNode
from langgraph.pregel.types import EXACT_MATCH, All, PregelExecutableTask, PregelTask
from langgraph.pregel.types import All, PregelExecutableTask, PregelTask
class WritesProtocol(Protocol):
@@ -357,10 +357,7 @@ def prepare_next_tasks(
),
triggers,
proc.retry_policy,
None
if task_checkpoint_ns
in configurable.get(CONFIG_KEY_CHECKPOINT_MAP, {})
else EXACT_MATCH,
None,
task_id,
)
)
@@ -466,10 +463,7 @@ def prepare_next_tasks(
),
triggers,
proc.retry_policy,
None
if task_checkpoint_ns
in configurable.get(CONFIG_KEY_CHECKPOINT_MAP, {})
else EXACT_MATCH,
None,
task_id,
)
)
libs/langgraph/langgraph/pregel/loop.py
+4 -9
View File
@@ -122,6 +122,7 @@ class PregelLoop:
output_keys: Union[str, Sequence[str]]
stream_keys: Union[str, Sequence[str]]
is_nested: bool
skip_done_tasks: bool
checkpointer_get_next_version: Callable[[Optional[V]], V]
checkpointer_put_writes: Optional[
@@ -180,6 +181,7 @@ class PregelLoop:
self.output_keys = output_keys
self.stream_keys = stream_keys
self.is_nested = CONFIG_KEY_TASK_ID in self.config.get("configurable", {})
self.skip_done_tasks = "checkpoint_id" not in config["configurable"]
if CONFIG_KEY_STREAM in config["configurable"]:
self.stream = DuplexStream(
self.stream, config["configurable"][CONFIG_KEY_STREAM]
@@ -340,18 +342,11 @@ class PregelLoop:
return False
# if there are pending writes from a previous loop, apply them
if self.checkpoint_pending_writes:
if self.skip_done_tasks and self.checkpoint_pending_writes:
for tid, k, v in self.checkpoint_pending_writes:
if k in (ERROR, INTERRUPT):
continue
if task := next(
(
t
for t in self.tasks
if t.id == tid and t.cache_policy is not None
),
None,
):
if task := next((t for t in self.tasks if t.id == tid), None):
task.writes.append((k, v))
# print output for any tasks we applied previous writes to
for task in self.tasks:
libs/langgraph/langgraph/pregel/types.py
-3
View File
@@ -63,9 +63,6 @@ class CachePolicy(NamedTuple):
pass
EXACT_MATCH = CachePolicy()
class PregelTask(NamedTuple):
id: str
name: str
libs/langgraph/tests/test_pregel.py
+9 -42
View File
@@ -798,27 +798,13 @@ def test_invoke_two_processes_in_out_interrupt(
),
]
# re-running from any previous checkpoint w/out forking should do nothing
# re-running from any previous checkpoint should re-run nodes
assert [c for c in app.stream(None, history[0].config, stream_mode="updates")] == []
assert [c for c in app.stream(None, history[1].config, stream_mode="updates")] == [
{"two": {"output": 5}, "__metadata__": {"cached": True}},
{"two": {"output": 5}},
]
assert [c for c in app.stream(None, history[2].config, stream_mode="updates")] == [
{"one": {"inbox": 4}, "__metadata__": {"cached": True}},
]
# forking and re-running from any prev checkpoint should re-run nodes
fork_config = app.update_state(history[0].config, None)
assert [c for c in app.stream(None, fork_config, stream_mode="updates")] == []
fork_config = app.update_state(history[1].config, None)
assert [c for c in app.stream(None, fork_config, stream_mode="updates")] == [
{"two": {"output": 5}}
]
fork_config = app.update_state(history[2].config, None)
assert [c for c in app.stream(None, fork_config, stream_mode="updates")] == [
{"one": {"inbox": 4}}
{"one": {"inbox": 4}},
]
@@ -987,25 +973,18 @@ def test_fork_always_re_runs_nodes(
),
]
# forking from any previous checkpoint w/out forking should do nothing
# forking from any previous checkpoint should re-run nodes
assert [
c for c in graph.stream(None, history[0].config, stream_mode="updates")
] == []
assert [
c for c in graph.stream(None, history[1].config, stream_mode="updates")
] == [{"add_one": 1, "__metadata__": {"cached": True}}]
# forking and re-running from any prev checkpoint should re-run nodes
fork_config = graph.update_state(history[0].config, None)
assert [c for c in graph.stream(None, fork_config, stream_mode="updates")] == []
fork_config = graph.update_state(history[1].config, None)
assert [c for c in graph.stream(None, fork_config, stream_mode="updates")] == [
{"add_one": 1}
] == [
{"add_one": 1},
]
fork_config = graph.update_state(history[2].config, None)
assert [c for c in graph.stream(None, fork_config, stream_mode="updates")] == [
assert [
c for c in graph.stream(None, history[2].config, stream_mode="updates")
] == [
{"add_one": 1},
{"add_one": 1},
]
@@ -9813,18 +9792,6 @@ def test_doubly_nested_graph_state(
assert [
c for c in app.stream(None, grandchild_history[2].config, subgraphs=True)
] == [
(
(AnyStr("child:"), AnyStr("child_1:")),
{
"grandchild_1": {"my_key": "hi my value here"},
"__metadata__": {"cached": True},
},
)
]
# fork and replay
fork = app.update_state(grandchild_history[2].config, None)
assert [c for c in app.stream(None, fork, subgraphs=True)] == [
(
(AnyStr("child:"), AnyStr("child_1:")),
{"grandchild_1": {"my_key": "hi my value here"}},
libs/langgraph/tests/test_pregel_async.py
+8 -47
View File
@@ -1020,35 +1020,19 @@ async def test_invoke_two_processes_in_out_interrupt(
),
]
# forking from any previous checkpoint w/out forking should do nothing
# forking from any previous checkpoint should re-run nodes
assert [
c async for c in app.astream(None, history[0].config, stream_mode="updates")
] == []
assert [
c async for c in app.astream(None, history[1].config, stream_mode="updates")
] == [
{"two": {"output": 5}, "__metadata__": {"cached": True}},
{"two": {"output": 5}},
]
assert [
c async for c in app.astream(None, history[2].config, stream_mode="updates")
] == [
{"one": {"inbox": 4}, "__metadata__": {"cached": True}},
]
# forking and re-running from any prev checkpoint should re-run nodes
fork_config = await app.aupdate_state(history[0].config, None)
assert [
c async for c in app.astream(None, fork_config, stream_mode="updates")
] == []
fork_config = await app.aupdate_state(history[1].config, None)
assert [c async for c in app.astream(None, fork_config, stream_mode="updates")] == [
{"two": {"output": 5}}
]
fork_config = await app.aupdate_state(history[2].config, None)
assert [c async for c in app.astream(None, fork_config, stream_mode="updates")] == [
{"one": {"inbox": 4}}
{"one": {"inbox": 4}},
]
@@ -1219,28 +1203,17 @@ async def test_fork_always_re_runs_nodes(
),
]
# forking from any previous checkpoint w/out forking should do nothing
# forking from any previous checkpoint should re-run nodes
assert [
c async for c in graph.astream(None, history[0].config, stream_mode="updates")
] == []
assert [
c async for c in graph.astream(None, history[1].config, stream_mode="updates")
] == [{"add_one": 1, "__metadata__": {"cached": True}}]
# forking and re-running from any prev checkpoint should re-run nodes
fork_config = await graph.aupdate_state(history[0].config, None)
] == [
{"add_one": 1},
]
assert [
c async for c in graph.astream(None, fork_config, stream_mode="updates")
] == []
fork_config = await graph.aupdate_state(history[1].config, None)
assert [
c async for c in graph.astream(None, fork_config, stream_mode="updates")
] == [{"add_one": 1}]
fork_config = await graph.aupdate_state(history[2].config, None)
assert [
c async for c in graph.astream(None, fork_config, stream_mode="updates")
c async for c in graph.astream(None, history[2].config, stream_mode="updates")
] == [
{"add_one": 1},
{"add_one": 1},
@@ -8340,18 +8313,6 @@ async def test_doubly_nested_graph_state(
assert [
c async for c in app.astream(None, grandchild_history[2].config, subgraphs=True)
] == [
(
(AnyStr("child:"), AnyStr("child_1:")),
{
"grandchild_1": {"my_key": "hi my value here"},
"__metadata__": {"cached": True},
},
)
]
# fork and replay
fork = await app.aupdate_state(grandchild_history[2].config, None)
assert [c async for c in app.astream(None, fork, subgraphs=True)] == [
(
(AnyStr("child:"), AnyStr("child_1:")),
{"grandchild_1": {"my_key": "hi my value here"}},