Better error messages for invalid update in all channel types (#1437)

* langgraph: support multiple edges for Topic channel annotations

* remove support_multiple_edges

* Better error messages for invalid update in all channel types

---------

Co-authored-by: vbarda <vadym@langchain.dev>

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

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:

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

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]

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

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:

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]

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))

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]:

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]

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()
             },

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)