diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py index b6b6fc9e5..59fa3b874 100644 --- a/langgraph/graph/graph.py +++ b/langgraph/graph/graph.py @@ -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(): diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py index d5ac32346..5e7bfe8b1 100644 --- a/langgraph/graph/state.py +++ b/langgraph/graph/state.py @@ -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: diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 552021e4d..c1d61e7e0 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -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 + ) + ) ) diff --git a/langgraph/pregel/types.py b/langgraph/pregel/types.py index 8ebe52e99..fd728a7f1 100644 --- a/langgraph/pregel/types.py +++ b/langgraph/pregel/types.py @@ -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["*"] diff --git a/langgraph/pregel/validate.py b/langgraph/pregel/validate.py index 0b1bfd16b..0a4e4cea5 100644 --- a/langgraph/pregel/validate.py +++ b/langgraph/pregel/validate.py @@ -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( diff --git a/tests/__snapshots__/test_pregel.ambr b/tests/__snapshots__/test_pregel.ambr index 19fd322df..b5bc325f6 100644 --- a/tests/__snapshots__/test_pregel.ambr +++ b/tests/__snapshots__/test_pregel.ambr @@ -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; diff --git a/tests/test_pregel.py b/tests/test_pregel.py index d12ce4898..84d5db4d5 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -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):