diff --git a/libs/langgraph/langgraph/channels/base.py b/libs/langgraph/langgraph/channels/base.py index fe47f0d8f..885698743 100644 --- a/libs/langgraph/langgraph/channels/base.py +++ b/libs/langgraph/langgraph/channels/base.py @@ -2,9 +2,9 @@ from abc import ABC, abstractmethod from contextlib import asynccontextmanager, contextmanager from typing import ( Any, - AsyncGenerator, - Generator, + AsyncIterator, Generic, + Iterator, Optional, Sequence, TypeVar, @@ -21,6 +21,8 @@ C = TypeVar("C") class BaseChannel(Generic[Value, Update, C], ABC): + key: str = "" + @property @abstractmethod def ValueType(self) -> Any: @@ -43,19 +45,35 @@ class BaseChannel(Generic[Value, Update, C], ABC): @abstractmethod def from_checkpoint( self, checkpoint: Optional[C], config: RunnableConfig - ) -> Generator[Self, None, None]: + ) -> Iterator[Self]: """Return a new identical channel, optionally initialized from a checkpoint. If the checkpoint contains complex data structures, they should be copied.""" + @contextmanager + def from_checkpoint_named( + self, checkpoint: Optional[C], config: RunnableConfig + ) -> Iterator[Self]: + with self.from_checkpoint(checkpoint, config) as value: + value.key = self.key + yield value + @asynccontextmanager async def afrom_checkpoint( self, checkpoint: Optional[C], config: RunnableConfig - ) -> AsyncGenerator[Self, None]: + ) -> AsyncIterator[Self]: """Return a new identical channel, optionally initialized from a checkpoint. If the checkpoint contains complex data structures, they should be copied.""" with self.from_checkpoint(checkpoint, config) as value: yield value + @asynccontextmanager + async def afrom_checkpoint_named( + self, checkpoint: Optional[C], config: RunnableConfig + ) -> AsyncIterator[Self]: + async with self.afrom_checkpoint(checkpoint, config) as value: + value.key = self.key + yield value + # state methods @abstractmethod diff --git a/libs/langgraph/langgraph/channels/context.py b/libs/langgraph/langgraph/channels/context.py index 914de9348..b48260b40 100644 --- a/libs/langgraph/langgraph/channels/context.py +++ b/libs/langgraph/langgraph/channels/context.py @@ -112,7 +112,9 @@ class Context(Generic[Value], BaseChannel[Value, None, None]): def update(self, values: Sequence[None]) -> bool: if values: - raise InvalidUpdateError("Context channel does not accept writes.") + raise InvalidUpdateError( + f"At key '{self.key}': Context channel does not accept writes." + ) return False def get(self) -> Value: diff --git a/libs/langgraph/langgraph/channels/dynamic_barrier_value.py b/libs/langgraph/langgraph/channels/dynamic_barrier_value.py index bb0d447fa..64406b8f8 100644 --- a/libs/langgraph/langgraph/channels/dynamic_barrier_value.py +++ b/libs/langgraph/langgraph/channels/dynamic_barrier_value.py @@ -69,7 +69,7 @@ class DynamicBarrierValue( if wait_for_names := [v for v in values if isinstance(v, WaitForNames)]: if len(wait_for_names) > 1: raise InvalidUpdateError( - "Received multiple WaitForNames updates in the same step." + f"At key '{self.key}': Received multiple WaitForNames updates in the same step." ) self.names = wait_for_names[0].names return True diff --git a/libs/langgraph/langgraph/channels/ephemeral_value.py b/libs/langgraph/langgraph/channels/ephemeral_value.py index 15e11550d..4e7f2ed63 100644 --- a/libs/langgraph/langgraph/channels/ephemeral_value.py +++ b/libs/langgraph/langgraph/channels/ephemeral_value.py @@ -58,7 +58,7 @@ class EphemeralValue(Generic[Value], BaseChannel[Value, Value, Value]): return False if len(values) != 1 and self.guard: raise InvalidUpdateError( - "EphemeralValue can only receive one value per step." + f"At key '{self.key}': EphemeralValue(guard=True) can receive only one value per step. Use guard=False if you want to store any one of multiple values." ) self.value = values[-1] diff --git a/libs/langgraph/langgraph/channels/last_value.py b/libs/langgraph/langgraph/channels/last_value.py index a207ebce3..e74580d6a 100644 --- a/libs/langgraph/langgraph/channels/last_value.py +++ b/libs/langgraph/langgraph/channels/last_value.py @@ -52,7 +52,9 @@ class LastValue(Generic[Value], BaseChannel[Value, Value, Value]): if len(values) == 0: return False if len(values) != 1: - raise InvalidUpdateError("LastValue can only receive one value per step.") + raise InvalidUpdateError( + f"At key '{self.key}': Can receive only one value per step. Use an Annotated key to handle multiple values." + ) self.value = values[-1] return True diff --git a/libs/langgraph/langgraph/channels/named_barrier_value.py b/libs/langgraph/langgraph/channels/named_barrier_value.py index bdfd4660b..023f54e6c 100644 --- a/libs/langgraph/langgraph/channels/named_barrier_value.py +++ b/libs/langgraph/langgraph/channels/named_barrier_value.py @@ -53,7 +53,9 @@ class NamedBarrierValue(Generic[Value], BaseChannel[Value, Value, set[Value]]): self.seen.add(value) updated = True else: - raise InvalidUpdateError(f"Value {value} not in {self.names}") + raise InvalidUpdateError( + f"At key '{self.key}': Value {value} not in {self.names}" + ) return updated def get(self) -> Value: diff --git a/libs/langgraph/langgraph/channels/untracked_value.py b/libs/langgraph/langgraph/channels/untracked_value.py index 989bba35e..a112b0e81 100644 --- a/libs/langgraph/langgraph/channels/untracked_value.py +++ b/libs/langgraph/langgraph/channels/untracked_value.py @@ -49,7 +49,7 @@ class UntrackedValue(Generic[Value], BaseChannel[Value, Value, Value]): return False if len(values) != 1 and self.guard: raise InvalidUpdateError( - "UntrackedValue can only receive one value per step." + f"At key '{self.key}': UntrackedValue(guard=True) can receive only one value per step. Use guard=False if you want to store any one of multiple values." ) self.value = values[-1] diff --git a/libs/langgraph/langgraph/graph/graph.py b/libs/langgraph/langgraph/graph/graph.py index dc927446d..d5eeb88e2 100644 --- a/libs/langgraph/langgraph/graph/graph.py +++ b/libs/langgraph/langgraph/graph/graph.py @@ -192,12 +192,14 @@ class Graph: raise ValueError("END cannot be a start node") if end_key == START: raise ValueError("START cannot be an end node") - if not self.support_multiple_edges and start_key in set( + + # run this validation only for non-StateGraph graphs + if not hasattr(self, "channels") and start_key in set( start for start, _ in self.edges ): raise ValueError( f"Already found path for node '{start_key}'.\n" - "For multiple edges, use StateGraph with an annotated state key." + "For multiple edges, use StateGraph with an Annotated state key." ) self.edges.add((start_key, end_key)) diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 3ff07b15c..7228c5fa2 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -193,10 +193,6 @@ class StateGraph(Graph): ) else: self.managed[key] = managed - if any( - isinstance(c, BinaryOperatorAggregate) for c in self.channels.values() - ): - self.support_multiple_edges = True @overload def add_node( @@ -723,10 +719,15 @@ def _get_channel( else: raise ValueError(f"This {annotation} not allowed in this position") elif channel := _is_field_channel(annotation): + channel.key = name return channel elif channel := _is_field_binop(annotation): + channel.key = name return channel - return LastValue(annotation) + + fallback = LastValue(annotation) + fallback.key = name + return fallback def _is_field_channel(typ: Type[Any]) -> Optional[BaseChannel]: diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 28a8dcd5d..1d3f57040 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -198,12 +198,7 @@ def apply_writes( updated_channels: set[str] = set() for chan, vals in pending_writes_by_channel.items(): if chan in channels: - try: - updated = channels[chan].update(vals) - except InvalidUpdateError as e: - raise InvalidUpdateError( - f"Invalid update for channel {chan} with values {vals}" - ) from e + updated = channels[chan].update(vals) if updated and get_next_version is not None: checkpoint["channel_versions"][chan] = get_next_version( max_version, channels[chan] diff --git a/libs/langgraph/langgraph/pregel/manager.py b/libs/langgraph/langgraph/pregel/manager.py index 437019113..849395c50 100644 --- a/libs/langgraph/langgraph/pregel/manager.py +++ b/libs/langgraph/langgraph/pregel/manager.py @@ -41,7 +41,7 @@ def ChannelsManager( yield ( { k: stack.enter_context( - v.from_checkpoint(checkpoint["channel_values"].get(k), config) + v.from_checkpoint_named(checkpoint["channel_values"].get(k), config) ) for k, v in channel_specs.items() }, @@ -95,7 +95,9 @@ async def AsyncChannelsManager( # channels: enter each channel with checkpoint { k: await stack.enter_async_context( - v.afrom_checkpoint(checkpoint["channel_values"].get(k), config) + v.afrom_checkpoint_named( + checkpoint["channel_values"].get(k), config + ) ) for k, v in channel_specs.items() }, diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index c1fc68c2a..3e40f31fc 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -198,6 +198,21 @@ def test_graph_validation() -> None: with pytest.raises(ValueError, match="Invalid reducer"): StateGraph(BadReducerState) + def node_b(state: State) -> State: + return {"hello": "world"} + + builder = StateGraph(State) + builder.add_node("a", node_b) + builder.add_node("b", node_b) + builder.add_node("c", node_b) + builder.set_entry_point("a") + builder.add_edge("a", "b") + builder.add_edge("a", "c") + graph = builder.compile() + + with pytest.raises(InvalidUpdateError, match="At key 'hello'"): + graph.invoke({"hello": "there"}) + def test_checkpoint_errors() -> None: class FaultyGetCheckpointer(MemorySaver): @@ -1197,6 +1212,21 @@ def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture) -> N # LastValue channels can only be updated once per iteration app.invoke(2) + class State(TypedDict): + hello: str + + def my_node(input: State) -> State: + return {"hello": "world"} + + builder = StateGraph(State) + builder.add_node("one", my_node) + builder.add_node("two", my_node) + builder.set_conditional_entry_point(lambda _: ["one", "two"]) + + graph = builder.compile() + with pytest.raises(InvalidUpdateError, match="At key 'hello'"): + graph.invoke({"hello": "there"}, debug=True) + def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1)