diff --git a/libs/langgraph/langgraph/channels/dynamic_barrier_value.py b/libs/langgraph/langgraph/channels/dynamic_barrier_value.py index 4f75f2a8c..48fb3c7db 100644 --- a/libs/langgraph/langgraph/channels/dynamic_barrier_value.py +++ b/libs/langgraph/langgraph/channels/dynamic_barrier_value.py @@ -1,3 +1,4 @@ +from collections.abc import Set from typing import Any, Generic, NamedTuple, Optional, Sequence, Type, Union from typing_extensions import Self @@ -8,7 +9,7 @@ from langgraph.errors import EmptyChannelError, InvalidUpdateError class WaitForNames(NamedTuple): - names: set[Any] + names: Set[Any] class DynamicBarrierValue( @@ -25,7 +26,7 @@ class DynamicBarrierValue( __slots__ = ("names", "seen") - names: Optional[set[Value]] + names: Optional[Set[Value]] seen: set[Value] def __init__(self, typ: Type[Value]) -> None: @@ -54,11 +55,11 @@ class DynamicBarrierValue( empty.seen = self.seen.copy() return empty - def checkpoint(self) -> tuple[Optional[set[Value]], set[Value]]: + def checkpoint(self) -> tuple[Optional[Set[Value]], set[Value]]: return (self.names, self.seen) def from_checkpoint( - self, checkpoint: tuple[Optional[set[Value]], set[Value]] + self, checkpoint: tuple[Optional[Set[Value]], set[Value]] ) -> Self: empty = self.__class__(self.typ) empty.key = self.key diff --git a/libs/langgraph/langgraph/graph/branch.py b/libs/langgraph/langgraph/graph/branch.py index ce0f8bd55..8e3b847f5 100644 --- a/libs/langgraph/langgraph/graph/branch.py +++ b/libs/langgraph/langgraph/graph/branch.py @@ -133,16 +133,6 @@ class Branch(NamedTuple): writer: Writer, reader: Optional[Callable[[RunnableConfig], Any]] = None, ) -> RunnableCallable: - print( - list( - zip_longest( - writer([e for e in self.ends.values() if e != END]), - [la for la, e in self.ends.items() if e != END], - ) - ) - if self.ends - else None - ) return ChannelWrite.register_writer( RunnableCallable( func=self._route, @@ -156,7 +146,7 @@ class Branch(NamedTuple): list( zip_longest( writer([e for e in self.ends.values() if e != END]), - [la for la, e in self.ends.items() if e != END], + [str(la) for la, e in self.ends.items() if e != END], ) ) if self.ends diff --git a/libs/langgraph/langgraph/pregel/draw.py b/libs/langgraph/langgraph/pregel/draw.py index d6092e978..6c6cff9fd 100644 --- a/libs/langgraph/langgraph/pregel/draw.py +++ b/libs/langgraph/langgraph/pregel/draw.py @@ -1,8 +1,8 @@ from collections import defaultdict -from typing import Any, Mapping, Optional, Sequence, Union +from typing import Any, Mapping, Optional, Sequence, Union, cast from langchain_core.runnables.config import RunnableConfig -from langchain_core.runnables.graph import Graph +from langchain_core.runnables.graph import Graph, Node from langgraph.channels.base import BaseChannel from langgraph.checkpoint.base import BaseCheckpointSaver @@ -45,7 +45,7 @@ def draw_graph( The graph for this Pregel instance. """ # (src, dest, is_conditional) - edges: set[tuple[str, str, bool]] = set() + edges: set[tuple[str, str, bool, Optional[str]]] = set() step = -1 checkpoint = empty_checkpoint() @@ -204,8 +204,8 @@ def draw_graph( first, last = graph.extend(subgraph, prefix=name) for idx, edge in enumerate(graph.edges): if edge.source == name: - graph.edges[idx] = edge.copy(source=last.id) + graph.edges[idx] = edge.copy(source=cast(Node, last).id) elif edge.target == name: - graph.edges[idx] = edge.copy(target=first.id) + graph.edges[idx] = edge.copy(target=cast(Node, first).id) return graph diff --git a/libs/langgraph/langgraph/pregel/write.py b/libs/langgraph/langgraph/pregel/write.py index 42637efc0..c1badb549 100644 --- a/libs/langgraph/langgraph/pregel/write.py +++ b/libs/langgraph/langgraph/pregel/write.py @@ -165,6 +165,10 @@ class ChannelWrite(RunnableCallable): ] or None elif writes := getattr(runnable, "_is_channel_writer", MISSING): if writes is not MISSING: + writes = cast( + Sequence[tuple[Union[ChannelWriteEntry, Send], Optional[str]]], + writes, + ) entries = [e for e, _ in writes] labels = [la for _, la in writes] return [(*t, la) for t, la in zip(_assemble_writes(entries), labels)]