diff --git a/libs/langgraph/langgraph/checkpoint/base.py b/libs/langgraph/langgraph/checkpoint/base.py index b6694db89..fc53833c1 100644 --- a/libs/langgraph/langgraph/checkpoint/base.py +++ b/libs/langgraph/langgraph/checkpoint/base.py @@ -231,6 +231,11 @@ class BaseCheckpointSaver(ABC): """ raise NotImplementedError + def list_subgraph_checkpoints( + self, config: RunnableConfig + ) -> Iterator[CheckpointTuple]: + raise NotImplementedError + def put( self, config: RunnableConfig, diff --git a/libs/langgraph/langgraph/checkpoint/memory.py b/libs/langgraph/langgraph/checkpoint/memory.py index 72b8c93db..55246be4b 100644 --- a/libs/langgraph/langgraph/checkpoint/memory.py +++ b/libs/langgraph/langgraph/checkpoint/memory.py @@ -171,6 +171,39 @@ class MemorySaver(BaseCheckpointSaver): else None, ) + def list_subgraph_checkpoints( + self, config: RunnableConfig + ) -> Iterator[CheckpointTuple]: + thread_id_prefix = config["configurable"]["thread_id"] + matching_thread_ids = [ + key for key in self.storage.keys() if key.startswith(thread_id_prefix) + ] + for thread_id in matching_thread_ids: + ts = config["configurable"].get("thread_ts") + if not ts: + if checkpoints := self.storage[thread_id]: + ts = max(checkpoints.keys()) + + if saved := self.storage[thread_id].get(ts): + checkpoint, metadata, parent_ts = saved + writes = self.writes[(thread_id, ts)] + yield CheckpointTuple( + config={"configurable": {"thread_id": thread_id, "thread_ts": ts}}, + checkpoint=self.serde.loads(checkpoint), + metadata=self.serde.loads(metadata), + pending_writes=[ + (id, c, self.serde.loads(v)) for id, c, v in writes + ], + parent_config={ + "configurable": { + "thread_id": thread_id, + "thread_ts": parent_ts, + } + } + if parent_ts + else None, + ) + def put( self, config: RunnableConfig, diff --git a/libs/langgraph/langgraph/checkpoint/sqlite.py b/libs/langgraph/langgraph/checkpoint/sqlite.py index eee6e05c7..f6bf4b857 100644 --- a/libs/langgraph/langgraph/checkpoint/sqlite.py +++ b/libs/langgraph/langgraph/checkpoint/sqlite.py @@ -357,6 +357,49 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager): ), ) + def list_subgraph_checkpoints( + self, config: RunnableConfig + ) -> Iterator[CheckpointTuple]: + with self.cursor(transaction=False) as cur: + if config["configurable"].get("thread_ts"): + cur.execute( + "SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id LIKE ? || '%' AND thread_ts = ?", + ( + str(config["configurable"]["thread_id"]), + str(config["configurable"]["thread_ts"]), + ), + ) + else: + cur.execute( + """SELECT checkpoints.thread_id, checkpoints.thread_ts, checkpoints.parent_ts, checkpoints.checkpoint, checkpoints.metadata + FROM checkpoints + INNER JOIN ( + SELECT thread_id, MAX(thread_ts) as thread_ts + FROM checkpoints + WHERE thread_id LIKE ? || '%' + GROUP BY thread_id + ) latest_checkpoints + ON checkpoints.thread_id = latest_checkpoints.thread_id AND checkpoints.thread_ts = latest_checkpoints.thread_ts + ORDER BY checkpoints.thread_id, checkpoints.thread_ts DESC""", + (str(config["configurable"]["thread_id"]),), + ) + for thread_id, thread_ts, parent_ts, value, metadata in cur: + yield CheckpointTuple( + {"configurable": {"thread_id": thread_id, "thread_ts": thread_ts}}, + self.serde.loads(value), + self.serde.loads(metadata) if metadata is not None else {}, + ( + { + "configurable": { + "thread_id": thread_id, + "thread_ts": parent_ts, + } + } + if parent_ts + else None + ), + ) + def put( self, config: RunnableConfig, diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index f3aeb6a2e..edab6d6b7 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -18,6 +18,8 @@ RESERVED = { } TAG_HIDDEN = "langsmith:hidden" +THREAD_ID_SEPARATOR = "__" + START = "__start__" END = "__end__" diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index c0edcacc3..c04f3e9e8 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -59,6 +59,7 @@ from langgraph.channels.manager import ( ) from langgraph.checkpoint.base import ( BaseCheckpointSaver, + CheckpointTuple, copy_checkpoint, empty_checkpoint, ) @@ -68,6 +69,7 @@ from langgraph.constants import ( CONFIG_KEY_RESUMING, CONFIG_KEY_SEND, INTERRUPT, + THREAD_ID_SEPARATOR, ) from langgraph.errors import GraphRecursionError, InvalidUpdateError from langgraph.managed.base import ( @@ -350,12 +352,9 @@ class Pregel( if is_managed_value(v) } - def get_state(self, config: RunnableConfig) -> StateSnapshot: - """Get the current state of the graph.""" - if not self.checkpointer: - raise ValueError("No checkpointer set") - - saved = self.checkpointer.get_tuple(config) + def _prepare_state_snapshot( + self, saved: CheckpointTuple, config: RunnableConfig + ) -> StateSnapshot: checkpoint = saved.checkpoint if saved else empty_checkpoint() config = saved.config if saved else config with ChannelsManager( @@ -373,14 +372,58 @@ class Pregel( for_execution=False, ) return StateSnapshot( - read_channels(channels, self.stream_channels_asis), - tuple(name for name, _ in next_tasks), - saved.config if saved else config, - saved.metadata if saved else None, - saved.checkpoint["ts"] if saved else None, - saved.parent_config if saved else None, + values=read_channels(channels, self.stream_channels_asis), + next=tuple(name for name, _ in next_tasks), + config=saved.config if saved else config, + metadata=saved.metadata if saved else None, + created_at=saved.checkpoint["ts"] if saved else None, + parent_config=saved.parent_config if saved else None, ) + @staticmethod + def _assemble_state_snapshot_hierarchy( + root_thread_id: str, subgraph_state_snapshots: dict[str, StateSnapshot] + ) -> StateSnapshot: + thread_ids_to_visit = sorted( + subgraph_state_snapshots.keys(), + key=lambda x: len(x.split(THREAD_ID_SEPARATOR)), + ) + while thread_ids_to_visit: + thread_id = thread_ids_to_visit.pop() + state_snapshot = subgraph_state_snapshots[thread_id] + *path, subgraph_node = thread_id.split(THREAD_ID_SEPARATOR) + parent_thread_id = THREAD_ID_SEPARATOR.join(path) + if parent_thread_id and THREAD_ID_SEPARATOR in parent_thread_id: + parent_subgraph_snapshots = ( + subgraph_state_snapshots[parent_thread_id].subgraph_state_snapshots + or {} + ) + parent_subgraph_snapshots[subgraph_node] = state_snapshot + subgraph_state_snapshots[parent_thread_id] = subgraph_state_snapshots[ + parent_thread_id + ]._replace(subgraph_state_snapshots=parent_subgraph_snapshots) + + state_snapshot = subgraph_state_snapshots.pop(root_thread_id) + return state_snapshot + + def get_state(self, config: RunnableConfig) -> StateSnapshot: + """Get the current state of the graph.""" + if not self.checkpointer: + raise ValueError("No checkpointer set") + + subgraph_state_snapshots: dict[str, StateSnapshot] = { + checkpoint.config["configurable"][ + "thread_id" + ]: self._prepare_state_snapshot(checkpoint, config) + for checkpoint in self.checkpointer.list_subgraph_checkpoints(config) + } + + thread_id = config["configurable"]["thread_id"] + state_snapshot = self._assemble_state_snapshot_hierarchy( + thread_id, subgraph_state_snapshots + ) + return state_snapshot + async def aget_state(self, config: RunnableConfig) -> StateSnapshot: """Get the current state of the graph.""" if not self.checkpointer: diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 13ce7e183..00d9e7809 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -36,6 +36,7 @@ from langgraph.constants import ( RESERVED, TAG_HIDDEN, TASKS, + THREAD_ID_SEPARATOR, Send, ) from langgraph.errors import EmptyChannelError, InvalidUpdateError @@ -345,7 +346,9 @@ def prepare_next_tasks( if parent_thread_id := config.get("configurable", {}).get( "thread_id" ): - thread_id: Optional[str] = f"{parent_thread_id}-{name}" + thread_id: Optional[ + str + ] = f"{parent_thread_id}{THREAD_ID_SEPARATOR}{name}" else: thread_id = None writes = deque() diff --git a/libs/langgraph/langgraph/pregel/types.py b/libs/langgraph/langgraph/pregel/types.py index 19d3c3301..d0648cf3b 100644 --- a/libs/langgraph/langgraph/pregel/types.py +++ b/libs/langgraph/langgraph/pregel/types.py @@ -85,6 +85,8 @@ class StateSnapshot(NamedTuple): """Timestamp of snapshot creation""" parent_config: Optional[RunnableConfig] = None """Config used to fetch the parent snapshot, if any""" + subgraph_state_snapshots: Optional[dict[str, "StateSnapshot"]] = None + """State snapshots of subgraphs represented as a mapping from thread ID suffix to snapshot.""" All = Literal["*"] diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index db902eca0..b218fabda 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -568,53 +568,53 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: ) # start execution, stop at inbox - assert app.invoke(2, {"configurable": {"thread_id": 1}}) is None + assert app.invoke(2, {"configurable": {"thread_id": "1"}}) is None # inbox == 3 - checkpoint = memory.get({"configurable": {"thread_id": 1}}) + checkpoint = memory.get({"configurable": {"thread_id": "1"}}) assert checkpoint is not None assert checkpoint["channel_values"]["inbox"] == 3 # resume execution, finish - assert app.invoke(None, {"configurable": {"thread_id": 1}}) == 4 + assert app.invoke(None, {"configurable": {"thread_id": "1"}}) == 4 # start execution again, stop at inbox - assert app.invoke(20, {"configurable": {"thread_id": 1}}) is None + assert app.invoke(20, {"configurable": {"thread_id": "1"}}) is None # inbox == 21 - checkpoint = memory.get({"configurable": {"thread_id": 1}}) + checkpoint = memory.get({"configurable": {"thread_id": "1"}}) assert checkpoint is not None assert checkpoint["channel_values"]["inbox"] == 21 # send a new value in, interrupting the previous execution - assert app.invoke(3, {"configurable": {"thread_id": 1}}) is None - assert app.invoke(None, {"configurable": {"thread_id": 1}}) == 5 + assert app.invoke(3, {"configurable": {"thread_id": "1"}}) is None + assert app.invoke(None, {"configurable": {"thread_id": "1"}}) == 5 # start execution again, stopping at inbox - assert app.invoke(20, {"configurable": {"thread_id": 2}}) is None + assert app.invoke(20, {"configurable": {"thread_id": "2"}}) is None # inbox == 21 - snapshot = app.get_state({"configurable": {"thread_id": 2}}) + snapshot = app.get_state({"configurable": {"thread_id": "2"}}) assert snapshot.values["inbox"] == 21 assert snapshot.next == ("two",) # update the state, resume - app.update_state({"configurable": {"thread_id": 2}}, 25, as_node="one") - assert app.invoke(None, {"configurable": {"thread_id": 2}}) == 26 + app.update_state({"configurable": {"thread_id": "2"}}, 25, as_node="one") + assert app.invoke(None, {"configurable": {"thread_id": "2"}}) == 26 # no pending tasks - snapshot = app.get_state({"configurable": {"thread_id": 2}}) + snapshot = app.get_state({"configurable": {"thread_id": "2"}}) assert snapshot.next == () # list history - thread1 = {"configurable": {"thread_id": 1}} + thread1 = {"configurable": {"thread_id": "1"}} assert [c for c in app.get_state_history(thread1)] == [ StateSnapshot( values={"inbox": 4, "output": 5, "input": 3}, next=(), config={ "configurable": { - "thread_id": 1, + "thread_id": "1", "thread_ts": AnyStr(), } }, @@ -627,7 +627,7 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: next=("two",), config={ "configurable": { - "thread_id": 1, + "thread_id": "1", "thread_ts": AnyStr(), } }, @@ -640,7 +640,7 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: next=("one",), config={ "configurable": { - "thread_id": 1, + "thread_id": "1", "thread_ts": AnyStr(), } }, @@ -653,7 +653,7 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: next=("two",), config={ "configurable": { - "thread_id": 1, + "thread_id": "1", "thread_ts": AnyStr(), } }, @@ -666,7 +666,7 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: next=("one",), config={ "configurable": { - "thread_id": 1, + "thread_id": "1", "thread_ts": AnyStr(), } }, @@ -679,7 +679,7 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: next=(), config={ "configurable": { - "thread_id": 1, + "thread_id": "1", "thread_ts": AnyStr(), } }, @@ -692,7 +692,7 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: next=("two",), config={ "configurable": { - "thread_id": 1, + "thread_id": "1", "thread_ts": AnyStr(), } }, @@ -705,7 +705,7 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: next=("one",), config={ "configurable": { - "thread_id": 1, + "thread_id": "1", "thread_ts": AnyStr(), } }, @@ -1080,7 +1080,7 @@ def test_pending_writes_resume(checkpointer: BaseCheckpointSaver) -> None: builder.add_edge(START, "two") graph = builder.compile(checkpointer=checkpointer) - thread1: RunnableConfig = {"configurable": {"thread_id": 1}} + thread1: RunnableConfig = {"configurable": {"thread_id": "1"}} with pytest.raises(ConnectionError, match="I'm not good"): graph.invoke({"value": 1}, thread1)