mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-11-07 14:51:15 +01:00
Compare commits
2
Commits
sdk==0.4.14
...
main
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
e5e37b5d70 | ||
|
|
22025bbdbe |
No files matched your search
@@ -884,11 +884,23 @@ class Pregel(
|
||||
)
|
||||
|
||||
def _subgraph_for_namespace(
|
||||
self, config: RunnableConfig, checkpointer: BaseCheckpointSaver
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
checkpointer: BaseCheckpointSaver,
|
||||
*,
|
||||
for_update: bool = False,
|
||||
) -> tuple[PregelProtocol, RunnableConfig] | None:
|
||||
"""The subgraph a state method's namespaced config addresses, and the
|
||||
config to call it with, lending it `checkpointer`. `None` when the
|
||||
config is for this graph."""
|
||||
config is for this graph, or names no subgraph this graph declares or
|
||||
detects (such as one a tool invokes), so a read uses this graph's
|
||||
channels.
|
||||
|
||||
Raises:
|
||||
ValueError: If `for_update` and the namespace names no such
|
||||
subgraph: this graph's nodes and channels would write the wrong
|
||||
state into it.
|
||||
"""
|
||||
checkpoint_ns = config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "")
|
||||
if not checkpoint_ns or CONFIG_KEY_CHECKPOINTER in config[CONF]:
|
||||
return None
|
||||
@@ -897,7 +909,9 @@ class Pregel(
|
||||
return pregel, patch_configurable(
|
||||
config, {CONFIG_KEY_CHECKPOINTER: checkpointer}
|
||||
)
|
||||
raise ValueError(f"Subgraph {recast} not found")
|
||||
if for_update:
|
||||
raise ValueError(f"Subgraph {recast} not found")
|
||||
return None
|
||||
|
||||
def _apply_checkpointer_allowlist(
|
||||
self, checkpointer: BaseCheckpointSaver | None
|
||||
@@ -1143,7 +1157,9 @@ class Pregel(
|
||||
"""Get the subgraphs of the graph.
|
||||
|
||||
Args:
|
||||
namespace: The namespace to filter the subgraphs by.
|
||||
namespace: Yield only the subgraph at this exact namespace, such
|
||||
as `parent|child` (without task IDs). A nested namespace needs
|
||||
`recurse=True`.
|
||||
recurse: Whether to recurse into the subgraphs.
|
||||
If `False`, only the immediate subgraphs will be returned.
|
||||
|
||||
@@ -1151,9 +1167,9 @@ class Pregel(
|
||||
An iterator of the `(namespace, subgraph)` pairs.
|
||||
"""
|
||||
for name, node in self.nodes.items():
|
||||
# filter by prefix
|
||||
# filter by the namespace's first segment, matched whole
|
||||
if namespace is not None:
|
||||
if not namespace.startswith(name):
|
||||
if namespace != name and not namespace.startswith(name + NS_SEP):
|
||||
continue
|
||||
|
||||
# find the subgraph, if any
|
||||
@@ -1167,12 +1183,13 @@ class Pregel(
|
||||
if namespace is None:
|
||||
yield name, graph
|
||||
if recurse and isinstance(graph, Pregel):
|
||||
if namespace is not None:
|
||||
namespace = namespace[len(name) + 1 :]
|
||||
child_namespace = (
|
||||
None if namespace is None else namespace[len(name) + 1 :]
|
||||
)
|
||||
yield from (
|
||||
(f"{name}{NS_SEP}{n}", s)
|
||||
for n, s in graph.get_subgraphs(
|
||||
namespace=namespace, recurse=recurse
|
||||
namespace=child_namespace, recurse=recurse
|
||||
)
|
||||
)
|
||||
|
||||
@@ -1182,7 +1199,9 @@ class Pregel(
|
||||
"""Get the subgraphs of the graph.
|
||||
|
||||
Args:
|
||||
namespace: The namespace to filter the subgraphs by.
|
||||
namespace: Yield only the subgraph at this exact namespace, such
|
||||
as `parent|child` (without task IDs). A nested namespace needs
|
||||
`recurse=True`.
|
||||
recurse: Whether to recurse into the subgraphs.
|
||||
If `False`, only the immediate subgraphs will be returned.
|
||||
|
||||
@@ -1641,7 +1660,9 @@ class Pregel(
|
||||
if any(len(u) == 0 for u in supersteps):
|
||||
raise ValueError("No updates provided")
|
||||
|
||||
if subgraph := self._subgraph_for_namespace(config, checkpointer):
|
||||
if subgraph := self._subgraph_for_namespace(
|
||||
config, checkpointer, for_update=True
|
||||
):
|
||||
pregel, subgraph_config = subgraph
|
||||
return pregel.bulk_update_state(subgraph_config, supersteps)
|
||||
|
||||
@@ -2142,7 +2163,9 @@ class Pregel(
|
||||
if any(len(u) == 0 for u in supersteps):
|
||||
raise ValueError("No updates provided")
|
||||
|
||||
if subgraph := self._subgraph_for_namespace(config, checkpointer):
|
||||
if subgraph := self._subgraph_for_namespace(
|
||||
config, checkpointer, for_update=True
|
||||
):
|
||||
pregel, subgraph_config = subgraph
|
||||
return await pregel.abulk_update_state(subgraph_config, supersteps)
|
||||
|
||||
|
||||
@@ -6,16 +6,24 @@ Covers three checkpointer settings for subgraph state:
|
||||
interrupt support, but state resets each invocation. This is the common case
|
||||
when an agent is invoked from inside a tool used by another agent.
|
||||
- checkpointer=True: "stateful" — state accumulates across invocations on the same thread id
|
||||
|
||||
Also covers reading a subgraph's namespace from the parent graph, including
|
||||
one where no subgraph is declared, such as a subagent a tool invokes, and
|
||||
finding the subgraph a namespace names.
|
||||
"""
|
||||
|
||||
from typing import Annotated, Any
|
||||
from uuid import uuid4
|
||||
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
import pytest
|
||||
from langchain_core.messages import AIMessage, AnyMessage, HumanMessage
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.checkpoint.base import BaseCheckpointSaver
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.graph.message import MessagesState
|
||||
from langgraph.graph.message import MessagesState, _messages_delta_reducer
|
||||
from langgraph.types import Command, Interrupt, interrupt
|
||||
from tests.any_str import AnyStr
|
||||
|
||||
@@ -685,3 +693,301 @@ def test_child_with_own_thread_id_keeps_namespace(
|
||||
parent.invoke({"result": ""}, config=parent_config)
|
||||
state2 = child.get_state(child_config)
|
||||
assert state2.values.get("count") == 2
|
||||
|
||||
|
||||
# -- reading a subgraph's namespace from the parent graph --
|
||||
|
||||
|
||||
class NotesState(MessagesState):
|
||||
notes: str # a key only the subgraphs have
|
||||
|
||||
|
||||
class DeltaMessagesState(TypedDict):
|
||||
# Stores `messages` as deltas between snapshots, as deepagents' parent does.
|
||||
messages: Annotated[list[AnyMessage], DeltaChannel(_messages_delta_reducer)]
|
||||
|
||||
|
||||
class DeltaNotesState(DeltaMessagesState):
|
||||
notes: str
|
||||
|
||||
|
||||
def _notes_graph(label: str, state_schema: type = NotesState) -> Any:
|
||||
# Typed `dict` so it fits either schema (see `_run_tool_subgraph`).
|
||||
def work(state: dict) -> dict:
|
||||
return {"messages": [AIMessage(content=f"{label} done")], "notes": label}
|
||||
|
||||
return (
|
||||
StateGraph(state_schema)
|
||||
.add_node("work", work)
|
||||
.add_edge(START, "work")
|
||||
.compile()
|
||||
)
|
||||
|
||||
|
||||
def _namespace_configs(parent: Any) -> dict[str, RunnableConfig]:
|
||||
"""Run `parent` on a new thread. Returns a config for each subgraph
|
||||
namespace it wrote, keyed by the node the namespace starts with."""
|
||||
thread_id = str(uuid4())
|
||||
return {
|
||||
ns[0].split(":")[0]: {
|
||||
"configurable": {"thread_id": thread_id, "checkpoint_ns": ns[0]}
|
||||
}
|
||||
for ns, _ in parent.stream(
|
||||
{"messages": [HumanMessage(content="go")]},
|
||||
{"configurable": {"thread_id": thread_id}},
|
||||
subgraphs=True,
|
||||
)
|
||||
if ns
|
||||
}
|
||||
|
||||
|
||||
def _run_tool_subgraph(
|
||||
checkpointer: BaseCheckpointSaver,
|
||||
*,
|
||||
parent_schema: type = MessagesState,
|
||||
subagent_schema: type = NotesState,
|
||||
) -> tuple[Any, RunnableConfig]:
|
||||
"""Run a parent whose `tools` node invokes a subgraph it doesn't declare,
|
||||
as deepagents runs subagents. Returns the parent and a config for the
|
||||
subgraph's `tools:<task_id>` namespace."""
|
||||
# Subgraph detection doesn't follow a subscript, so `tools` declares none.
|
||||
subagents = {"researcher": _notes_graph("researcher", subagent_schema)}
|
||||
|
||||
# Typed `dict`: a schema here would add its `messages` channel to the
|
||||
# parent's, and a plain one clashes with a `DeltaChannel` one.
|
||||
def tools(state: dict) -> dict:
|
||||
subagents["researcher"].invoke(
|
||||
{"messages": [HumanMessage(content="research apples")]}
|
||||
)
|
||||
return {"messages": [AIMessage(content="tool done")]}
|
||||
|
||||
parent = (
|
||||
StateGraph(parent_schema)
|
||||
.add_node("tools", tools)
|
||||
.add_edge(START, "tools")
|
||||
.compile(checkpointer=checkpointer)
|
||||
)
|
||||
return parent, _namespace_configs(parent)["tools"]
|
||||
|
||||
|
||||
def test_undeclared_subgraph_get_state(sync_checkpointer: BaseCheckpointSaver) -> None:
|
||||
"""A namespace with no declared subgraph is read with the parent graph."""
|
||||
parent, ns_config = _run_tool_subgraph(sync_checkpointer)
|
||||
|
||||
state = parent.get_state(ns_config)
|
||||
|
||||
assert [m.text for m in state.values["messages"]] == [
|
||||
"research apples",
|
||||
"researcher done",
|
||||
]
|
||||
assert state.config["configurable"]["checkpoint_ns"] == AnyStr("tools:")
|
||||
|
||||
|
||||
def test_undeclared_subgraph_get_state_history(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
"""History of a namespace with no declared subgraph is read with the parent
|
||||
graph, and only holds that namespace's checkpoints."""
|
||||
parent, ns_config = _run_tool_subgraph(sync_checkpointer)
|
||||
|
||||
history = list(parent.get_state_history(ns_config))
|
||||
|
||||
assert [[m.text for m in s.values.get("messages", [])] for s in history] == [
|
||||
["research apples", "researcher done"],
|
||||
["research apples"],
|
||||
[],
|
||||
]
|
||||
assert {s.config["configurable"]["checkpoint_ns"] for s in history} == {
|
||||
ns_config["configurable"]["checkpoint_ns"]
|
||||
}
|
||||
|
||||
|
||||
def test_undeclared_subgraph_update_state_raises(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
"""Unlike a read, an update can't fall back to the parent graph: its nodes
|
||||
and channels would write the wrong state into the subgraph's namespace."""
|
||||
parent, ns_config = _run_tool_subgraph(sync_checkpointer)
|
||||
|
||||
with pytest.raises(ValueError, match="Subgraph tools not found"):
|
||||
parent.update_state(ns_config, {"messages": [AIMessage(content="edit")]})
|
||||
|
||||
assert len(list(parent.get_state_history(ns_config))) == 3
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("parent_schema", "subagent_schema"),
|
||||
[(DeltaMessagesState, NotesState), (MessagesState, DeltaNotesState)],
|
||||
ids=["delta_parent", "delta_subagent"],
|
||||
)
|
||||
def test_undeclared_subgraph_read_across_messages_channel_types(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
parent_schema: type,
|
||||
subagent_schema: type,
|
||||
) -> None:
|
||||
"""deepagents' parent stores `messages` in a `DeltaChannel`, and its
|
||||
subagents in a plain `add_messages` channel. Reading a subagent's namespace
|
||||
with the parent's channel restores its messages, either way round."""
|
||||
parent, ns_config = _run_tool_subgraph(
|
||||
sync_checkpointer,
|
||||
parent_schema=parent_schema,
|
||||
subagent_schema=subagent_schema,
|
||||
)
|
||||
|
||||
state = parent.get_state(ns_config)
|
||||
history = list(parent.get_state_history(ns_config))
|
||||
|
||||
assert [m.text for m in state.values["messages"]] == [
|
||||
"research apples",
|
||||
"researcher done",
|
||||
]
|
||||
assert [[m.text for m in s.values.get("messages", [])] for s in history] == [
|
||||
["research apples", "researcher done"],
|
||||
["research apples"],
|
||||
[],
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("review_declared", [False, True])
|
||||
def test_prefix_named_namespaces_resolve_to_their_own_graph(
|
||||
sync_checkpointer: BaseCheckpointSaver, review_declared: bool
|
||||
) -> None:
|
||||
"""`research_review` starts with `research`, whose subgraph has a `review`
|
||||
node with its own subgraph. Each namespace still reads with its own graph:
|
||||
a declared subgraph's channels, or the parent's when none is declared."""
|
||||
research = (
|
||||
StateGraph(NotesState)
|
||||
.add_node("review", _notes_graph("research"))
|
||||
.add_edge(START, "review")
|
||||
.compile()
|
||||
)
|
||||
|
||||
class ReviewState(NotesState):
|
||||
verdict: str # a key `research|review`'s graph doesn't have
|
||||
|
||||
def review(state: ReviewState) -> dict:
|
||||
return {
|
||||
"messages": [AIMessage(content="review done")],
|
||||
"notes": "review",
|
||||
"verdict": "approved",
|
||||
}
|
||||
|
||||
# Subgraph detection doesn't follow a subscript, so `call_review` declares
|
||||
# none.
|
||||
subgraphs = {
|
||||
"review": StateGraph(ReviewState)
|
||||
.add_node("work", review)
|
||||
.add_edge(START, "work")
|
||||
.compile()
|
||||
}
|
||||
|
||||
def call_review(state: MessagesState) -> dict:
|
||||
subgraphs["review"].invoke({"messages": [HumanMessage(content="check it")]})
|
||||
return {}
|
||||
|
||||
parent = (
|
||||
StateGraph(MessagesState)
|
||||
.add_node("research", research)
|
||||
.add_node(
|
||||
"research_review",
|
||||
subgraphs["review"] if review_declared else call_review,
|
||||
)
|
||||
.add_edge(START, "research")
|
||||
.add_edge("research", "research_review")
|
||||
.compile(checkpointer=sync_checkpointer)
|
||||
)
|
||||
configs = _namespace_configs(parent)
|
||||
|
||||
def values(name: str) -> dict:
|
||||
return {
|
||||
k: [m.text for m in v] if k == "messages" else v
|
||||
for k, v in parent.get_state(configs[name]).values.items()
|
||||
}
|
||||
|
||||
assert values("research") == {
|
||||
"messages": ["go", "research done"],
|
||||
"notes": "research",
|
||||
}
|
||||
if review_declared:
|
||||
assert values("research_review") == {
|
||||
"messages": ["go", "research done", "review done"],
|
||||
"notes": "review",
|
||||
"verdict": "approved",
|
||||
}
|
||||
else:
|
||||
# The parent's channels have no `notes` or `verdict`.
|
||||
assert values("research_review") == {"messages": ["check it", "review done"]}
|
||||
|
||||
|
||||
def _prefix_named_subgraphs() -> tuple[Any, dict[str, Any]]:
|
||||
"""A parent whose subgraph namespaces start with one another's names:
|
||||
`research` has a `review` subgraph, and the parent declares
|
||||
`research_review` and `review` too. Returns the parent and each subgraph,
|
||||
keyed by its namespace."""
|
||||
graphs = {
|
||||
"research|review": _notes_graph("research's review"),
|
||||
"research_review": _notes_graph("research_review"),
|
||||
"review": _notes_graph("review"),
|
||||
}
|
||||
graphs["research"] = (
|
||||
StateGraph(NotesState)
|
||||
.add_node("review", graphs["research|review"])
|
||||
.add_edge(START, "review")
|
||||
.compile()
|
||||
)
|
||||
parent = (
|
||||
StateGraph(NotesState)
|
||||
.add_node("research", graphs["research"])
|
||||
.add_node("research_review", graphs["research_review"])
|
||||
.add_node("review", graphs["review"])
|
||||
.add_edge(START, "research")
|
||||
.add_edge("research", "research_review")
|
||||
.add_edge("research_review", "review")
|
||||
.compile()
|
||||
)
|
||||
return parent, graphs
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"namespace", ["research", "research|review", "research_review", "review"]
|
||||
)
|
||||
def test_get_subgraphs_matches_whole_namespace_segments(namespace: str) -> None:
|
||||
"""`get_subgraphs(namespace=...)` yields only the subgraph at that
|
||||
namespace. `research_review` isn't `research` then `review`, and
|
||||
`research|review` doesn't also find the parent's own `review`."""
|
||||
parent, graphs = _prefix_named_subgraphs()
|
||||
|
||||
assert list(parent.get_subgraphs(namespace=namespace, recurse=True)) == [
|
||||
(namespace, graphs[namespace])
|
||||
]
|
||||
|
||||
|
||||
def test_nested_namespace_update_raises_despite_top_level_namesake(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
"""`research|review` names a `review` inside `research`'s subgraph, which
|
||||
declares no subgraph there. An update to it raises rather than writing
|
||||
through the parent's own `review`."""
|
||||
research = (
|
||||
StateGraph(NotesState)
|
||||
.add_node("review", lambda state: {})
|
||||
.add_edge(START, "review")
|
||||
.compile()
|
||||
)
|
||||
parent = (
|
||||
StateGraph(NotesState)
|
||||
.add_node("research", research)
|
||||
.add_node("review", _notes_graph("review"))
|
||||
.add_edge(START, "research")
|
||||
.add_edge("research", "review")
|
||||
.compile(checkpointer=sync_checkpointer)
|
||||
)
|
||||
config: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": str(uuid4()),
|
||||
"checkpoint_ns": "research:a|review:b",
|
||||
}
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match=r"Subgraph research\|review not found"):
|
||||
parent.update_state(config, {"notes": "edit"})
|
||||
@@ -6,18 +6,25 @@ Covers three checkpointer settings for subgraph state:
|
||||
interrupt support, but state resets each invocation. This is the common case
|
||||
when an agent is invoked from inside a tool used by another agent.
|
||||
- checkpointer=True: "stateful" — state accumulates across invocations on the same thread id
|
||||
|
||||
Also covers reading a subgraph's namespace from the parent graph, including
|
||||
one where no subgraph is declared, such as a subagent a tool invokes, and
|
||||
finding the subgraph a namespace names.
|
||||
"""
|
||||
|
||||
import sys
|
||||
from typing import Annotated, Any
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langchain_core.messages import AIMessage, AnyMessage, HumanMessage
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.checkpoint.base import BaseCheckpointSaver
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.graph.message import MessagesState
|
||||
from langgraph.graph.message import MessagesState, _messages_delta_reducer
|
||||
from langgraph.types import Command, Interrupt, interrupt
|
||||
from tests.any_str import AnyStr
|
||||
|
||||
@@ -707,3 +714,315 @@ async def test_child_with_own_thread_id_keeps_namespace_async(
|
||||
await parent.ainvoke({"result": ""}, config=parent_config)
|
||||
state2 = await child.aget_state(child_config)
|
||||
assert state2.values.get("count") == 2
|
||||
|
||||
|
||||
# -- reading a subgraph's namespace from the parent graph --
|
||||
|
||||
|
||||
class NotesState(MessagesState):
|
||||
notes: str # a key only the subgraphs have
|
||||
|
||||
|
||||
class DeltaMessagesState(TypedDict):
|
||||
# Stores `messages` as deltas between snapshots, as deepagents' parent does.
|
||||
messages: Annotated[list[AnyMessage], DeltaChannel(_messages_delta_reducer)]
|
||||
|
||||
|
||||
class DeltaNotesState(DeltaMessagesState):
|
||||
notes: str
|
||||
|
||||
|
||||
def _notes_graph(label: str, state_schema: type = NotesState) -> Any:
|
||||
# Typed `dict` so it fits either schema (see `_run_tool_subgraph`).
|
||||
def work(state: dict) -> dict:
|
||||
return {"messages": [AIMessage(content=f"{label} done")], "notes": label}
|
||||
|
||||
return (
|
||||
StateGraph(state_schema)
|
||||
.add_node("work", work)
|
||||
.add_edge(START, "work")
|
||||
.compile()
|
||||
)
|
||||
|
||||
|
||||
async def _namespace_configs(parent: Any) -> dict[str, RunnableConfig]:
|
||||
"""Run `parent` on a new thread. Returns a config for each subgraph
|
||||
namespace it wrote, keyed by the node the namespace starts with."""
|
||||
thread_id = str(uuid4())
|
||||
return {
|
||||
ns[0].split(":")[0]: {
|
||||
"configurable": {"thread_id": thread_id, "checkpoint_ns": ns[0]}
|
||||
}
|
||||
async for ns, _ in parent.astream(
|
||||
{"messages": [HumanMessage(content="go")]},
|
||||
{"configurable": {"thread_id": thread_id}},
|
||||
subgraphs=True,
|
||||
)
|
||||
if ns
|
||||
}
|
||||
|
||||
|
||||
async def _run_tool_subgraph(
|
||||
checkpointer: BaseCheckpointSaver,
|
||||
*,
|
||||
parent_schema: type = MessagesState,
|
||||
subagent_schema: type = NotesState,
|
||||
) -> tuple[Any, RunnableConfig]:
|
||||
"""Run a parent whose `tools` node invokes a subgraph it doesn't declare,
|
||||
as deepagents runs subagents. Returns the parent and a config for the
|
||||
subgraph's `tools:<task_id>` namespace."""
|
||||
# Subgraph detection doesn't follow a subscript, so `tools` declares none.
|
||||
subagents = {"researcher": _notes_graph("researcher", subagent_schema)}
|
||||
|
||||
# Typed `dict`: a schema here would add its `messages` channel to the
|
||||
# parent's, and a plain one clashes with a `DeltaChannel` one.
|
||||
async def tools(state: dict) -> dict:
|
||||
await subagents["researcher"].ainvoke(
|
||||
{"messages": [HumanMessage(content="research apples")]}
|
||||
)
|
||||
return {"messages": [AIMessage(content="tool done")]}
|
||||
|
||||
parent = (
|
||||
StateGraph(parent_schema)
|
||||
.add_node("tools", tools)
|
||||
.add_edge(START, "tools")
|
||||
.compile(checkpointer=checkpointer)
|
||||
)
|
||||
return parent, (await _namespace_configs(parent))["tools"]
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
async def test_undeclared_subgraph_get_state_async(
|
||||
async_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
"""A namespace with no declared subgraph is read with the parent graph."""
|
||||
parent, ns_config = await _run_tool_subgraph(async_checkpointer)
|
||||
|
||||
state = await parent.aget_state(ns_config)
|
||||
|
||||
assert [m.text for m in state.values["messages"]] == [
|
||||
"research apples",
|
||||
"researcher done",
|
||||
]
|
||||
assert state.config["configurable"]["checkpoint_ns"] == AnyStr("tools:")
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
async def test_undeclared_subgraph_get_state_history_async(
|
||||
async_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
"""History of a namespace with no declared subgraph is read with the parent
|
||||
graph, and only holds that namespace's checkpoints."""
|
||||
parent, ns_config = await _run_tool_subgraph(async_checkpointer)
|
||||
|
||||
history = [s async for s in parent.aget_state_history(ns_config)]
|
||||
|
||||
assert [[m.text for m in s.values.get("messages", [])] for s in history] == [
|
||||
["research apples", "researcher done"],
|
||||
["research apples"],
|
||||
[],
|
||||
]
|
||||
assert {s.config["configurable"]["checkpoint_ns"] for s in history} == {
|
||||
ns_config["configurable"]["checkpoint_ns"]
|
||||
}
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
async def test_undeclared_subgraph_update_state_raises_async(
|
||||
async_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
"""Unlike a read, an update can't fall back to the parent graph: its nodes
|
||||
and channels would write the wrong state into the subgraph's namespace."""
|
||||
parent, ns_config = await _run_tool_subgraph(async_checkpointer)
|
||||
|
||||
with pytest.raises(ValueError, match="Subgraph tools not found"):
|
||||
await parent.aupdate_state(ns_config, {"messages": [AIMessage(content="edit")]})
|
||||
|
||||
assert len([s async for s in parent.aget_state_history(ns_config)]) == 3
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
@pytest.mark.parametrize(
|
||||
("parent_schema", "subagent_schema"),
|
||||
[(DeltaMessagesState, NotesState), (MessagesState, DeltaNotesState)],
|
||||
ids=["delta_parent", "delta_subagent"],
|
||||
)
|
||||
async def test_undeclared_subgraph_read_across_messages_channel_types_async(
|
||||
async_checkpointer: BaseCheckpointSaver,
|
||||
parent_schema: type,
|
||||
subagent_schema: type,
|
||||
) -> None:
|
||||
"""deepagents' parent stores `messages` in a `DeltaChannel`, and its
|
||||
subagents in a plain `add_messages` channel. Reading a subagent's namespace
|
||||
with the parent's channel restores its messages, either way round."""
|
||||
parent, ns_config = await _run_tool_subgraph(
|
||||
async_checkpointer,
|
||||
parent_schema=parent_schema,
|
||||
subagent_schema=subagent_schema,
|
||||
)
|
||||
|
||||
state = await parent.aget_state(ns_config)
|
||||
history = [s async for s in parent.aget_state_history(ns_config)]
|
||||
|
||||
assert [m.text for m in state.values["messages"]] == [
|
||||
"research apples",
|
||||
"researcher done",
|
||||
]
|
||||
assert [[m.text for m in s.values.get("messages", [])] for s in history] == [
|
||||
["research apples", "researcher done"],
|
||||
["research apples"],
|
||||
[],
|
||||
]
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
@pytest.mark.parametrize("review_declared", [False, True])
|
||||
async def test_prefix_named_namespaces_resolve_to_their_own_graph_async(
|
||||
async_checkpointer: BaseCheckpointSaver, review_declared: bool
|
||||
) -> None:
|
||||
"""`research_review` starts with `research`, whose subgraph has a `review`
|
||||
node with its own subgraph. Each namespace still reads with its own graph:
|
||||
a declared subgraph's channels, or the parent's when none is declared."""
|
||||
research = (
|
||||
StateGraph(NotesState)
|
||||
.add_node("review", _notes_graph("research"))
|
||||
.add_edge(START, "review")
|
||||
.compile()
|
||||
)
|
||||
|
||||
class ReviewState(NotesState):
|
||||
verdict: str # a key `research|review`'s graph doesn't have
|
||||
|
||||
def review(state: ReviewState) -> dict:
|
||||
return {
|
||||
"messages": [AIMessage(content="review done")],
|
||||
"notes": "review",
|
||||
"verdict": "approved",
|
||||
}
|
||||
|
||||
# Subgraph detection doesn't follow a subscript, so `call_review` declares
|
||||
# none.
|
||||
subgraphs = {
|
||||
"review": StateGraph(ReviewState)
|
||||
.add_node("work", review)
|
||||
.add_edge(START, "work")
|
||||
.compile()
|
||||
}
|
||||
|
||||
async def call_review(state: MessagesState) -> dict:
|
||||
await subgraphs["review"].ainvoke(
|
||||
{"messages": [HumanMessage(content="check it")]}
|
||||
)
|
||||
return {}
|
||||
|
||||
parent = (
|
||||
StateGraph(MessagesState)
|
||||
.add_node("research", research)
|
||||
.add_node(
|
||||
"research_review",
|
||||
subgraphs["review"] if review_declared else call_review,
|
||||
)
|
||||
.add_edge(START, "research")
|
||||
.add_edge("research", "research_review")
|
||||
.compile(checkpointer=async_checkpointer)
|
||||
)
|
||||
configs = await _namespace_configs(parent)
|
||||
|
||||
async def values(name: str) -> dict:
|
||||
state = await parent.aget_state(configs[name])
|
||||
return {
|
||||
k: [m.text for m in v] if k == "messages" else v
|
||||
for k, v in state.values.items()
|
||||
}
|
||||
|
||||
assert await values("research") == {
|
||||
"messages": ["go", "research done"],
|
||||
"notes": "research",
|
||||
}
|
||||
if review_declared:
|
||||
assert await values("research_review") == {
|
||||
"messages": ["go", "research done", "review done"],
|
||||
"notes": "review",
|
||||
"verdict": "approved",
|
||||
}
|
||||
else:
|
||||
# The parent's channels have no `notes` or `verdict`.
|
||||
assert await values("research_review") == {
|
||||
"messages": ["check it", "review done"]
|
||||
}
|
||||
|
||||
|
||||
def _prefix_named_subgraphs() -> tuple[Any, dict[str, Any]]:
|
||||
"""A parent whose subgraph namespaces start with one another's names:
|
||||
`research` has a `review` subgraph, and the parent declares
|
||||
`research_review` and `review` too. Returns the parent and each subgraph,
|
||||
keyed by its namespace."""
|
||||
graphs = {
|
||||
"research|review": _notes_graph("research's review"),
|
||||
"research_review": _notes_graph("research_review"),
|
||||
"review": _notes_graph("review"),
|
||||
}
|
||||
graphs["research"] = (
|
||||
StateGraph(NotesState)
|
||||
.add_node("review", graphs["research|review"])
|
||||
.add_edge(START, "review")
|
||||
.compile()
|
||||
)
|
||||
parent = (
|
||||
StateGraph(NotesState)
|
||||
.add_node("research", graphs["research"])
|
||||
.add_node("research_review", graphs["research_review"])
|
||||
.add_node("review", graphs["review"])
|
||||
.add_edge(START, "research")
|
||||
.add_edge("research", "research_review")
|
||||
.add_edge("research_review", "review")
|
||||
.compile()
|
||||
)
|
||||
return parent, graphs
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"namespace", ["research", "research|review", "research_review", "review"]
|
||||
)
|
||||
async def test_get_subgraphs_matches_whole_namespace_segments_async(
|
||||
namespace: str,
|
||||
) -> None:
|
||||
"""`aget_subgraphs(namespace=...)` yields only the subgraph at that
|
||||
namespace. `research_review` isn't `research` then `review`, and
|
||||
`research|review` doesn't also find the parent's own `review`."""
|
||||
parent, graphs = _prefix_named_subgraphs()
|
||||
|
||||
assert [
|
||||
s async for s in parent.aget_subgraphs(namespace=namespace, recurse=True)
|
||||
] == [(namespace, graphs[namespace])]
|
||||
|
||||
|
||||
async def test_nested_namespace_update_raises_despite_top_level_namesake_async(
|
||||
async_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
"""`research|review` names a `review` inside `research`'s subgraph, which
|
||||
declares no subgraph there. An update to it raises rather than writing
|
||||
through the parent's own `review`."""
|
||||
research = (
|
||||
StateGraph(NotesState)
|
||||
.add_node("review", lambda state: {})
|
||||
.add_edge(START, "review")
|
||||
.compile()
|
||||
)
|
||||
parent = (
|
||||
StateGraph(NotesState)
|
||||
.add_node("research", research)
|
||||
.add_node("review", _notes_graph("review"))
|
||||
.add_edge(START, "research")
|
||||
.add_edge("research", "review")
|
||||
.compile(checkpointer=async_checkpointer)
|
||||
)
|
||||
config: RunnableConfig = {
|
||||
"configurable": {
|
||||
"thread_id": str(uuid4()),
|
||||
"checkpoint_ns": "research:a|review:b",
|
||||
}
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match=r"Subgraph research\|review not found"):
|
||||
await parent.aupdate_state(config, {"notes": "edit"})
|
||||
Reference in new issue
Block a user