Merge branch 'master' into wfh/update-nb

langgraph/graph/state.py

@@ -5,7 +5,7 @@ from typing import Any, Optional, Type
 
 from langchain_core.runnables import RunnableConfig, RunnableLambda, RunnablePassthrough
 
-from langgraph.channels.base import BaseChannel
+from langgraph.channels.base import BaseChannel, InvalidUpdateError
 from langgraph.channels.binop import BinaryOperatorAggregate
 from langgraph.channels.last_value import LastValue
 from langgraph.checkpoint import BaseCheckpointSaver
@@ -34,7 +34,7 @@ class StateGraph(Graph):
         state_keys = list(self.channels)
         state_keys_read = state_keys[0] if state_keys == ["__root__"] else state_keys
         update_state = (
-            _update_state_dict
+            partial(_update_state_dict, state_keys)
             if isinstance(state_keys_read, list)
             else _update_state_root
         )
@@ -53,7 +53,7 @@ class StateGraph(Graph):
                 Channel.subscribe_to(f"{key}:inbox")
                 | coerce_state  # coerce/validate using schema
                 | node
-                | update_state
+                | partial(update_state, key)
                 | Channel.write_to(key)
             )
             for key, node in self.nodes.items()
@@ -76,7 +76,7 @@ class StateGraph(Graph):
 
         nodes[START] = (
             Channel.subscribe_to(f"{START}:inbox", tags=["langsmith:hidden"])
-            | update_state
+            | partial(update_state, START)
             | Channel.write_to(START)
         )
         nodes[f"{START}:edges"] = (
@@ -99,13 +99,22 @@ def _coerce_state(schema: Type[Any], input: dict[str, Any]) -> dict[str, Any]:
     return schema(**input)
 
 
-def _update_state_dict(input: dict[str, Any], config: RunnableConfig) -> dict[str, Any]:
+def _update_state_dict(
+    state_keys: list[str], node_name: str, input: dict[str, Any], config: RunnableConfig
+) -> dict[str, Any]:
     if input is not None:
+        if not isinstance(input, dict) or any(key not in state_keys for key in input):
+            raise InvalidUpdateError(
+                f"Invalid state update from node {node_name},"
+                f" expected dict with one or more of {state_keys}, got {input}"
+            )
         ChannelWrite.do_write(config, **input)
     return input
 
 
-def _update_state_root(input: Any, config: RunnableConfig) -> dict[str, Any]:
+def _update_state_root(
+    node_name: str, input: Any, config: RunnableConfig
+) -> dict[str, Any]:
     if input is not None:
         ChannelWrite.do_write(config, __root__=input)
     return input

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph"
-version = "0.0.16"
+version = "0.0.17"
 description = "langgraph"
 authors = []
 license = "LangGraph License"