From 8ea24435feb49edbf6eceb9310ea2f6b319a5a5e Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Tue, 16 Jan 2024 10:46:21 -0800
Subject: [PATCH] Add more validation

---
 langgraph/pregel/__init__.py | 17 +++++++++++++++--
 langgraph/pregel/validate.py | 18 ++++++++++++++++++
 2 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py
index 49187d1f4..4ec34ae32 100644
--- a/langgraph/pregel/__init__.py
+++ b/langgraph/pregel/__init__.py
@@ -61,7 +61,7 @@ from langgraph.pregel.io import map_input, map_output
 from langgraph.pregel.log import logger
 from langgraph.pregel.read import ChannelBatch, ChannelInvoke
 from langgraph.pregel.reserved import ReservedChannels
-from langgraph.pregel.validate import validate_graph
+from langgraph.pregel.validate import validate_graph, validate_keys
 from langgraph.pregel.write import ChannelWrite
 
 WriteValue = Union[
@@ -179,7 +179,12 @@ class Pregel(
     @root_validator(skip_on_failure=True)
     def validate_pregel(cls, values: dict[str, Any]) -> dict[str, Any]:
         validate_graph(
-            values["nodes"], values["channels"], values["input"], values["output"]
+            values["nodes"],
+            values["channels"],
+            values["input"],
+            values["output"],
+            values["hidden"],
+            values["interrupt"],
         )
         return values
 
@@ -239,8 +244,12 @@ class Pregel(
         # assign defaults
         if output_keys is None:
             output_keys = [chan for chan in self.channels if chan not in self.hidden]
+        else:
+            validate_keys(output_keys, self.channels)
         if input_keys is None:
             input_keys = self.input
+        else:
+            validate_keys(input_keys, self.channels)
         # copy nodes to ignore mutations during execution
         processes = {**self.nodes}
         # get checkpoint from saver, or create an empty one
@@ -369,8 +378,12 @@ class Pregel(
         # assign defaults
         if output_keys is None:
             output_keys = [chan for chan in self.channels if chan not in self.hidden]
+        else:
+            validate_keys(output_keys, self.channels)
         if input_keys is None:
             input_keys = self.input
+        else:
+            validate_keys(input_keys, self.channels)
         # copy nodes to ignore mutations during execution
         processes = {**self.nodes}
         # get checkpoint from saver, or create an empty one
diff --git a/langgraph/pregel/validate.py b/langgraph/pregel/validate.py
index 096061665..1186bb8be 100644
--- a/langgraph/pregel/validate.py
+++ b/langgraph/pregel/validate.py
@@ -11,6 +11,8 @@ def validate_graph(
     channels: dict[str, BaseChannel],
     input: Union[str, Sequence[str]],
     output: Union[str, Sequence[str]],
+    hidden: Sequence[str],
+    interrupt: Sequence[str],
 ) -> None:
     subscribed_channels = set[str]()
     for node in nodes.values():
@@ -52,3 +54,19 @@ def validate_graph(
     for chan in ReservedChannels:
         if chan not in channels:
             channels[chan] = LastValue(Any)  # type: ignore[arg-type]
+
+    validate_keys(hidden, channels)
+    validate_keys(interrupt, channels)
+
+
+def validate_keys(
+    keys: Union[str, Sequence[str]],
+    channels: dict[str, BaseChannel],
+) -> None:
+    if isinstance(keys, str):
+        if keys not in channels:
+            raise ValueError(f"Key {keys} not in channels")
+    else:
+        for chan in keys:
+            if chan not in channels:
+                raise ValueError(f"Key {chan} not in channels")