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

Implement interrupt_before=* and interrupt_after=*

Browse Source 
- This interrupts execution before/after all nodes
This commit is contained in:
Nuno Campos
2024-04-30 10:05:53 -07:00
parent 9a79eb54f9
commit afa6a91934
7 changed files with 262 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files
langgraph/graph/graph.py
+15 -9
View File
@@ -28,6 +28,7 @@ from langgraph.checkpoint import BaseCheckpointSaver
from langgraph.constants import TAG_HIDDEN
from langgraph.pregel import Channel, Pregel
from langgraph.pregel.read import PregelNode
from langgraph.pregel.types import All
from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
from langgraph.utils import DrawableGraph, RunnableCallable, coerce_to_runnable
@@ -295,8 +296,8 @@ class Graph:
def compile(
self,
checkpointer: Optional[BaseCheckpointSaver] = None,
interrupt_before: Optional[Sequence[str]] = None,
interrupt_after: Optional[Sequence[str]] = None,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
debug: bool = False,
) -> "CompiledGraph":
# assign default values
@@ -304,11 +305,16 @@ class Graph:
interrupt_after = interrupt_after or []
# validate the graph
self.validate(interrupt=interrupt_before + interrupt_after)
self.validate(
interrupt=(interrupt_before if interrupt_before != "*" else [])
+ interrupt_after
if interrupt_after != "*"
else []
)
# create empty compiled graph
compiled = CompiledGraph(
graph=self,
builder=self,
nodes={},
channels={START: EphemeralValue(Any), END: EphemeralValue(Any)},
input_channels=START,
@@ -338,7 +344,7 @@ class Graph:
class CompiledGraph(Pregel):
graph: Graph
builder: Graph
def attach_node(self, key: str, node: Runnable) -> None:
self.channels[key] = EphemeralValue(Any)
@@ -400,7 +406,7 @@ class CompiledGraph(Pregel):
END: graph.add_node(self.get_output_schema(config), END)
}
for key, node in self.graph.nodes.items():
for key, node in self.builder.nodes.items():
if xray:
subgraph = (
node.get_graph(
@@ -424,11 +430,11 @@ class CompiledGraph(Pregel):
n = graph.add_node(node, key)
start_nodes[key] = n
end_nodes[key] = n
for start, end in sorted(self.graph._all_edges):
for start, end in sorted(self.builder._all_edges):
graph.add_edge(start_nodes[start], end_nodes[end])
for start, branches in self.graph.branches.items():
for start, branches in self.builder.branches.items():
default_ends = {
**{k: k for k in self.graph.nodes if k != start},
**{k: k for k in self.builder.nodes if k != start},
END: END,
}
for _, branch in branches.items():
langgraph/graph/state.py
+15 -9
View File
@@ -16,6 +16,7 @@ from langgraph.checkpoint import BaseCheckpointSaver
from langgraph.constants import TAG_HIDDEN
from langgraph.graph.graph import END, START, Branch, CompiledGraph, Graph
from langgraph.pregel.read import ChannelRead, PregelNode
from langgraph.pregel.types import All
from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry
from langgraph.utils import RunnableCallable
@@ -100,8 +101,8 @@ class StateGraph(Graph):
def compile(
self,
checkpointer: Optional[BaseCheckpointSaver] = None,
interrupt_before: Optional[Sequence[str]] = None,
interrupt_after: Optional[Sequence[str]] = None,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
debug: bool = False,
) -> CompiledGraph:
"""Compiles the state graph into a `CompiledGraph` object.
@@ -120,14 +121,19 @@ class StateGraph(Graph):
interrupt_after = interrupt_after or []
# validate the graph
self.validate(interrupt=interrupt_before + interrupt_after)
self.validate(
interrupt=(interrupt_before if interrupt_before != "*" else [])
+ interrupt_after
if interrupt_after != "*"
else []
)
# prepare output channels
state_keys = list(self.channels)
output_channels = state_keys[0] if state_keys == ["__root__"] else state_keys
compiled = CompiledStateGraph(
graph=self,
builder=self,
nodes={},
channels={**self.channels, START: EphemeralValue(self.schema)},
input_channels=START,
@@ -159,7 +165,7 @@ class StateGraph(Graph):
class CompiledStateGraph(CompiledGraph):
graph: StateGraph
builder: StateGraph
def attach_node(self, key: str, node: Optional[Runnable]) -> None:
def _get_state_key(input: dict, config: RunnableConfig, *, key: str) -> Any:
@@ -170,7 +176,7 @@ class CompiledStateGraph(CompiledGraph):
else:
return input.get(key, SKIP_WRITE)
state_keys = list(self.graph.channels)
state_keys = list(self.builder.channels)
# state updaters
state_write_entries = [
(
@@ -210,7 +216,7 @@ class CompiledStateGraph(CompiledGraph):
mapper=(
None
if state_keys == ["__root__"]
else partial(_coerce_state, self.graph.schema)
else partial(_coerce_state, self.builder.schema)
),
writers=[
# publish to this channel and state keys
@@ -265,13 +271,13 @@ class CompiledStateGraph(CompiledGraph):
return ChannelWrite(writes, tags=[TAG_HIDDEN])
# attach branch publisher
self.nodes[start] |= branch.run(branch_writer, _get_state_reader(self.graph))
self.nodes[start] |= branch.run(branch_writer, _get_state_reader(self.builder))
# attach branch subscribers
ends = (
branch.ends.values()
if branch.ends
else [node for node in self.graph.nodes if node != branch.then]
else [node for node in self.builder.nodes if node != branch.then]
)
for end in ends:
if end != END:
langgraph/pregel/__init__.py
+32 -22
View File
@@ -65,6 +65,7 @@ from langgraph.constants import (
CONFIG_KEY_READ,
CONFIG_KEY_SEND,
INTERRUPT,
TAG_HIDDEN,
)
from langgraph.pregel.debug import (
print_step_checkpoint,
@@ -81,6 +82,7 @@ from langgraph.pregel.io import (
from langgraph.pregel.log import logger
from langgraph.pregel.read import PregelNode
from langgraph.pregel.types import (
All,
PregelExecutableTask,
PregelTaskDescription,
StateSnapshot,
@@ -194,9 +196,9 @@ class Pregel(
stream_channels: Optional[Union[str, Sequence[str]]] = None
"""Channels to stream, defaults to all channels not in reserved channels"""
interrupt_after_nodes: Sequence[str] = Field(default_factory=list)
interrupt_after_nodes: Union[All, Sequence[str]] = Field(default_factory=list)
interrupt_before_nodes: Sequence[str] = Field(default_factory=list)
interrupt_before_nodes: Union[All, Sequence[str]] = Field(default_factory=list)
input_channels: Union[str, Sequence[str]]
@@ -518,8 +520,8 @@ class Pregel(
stream_mode: Optional[StreamMode] = None,
input_keys: Optional[Union[str, Sequence[str]]] = None,
output_keys: Optional[Union[str, Sequence[str]]] = None,
interrupt_before: Optional[Sequence[str]] = None,
interrupt_after: Optional[Sequence[str]] = None,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
debug: Optional[bool] = None,
) -> tuple[
bool,
@@ -565,8 +567,8 @@ class Pregel(
stream_mode: Optional[StreamMode] = None,
output_keys: Optional[Union[str, Sequence[str]]] = None,
input_keys: Optional[Union[str, Sequence[str]]] = None,
interrupt_before: Optional[Sequence[str]] = None,
interrupt_after: Optional[Sequence[str]] = None,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
debug: Optional[bool] = None,
) -> Iterator[Union[dict[str, Any], Any]]:
"""Stream graph steps for a single input."""
@@ -775,8 +777,8 @@ class Pregel(
stream_mode: Optional[StreamMode] = None,
output_keys: Optional[Union[str, Sequence[str]]] = None,
input_keys: Optional[Union[str, Sequence[str]]] = None,
interrupt_before: Optional[Sequence[str]] = None,
interrupt_after: Optional[Sequence[str]] = None,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
debug: Optional[bool] = None,
) -> AsyncIterator[Union[dict[str, Any], Any]]:
config = ensure_config(config)
@@ -1002,8 +1004,8 @@ class Pregel(
stream_mode: StreamMode = "values",
output_keys: Optional[Union[str, Sequence[str]]] = None,
input_keys: Optional[Union[str, Sequence[str]]] = None,
interrupt_before_nodes: Optional[Sequence[str]] = None,
interrupt_after_nodes: Optional[Sequence[str]] = None,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
debug: Optional[bool] = None,
**kwargs: Any,
) -> Union[dict[str, Any], Any]:
@@ -1015,8 +1017,8 @@ class Pregel(
stream_mode: Optional[str]. The stream mode for the graph run. Default is "values".
output_keys: Optional. The output keys to retrieve from the graph run.
input_keys: Optional. The input keys to provide for the graph run.
interrupt_before_nodes: Optional. The nodes to interrupt the graph run before.
interrupt_after_nodes: Optional. The nodes to interrupt the graph run after.
interrupt_before: Optional. The nodes to interrupt the graph run before.
interrupt_after: Optional. The nodes to interrupt the graph run after.
debug: Optional. Enable debug mode for the graph run.
**kwargs: Additional keyword arguments to pass to the graph run.
@@ -1035,8 +1037,8 @@ class Pregel(
stream_mode=stream_mode,
output_keys=output_keys,
input_keys=input_keys,
interrupt_before=interrupt_before_nodes,
interrupt_after=interrupt_after_nodes,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
debug=debug,
**kwargs,
):
@@ -1057,8 +1059,8 @@ class Pregel(
stream_mode: StreamMode = "values",
output_keys: Optional[Union[str, Sequence[str]]] = None,
input_keys: Optional[Union[str, Sequence[str]]] = None,
interrupt_before_nodes: Optional[Sequence[str]] = None,
interrupt_after_nodes: Optional[Sequence[str]] = None,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
debug: Optional[bool] = None,
**kwargs: Any,
) -> Union[dict[str, Any], Any]:
@@ -1070,8 +1072,8 @@ class Pregel(
stream_mode: Optional. The stream mode for the computation. Default is "values".
output_keys: Optional. The output keys to include in the result. Default is None.
input_keys: Optional. The input keys to include in the result. Default is None.
interrupt_before_nodes: Optional. The nodes to interrupt before. Default is None.
interrupt_after_nodes: Optional. The nodes to interrupt after. Default is None.
interrupt_before: Optional. The nodes to interrupt before. Default is None.
interrupt_after: Optional. The nodes to interrupt after. Default is None.
debug: Optional. Whether to enable debug mode. Default is None.
**kwargs: Additional keyword arguments.
@@ -1091,8 +1093,8 @@ class Pregel(
stream_mode=stream_mode,
output_keys=output_keys,
input_keys=input_keys,
interrupt_before=interrupt_before_nodes,
interrupt_after=interrupt_after_nodes,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
debug=debug,
**kwargs,
):
@@ -1132,7 +1134,7 @@ def _panic_or_proceed(
def _should_interrupt(
checkpoint: Checkpoint,
interrupt_nodes: Sequence[str],
interrupt_nodes: Union[All, Sequence[str]],
snapshot_channels: Sequence[str],
tasks: list[PregelExecutableTask],
) -> bool:
@@ -1145,7 +1147,15 @@ def _should_interrupt(
for chan in snapshot_channels
)
# and any channel written to is in interrupt_nodes list
and any(node for node, _, _, _, _ in tasks if node in interrupt_nodes)
and any(
node
for node, _, _, _, config in tasks
if (
(not config or TAG_HIDDEN not in config.get("tags"))
if interrupt_nodes == "*"
else node in interrupt_nodes
)
)
)
langgraph/pregel/types.py
+4 -1
View File
@@ -1,5 +1,5 @@
from collections import deque
from typing import Any, NamedTuple, Optional, Union
from typing import Any, Literal, NamedTuple, Optional, Union
from langchain_core.runnables import Runnable, RunnableConfig
@@ -26,3 +26,6 @@ class StateSnapshot(NamedTuple):
"""Config used to fetch this snapshot"""
parent_config: Optional[RunnableConfig] = None
"""Config used to fetch the parent snapshot, if any"""
All = Literal["*"]
langgraph/pregel/validate.py
+11 -8
View File
@@ -3,6 +3,7 @@ from typing import Mapping, Optional, Sequence, Union
from langgraph.channels.base import BaseChannel
from langgraph.constants import INTERRUPT
from langgraph.pregel.read import PregelNode
from langgraph.pregel.types import All
def validate_graph(
@@ -11,8 +12,8 @@ def validate_graph(
input_channels: Union[str, Sequence[str]],
output_channels: Union[str, Sequence[str]],
stream_channels: Optional[Union[str, Sequence[str]]],
interrupt_after_nodes: Sequence[str],
interrupt_before_nodes: Sequence[str],
interrupt_after_nodes: Union[All, Sequence[str]],
interrupt_before_nodes: Union[All, Sequence[str]],
) -> None:
subscribed_channels = set[str]()
for name, node in nodes.items():
@@ -59,12 +60,14 @@ def validate_graph(
if chan not in channels:
raise ValueError(f"Output channel '{chan}' not in 'channels'")
for node in interrupt_after_nodes:
if node not in nodes:
raise ValueError(f"Node {node} not in nodes")
for node in interrupt_before_nodes:
if node not in nodes:
raise ValueError(f"Node {node} not in nodes")
if interrupt_after_nodes != "*":
for node in interrupt_after_nodes:
if node not in nodes:
raise ValueError(f"Node {node} not in nodes")
if interrupt_before_nodes != "*":
for node in interrupt_before_nodes:
if node not in nodes:
raise ValueError(f"Node {node} not in nodes")
def validate_keys(
tests/__snapshots__/test_pregel.ambr
+8 -8
View File
@@ -687,13 +687,13 @@
'''
# ---
# name: test_conditional_graph_state[end_of_run]
# name: test_conditional_state_graph[end_of_run]
'{"title": "LangGraphInput", "$ref": "#/definitions/AgentState", "definitions": {"AgentAction": {"title": "AgentAction", "description": "A full description of an action for an ActionAgent to execute.", "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": "The final return value of an ActionAgent.", "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"]}, "AgentState": {"title": "AgentState", "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"}]}}}}}}'
# ---
# name: test_conditional_graph_state[end_of_run].1
# name: test_conditional_state_graph[end_of_run].1
'{"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": "A full description of an action for an ActionAgent to execute.", "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": "The final return value of an ActionAgent.", "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_conditional_graph_state[end_of_run].2
# name: test_conditional_state_graph[end_of_run].2
'''
{
"nodes": [
@@ -758,7 +758,7 @@
}
'''
# ---
# name: test_conditional_graph_state[end_of_run].3
# name: test_conditional_state_graph[end_of_run].3
'''
graph TD;
__start__ --> agent;
@@ -768,13 +768,13 @@
'''
# ---
# name: test_conditional_graph_state[end_of_step]
# name: test_conditional_state_graph[end_of_step]
'{"title": "LangGraphInput", "$ref": "#/definitions/AgentState", "definitions": {"AgentAction": {"title": "AgentAction", "description": "A full description of an action for an ActionAgent to execute.", "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": "The final return value of an ActionAgent.", "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"]}, "AgentState": {"title": "AgentState", "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"}]}}}}}}'
# ---
# name: test_conditional_graph_state[end_of_step].1
# name: test_conditional_state_graph[end_of_step].1
'{"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": "A full description of an action for an ActionAgent to execute.", "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": "The final return value of an ActionAgent.", "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_conditional_graph_state[end_of_step].2
# name: test_conditional_state_graph[end_of_step].2
'''
{
"nodes": [
@@ -839,7 +839,7 @@
}
'''
# ---
# name: test_conditional_graph_state[end_of_step].3
# name: test_conditional_state_graph[end_of_step].3
'''
graph TD;
__start__ --> agent;
tests/test_pregel.py
+177 -1
View File
@@ -1574,7 +1574,7 @@ def test_conditional_entrypoint_graph(snapshot: SnapshotAssertion) -> None:
@pytest.mark.parametrize(
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
)
def test_conditional_graph_state(
def test_conditional_state_graph(
snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt
) -> None:
from langchain.llms.fake import FakeStreamingListLLM
@@ -1972,6 +1972,182 @@ def test_conditional_graph_state(
config=app_w_interrupt.checkpointer.get_tuple(config).config,
)
# test w interrupt before all
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(at=checkpoint_at),
interrupt_before="*",
debug=True,
)
config = {"configurable": {"thread_id": "3"}}
llm.i = 0 # reset the llm
assert [
c for c in app_w_interrupt.stream({"input": "what is weather in sf"}, config)
] == []
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"input": "what is weather in sf",
"intermediate_steps": [],
},
next=("agent",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
)
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:query"
),
}
},
]
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
"intermediate_steps": [],
},
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
)
assert [c for c in app_w_interrupt.stream(None, config)] == [
{
"tools": {
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"result for query",
)
],
}
},
]
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"result for query",
)
],
},
next=("agent",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
)
assert [c for c in app_w_interrupt.stream(None, config)] == [
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
}
},
]
# test w interrupt after all
app_w_interrupt = workflow.compile(
checkpointer=MemorySaverAssertImmutable(at=checkpoint_at),
interrupt_after="*",
)
config = {"configurable": {"thread_id": "4"}}
llm.i = 0 # reset the llm
assert [
c for c in app_w_interrupt.stream({"input": "what is weather in sf"}, config)
] == [
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
}
},
]
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
"intermediate_steps": [],
},
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
)
assert [c for c in app_w_interrupt.stream(None, config)] == [
{
"tools": {
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"result for query",
)
],
}
},
]
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"result for query",
)
],
},
next=("agent",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
)
assert [c for c in app_w_interrupt.stream(None, config)] == [
{
"agent": {
"agent_outcome": AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
}
},
]
def test_conditional_entrypoint_graph_state(snapshot: SnapshotAssertion) -> None:
class AgentState(TypedDict, total=False):