From 95c1ca3adcd1b14f304aecd01140a9dd636f1e8f Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Thu, 19 Sep 2024 13:20:45 -0700
Subject: [PATCH] Fix

---
 libs/langgraph/Makefile                        |  2 +-
 .../channels/dynamic_barrier_value.py          |  6 +++---
 libs/langgraph/langgraph/graph/graph.py        | 18 ++++++++++++------
 libs/langgraph/langgraph/graph/state.py        |  2 +-
 4 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/libs/langgraph/Makefile b/libs/langgraph/Makefile
index 7ec12f80a..1e249a0cd 100644
--- a/libs/langgraph/Makefile
+++ b/libs/langgraph/Makefile
@@ -75,7 +75,7 @@ lint lint_diff lint_package lint_tests:
 	[ "$(PYTHON_FILES)" = "" ] || poetry run ruff format $(PYTHON_FILES) --diff
 	[ "$(PYTHON_FILES)" = "" ] || poetry run ruff check --select I $(PYTHON_FILES)
 	[ "$(PYTHON_FILES)" = "" ] || mkdir -p $(MYPY_CACHE)
-	[ "$(PYTHON_FILES)" != "langgraph" ] || poetry run mypy $(PYTHON_FILES) --cache-dir $(MYPY_CACHE)
+	[ "$(PYTHON_FILES)" = "" ] || poetry run mypy langgraph --cache-dir $(MYPY_CACHE)
 
 format format_diff:
 	poetry run ruff format $(PYTHON_FILES)
diff --git a/libs/langgraph/langgraph/channels/dynamic_barrier_value.py b/libs/langgraph/langgraph/channels/dynamic_barrier_value.py
index dfa77f350..f64191e86 100644
--- a/libs/langgraph/langgraph/channels/dynamic_barrier_value.py
+++ b/libs/langgraph/langgraph/channels/dynamic_barrier_value.py
@@ -1,4 +1,4 @@
-from typing import Generic, NamedTuple, Optional, Sequence, Type, Union
+from typing import Any, Generic, NamedTuple, Optional, Sequence, Type, Union
 
 from typing_extensions import Self
 
@@ -6,8 +6,8 @@ from langgraph.channels.base import BaseChannel, Value
 from langgraph.errors import EmptyChannelError, InvalidUpdateError
 
 
-class WaitForNames(NamedTuple, Generic[Value]):
-    names: set[Value]
+class WaitForNames(NamedTuple):
+    names: set[Any]
 
 
 class DynamicBarrierValue(
diff --git a/libs/langgraph/langgraph/graph/graph.py b/libs/langgraph/langgraph/graph/graph.py
index 93fe28b3d..12f9cf943 100644
--- a/libs/langgraph/langgraph/graph/graph.py
+++ b/libs/langgraph/langgraph/graph/graph.py
@@ -268,13 +268,15 @@ class Graph:
                 path_map_ = path_map.copy()
             elif isinstance(path_map, list):
                 path_map_ = {name: name for name in path_map}
-            elif callable(path) and (
-                rtn_type := get_type_hints(path.__call__).get("return")
-                if hasattr(path, "__call__")
-                else get_type_hints(path).get("return")
-            ):
+            elif isinstance(path, Runnable):
+                path_map_ = None
+            elif rtn_type := get_type_hints(path.__call__).get(  # type: ignore[operator]
+                "return"
+            ) or get_type_hints(path).get("return"):
                 if get_origin(rtn_type) is Literal:
                     path_map_ = {name: name for name in get_args(rtn_type)}
+                else:
+                    path_map_ = None
             else:
                 path_map_ = None
         except Exception:
@@ -417,6 +419,7 @@ class Graph:
 
         # create empty compiled graph
         compiled = CompiledGraph(
+            builder=self,
             nodes={},
             channels={START: EphemeralValue(Any), END: EphemeralValue(Any)},
             input_channels=START,
@@ -429,7 +432,6 @@ class Graph:
             auto_validate=False,
             debug=debug,
         )
-        compiled.builder = self
 
         # attach nodes, edges, and branches
         for key, node in self.nodes.items():
@@ -449,6 +451,10 @@ class Graph:
 class CompiledGraph(Pregel):
     builder: Graph
 
+    def __init__(self, *, builder: Graph, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.builder = builder
+
     def attach_node(self, key: str, node: NodeSpec) -> None:
         self.channels[key] = EphemeralValue(Any)
         self.nodes[key] = (
diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py
index 803434316..ace4d2553 100644
--- a/libs/langgraph/langgraph/graph/state.py
+++ b/libs/langgraph/langgraph/graph/state.py
@@ -453,6 +453,7 @@ class StateGraph(Graph):
         )
 
         compiled = CompiledStateGraph(
+            builder=self,
             config_type=self.config_schema,
             nodes={},
             channels={
@@ -471,7 +472,6 @@ class StateGraph(Graph):
             debug=debug,
             store=store,
         )
-        compiled.builder = self
 
         compiled.attach_node(START, None)
         for key, node in self.nodes.items():