From 08b505e017add4a0c9f6e4b031de52eadf58a304 Mon Sep 17 00:00:00 2001 From: Andrew Nguonly Date: Tue, 21 May 2024 11:03:25 -0700 Subject: [PATCH] Merge configurable fields with previous checkpoint config before each run (#510) * Merge configurable fields with previous checkpoint config. * Update update_state() and aupdate_state() to merge configurable fields with previous checkpoint config. * Update tests to verify that all checkpoint metadata contain the expected configurable field keys. This assertion is needed because a run can have an arbitrary number of steps based on the construction of the graph. --- langgraph/pregel/__init__.py | 50 ++++++++++- tests/memory_assert.py | 63 ++++++++++++++ tests/test_pregel.py | 163 ++++++++++++++++++++++++++++++++++- tests/test_pregel_async.py | 163 ++++++++++++++++++++++++++++++++++- 4 files changed, 433 insertions(+), 6 deletions(-) diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index e535ed735..00418f866 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -550,8 +550,19 @@ class Pregel( # apply to checkpoint and save _apply_writes(checkpoint, channels, task.writes) step = saved.metadata.get("step", -2) + 1 if saved else -1 + + # merge configurable fields with previous checkpoint config + checkpoint_config = config + if saved: + checkpoint_config = { + "configurable": { + **config.get("configurable", {}), + **saved.config["configurable"], + } + } + return self.checkpointer.put( - saved.config if saved else config, + checkpoint_config, create_checkpoint(checkpoint, channels, step), { "source": "update", @@ -625,8 +636,19 @@ class Pregel( # apply to checkpoint and save _apply_writes(checkpoint, channels, task.writes) step = saved.metadata.get("step", -2) + 1 if saved else -1 + + # merge configurable fields with previous checkpoint config + checkpoint_config = config + if saved: + checkpoint_config = { + "configurable": { + **config.get("configurable", {}), + **saved.config["configurable"], + } + } + return await self.checkpointer.aput( - saved.config if saved else config, + checkpoint_config, create_checkpoint(checkpoint, channels, step), { "source": "update", @@ -728,7 +750,17 @@ class Pregel( # get checkpoint from saver, or create an empty one saved = self.checkpointer.get_tuple(config) if self.checkpointer else None checkpoint = saved.checkpoint if saved else empty_checkpoint() - checkpoint_config = saved.config if saved else config + + # merge configurable fields with previous checkpoint config + checkpoint_config = config + if saved: + checkpoint_config = { + "configurable": { + **config.get("configurable", {}), + **saved.config["configurable"], + } + } + start = saved.metadata.get("step", -2) + 1 if saved else -1 # create channels from checkpoint with ChannelsManager( @@ -1002,7 +1034,17 @@ class Pregel( else None ) checkpoint = saved.checkpoint if saved else empty_checkpoint() - checkpoint_config = saved.config if saved else config + + # merge configurable fields with previous checkpoint config + checkpoint_config = config + if saved: + checkpoint_config = { + "configurable": { + **config.get("configurable", {}), + **saved.config["configurable"], + } + } + start = saved.metadata.get("step", -2) + 1 if saved else -1 # create channels from checkpoint async with AsyncChannelsManager( diff --git a/tests/memory_assert.py b/tests/memory_assert.py index 99a438538..576fde89a 100644 --- a/tests/memory_assert.py +++ b/tests/memory_assert.py @@ -1,6 +1,9 @@ +import asyncio from collections import defaultdict from typing import Any, Optional +from langchain_core.runnables import RunnableConfig + from langgraph.checkpoint.base import ( Checkpoint, CheckpointMetadata, @@ -46,3 +49,63 @@ class MemorySaverAssertImmutable(MemorySaver): ) # call super to write checkpoint return super().put(config, checkpoint, metadata) + + +class MemorySaverAssertCheckpointMetadata(MemorySaver): + """This custom checkpointer is for verifying that a run's configurable + fields are merged with the previous checkpoint config for each step in + the run. This is the desired behavior. Because the checkpointer's (a)put() + method is called for each step, the implementation of this checkpointer + should produce a side effect that can be asserted. + """ + + serde = NoopSerializer() + + def __init__( + self, + *, + serde: Optional[SerializerProtocol] = None, + ) -> None: + super().__init__(serde=serde) + + def put( + self, + config: RunnableConfig, + checkpoint: Checkpoint, + metadata: Optional[CheckpointMetadata] = None, + ) -> None: + """The implementation of put() merges config["configurable"] (a run's + configurable fields) with the metadata field. The state of the + checkpoint metadata can be asserted to confirm that the run's + configurable fields were merged with the previous checkpoint config. + """ + configurable = config["configurable"].copy() + + # remove thread_ts to make testing simpler + configurable.pop("thread_ts", None) + + self.storage[config["configurable"]["thread_id"]].update( + { + checkpoint["id"]: ( + self.serde.dumps(checkpoint), + # merge configurable fields and metadata + self.serde.dumps({**configurable, **metadata}), + ) + } + ) + return { + "configurable": { + "thread_id": config["configurable"]["thread_id"], + "thread_ts": checkpoint["id"], + } + } + + async def aput( + self, + config: RunnableConfig, + checkpoint: Checkpoint, + metadata: CheckpointMetadata, + ) -> RunnableConfig: + return await asyncio.get_running_loop().run_in_executor( + None, self.put, config, checkpoint, metadata + ) diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 6800f3a28..f8620ada7 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -41,7 +41,10 @@ from langgraph.prebuilt.tool_node import ToolNode from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot from langgraph.pregel.retry import RetryPolicy from tests.any_str import AnyStr -from tests.memory_assert import MemorySaverAssertImmutable +from tests.memory_assert import ( + MemorySaverAssertCheckpointMetadata, + MemorySaverAssertImmutable, +) def test_graph_validation() -> None: @@ -5492,3 +5495,161 @@ def test_repeat_condition(snapshot: SnapshotAssertion) -> None: app = workflow.compile() assert app.get_graph().draw_mermaid(with_styles=False) == snapshot + + +def test_checkpoint_metadata() -> None: + """This test verifies that a run's configurable fields are merged with the + previous checkpoint config for each step in the run. + """ + # set up test + from langchain_core.language_models.fake_chat_models import ( + FakeMessagesListChatModel, + ) + from langchain_core.messages import AIMessage, AnyMessage + from langchain_core.prompts import ChatPromptTemplate + from langchain_core.tools import tool + + # graph state + class BaseState(TypedDict): + messages: Annotated[list[AnyMessage], add_messages] + + # initialize graph nodes + @tool() + def search_api(query: str) -> str: + """Searches the API for the query.""" + return f"result for {query}" + + tools = [search_api] + + prompt = ChatPromptTemplate.from_messages( + [ + ("system", "You are a nice assistant."), + ("placeholder", "{messages}"), + ] + ) + + model = FakeMessagesListChatModel( + responses=[ + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ), + AIMessage(content="answer"), + ] + ) + + def agent(state: BaseState, config: RunnableConfig) -> BaseState: + formatted = prompt.invoke(state) + response = model.invoke(formatted) + return {"messages": response} + + def should_continue(data: BaseState) -> str: + # Logic to decide whether to continue in the loop or exit + if not data["messages"][-1].tool_calls: + return "exit" + else: + return "continue" + + # define graphs w/ and w/o interrupt + workflow = StateGraph(BaseState) + workflow.add_node("agent", agent) + workflow.add_node("tools", ToolNode(tools)) + workflow.set_entry_point("agent") + workflow.add_conditional_edges( + "agent", should_continue, {"continue": "tools", "exit": END} + ) + workflow.add_edge("tools", "agent") + + # graph w/o interrupt + checkpointer_1 = MemorySaverAssertCheckpointMetadata() + app = workflow.compile(checkpointer=checkpointer_1) + + # graph w/ interrupt + checkpointer_2 = MemorySaverAssertCheckpointMetadata() + app_w_interrupt = workflow.compile( + checkpointer=checkpointer_2, interrupt_before=["tools"] + ) + + # assertions + + # invoke graph w/o interrupt + app.invoke( + {"messages": ["what is weather in sf"]}, + { + "configurable": { + "thread_id": "1", + "test_config_1": "foo", + "test_config_2": "bar", + }, + }, + ) + + config = {"configurable": {"thread_id": "1"}} + + # assert that checkpoint metadata contains the run's configurable fields + chkpnt_metadata_1 = checkpointer_1.get_tuple(config).metadata + assert chkpnt_metadata_1["thread_id"] == "1" + assert chkpnt_metadata_1["test_config_1"] == "foo" + assert chkpnt_metadata_1["test_config_2"] == "bar" + + # Verify that all checkpoint metadata have the expected keys. This check + # is needed because a run may have an arbitrary number of steps depending + # on how the graph is constructed. + chkpnt_tuples_1 = checkpointer_1.list(config) + for chkpnt_tuple in chkpnt_tuples_1: + assert chkpnt_tuple.metadata["thread_id"] == "1" + assert chkpnt_tuple.metadata["test_config_1"] == "foo" + assert chkpnt_tuple.metadata["test_config_2"] == "bar" + + # invoke graph, but interrupt before tool call + app_w_interrupt.invoke( + {"messages": ["what is weather in sf"]}, + { + "configurable": { + "thread_id": "2", + "test_config_3": "foo", + "test_config_4": "bar", + }, + }, + ) + + config = {"configurable": {"thread_id": "2"}} + + # assert that checkpoint metadata contains the run's configurable fields + chkpnt_metadata_2 = checkpointer_2.get_tuple(config).metadata + assert chkpnt_metadata_2["thread_id"] == "2" + assert chkpnt_metadata_2["test_config_3"] == "foo" + assert chkpnt_metadata_2["test_config_4"] == "bar" + + # resume graph execution + app_w_interrupt.invoke( + input=None, + config={ + "configurable": { + "thread_id": "2", + "test_config_3": "foo", + "test_config_4": "bar", + } + }, + ) + + # assert that checkpoint metadata contains the run's configurable fields + chkpnt_metadata_3 = checkpointer_2.get_tuple(config).metadata + assert chkpnt_metadata_3["thread_id"] == "2" + assert chkpnt_metadata_3["test_config_3"] == "foo" + assert chkpnt_metadata_3["test_config_4"] == "bar" + + # Verify that all checkpoint metadata have the expected keys. This check + # is needed because a run may have an arbitrary number of steps depending + # on how the graph is constructed. + chkpnt_tuples_2 = checkpointer_2.list(config) + for chkpnt_tuple in chkpnt_tuples_2: + assert chkpnt_tuple.metadata["thread_id"] == "2" + assert chkpnt_tuple.metadata["test_config_3"] == "foo" + assert chkpnt_tuple.metadata["test_config_4"] == "bar" diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 6ca1b4846..1fbda625d 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -41,7 +41,10 @@ from langgraph.prebuilt.tool_node import ToolNode from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot from langgraph.pregel.retry import RetryPolicy from tests.any_str import AnyStr -from tests.memory_assert import MemorySaverAssertImmutable +from tests.memory_assert import ( + MemorySaverAssertCheckpointMetadata, + MemorySaverAssertImmutable, +) async def test_node_cancellation_on_external_cancel() -> None: @@ -4779,3 +4782,161 @@ async def test_nested_graph(snapshot: SnapshotAssertion) -> None: ] } assert times_called == 1 + + +async def test_checkpoint_metadata() -> None: + """This test verifies that a run's configurable fields are merged with the + previous checkpoint config for each step in the run. + """ + # set up test + from langchain_core.language_models.fake_chat_models import ( + FakeMessagesListChatModel, + ) + from langchain_core.messages import AIMessage, AnyMessage + from langchain_core.prompts import ChatPromptTemplate + from langchain_core.tools import tool + + # graph state + class BaseState(TypedDict): + messages: Annotated[list[AnyMessage], add_messages] + + # initialize graph nodes + @tool() + def search_api(query: str) -> str: + """Searches the API for the query.""" + return f"result for {query}" + + tools = [search_api] + + prompt = ChatPromptTemplate.from_messages( + [ + ("system", "You are a nice assistant."), + ("placeholder", "{messages}"), + ] + ) + + model = FakeMessagesListChatModel( + responses=[ + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ), + AIMessage(content="answer"), + ] + ) + + def agent(state: BaseState, config: RunnableConfig) -> BaseState: + formatted = prompt.invoke(state) + response = model.invoke(formatted) + return {"messages": response} + + def should_continue(data: BaseState) -> str: + # Logic to decide whether to continue in the loop or exit + if not data["messages"][-1].tool_calls: + return "exit" + else: + return "continue" + + # define graphs w/ and w/o interrupt + workflow = StateGraph(BaseState) + workflow.add_node("agent", agent) + workflow.add_node("tools", ToolNode(tools)) + workflow.set_entry_point("agent") + workflow.add_conditional_edges( + "agent", should_continue, {"continue": "tools", "exit": END} + ) + workflow.add_edge("tools", "agent") + + # graph w/o interrupt + checkpointer_1 = MemorySaverAssertCheckpointMetadata() + app = workflow.compile(checkpointer=checkpointer_1) + + # graph w/ interrupt + checkpointer_2 = MemorySaverAssertCheckpointMetadata() + app_w_interrupt = workflow.compile( + checkpointer=checkpointer_2, interrupt_before=["tools"] + ) + + # assertions + + # invoke graph w/o interrupt + await app.ainvoke( + {"messages": ["what is weather in sf"]}, + { + "configurable": { + "thread_id": "1", + "test_config_1": "foo", + "test_config_2": "bar", + }, + }, + ) + + config = {"configurable": {"thread_id": "1"}} + + # assert that checkpoint metadata contains the run's configurable fields + chkpnt_metadata_1 = (await checkpointer_1.aget_tuple(config)).metadata + assert chkpnt_metadata_1["thread_id"] == "1" + assert chkpnt_metadata_1["test_config_1"] == "foo" + assert chkpnt_metadata_1["test_config_2"] == "bar" + + # Verify that all checkpoint metadata have the expected keys. This check + # is needed because a run may have an arbitrary number of steps depending + # on how the graph is constructed. + chkpnt_tuples_1 = checkpointer_1.alist(config) + async for chkpnt_tuple in chkpnt_tuples_1: + assert chkpnt_tuple.metadata["thread_id"] == "1" + assert chkpnt_tuple.metadata["test_config_1"] == "foo" + assert chkpnt_tuple.metadata["test_config_2"] == "bar" + + # invoke graph, but interrupt before tool call + await app_w_interrupt.ainvoke( + {"messages": ["what is weather in sf"]}, + { + "configurable": { + "thread_id": "2", + "test_config_3": "foo", + "test_config_4": "bar", + }, + }, + ) + + config = {"configurable": {"thread_id": "2"}} + + # assert that checkpoint metadata contains the run's configurable fields + chkpnt_metadata_2 = (await checkpointer_2.aget_tuple(config)).metadata + assert chkpnt_metadata_2["thread_id"] == "2" + assert chkpnt_metadata_2["test_config_3"] == "foo" + assert chkpnt_metadata_2["test_config_4"] == "bar" + + # resume graph execution + await app_w_interrupt.ainvoke( + input=None, + config={ + "configurable": { + "thread_id": "2", + "test_config_3": "foo", + "test_config_4": "bar", + } + }, + ) + + # assert that checkpoint metadata contains the run's configurable fields + chkpnt_metadata_3 = (await checkpointer_2.aget_tuple(config)).metadata + assert chkpnt_metadata_3["thread_id"] == "2" + assert chkpnt_metadata_3["test_config_3"] == "foo" + assert chkpnt_metadata_3["test_config_4"] == "bar" + + # Verify that all checkpoint metadata have the expected keys. This check + # is needed because a run may have an arbitrary number of steps depending + # on how the graph is constructed. + chkpnt_tuples_2 = checkpointer_2.alist(config) + async for chkpnt_tuple in chkpnt_tuples_2: + assert chkpnt_tuple.metadata["thread_id"] == "2" + assert chkpnt_tuple.metadata["test_config_3"] == "foo" + assert chkpnt_tuple.metadata["test_config_4"] == "bar"