From f1523f3e8dc35bed5128187456379cee6524afe4 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Sat, 10 Feb 2024 11:39:38 -0800
Subject: [PATCH 1/2] Add interrupt_after arg to .compile()

---
 langgraph/graph/graph.py | 12 ++++++++----
 langgraph/graph/state.py | 12 ++++++++----
 2 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py
index f3fa36085..cfc1401f8 100644
--- a/langgraph/graph/graph.py
+++ b/langgraph/graph/graph.py
@@ -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]
+            ),
         )
diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py
index fcfab0425..f011652d6 100644
--- a/langgraph/graph/state.py
+++ b/langgraph/graph/state.py
@@ -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]
+            ),
         )
 
 

From 29edaaead8401f8a48e7d4d3a892c9e67ada3f38 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Sat, 10 Feb 2024 13:39:46 -0800
Subject: [PATCH 2/2] Warn if graph is mutated after being compiled

---
 langgraph/graph/graph.py | 26 ++++++++++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py
index cfc1401f8..22e075575 100644
--- a/langgraph/graph/graph.py
+++ b/langgraph/graph/graph.py
@@ -1,3 +1,4 @@
+import logging
 from asyncio import iscoroutinefunction
 from collections import defaultdict
 from typing import Any, Callable, Dict, NamedTuple, Optional, Sequence
@@ -12,6 +13,8 @@ from langchain_core.runnables.base import (
 from langgraph.checkpoint import BaseCheckpointSaver
 from langgraph.pregel import Channel, Pregel
 
+logger = logging.getLogger(__name__)
+
 END = "__end__"
 
 
@@ -34,8 +37,14 @@ class Graph:
         self.edges = set[tuple[str, str]]()
         self.branches: defaultdict[str, list[Branch]] = defaultdict(list)
         self.support_multiple_edges = False
+        self.compiled = False
 
     def add_node(self, key: str, action: RunnableLike) -> None:
+        if self.compiled:
+            logger.warning(
+                "Adding a node to a graph that has already been compiled. This will "
+                "not be reflected in the compiled graph."
+            )
         if key in self.nodes:
             raise ValueError(f"Node `{key}` already present.")
         if key == END:
@@ -44,6 +53,11 @@ class Graph:
         self.nodes[key] = coerce_to_runnable(action)
 
     def add_edge(self, start_key: str, end_key: str) -> None:
+        if self.compiled:
+            logger.warning(
+                "Adding an edge to a graph that has already been compiled. This will "
+                "not be reflected in the compiled graph."
+            )
         if start_key == END:
             raise ValueError("END cannot be a start node")
         if start_key not in self.nodes:
@@ -64,6 +78,11 @@ class Graph:
         condition: Callable[..., str],
         conditional_edge_mapping: Optional[Dict[str, str]] = None,
     ) -> None:
+        if self.compiled:
+            logger.warning(
+                "Adding an edge to a graph that has already been compiled. This will "
+                "not be reflected in the compiled graph."
+            )
         if start_key not in self.nodes:
             raise ValueError(f"Need to add_node `{start_key}` first")
         if iscoroutinefunction(condition):
@@ -81,6 +100,11 @@ class Graph:
         self.branches[start_key].append(Branch(condition, conditional_edge_mapping))
 
     def set_entry_point(self, key: str) -> None:
+        if self.compiled:
+            logger.warning(
+                "Setting the entry point of a graph that has already been compiled. "
+                "This will not be reflected in the compiled graph."
+            )
         if key not in self.nodes:
             raise ValueError(f"Need to add_node `{key}` first")
         self.entry_point = key
@@ -119,6 +143,8 @@ class Graph:
                 if node not in self.nodes:
                     raise ValueError(f"Node `{node}` is not present")
 
+        self.compiled = True
+
     def compile(
         self,
         checkpointer: Optional[BaseCheckpointSaver] = None,