diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py index 7f3fcdb94..164ecaac4 100644 --- a/langgraph/graph/state.py +++ b/langgraph/graph/state.py @@ -27,20 +27,17 @@ class StateGraph(Graph): self.channels = _get_channels(schema) if any(isinstance(c, BinaryOperatorAggregate) for c in self.channels.values()): self.support_multiple_edges = True - self.w_edges: set[tuple[tuple[str, ...], str]] = set() + self.waiting_edges: set[tuple[tuple[str, ...], str]] = set() @property def _all_edges(self) -> set[tuple[str, str]]: return self.edges | { - (start, end) for starts, end in self.w_edges for start in starts + (start, end) for starts, end in self.waiting_edges for start in starts } def add_node(self, key: str, action: RunnableLike) -> None: if key in self.channels: - raise ValueError( - f"'{key}' is already being used as a state attribute " - "(a.k.a. a channel), cannot also be used as a node name." - ) + raise ValueError(f"'{key}' is already being used as a state key") return super().add_node(key, action) def add_edge(self, start_key: Union[str, list[str]], end_key: str) -> None: @@ -62,7 +59,7 @@ class StateGraph(Graph): if end_key not in self.nodes: raise ValueError(f"Need to add_node `{end_key}` first") - self.w_edges.add((tuple(start_key), end_key)) + self.waiting_edges.add((tuple(start_key), end_key)) def compile( self, @@ -71,10 +68,6 @@ class StateGraph(Graph): interrupt_after: Optional[Sequence[str]] = None, debug: bool = False, ) -> CompiledGraph: - # prepare state channels - state_keys = list(self.channels) - state_keys_read = state_keys[0] if state_keys == ["__root__"] else state_keys - # assign default values interrupt_before = interrupt_before or [] interrupt_after = interrupt_after or [] @@ -82,14 +75,18 @@ class StateGraph(Graph): # validate the graph self.validate(interrupt=interrupt_before + interrupt_after) + # prepare output channels + state_keys = list(self.channels) + output_channels = state_keys[0] if state_keys == ["__root__"] else state_keys + compiled = CompiledStateGraph( graph=self, nodes={}, channels={**self.channels, START: EphemeralValue(self.schema)}, input_channels=START, stream_mode="updates", - output_channels=state_keys_read, - stream_channels=state_keys_read, + output_channels=output_channels, + stream_channels=output_channels, checkpointer=checkpointer, interrupt_before_nodes=interrupt_before, interrupt_after_nodes=interrupt_after, @@ -104,7 +101,7 @@ class StateGraph(Graph): for start, end in self.edges: compiled.attach_edge(start, end) - for starts, end in self.w_edges: + for starts, end in self.waiting_edges: compiled.attach_edge(starts, end) for start, branches in self.branches.items(): @@ -118,17 +115,20 @@ class CompiledStateGraph(CompiledGraph): graph: StateGraph def attach_node(self, key: str, node: Optional[Runnable]) -> None: - # prepare state + def _get_state_key(key: str, input: dict) -> Any: + if input is None: + return SKIP_WRITE + elif not isinstance(input, dict): + raise InvalidUpdateError(f"Expected dict, got {input}") + else: + return input.get(key, SKIP_WRITE) + state_keys = list(self.graph.channels) - state_keys_read = state_keys[0] if state_keys == ["__root__"] else state_keys + # state updaters state_write_entries = [ - ChannelWriteEntry( - key, - RunnableLambda(partial(_dict_getter, state_keys, key)), - False, - ) - if key != "__root__" - else ChannelWriteEntry(key, None, True) + ChannelWriteEntry(key, None, skip_none=True) + if key == "__root__" + else ChannelWriteEntry(key, RunnableLambda(partial(_get_state_key, key))) for key in state_keys ] @@ -143,23 +143,24 @@ class CompiledStateGraph(CompiledGraph): triggers=[], # read state keys channels=( - {chan: chan for chan in state_keys} - if isinstance(state_keys_read, list) - else [state_keys_read] + state_keys + if state_keys == ["__root__"] + else {chan: chan for chan in state_keys} ), # coerce state dict to schema class (eg. pydantic model) mapper=( - partial(_coerce_state, self.graph.schema) - if isinstance(state_keys_read, list) - else None + None + if state_keys == ["__root__"] + else partial(_coerce_state, self.graph.schema) ), # publish to this channel and state keys writers=[ - ChannelWrite( - [ChannelWriteEntry(key, None, False)] + state_write_entries - ), + ChannelWrite([ChannelWriteEntry(key)] + state_write_entries), # read back state with updates applied - ChannelRead(state_keys_read, fresh=True), + ChannelRead( + state_keys[0] if state_keys == ["__root__"] else state_keys, + fresh=True, + ), ], ).pipe(node) @@ -173,7 +174,7 @@ class CompiledStateGraph(CompiledGraph): self.nodes[end].triggers.append(channel_name) # publish to channel self.nodes[START] |= ChannelWrite( - [ChannelWriteEntry(channel_name, START, False)] + [ChannelWriteEntry(channel_name, START)] ) elif end != END: # subscribe to start channel @@ -187,14 +188,14 @@ class CompiledStateGraph(CompiledGraph): # publish to channel for start in starts: self.nodes[start] |= ChannelWrite( - [ChannelWriteEntry(channel_name, start, False)] + [ChannelWriteEntry(channel_name, start)] ) def attach_branch(self, start: str, name: str, branch: Branch) -> None: def branch_writer(end: str) -> Optional[ChannelWrite]: if end != END: return ChannelWrite( - [ChannelWriteEntry(f"branch:{start}:{name}:{end}", start, False)] + [ChannelWriteEntry(f"branch:{start}:{name}:{end}", start)] ) # attach branch publisher @@ -213,18 +214,6 @@ def _coerce_state(schema: Type[Any], input: dict[str, Any]) -> dict[str, Any]: return schema(**input) -def _dict_getter(allowed_keys: list[str], key: str, input: dict) -> Any: - if input is not None: - if not isinstance(input, dict) or any(key not in allowed_keys for key in input): - raise InvalidUpdateError( - f"Invalid state update," - f" expected dict with one or more of {allowed_keys}, got {input}" - ) - return input.get(key, SKIP_WRITE) - else: - return SKIP_WRITE - - def _get_channels(schema: Type[dict]) -> dict[str, BaseChannel]: if not hasattr(schema, "__annotations__"): return { diff --git a/langgraph/pregel/write.py b/langgraph/pregel/write.py index f5340dba7..03e95fb28 100644 --- a/langgraph/pregel/write.py +++ b/langgraph/pregel/write.py @@ -21,8 +21,8 @@ SKIP_WRITE = object() class ChannelWriteEntry(NamedTuple): channel: str - value: Optional[Union[Any, Runnable]] - skip_none: bool + value: Optional[Union[Any, Runnable]] = None + skip_none: bool = False class ChannelWrite(RunnablePassthrough):