Merge pull request #102 from langchain-ai/nc/10feb/interrupt-after

Add interrupt_after arg to .compile()
This commit is contained in:
Nuno Campos
2024-02-10 11:58:04 -08:00
committed by GitHub
2 changed files with 16 additions and 8 deletions
+8 -4
View File
@@ -123,8 +123,11 @@ class Graph:
self,
checkpointer: Optional[BaseCheckpointSaver] = None,
interrupt_before: Optional[Sequence[str]] = None,
interrupt_after: Optional[Sequence[str]] = None,
) -> Pregel:
self.validate(interrupt=interrupt_before)
interrupt_before = interrupt_before or []
interrupt_after = interrupt_after or []
self.validate(interrupt=interrupt_before + interrupt_after)
outgoing_edges = defaultdict(list)
for start, end in self.edges:
@@ -154,7 +157,8 @@ class Graph:
output=END,
hidden=[f"{node}:inbox" for node in self.nodes],
checkpointer=checkpointer,
interrupt=[f"{node}:inbox" for node in interrupt_before]
if interrupt_before
else [],
interrupt=(
[f"{node}:inbox" for node in interrupt_before]
+ [node for node in interrupt_after]
),
)
+8 -4
View File
@@ -38,8 +38,11 @@ class StateGraph(Graph):
self,
checkpointer: Optional[BaseCheckpointSaver] = None,
interrupt_before: Optional[Sequence[str]] = None,
interrupt_after: Optional[Sequence[str]] = None,
) -> Pregel:
self.validate(interrupt=interrupt_before)
interrupt_before = interrupt_before or []
interrupt_after = interrupt_after or []
self.validate(interrupt=interrupt_before + interrupt_after)
state_keys = list(self.channels)
state_keys_read = state_keys[0] if state_keys == ["__root__"] else state_keys
@@ -102,9 +105,10 @@ class StateGraph(Graph):
output=END,
hidden=[f"{node}:inbox" for node in self.nodes] + [START] + state_keys,
checkpointer=checkpointer,
interrupt=[f"{node}:inbox" for node in interrupt_before]
if interrupt_before
else [],
interrupt=(
[f"{node}:inbox" for node in interrupt_before]
+ [node for node in interrupt_after]
),
)