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

Merge pull request #1539 from langchain-ai/nc/29aug/nested-update-state

Browse Source 
Implement update_state for nested graphs
This commit is contained in:
Nuno Campos
2024-08-30 10:30:40 -07:00
committed by GitHub
parent 24a7e6403f 9fd64cd37b
commit 15a204b2be
8 changed files with 902 additions and 160 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/Makefile
+3 -5
View File
@@ -18,20 +18,18 @@ start-postgres:
docker compose -f tests/compose-postgres.yml up -V --force-recreate --wait
stop-postgres:
docker compose -f tests/compose-postgres.yml down
docker compose -f tests/compose-postgres.yml down -v
TEST_PATH ?= .
test:
make start-postgres; \
poetry run pytest $(TEST_PATH); \
make start-postgres && poetry run pytest $(TEST_PATH); \
EXIT_CODE=$$?; \
make stop-postgres; \
exit $$EXIT_CODE
test_watch:
make start-postgres; \
poetry run ptw . -- --ff -v -x -n auto --dist worksteal --snapshot-update --tb short $(TEST_PATH); \
make start-postgres && poetry run ptw . -- --ff -v -x -n auto --dist worksteal --snapshot-update --tb short $(TEST_PATH); \
EXIT_CODE=$$?; \
make stop-postgres; \
exit $$EXIT_CODE
libs/langgraph/langgraph/pregel/__init__.py
+43 -41
View File
@@ -59,7 +59,6 @@ from langgraph.checkpoint.base import (
empty_checkpoint,
)
from langgraph.constants import (
CONFIG_KEY_CHECKPOINT_MAP,
CONFIG_KEY_CHECKPOINTER,
CONFIG_KEY_READ,
CONFIG_KEY_RESUMING,
@@ -79,7 +78,7 @@ from langgraph.pregel.algo import (
local_write,
prepare_next_tasks,
)
from langgraph.pregel.config import patch_configurable
from langgraph.pregel.config import patch_checkpoint_map, patch_configurable
from langgraph.pregel.debug import (
print_step_checkpoint,
print_step_tasks,
@@ -445,19 +444,7 @@ class Pregel(
return StateSnapshot(
read_channels(channels, self.stream_channels_asis),
tuple(t.name for t in next_tasks),
patch_configurable(
saved.config,
{
CONFIG_KEY_CHECKPOINT_MAP: {
**saved.metadata["parents"],
saved.config["configurable"][
"checkpoint_ns"
]: saved.checkpoint["id"],
}
},
)
if saved.metadata.get("parents")
else saved.config,
patch_checkpoint_map(saved.config, saved.metadata),
saved.metadata,
saved.checkpoint["ts"],
saved.parent_config,
@@ -533,19 +520,7 @@ class Pregel(
return StateSnapshot(
read_channels(channels, self.stream_channels_asis),
tuple(t.name for t in next_tasks),
patch_configurable(
saved.config,
{
CONFIG_KEY_CHECKPOINT_MAP: {
**saved.metadata["parents"],
saved.config["configurable"][
"checkpoint_ns"
]: saved.checkpoint["id"],
}
},
)
if saved.metadata.get("parents")
else saved.config,
patch_checkpoint_map(saved.config, saved.metadata),
saved.metadata,
saved.checkpoint["ts"],
saved.parent_config,
@@ -738,6 +713,7 @@ class Pregel(
if not checkpointer:
raise ValueError("No checkpointer set")
# delegate to subgraph
if (
checkpoint_ns := config["configurable"].get("checkpoint_ns", "")
) and CONFIG_KEY_CHECKPOINTER not in config["configurable"]:
@@ -760,7 +736,7 @@ class Pregel(
# get last checkpoint
config = merge_configs(self.config, config) if self.config else config
saved = self.checkpointer.get_tuple(config)
saved = checkpointer.get_tuple(config)
checkpoint = copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint()
checkpoint_previous_versions = (
saved.checkpoint["channel_versions"].copy() if saved else {}
@@ -775,7 +751,7 @@ class Pregel(
checkpoint_config = patch_configurable(config, saved.config["configurable"])
# find last node that updated the state, if not provided
if values is None and as_node is None:
return self.checkpointer.put(
next_config = checkpointer.put(
checkpoint_config,
create_checkpoint(checkpoint, None, step),
{
@@ -786,6 +762,7 @@ class Pregel(
},
{},
)
return patch_checkpoint_map(next_config, saved.metadata if saved else None)
elif as_node is None and not any(
v for vv in checkpoint["versions_seen"].values() for v in vv.values()
):
@@ -860,13 +837,13 @@ class Pregel(
)
# save task writes
if saved:
self.checkpointer.put_writes(checkpoint_config, task.writes, task.id)
checkpointer.put_writes(checkpoint_config, task.writes, task.id)
# apply to checkpoint and save
assert not apply_writes(
checkpoint, channels, [task], self.checkpointer.get_next_version
checkpoint, channels, [task], checkpointer.get_next_version
), "Can't write to SharedValues from update_state"
checkpoint = create_checkpoint(checkpoint, channels, step + 1)
return self.checkpointer.put(
next_config = checkpointer.put(
checkpoint_config,
checkpoint,
{
@@ -879,6 +856,7 @@ class Pregel(
checkpoint_previous_versions, checkpoint["channel_versions"]
),
)
return patch_checkpoint_map(next_config, saved.metadata if saved else None)
async def aupdate_state(
self,
@@ -886,12 +864,36 @@ class Pregel(
values: dict[str, Any] | Any,
as_node: Optional[str] = None,
) -> RunnableConfig:
if not self.checkpointer:
checkpointer: Optional[BaseCheckpointSaver] = config["configurable"].get(
CONFIG_KEY_CHECKPOINTER, self.checkpointer
)
if not checkpointer:
raise ValueError("No checkpointer set")
# delegate to subgraph
if (
checkpoint_ns := config["configurable"].get("checkpoint_ns", "")
) and CONFIG_KEY_CHECKPOINTER not in config["configurable"]:
# remove task_ids from checkpoint_ns
recast_checkpoint_ns = NS_SEP.join(
part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP)
)
# find the subgraph with the matching name
async for name, pregel in self.aget_subgraphs(recurse=True):
if name == recast_checkpoint_ns:
return await pregel.aupdate_state(
patch_configurable(
config, {CONFIG_KEY_CHECKPOINTER: checkpointer}
),
values,
as_node,
)
else:
raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
# get last checkpoint
config = merge_configs(self.config, config) if self.config else config
saved = await self.checkpointer.aget_tuple(config)
saved = await checkpointer.aget_tuple(config)
checkpoint = copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint()
checkpoint_previous_versions = (
saved.checkpoint["channel_versions"].copy() if saved else {}
@@ -915,7 +917,7 @@ class Pregel(
}
# find last node that updated the state, if not provided
if values is None and as_node is None:
return await self.checkpointer.aput(
next_config = await checkpointer.aput(
checkpoint_config,
create_checkpoint(checkpoint, None, step),
{
@@ -926,6 +928,7 @@ class Pregel(
},
{},
)
return patch_checkpoint_map(next_config, saved.metadata if saved else None)
elif as_node is None and not saved:
if (
isinstance(self.input_channels, str)
@@ -998,15 +1001,13 @@ class Pregel(
)
# save task writes
if saved:
await self.checkpointer.aput_writes(
checkpoint_config, task.writes, task.id
)
await checkpointer.aput_writes(checkpoint_config, task.writes, task.id)
# apply to checkpoint and save
assert not apply_writes(
checkpoint, channels, [task], self.checkpointer.get_next_version
checkpoint, channels, [task], checkpointer.get_next_version
), "Can't write to SharedValues from update_state"
checkpoint = create_checkpoint(checkpoint, channels, step + 1)
return await self.checkpointer.aput(
next_config = await checkpointer.aput(
checkpoint_config,
checkpoint,
{
@@ -1019,6 +1020,7 @@ class Pregel(
checkpoint_previous_versions, checkpoint["channel_versions"]
),
)
return patch_checkpoint_map(next_config, saved.metadata if saved else None)
def _defaults(
self,
libs/langgraph/langgraph/pregel/config.py
+22
View File
@@ -2,6 +2,9 @@ from typing import Any, Optional
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import CheckpointMetadata
from langgraph.constants import CONFIG_KEY_CHECKPOINT_MAP
def patch_configurable(
config: Optional[RunnableConfig], patch: dict[str, Any]
@@ -10,3 +13,22 @@ def patch_configurable(
return {"configurable": patch}
else:
return {**config, "configurable": {**config["configurable"], **patch}}
def patch_checkpoint_map(
config: RunnableConfig, metadata: Optional[CheckpointMetadata]
) -> RunnableConfig:
if parents := (metadata.get("parents") if metadata else None):
return patch_configurable(
config,
{
CONFIG_KEY_CHECKPOINT_MAP: {
**parents,
config["configurable"]["checkpoint_ns"]: config["configurable"][
"checkpoint_id"
],
},
},
)
else:
return config
libs/langgraph/langgraph/pregel/io.py
+8 -4
View File
@@ -1,4 +1,4 @@
from typing import Any, Iterator, Mapping, Optional, Sequence, TypeVar, Union
from typing import Any, Iterator, Literal, Mapping, Optional, Sequence, TypeVar, Union
from langchain_core.runnables.utils import AddableDict
@@ -73,15 +73,19 @@ class AddableValuesDict(AddableDict):
def map_output_values(
output_channels: Union[str, Sequence[str]],
pending_writes: Sequence[tuple[str, Any]],
pending_writes: Union[Literal[True], Sequence[tuple[str, Any]]],
channels: Mapping[str, BaseChannel],
) -> Iterator[Union[dict[str, Any], Any]]:
"""Map pending writes (a sequence of tuples (channel, value)) to output chunk."""
if isinstance(output_channels, str):
if any(chan == output_channels for chan, _ in pending_writes):
if pending_writes is True or any(
chan == output_channels for chan, _ in pending_writes
):
yield read_channel(channels, output_channels)
else:
if {c for c, _ in pending_writes if c in output_channels}:
if pending_writes is True or {
c for c, _ in pending_writes if c in output_channels
}:
yield AddableValuesDict(read_channels(channels, output_channels))
libs/langgraph/langgraph/pregel/loop.py
+5
View File
@@ -401,6 +401,11 @@ class PregelLoop:
if k in self.checkpoint["channel_versions"]:
version = self.checkpoint["channel_versions"][k]
self.checkpoint["versions_seen"][INTERRUPT][k] = version
# produce values output
self.stream.extend(
(self.config["configurable"].get("checkpoint_ns", ""), "values", v)
for v in map_output_values(self.output_keys, True, self.channels)
)
# map inputs to channel updates
elif input_writes := deque(map_input(input_keys, self.input)):
# discard any unfinished tasks from previous checkpoint
libs/langgraph/tests/__snapshots__/test_pregel.ambr
+125
View File
@@ -5099,6 +5099,131 @@
# name: test_state_graph_w_config_inherited_state_keys.2
'{"title": "LangGraphOutput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "Represents a request to execute an action by an agent.\\n\\nThe action consists of the name of the tool to execute and the input to pass\\nto the tool. The log is used to pass along extra information about the action.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}'
# ---
# name: test_weather_subgraph[memory]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([__start__]):::first
router_node(router_node)
normal_llm_node(normal_llm_node)
weather_graph_model_node(model_node)
weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
__end__([__end__]):::last
subgraph weather_graph
weather_graph_model_node --> weather_graph_weather_node;
end
__start__ --> router_node;
normal_llm_node --> __end__;
weather_graph_weather_node --> __end__;
router_node -.-> normal_llm_node;
router_node -.-> weather_graph_model_node;
router_node -.-> __end__;
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_weather_subgraph[postgres]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([__start__]):::first
router_node(router_node)
normal_llm_node(normal_llm_node)
weather_graph_model_node(model_node)
weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
__end__([__end__]):::last
subgraph weather_graph
weather_graph_model_node --> weather_graph_weather_node;
end
__start__ --> router_node;
normal_llm_node --> __end__;
weather_graph_weather_node --> __end__;
router_node -.-> normal_llm_node;
router_node -.-> weather_graph_model_node;
router_node -.-> __end__;
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_weather_subgraph[postgres_pipe]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([__start__]):::first
router_node(router_node)
normal_llm_node(normal_llm_node)
weather_graph_model_node(model_node)
weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
__end__([__end__]):::last
subgraph weather_graph
weather_graph_model_node --> weather_graph_weather_node;
end
__start__ --> router_node;
normal_llm_node --> __end__;
weather_graph_weather_node --> __end__;
router_node -.-> normal_llm_node;
router_node -.-> weather_graph_model_node;
router_node -.-> __end__;
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_weather_subgraph[postgres_pool]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([__start__]):::first
router_node(router_node)
normal_llm_node(normal_llm_node)
weather_graph_model_node(model_node)
weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
__end__([__end__]):::last
subgraph weather_graph
weather_graph_model_node --> weather_graph_weather_node;
end
__start__ --> router_node;
normal_llm_node --> __end__;
weather_graph_weather_node --> __end__;
router_node -.-> normal_llm_node;
router_node -.-> weather_graph_model_node;
router_node -.-> __end__;
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_weather_subgraph[sqlite]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([__start__]):::first
router_node(router_node)
normal_llm_node(normal_llm_node)
weather_graph_model_node(model_node)
weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
__end__([__end__]):::last
subgraph weather_graph
weather_graph_model_node --> weather_graph_weather_node;
end
__start__ --> router_node;
normal_llm_node --> __end__;
weather_graph_weather_node --> __end__;
router_node -.-> normal_llm_node;
router_node -.-> weather_graph_model_node;
router_node -.-> __end__;
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_xray_issue
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
libs/langgraph/tests/test_pregel.py
+573 -64
View File
@@ -20,6 +20,7 @@ from typing import (
Tuple,
TypedDict,
Union,
cast,
get_type_hints,
)
@@ -2293,6 +2294,16 @@ def test_conditional_graph(
)
assert [c for c in app_w_interrupt.stream(None, config)] == [
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"input": "what is weather in sf",
}
},
{
"tools": {
"input": "what is weather in sf",
@@ -2503,6 +2514,16 @@ def test_conditional_graph(
)
assert [c for c in app_w_interrupt.stream(None, config)] == [
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"input": "what is weather in sf",
},
},
{
"tools": {
"input": "what is weather in sf",
@@ -2668,6 +2689,14 @@ def test_conditional_graph(
)
assert [c for c in app_w_interrupt.stream(None, config)] == [
{
"agent": {
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
},
},
{
"tools": {
"input": "what is weather in sf",
@@ -2706,6 +2735,26 @@ def test_conditional_graph(
]
assert [c for c in app_w_interrupt.stream(None, config)] == [
{
"agent": {
"input": "what is weather in sf",
"intermediate_steps": [
[
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"result for query",
]
],
"agent_outcome": AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
}
},
{
"tools": {
"input": "what is weather in sf",
@@ -8396,17 +8445,12 @@ def test_nested_graph_interrupts_parallel(
# test stream values w/ nested interrupt
config = {"configurable": {"thread_id": "3"}}
assert [*app.stream({"my_key": ""}, config, stream_mode="values")] == [
{
"my_key": "",
},
{"my_key": ""},
]
assert [*app.stream(None, config, stream_mode="values")] == [
{
"my_key": "got here and there and parallel",
},
{
"my_key": "got here and there and parallel and back again",
},
{"my_key": ""},
{"my_key": "got here and there and parallel"},
{"my_key": "got here and there and parallel and back again"},
]
# test interrupts BEFORE the parallel node
@@ -8416,14 +8460,13 @@ def test_nested_graph_interrupts_parallel(
{"my_key": ""}
]
# while we're waiting for the node w/ interrupt inside to finish
assert [*app.stream(None, config, stream_mode="values")] == []
assert [*app.stream(None, config, stream_mode="values")] == [
{
"my_key": "got here and there and parallel",
},
{
"my_key": "got here and there and parallel and back again",
},
{"my_key": ""},
]
assert [*app.stream(None, config, stream_mode="values")] == [
{"my_key": ""},
{"my_key": "got here and there and parallel"},
{"my_key": "got here and there and parallel and back again"},
]
# test interrupts AFTER the parallel node
@@ -8433,12 +8476,12 @@ def test_nested_graph_interrupts_parallel(
{"my_key": ""}
]
assert [*app.stream(None, config, stream_mode="values")] == [
{"my_key": ""},
{"my_key": "got here and there and parallel"},
]
assert [*app.stream(None, config, stream_mode="values")] == [
{
"my_key": "got here and there and parallel and back again",
},
{"my_key": "got here and there and parallel"},
{"my_key": "got here and there and parallel and back again"},
]
@@ -8520,20 +8563,13 @@ def test_doubly_nested_graph_interrupts(
# test stream values w/ nested interrupt
config = {"configurable": {"thread_id": "3"}}
assert [*app.stream({"my_key": "my value"}, config, stream_mode="values")] == [
{
"my_key": "my value",
},
{
"my_key": "hi my value",
},
{"my_key": "my value"},
{"my_key": "hi my value"},
]
assert [*app.stream(None, config, stream_mode="values")] == [
{
"my_key": "hi my value here and there",
},
{
"my_key": "hi my value here and there and back again",
},
{"my_key": "hi my value"},
{"my_key": "hi my value here and there"},
{"my_key": "hi my value here and there and back again"},
]
@@ -9786,8 +9822,16 @@ def test_doubly_nested_graph_state(
)
]
# 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"}},
)
]
@pytest.mark.skip("TODO")
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
def test_send_to_nested_graphs(
request: pytest.FixtureRequest, checkpointer_name: str
@@ -9835,24 +9879,30 @@ def test_send_to_nested_graphs(
"subjects": ["cats", "dogs"],
"jokes": [],
}
actual_snapshot = graph.get_state(config)
subgraph_nodes = list(actual_snapshot.subgraphs.keys())
assert len(subgraph_nodes) == 2
for subgraph_node in subgraph_nodes:
assert subgraph_node.split(":")[0] == "generate_joke"
subgraphs = {
subgraph_node: graph.get_state(
{"configurable": {"thread_id": "1", "checkpoint_ns": subgraph_node}}
)
for subgraph_node in subgraph_nodes
}
expected_snapshot = StateSnapshot(
outer_state = graph.get_state(config)
assert outer_state == StateSnapshot(
values={"subjects": ["cats", "dogs"], "jokes": []},
tasks=(
PregelTask(AnyStr(), "generate_joke"),
PregelTask(AnyStr(), "generate_joke"),
PregelTask(
AnyStr(),
"generate_joke",
state={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
}
},
),
PregelTask(
AnyStr(),
"generate_joke",
state={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
}
},
),
),
next=("generate_joke", "generate_joke"),
config={
@@ -9871,21 +9921,86 @@ def test_send_to_nested_graphs(
"checkpoint_id": AnyStr(),
}
},
subgraphs=subgraphs,
)
assert actual_snapshot == expected_snapshot
# check state of each of the inner tasks
assert graph.get_state(outer_state.tasks[0].state) == StateSnapshot(
values={"subject": "cats - hohoho", "jokes": []},
next=("generate",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{
"": AnyStr(),
AnyStr("generate_joke:"): AnyStr(),
}
),
}
},
metadata={
"step": 1,
"source": "loop",
"writes": {"edit": None},
"parents": {"": AnyStr()},
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
"checkpoint_id": AnyStr(),
}
},
tasks=(PregelTask(id=AnyStr(""), name="generate"),),
)
assert graph.get_state(outer_state.tasks[1].state) == StateSnapshot(
values={"subject": "dogs - hohoho", "jokes": []},
next=("generate",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{
"": AnyStr(),
AnyStr("generate_joke:"): AnyStr(),
}
),
}
},
metadata={
"step": 1,
"source": "loop",
"writes": {"edit": None},
"parents": {"": AnyStr()},
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
"checkpoint_id": AnyStr(),
}
},
tasks=(PregelTask(id=AnyStr(""), name="generate"),),
)
# update state of dogs joke graph
graph.update_state(outer_state.tasks[1].state, {"subject": "turtles - hohoho"})
# continue past interrupt
assert graph.invoke(None, config=config) == {
"subjects": ["cats", "dogs"],
"jokes": ["Joke about cats - hohoho", "Joke about dogs - hohoho"],
"jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"],
}
actual_snapshot = graph.get_state(config)
expected_snapshot = StateSnapshot(
values={
"subjects": ["cats", "dogs"],
"jokes": ["Joke about cats - hohoho", "Joke about dogs - hohoho"],
"jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"],
},
tasks=(),
next=(),
@@ -9902,7 +10017,7 @@ def test_send_to_nested_graphs(
"writes": {
"generate_joke": [
{"jokes": ["Joke about cats - hohoho"]},
{"jokes": ["Joke about dogs - hohoho"]},
{"jokes": ["Joke about turtles - hohoho"]},
]
},
"step": 1,
@@ -9922,17 +10037,11 @@ def test_send_to_nested_graphs(
actual_history = list(graph.get_state_history(config))
# get subgraph node state for expected history
subgraphs = {
subgraph_node: graph.get_state(
{"configurable": {"thread_id": "1", "checkpoint_ns": subgraph_node}}
)
for subgraph_node in subgraph_nodes
}
expected_history = [
StateSnapshot(
values={
"subjects": ["cats", "dogs"],
"jokes": ["Joke about cats - hohoho", "Joke about dogs - hohoho"],
"jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"],
},
tasks=(),
next=(),
@@ -9949,7 +10058,7 @@ def test_send_to_nested_graphs(
"writes": {
"generate_joke": [
{"jokes": ["Joke about cats - hohoho"]},
{"jokes": ["Joke about dogs - hohoho"]},
{"jokes": ["Joke about turtles - hohoho"]},
]
},
"step": 1,
@@ -9966,8 +10075,26 @@ def test_send_to_nested_graphs(
StateSnapshot(
values={"subjects": ["cats", "dogs"], "jokes": []},
tasks=(
PregelTask(AnyStr(), "generate_joke"),
PregelTask(AnyStr(), "generate_joke"),
PregelTask(
AnyStr(),
"generate_joke",
state={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
}
},
),
PregelTask(
AnyStr(),
"generate_joke",
state={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
}
},
),
),
next=("generate_joke", "generate_joke"),
config={
@@ -9986,7 +10113,6 @@ def test_send_to_nested_graphs(
"checkpoint_id": AnyStr(),
}
},
subgraphs=subgraphs,
),
StateSnapshot(
values={"jokes": []},
@@ -10012,6 +10138,389 @@ def test_send_to_nested_graphs(
assert actual_history == expected_history
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
def test_weather_subgraph(
request: pytest.FixtureRequest, checkpointer_name: str, snapshot: SnapshotAssertion
) -> None:
from langchain_core.language_models.fake_chat_models import (
FakeMessagesListChatModel,
)
from langchain_core.messages import AIMessage, HumanMessage, ToolCall
from langchain_core.tools import tool
from langgraph.graph import MessagesState
checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name)
# setup subgraph
@tool
def get_weather(city: str):
"""Get the weather for a specific city"""
return f"I'ts sunny in {city}!"
weather_model = FakeMessagesListChatModel(
responses=[
AIMessage(
content="",
tool_calls=[
ToolCall(
id="tool_call123",
name="get_weather",
args={"city": "San Francisco"},
)
],
)
]
)
class SubGraphState(MessagesState):
city: str
def model_node(state: SubGraphState):
result = weather_model.invoke(state["messages"])
return {"city": cast(AIMessage, result).tool_calls[0]["args"]["city"]}
def weather_node(state: SubGraphState):
result = get_weather.invoke({"city": state["city"]})
return {"messages": [{"role": "assistant", "content": result}]}
subgraph = StateGraph(SubGraphState)
subgraph.add_node(model_node)
subgraph.add_node(weather_node)
subgraph.add_edge(START, "model_node")
subgraph.add_edge("model_node", "weather_node")
subgraph.add_edge("weather_node", END)
subgraph = subgraph.compile(interrupt_before=["weather_node"])
# setup main graph
class RouterState(MessagesState):
route: Literal["weather", "other"]
class Router(TypedDict):
route: Literal["weather", "other"]
router_model = FakeMessagesListChatModel(
responses=[
AIMessage(
content="",
tool_calls=[
ToolCall(
id="tool_call123",
name="router",
args={"dest": "weather"},
)
],
)
]
)
def router_node(state: RouterState):
system_message = "Classify the incoming query as either about weather or not."
messages = [{"role": "system", "content": system_message}] + state["messages"]
route = router_model.invoke(messages)
return {"route": cast(AIMessage, route).tool_calls[0]["args"]["dest"]}
def normal_llm_node(state: RouterState):
return {"messages": [AIMessage("Hello!")]}
def route_after_prediction(state: RouterState):
if state["route"] == "weather":
return "weather_graph"
else:
return "normal_llm_node"
graph = StateGraph(RouterState)
graph.add_node(router_node)
graph.add_node(normal_llm_node)
graph.add_node("weather_graph", subgraph)
graph.add_edge(START, "router_node")
graph.add_conditional_edges("router_node", route_after_prediction)
graph.add_edge("normal_llm_node", END)
graph.add_edge("weather_graph", END)
graph = graph.compile(checkpointer=checkpointer)
assert graph.get_graph(xray=1).draw_mermaid() == snapshot
config = {"configurable": {"thread_id": "1"}}
inputs = {"messages": [{"role": "user", "content": "what's the weather in sf"}]}
# run until interrupt
assert [
c
for c in graph.stream(
inputs, config=config, stream_mode="updates", subgraphs=True
)
] == [
((), {"router_node": {"route": "weather"}}),
((AnyStr("weather_graph:"),), {"model_node": {"city": "San Francisco"}}),
]
# check current state
state = graph.get_state(config)
assert state == StateSnapshot(
values={
"messages": [HumanMessage(content="what's the weather in sf", id=AnyStr())],
"route": "weather",
},
next=("weather_graph",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {"router_node": {"route": "weather"}},
"step": 1,
"parents": {},
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
tasks=(
PregelTask(
id=AnyStr(),
name="weather_graph",
state={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("weather_graph:"),
}
},
),
),
)
# update
graph.update_state(state.tasks[0].state, {"city": "la"})
# run after update
assert [
c
for c in graph.stream(
None, config=config, stream_mode="updates", subgraphs=True
)
] == [
(
(AnyStr("weather_graph:"),),
{
"weather_node": {
"messages": [{"role": "assistant", "content": "I'ts sunny in la!"}]
}
},
),
(
(),
{
"weather_graph": {
"messages": [
HumanMessage(content="what's the weather in sf", id=AnyStr()),
AIMessage(content="I'ts sunny in la!", id=AnyStr()),
]
}
},
),
]
# try updating acting as weather node
config = {"configurable": {"thread_id": "14"}}
inputs = {"messages": [{"role": "user", "content": "what's the weather in sf"}]}
assert [
c
for c in graph.stream(
inputs, config=config, stream_mode="updates", subgraphs=True
)
] == [
((), {"router_node": {"route": "weather"}}),
((AnyStr("weather_graph:"),), {"model_node": {"city": "San Francisco"}}),
]
state = graph.get_state(config, subgraphs=True)
assert state == StateSnapshot(
values={
"messages": [HumanMessage(content="what's the weather in sf", id=AnyStr())],
"route": "weather",
},
next=("weather_graph",),
config={
"configurable": {
"thread_id": "14",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {"router_node": {"route": "weather"}},
"step": 1,
"parents": {},
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "14",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
tasks=(
PregelTask(
id=AnyStr(),
name="weather_graph",
state=StateSnapshot(
values={
"messages": [
HumanMessage(
content="what's the weather in sf", id=AnyStr()
)
],
"city": "San Francisco",
},
next=("weather_node",),
config={
"configurable": {
"thread_id": "14",
"checkpoint_ns": AnyStr("weather_graph:"),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{
"": AnyStr(),
AnyStr("weather_graph:"): AnyStr(),
}
),
}
},
metadata={
"source": "loop",
"writes": {"model_node": {"city": "San Francisco"}},
"step": 1,
"parents": {"": AnyStr()},
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "14",
"checkpoint_ns": AnyStr("weather_graph:"),
"checkpoint_id": AnyStr(),
}
},
tasks=(PregelTask(id=AnyStr(), name="weather_node"),),
),
),
),
)
graph.update_state(
state.tasks[0].state.config,
{"messages": [{"role": "assistant", "content": "rainy"}]},
as_node="weather_node",
)
state = graph.get_state(config, subgraphs=True)
assert state == StateSnapshot(
values={
"messages": [HumanMessage(content="what's the weather in sf", id=AnyStr())],
"route": "weather",
},
next=("weather_graph",),
config={
"configurable": {
"thread_id": "14",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {"router_node": {"route": "weather"}},
"step": 1,
"parents": {},
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "14",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
tasks=(
PregelTask(
id=AnyStr(),
name="weather_graph",
state=StateSnapshot(
values={
"messages": [
HumanMessage(
content="what's the weather in sf", id=AnyStr()
),
AIMessage(content="rainy", id=AnyStr()),
],
"city": "San Francisco",
},
next=(),
config={
"configurable": {
"thread_id": "14",
"checkpoint_ns": AnyStr("weather_graph:"),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{
"": AnyStr(),
AnyStr("weather_graph:"): AnyStr(),
}
),
}
},
metadata={
"source": "update",
"step": 2,
"writes": {
"weather_node": {
"messages": [{"role": "assistant", "content": "rainy"}]
}
},
"parents": {"": AnyStr()},
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "14",
"checkpoint_ns": AnyStr("weather_graph:"),
"checkpoint_id": AnyStr(),
}
},
tasks=(),
),
),
),
)
assert [
c
for c in graph.stream(
None, config=config, stream_mode="updates", subgraphs=True
)
] == [
(
(),
{
"weather_graph": {
"messages": [
HumanMessage(content="what's the weather in sf", id=AnyStr()),
AIMessage(content="rainy", id=AnyStr()),
]
}
},
),
]
def test_repeat_condition(snapshot: SnapshotAssertion) -> None:
class AgentState(TypedDict):
hello: str
libs/langgraph/tests/test_pregel_async.py
+123 -46
View File
@@ -347,7 +347,7 @@ async def test_node_not_cancelled_on_other_node_interrupted(
assert not inner_task_cancelled
assert awhiles == 1
assert await graph.ainvoke(None, thread, debug=True) is None
assert await graph.ainvoke(None, thread, debug=True) == {"hello": "world"}
assert not inner_task_cancelled
assert awhiles == 1
@@ -2572,6 +2572,16 @@ async def test_conditional_graph(
)
assert [c async for c in app_w_interrupt.astream(None, config)] == [
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"input": "what is weather in sf",
},
},
{
"tools": {
"input": "what is weather in sf",
@@ -2797,6 +2807,16 @@ async def test_conditional_graph(
)
assert [c async for c in app_w_interrupt.astream(None, config)] == [
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:a different query",
),
"input": "what is weather in sf",
},
},
{
"tools": {
"input": "what is weather in sf",
@@ -2973,6 +2993,14 @@ async def test_conditional_graph(
)
assert [c async for c in app_w_interrupt.astream(None, config)] == [
{
"agent": {
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
},
},
{
"tools": {
"input": "what is weather in sf",
@@ -3011,6 +3039,26 @@ async def test_conditional_graph(
]
assert [c async for c in app_w_interrupt.astream(None, config)] == [
{
"agent": {
"input": "what is weather in sf",
"intermediate_steps": [
[
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"result for query",
]
],
"agent_outcome": AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
}
},
{
"tools": {
"input": "what is weather in sf",
@@ -6891,6 +6939,7 @@ async def test_nested_graph_interrupts_parallel(
{"my_key": ""},
]
assert [c async for c in app.astream(None, config, stream_mode="values")] == [
{"my_key": ""},
{"my_key": "got here and there and parallel"},
{"my_key": "got here and there and parallel and back again"},
]
@@ -6904,8 +6953,11 @@ async def test_nested_graph_interrupts_parallel(
{"my_key": ""},
]
# while we're waiting for the node w/ interrupt inside to finish
assert [c async for c in app.astream(None, config, stream_mode="values")] == []
assert [c async for c in app.astream(None, config, stream_mode="values")] == [
{"my_key": ""},
]
assert [c async for c in app.astream(None, config, stream_mode="values")] == [
{"my_key": ""},
{"my_key": "got here and there and parallel"},
{"my_key": "got here and there and parallel and back again"},
]
@@ -6919,9 +6971,11 @@ async def test_nested_graph_interrupts_parallel(
{"my_key": ""},
]
assert [c async for c in app.astream(None, config, stream_mode="values")] == [
{"my_key": ""},
{"my_key": "got here and there and parallel"},
]
assert [c async for c in app.astream(None, config, stream_mode="values")] == [
{"my_key": "got here and there and parallel"},
{"my_key": "got here and there and parallel and back again"},
]
@@ -7007,20 +7061,13 @@ async def test_doubly_nested_graph_interrupts(
c
async for c in app.astream({"my_key": "my value"}, config, stream_mode="values")
] == [
{
"my_key": "my value",
},
{
"my_key": "hi my value",
},
{"my_key": "my value"},
{"my_key": "hi my value"},
]
assert [c async for c in app.astream(None, config, stream_mode="values")] == [
{
"my_key": "hi my value here and there",
},
{
"my_key": "hi my value here and there and back again",
},
{"my_key": "hi my value"},
{"my_key": "hi my value here and there"},
{"my_key": "hi my value here and there and back again"},
]
@@ -8302,8 +8349,16 @@ async def test_doubly_nested_graph_state(
)
]
# 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"}},
)
]
@pytest.mark.skip("TODO")
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_send_to_nested_graphs(
request: pytest.FixtureRequest, checkpointer_name: str
@@ -8351,23 +8406,32 @@ async def test_send_to_nested_graphs(
"subjects": ["cats", "dogs"],
"jokes": [],
}
actual_snapshot = await graph.aget_state(config)
subgraph_nodes = list(actual_snapshot.subgraphs.keys())
assert len(subgraph_nodes) == 2
for subgraph_node in subgraph_nodes:
assert subgraph_node.split(":")[0] == "generate_joke"
# check state
outer_state = await graph.aget_state(config)
subgraphs = {
subgraph_node: await graph.aget_state(
{"configurable": {"thread_id": "1", "checkpoint_ns": subgraph_node}}
)
for subgraph_node in subgraph_nodes
}
expected_snapshot = StateSnapshot(
assert outer_state == StateSnapshot(
values={"subjects": ["cats", "dogs"], "jokes": []},
tasks=(
PregelTask(AnyStr(), "generate_joke"),
PregelTask(AnyStr(), "generate_joke"),
PregelTask(
AnyStr(),
"generate_joke",
state={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
}
},
),
PregelTask(
AnyStr(),
"generate_joke",
state={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
}
},
),
),
next=("generate_joke", "generate_joke"),
config={
@@ -8386,21 +8450,24 @@ async def test_send_to_nested_graphs(
"checkpoint_id": AnyStr(),
}
},
subgraphs=subgraphs,
)
assert actual_snapshot == expected_snapshot
# update state of dogs joke graph
await graph.aupdate_state(
outer_state.tasks[1].state, {"subject": "turtles - hohoho"}
)
# continue past interrupt
assert await graph.ainvoke(None, config=config) == {
"subjects": ["cats", "dogs"],
"jokes": ["Joke about cats - hohoho", "Joke about dogs - hohoho"],
"jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"],
}
actual_snapshot = await graph.aget_state(config)
expected_snapshot = StateSnapshot(
values={
"subjects": ["cats", "dogs"],
"jokes": ["Joke about cats - hohoho", "Joke about dogs - hohoho"],
"jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"],
},
tasks=(),
next=(),
@@ -8417,7 +8484,7 @@ async def test_send_to_nested_graphs(
"writes": {
"generate_joke": [
{"jokes": ["Joke about cats - hohoho"]},
{"jokes": ["Joke about dogs - hohoho"]},
{"jokes": ["Joke about turtles - hohoho"]},
]
},
"step": 1,
@@ -8435,18 +8502,11 @@ async def test_send_to_nested_graphs(
# test full history
actual_history = [c async for c in graph.aget_state_history(config)]
# get subgraph node state for expected history
subgraphs = {
subgraph_node: await graph.aget_state(
{"configurable": {"thread_id": "1", "checkpoint_ns": subgraph_node}}
)
for subgraph_node in subgraph_nodes
}
expected_history = [
StateSnapshot(
values={
"subjects": ["cats", "dogs"],
"jokes": ["Joke about cats - hohoho", "Joke about dogs - hohoho"],
"jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"],
},
tasks=(),
next=(),
@@ -8463,7 +8523,7 @@ async def test_send_to_nested_graphs(
"writes": {
"generate_joke": [
{"jokes": ["Joke about cats - hohoho"]},
{"jokes": ["Joke about dogs - hohoho"]},
{"jokes": ["Joke about turtles - hohoho"]},
]
},
"step": 1,
@@ -8481,8 +8541,26 @@ async def test_send_to_nested_graphs(
values={"subjects": ["cats", "dogs"], "jokes": []},
next=("generate_joke", "generate_joke"),
tasks=(
PregelTask(AnyStr(), "generate_joke"),
PregelTask(AnyStr(), "generate_joke"),
PregelTask(
AnyStr(),
"generate_joke",
state={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
}
},
),
PregelTask(
AnyStr(),
"generate_joke",
state={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
}
},
),
),
config={
"configurable": {
@@ -8500,7 +8578,6 @@ async def test_send_to_nested_graphs(
"checkpoint_id": AnyStr(),
}
},
subgraphs=subgraphs,
),
StateSnapshot(
values={"jokes": []},