diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py index a11288d91..fc99ed3fe 100644 --- a/langgraph/graph/state.py +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index c4802162d..2dd6bc6df 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langgraph" -version = "0.0.16" +version = "0.0.17" description = "langgraph" authors = [] license = "LangGraph License"