From 158f84c826fc8686e7d00a4276d9b3e880540613 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 29 Feb 2024 18:39:09 -0800 Subject: [PATCH 1/7] WIP Waiting edge --- langgraph/channels/named_barrier_value.py | 60 +++++++++++++++++ langgraph/graph/graph.py | 11 +++- langgraph/graph/state.py | 46 ++++++++++++- langgraph/pregel/io.py | 18 ++++- tests/test_pregel.py | 80 +++++++++++++++++++++++ 5 files changed, 207 insertions(+), 8 deletions(-) create mode 100644 langgraph/channels/named_barrier_value.py diff --git a/langgraph/channels/named_barrier_value.py b/langgraph/channels/named_barrier_value.py new file mode 100644 index 000000000..ccc86a541 --- /dev/null +++ b/langgraph/channels/named_barrier_value.py @@ -0,0 +1,60 @@ +from contextlib import contextmanager +from typing import Generator, Generic, Optional, Sequence, Type + +from typing_extensions import Self + +from langgraph.channels.base import ( + BaseChannel, + EmptyChannelError, + InvalidUpdateError, + Value, +) + + +class NamedBarrierValue(Generic[Value], BaseChannel[Value, Value, Value]): + """A channel that waits until all named values are received before making the value available.""" + + def __init__(self, typ: Type[Value], names: set[str]) -> None: + self.typ = typ + self.names = names + self.seen = set() + + @property + def ValueType(self) -> Type[Value]: + """The type of the value stored in the channel.""" + return self.typ + + @property + def UpdateType(self) -> Type[Value]: + """The type of the update received by the channel.""" + return self.typ + + @contextmanager + def empty(self, checkpoint: Optional[Value] = None) -> Generator[Self, None, None]: + empty = self.__class__(self.typ, self.names) + if checkpoint is not None: + empty.value = checkpoint + try: + yield empty + finally: + try: + del empty.value + except AttributeError: + pass + + def update(self, values: Sequence[Value]) -> None: + if self.seen == self.names: + self.seen = set() + for value in values: + if value in self.names: + self.seen.add(value) + else: + raise InvalidUpdateError(f"Value {value} not in {self.names}") + + def get(self) -> Value: + if self.seen != self.names: + raise EmptyChannelError() + return None + + def checkpoint(self) -> Value: + return self.seen diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py index 72ead903b..c7b6a2074 100644 --- a/langgraph/graph/graph.py +++ b/langgraph/graph/graph.py @@ -146,8 +146,13 @@ class Graph: def set_finish_point(self, key: str) -> None: return self.add_edge(key, END) - def validate(self, interrupt: Optional[Sequence[str]] = None) -> None: - all_starts = {src for src, _ in self.edges} | {src for src in self.branches} + def validate( + self, + interrupt: Optional[Sequence[str]] = None, + additional_edges: Optional[set[tuple[str, str]]] = None, + ) -> None: + edges = self.edges.union(additional_edges or []) + all_starts = {src for src, _ in edges} | {src for src in self.branches} for node in self.nodes: if node not in all_starts: raise ValueError(f"Node `{node}` is a dead-end") @@ -158,7 +163,7 @@ class Graph: if self.entry_point_branch is not None: branches.append(self.entry_point_branch) - all_hard_ends = {end for _, end in self.edges} + all_hard_ends = {end for _, end in edges} if self.entry_point is not None: all_hard_ends.add(self.entry_point) diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py index 76362979d..97c0b2abf 100644 --- a/langgraph/graph/state.py +++ b/langgraph/graph/state.py @@ -1,3 +1,4 @@ +import logging from collections import defaultdict from functools import partial from inspect import signature @@ -11,12 +12,15 @@ from langgraph.channels.base import BaseChannel, InvalidUpdateError from langgraph.channels.binop import BinaryOperatorAggregate from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue +from langgraph.channels.named_barrier_value import NamedBarrierValue from langgraph.checkpoint import BaseCheckpointSaver from langgraph.graph.graph import END, START, CompiledGraph, Graph from langgraph.pregel import Channel from langgraph.pregel.read import ChannelInvoke from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry +logger = logging.getLogger(__name__) + class StateGraph(Graph): def __init__(self, schema: Type[Any]) -> None: @@ -25,6 +29,7 @@ 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() def add_node(self, key: str, action: RunnableLike) -> None: if key in self.channels: @@ -34,6 +39,24 @@ class StateGraph(Graph): ) return super().add_node(key, action) + def add_waiting_edge(self, starts: Sequence[str], end: str) -> None: + if self.compiled: + logger.warning( + "Adding an edge to a graph that has already been compiled. This will " + "not be reflected in the compiled graph." + ) + for start in starts: + if start == END: + raise ValueError("END cannot be a start node") + if start not in self.nodes: + raise ValueError(f"Need to add_node `{start}` first") + if end == END: + raise ValueError("END cannot be an end node") + if end not in self.nodes: + raise ValueError(f"Need to add_node `{end}` first") + + self.w_edges.add((tuple(starts), end)) + def compile( self, checkpointer: Optional[BaseCheckpointSaver] = None, @@ -43,7 +66,12 @@ class StateGraph(Graph): ) -> CompiledGraph: interrupt_before = interrupt_before or [] interrupt_after = interrupt_after or [] - self.validate(interrupt=interrupt_before + interrupt_after) + self.validate( + interrupt=interrupt_before + interrupt_after, + additional_edges={ + (start, end) for starts, end in self.w_edges for start in starts + }, + ) state_keys = list(self.channels) state_keys_read = state_keys[0] if state_keys == ["__root__"] else state_keys @@ -68,14 +96,27 @@ class StateGraph(Graph): else None ) + waiting_edges = { + (f"{starts}:{end}", starts, end) for starts, end in self.w_edges + } + waiting_edge_channels = { + key: NamedBarrierValue(str, set(starts)) for key, starts, _ in waiting_edges + } + outgoing_edges = defaultdict(list) for start, end in self.edges: outgoing_edges[start].append(f"{end}:inbox" if end != END else END) + for key, starts, end in waiting_edges: + for start in starts: + outgoing_edges[start].append(key) nodes = { key: ( ChannelInvoke( - triggers=[f"{key}:inbox"], + triggers=[ + f"{key}:inbox", + *[chan for chan, _, end in waiting_edges if end == key], + ], channels=state_channels, mapper=coerce_state, ) @@ -142,6 +183,7 @@ class StateGraph(Graph): **self.channels, **node_inboxes, **node_outboxes, + **waiting_edge_channels, END: LastValue(self.schema), }, input=f"{START}:inbox", diff --git a/langgraph/pregel/io.py b/langgraph/pregel/io.py index 4dfc7b967..829cb9242 100644 --- a/langgraph/pregel/io.py +++ b/langgraph/pregel/io.py @@ -1,9 +1,21 @@ from typing import Any, Iterator, Mapping, Optional, Sequence, Union -from langgraph.channels.base import BaseChannel +from langgraph.channels.base import BaseChannel, EmptyChannelError from langgraph.pregel.log import logger +def _read_channel( + channels: Mapping[str, BaseChannel], chan: str, catch: bool = True +) -> Any: + try: + return channels[chan].get() + except EmptyChannelError: + if catch: + return None + else: + raise + + def map_input( input_channels: Union[str, Sequence[str]], chunk: Optional[Union[dict[str, Any], Any]], @@ -31,8 +43,8 @@ def map_output( """Map pending writes (a sequence of tuples (channel, value)) to output chunk.""" if isinstance(output_channels, str): if any(chan == output_channels for chan, _ in pending_writes): - return channels[output_channels].get() + return _read_channel(channels, output_channels) else: if updated := {c for c, _ in pending_writes if c in output_channels}: - return {chan: channels[chan].get() for chan in updated} + return {chan: _read_channel(channels, chan) for chan in updated} return None diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 37a57d90b..a9eb983da 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -2783,3 +2783,83 @@ def test_in_one_fan_out_out_one_graph_state() -> None: } }, ] + + +def test_in_one_fan_out_waiting_edge_graph_state() -> None: + def sorted_add( + x: list[str], y: Union[list[str], list[tuple[str, str]]] + ) -> list[str]: + if isinstance(y[0], tuple): + for rem, _ in y: + x.remove(rem) + y = [t[1] for t in y] + return sorted(operator.add(x, y)) + + class State(TypedDict, total=False): + query: str + answer: str + docs: Annotated[list[str], sorted_add] + + def rewrite_query(data: State) -> State: + return {"query": f'query: {data["query"]}'} + + def analyzer_one(data: State) -> State: + return {"query": f'analyzed: {data["query"]}'} + + def retriever_one(data: State) -> State: + return {"docs": ["doc1", "doc2"]} + + def retriever_two(data: State) -> State: + return {"docs": ["doc3", "doc4"]} + + def qa(data: State) -> State: + return {"answer": ",".join(data["docs"])} + + workflow = StateGraph(State) + + workflow.add_node("rewrite_query", rewrite_query) + workflow.add_node("analyzer_one", analyzer_one) + workflow.add_node("retriever_one", retriever_one) + workflow.add_node("retriever_two", retriever_two) + workflow.add_node("qa", qa) + + workflow.set_entry_point("rewrite_query") + workflow.add_edge("rewrite_query", "analyzer_one") + workflow.add_edge("analyzer_one", "retriever_one") + workflow.add_edge("rewrite_query", "retriever_two") + workflow.add_waiting_edge(["retriever_one", "retriever_two"], "qa") + workflow.set_finish_point("qa") + + app = workflow.compile() + + assert app.invoke({"query": "what is weather in sf"}, debug=True) == { + "query": "analyzed: query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + "answer": "doc1,doc2,doc3,doc4", + } + + assert [*app.stream({"query": "what is weather in sf"})] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + { + "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, + "retriever_two": {"docs": ["doc3", "doc4"]}, + }, + {"('retriever_one', 'retriever_two'):qa": None}, + {"retriever_one": {"docs": ["doc1", "doc2"]}, "qa": {"answer": "doc3,doc4"}}, + { + "('retriever_on', 'retriever_two'):qa": None, + "__end__": { + "query": "analyzed: query: what is weather in sf", + "answer": "doc3,doc4", + "docs": ["doc1", "doc2", "doc3", "doc4"], + }, + }, + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + { + "__end__": { + "query": "analyzed: query: what is weather in sf", + "answer": "doc1,doc2,doc3,doc4", + "docs": ["doc1", "doc2", "doc3", "doc4"], + } + }, + ] From 6cb5062c1ad50daf45c83aa0431c88eddb52068f Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 21 Mar 2024 14:30:51 -0700 Subject: [PATCH 2/7] Lol --- tests/test_pregel.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_pregel.py b/tests/test_pregel.py index a9eb983da..b97a9ff04 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -2847,7 +2847,7 @@ def test_in_one_fan_out_waiting_edge_graph_state() -> None: {"('retriever_one', 'retriever_two'):qa": None}, {"retriever_one": {"docs": ["doc1", "doc2"]}, "qa": {"answer": "doc3,doc4"}}, { - "('retriever_on', 'retriever_two'):qa": None, + "('retriever_one', 'retriever_two'):qa": None, "__end__": { "query": "analyzed: query: what is weather in sf", "answer": "doc3,doc4", From 01c32c2464d32ecd93e9aa16bfad0a8b6ef34210 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 21 Mar 2024 15:54:03 -0700 Subject: [PATCH 3/7] Finish --- langgraph/channels/any_value.py | 4 + langgraph/channels/named_barrier_value.py | 8 +- langgraph/graph/state.py | 5 +- langgraph/pregel/__init__.py | 20 +- tests/conftest.py | 3 + tests/test_pregel.py | 263 ++++++++++++++++- tests/test_pregel_async.py | 325 ++++++++++++++++++++++ 7 files changed, 605 insertions(+), 23 deletions(-) diff --git a/langgraph/channels/any_value.py b/langgraph/channels/any_value.py index 2cbb2c25a..73edb3494 100644 --- a/langgraph/channels/any_value.py +++ b/langgraph/channels/any_value.py @@ -38,6 +38,10 @@ class AnyValue(Generic[Value], BaseChannel[Value, Value, Value]): def update(self, values: Sequence[Value]) -> None: if len(values) == 0: + try: + del self.value + except AttributeError: + pass return self.value = values[-1] diff --git a/langgraph/channels/named_barrier_value.py b/langgraph/channels/named_barrier_value.py index ccc86a541..7f55de31b 100644 --- a/langgraph/channels/named_barrier_value.py +++ b/langgraph/channels/named_barrier_value.py @@ -33,14 +33,12 @@ class NamedBarrierValue(Generic[Value], BaseChannel[Value, Value, Value]): def empty(self, checkpoint: Optional[Value] = None) -> Generator[Self, None, None]: empty = self.__class__(self.typ, self.names) if checkpoint is not None: - empty.value = checkpoint + empty.seen = checkpoint + try: yield empty finally: - try: - del empty.value - except AttributeError: - pass + pass def update(self, values: Sequence[Value]) -> None: if self.seen == self.names: diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py index 97c0b2abf..c12f2f9d7 100644 --- a/langgraph/graph/state.py +++ b/langgraph/graph/state.py @@ -188,7 +188,10 @@ class StateGraph(Graph): }, input=f"{START}:inbox", output=END, - hidden=[f"{node}:inbox" for node in self.nodes] + [START] + state_keys, + hidden=[f"{node}:inbox" for node in self.nodes] + + [START] + + state_keys + + [key for key, _, _ in waiting_edges], snapshot_channels=state_keys_read, checkpointer=checkpointer, interrupt_before_nodes=[f"{node}:inbox" for node in interrupt_before], diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index e7a45a38b..f170dc746 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -1018,12 +1018,18 @@ def _should_interrupt( def _read_channel( - channels: Mapping[str, BaseChannel], chan: str, catch: bool = True + channels: Mapping[str, BaseChannel], + chan: str, + *, + catch: bool = True, + return_exception: bool = False, ) -> Any: try: return channels[chan].get() - except EmptyChannelError: - if catch: + except EmptyChannelError as exc: + if return_exception: + return exc + elif catch: return None else: raise @@ -1090,7 +1096,7 @@ def _prepare_next_tasks( channels: Mapping[str, BaseChannel], update_seen: bool = True, ) -> tuple[Checkpoint, list[tuple[Runnable, Any, str]]]: - checkpoint = copy_checkpoint(checkpoint) if update_seen else checkpoint + checkpoint = copy_checkpoint(checkpoint) tasks: list[tuple[Runnable, Any, str]] = [] # Check if any processes should be run in next step # If so, prepare the values to be passed to them @@ -1098,7 +1104,11 @@ def _prepare_next_tasks( seen = checkpoint["versions_seen"][name] # If any of the channels read by this process were updated if any( - checkpoint["channel_versions"][chan] > seen[chan] for chan in proc.triggers + checkpoint["channel_versions"][chan] > seen[chan] + for chan in proc.triggers + if not isinstance( + _read_channel(channels, chan, return_exception=True), EmptyChannelError + ) ): # If all trigger channels subscribed by this process are not empty # then invoke the process with the values of all non-empty channels diff --git a/tests/conftest.py b/tests/conftest.py index 9b6b04c01..42278f79a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -10,3 +10,6 @@ def deterministic_uuids(mocker: MockerFixture) -> MockerFixture: UUID(f"00000000-0000-4000-8000-{i:012}", version=4) for i in range(10000) ) return mocker.patch("uuid.uuid4", side_effect=side_effect) + + +pytest.register_assert_rewrite("tests.memory_assert") diff --git a/tests/test_pregel.py b/tests/test_pregel.py index b97a9ff04..ab190751f 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -2785,7 +2785,7 @@ def test_in_one_fan_out_out_one_graph_state() -> None: ] -def test_in_one_fan_out_waiting_edge_graph_state() -> None: +def test_in_one_fan_out_state_graph_waiting_edge() -> None: def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] ) -> list[str]: @@ -2832,7 +2832,7 @@ def test_in_one_fan_out_waiting_edge_graph_state() -> None: app = workflow.compile() - assert app.invoke({"query": "what is weather in sf"}, debug=True) == { + assert app.invoke({"query": "what is weather in sf"}) == { "query": "analyzed: query: what is weather in sf", "docs": ["doc1", "doc2", "doc3", "doc4"], "answer": "doc1,doc2,doc3,doc4", @@ -2844,16 +2844,7 @@ def test_in_one_fan_out_waiting_edge_graph_state() -> None: "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, "retriever_two": {"docs": ["doc3", "doc4"]}, }, - {"('retriever_one', 'retriever_two'):qa": None}, - {"retriever_one": {"docs": ["doc1", "doc2"]}, "qa": {"answer": "doc3,doc4"}}, - { - "('retriever_one', 'retriever_two'):qa": None, - "__end__": { - "query": "analyzed: query: what is weather in sf", - "answer": "doc3,doc4", - "docs": ["doc1", "doc2", "doc3", "doc4"], - }, - }, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, { "__end__": { @@ -2863,3 +2854,251 @@ def test_in_one_fan_out_waiting_edge_graph_state() -> None: } }, ] + + app_w_interrupt = workflow.compile( + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["retriever_one"] + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c for c in app_w_interrupt.stream({"query": "what is weather in sf"}, config) + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + { + "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, + "retriever_two": {"docs": ["doc3", "doc4"]}, + }, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + ] + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + { + "__end__": { + "query": "analyzed: query: what is weather in sf", + "answer": "doc1,doc2,doc3,doc4", + "docs": ["doc1", "doc2", "doc3", "doc4"], + } + }, + ] + + +def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> None: + def sorted_add( + x: list[str], y: Union[list[str], list[tuple[str, str]]] + ) -> list[str]: + if isinstance(y[0], tuple): + for rem, _ in y: + x.remove(rem) + y = [t[1] for t in y] + return sorted(operator.add(x, y)) + + class State(TypedDict, total=False): + query: str + answer: str + docs: Annotated[list[str], sorted_add] + + def rewrite_query(data: State) -> State: + return {"query": f'query: {data["query"]}'} + + def analyzer_one(data: State) -> State: + return {"query": f'analyzed: {data["query"]}'} + + def retriever_one(data: State) -> State: + return {"docs": ["doc1", "doc2"]} + + def retriever_two(data: State) -> State: + return {"docs": ["doc3", "doc4"]} + + def qa(data: State) -> State: + return {"answer": ",".join(data["docs"])} + + workflow = StateGraph(State) + + workflow.add_node("rewrite_query", rewrite_query) + workflow.add_node("analyzer_one", analyzer_one) + workflow.add_node("retriever_one", retriever_one) + workflow.add_node("retriever_two", retriever_two) + workflow.add_node("qa", qa) + + workflow.set_entry_point("rewrite_query") + workflow.add_edge("rewrite_query", "analyzer_one") + workflow.add_edge("analyzer_one", "retriever_one") + workflow.add_edge("rewrite_query", "retriever_two") + workflow.add_waiting_edge(["retriever_one", "retriever_two"], "qa") + workflow.set_finish_point("qa") + + # silly edge, to make sure having been triggered before doesn't break + # semantics of named barrier (== waiting edges) + workflow.add_edge("rewrite_query", "qa") + + app = workflow.compile() + + assert app.invoke({"query": "what is weather in sf"}) == { + "query": "analyzed: query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + "answer": "doc1,doc2,doc3,doc4", + } + + assert [*app.stream({"query": "what is weather in sf"})] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + { + "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, + "retriever_two": {"docs": ["doc3", "doc4"]}, + "qa": {"answer": ""}, + }, + { + "__end__": { + "answer": "", + "docs": ["doc3", "doc4"], + "query": "analyzed: query: what is weather in sf", + } + }, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + { + "__end__": { + "query": "analyzed: query: what is weather in sf", + "answer": "doc1,doc2,doc3,doc4", + "docs": ["doc1", "doc2", "doc3", "doc4"], + } + }, + ] + + app_w_interrupt = workflow.compile( + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["retriever_one"] + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c for c in app_w_interrupt.stream({"query": "what is weather in sf"}, config) + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + { + "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, + "retriever_two": {"docs": ["doc3", "doc4"]}, + "qa": {"answer": ""}, + }, + { + "__end__": { + "answer": "", + "docs": ["doc3", "doc4"], + "query": "analyzed: query: what is weather in sf", + } + }, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + ] + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + { + "__end__": { + "query": "analyzed: query: what is weather in sf", + "answer": "doc1,doc2,doc3,doc4", + "docs": ["doc1", "doc2", "doc3", "doc4"], + } + }, + ] + + +def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None: + def sorted_add( + x: list[str], y: Union[list[str], list[tuple[str, str]]] + ) -> list[str]: + if isinstance(y[0], tuple): + for rem, _ in y: + x.remove(rem) + y = [t[1] for t in y] + return sorted(operator.add(x, y)) + + class State(TypedDict, total=False): + query: str + answer: str + docs: Annotated[list[str], sorted_add] + + def rewrite_query(data: State) -> State: + return {"query": f'query: {data["query"]}'} + + def analyzer_one(data: State) -> State: + return {"query": f'analyzed: {data["query"]}'} + + def retriever_one(data: State) -> State: + return {"docs": ["doc1", "doc2"]} + + def retriever_two(data: State) -> State: + return {"docs": ["doc3", "doc4"]} + + def qa(data: State) -> State: + return {"answer": ",".join(data["docs"])} + + def decider(data: State) -> None: + return None + + def decider_cond(data: State) -> str: + if data["query"].count("analyzed") > 1: + return "qa" + else: + return "rewrite_query" + + workflow = StateGraph(State) + + workflow.add_node("rewrite_query", rewrite_query) + workflow.add_node("analyzer_one", analyzer_one) + workflow.add_node("retriever_one", retriever_one) + workflow.add_node("retriever_two", retriever_two) + workflow.add_node("decider", decider) + workflow.add_node("qa", qa) + + workflow.set_entry_point("rewrite_query") + workflow.add_edge("rewrite_query", "analyzer_one") + workflow.add_edge("analyzer_one", "retriever_one") + workflow.add_edge("rewrite_query", "retriever_two") + workflow.add_waiting_edge(["retriever_one", "retriever_two"], "decider") + workflow.add_conditional_edges("decider", decider_cond) + workflow.set_finish_point("qa") + + app = workflow.compile() + + assert app.invoke({"query": "what is weather in sf"}) == { + "query": "analyzed: query: analyzed: query: what is weather in sf", + "answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4", + "docs": ["doc1", "doc1", "doc2", "doc2", "doc3", "doc3", "doc4", "doc4"], + } + + assert [*app.stream({"query": "what is weather in sf"})] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + { + "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, + "retriever_two": {"docs": ["doc3", "doc4"]}, + }, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"decider": None}, + {"rewrite_query": {"query": "query: analyzed: query: what is weather in sf"}}, + { + "analyzer_one": { + "query": "analyzed: query: analyzed: query: what is weather in sf" + }, + "retriever_two": {"docs": ["doc3", "doc4"]}, + }, + { + "retriever_one": {"docs": ["doc1", "doc2"]}, + }, + {"decider": None}, + {"qa": {"answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4"}}, + { + "__end__": { + "query": "analyzed: query: analyzed: query: what is weather in sf", + "answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4", + "docs": [ + "doc1", + "doc1", + "doc2", + "doc2", + "doc3", + "doc3", + "doc4", + "doc4", + ], + } + }, + ] diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 3ebb79910..b4d8cb1cf 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -2812,3 +2812,328 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: } }, ] + + +async def test_in_one_fan_out_state_graph_waiting_edge() -> None: + def sorted_add( + x: list[str], y: Union[list[str], list[tuple[str, str]]] + ) -> list[str]: + if isinstance(y[0], tuple): + for rem, _ in y: + x.remove(rem) + y = [t[1] for t in y] + return sorted(operator.add(x, y)) + + class State(TypedDict, total=False): + query: str + answer: str + docs: Annotated[list[str], sorted_add] + + async def rewrite_query(data: State) -> State: + return {"query": f'query: {data["query"]}'} + + async def analyzer_one(data: State) -> State: + return {"query": f'analyzed: {data["query"]}'} + + async def retriever_one(data: State) -> State: + return {"docs": ["doc1", "doc2"]} + + async def retriever_two(data: State) -> State: + return {"docs": ["doc3", "doc4"]} + + async def qa(data: State) -> State: + return {"answer": ",".join(data["docs"])} + + workflow = StateGraph(State) + + workflow.add_node("rewrite_query", rewrite_query) + workflow.add_node("analyzer_one", analyzer_one) + workflow.add_node("retriever_one", retriever_one) + workflow.add_node("retriever_two", retriever_two) + workflow.add_node("qa", qa) + + workflow.set_entry_point("rewrite_query") + workflow.add_edge("rewrite_query", "analyzer_one") + workflow.add_edge("analyzer_one", "retriever_one") + workflow.add_edge("rewrite_query", "retriever_two") + workflow.add_waiting_edge(["retriever_one", "retriever_two"], "qa") + workflow.set_finish_point("qa") + + app = workflow.compile() + + assert await app.ainvoke({"query": "what is weather in sf"}) == { + "query": "analyzed: query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + "answer": "doc1,doc2,doc3,doc4", + } + + assert [c async for c in app.astream({"query": "what is weather in sf"})] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + { + "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, + "retriever_two": {"docs": ["doc3", "doc4"]}, + }, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + { + "__end__": { + "query": "analyzed: query: what is weather in sf", + "answer": "doc1,doc2,doc3,doc4", + "docs": ["doc1", "doc2", "doc3", "doc4"], + } + }, + ] + + app_w_interrupt = workflow.compile( + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["retriever_one"] + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + async for c in app_w_interrupt.astream( + {"query": "what is weather in sf"}, config + ) + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + { + "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, + "retriever_two": {"docs": ["doc3", "doc4"]}, + }, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + ] + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + { + "__end__": { + "query": "analyzed: query: what is weather in sf", + "answer": "doc1,doc2,doc3,doc4", + "docs": ["doc1", "doc2", "doc3", "doc4"], + } + }, + ] + + +async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> None: + def sorted_add( + x: list[str], y: Union[list[str], list[tuple[str, str]]] + ) -> list[str]: + if isinstance(y[0], tuple): + for rem, _ in y: + x.remove(rem) + y = [t[1] for t in y] + return sorted(operator.add(x, y)) + + class State(TypedDict, total=False): + query: str + answer: str + docs: Annotated[list[str], sorted_add] + + async def rewrite_query(data: State) -> State: + return {"query": f'query: {data["query"]}'} + + async def analyzer_one(data: State) -> State: + return {"query": f'analyzed: {data["query"]}'} + + async def retriever_one(data: State) -> State: + return {"docs": ["doc1", "doc2"]} + + async def retriever_two(data: State) -> State: + return {"docs": ["doc3", "doc4"]} + + async def qa(data: State) -> State: + return {"answer": ",".join(data["docs"])} + + workflow = StateGraph(State) + + workflow.add_node("rewrite_query", rewrite_query) + workflow.add_node("analyzer_one", analyzer_one) + workflow.add_node("retriever_one", retriever_one) + workflow.add_node("retriever_two", retriever_two) + workflow.add_node("qa", qa) + + workflow.set_entry_point("rewrite_query") + workflow.add_edge("rewrite_query", "analyzer_one") + workflow.add_edge("analyzer_one", "retriever_one") + workflow.add_edge("rewrite_query", "retriever_two") + workflow.add_waiting_edge(["retriever_one", "retriever_two"], "qa") + workflow.set_finish_point("qa") + + # silly edge, to make sure having been triggered before doesn't break + # semantics of named barrier (== waiting edges) + workflow.add_edge("rewrite_query", "qa") + + app = workflow.compile() + + assert await app.ainvoke({"query": "what is weather in sf"}) == { + "query": "analyzed: query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + "answer": "doc1,doc2,doc3,doc4", + } + + assert [c async for c in app.astream({"query": "what is weather in sf"})] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + { + "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, + "retriever_two": {"docs": ["doc3", "doc4"]}, + "qa": {"answer": ""}, + }, + { + "__end__": { + "answer": "", + "docs": ["doc3", "doc4"], + "query": "analyzed: query: what is weather in sf", + } + }, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + { + "__end__": { + "query": "analyzed: query: what is weather in sf", + "answer": "doc1,doc2,doc3,doc4", + "docs": ["doc1", "doc2", "doc3", "doc4"], + } + }, + ] + + app_w_interrupt = workflow.compile( + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["retriever_one"] + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + async for c in app_w_interrupt.astream( + {"query": "what is weather in sf"}, config + ) + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + { + "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, + "retriever_two": {"docs": ["doc3", "doc4"]}, + "qa": {"answer": ""}, + }, + { + "__end__": { + "answer": "", + "docs": ["doc3", "doc4"], + "query": "analyzed: query: what is weather in sf", + } + }, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + ] + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + { + "__end__": { + "query": "analyzed: query: what is weather in sf", + "answer": "doc1,doc2,doc3,doc4", + "docs": ["doc1", "doc2", "doc3", "doc4"], + } + }, + ] + + +async def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None: + def sorted_add( + x: list[str], y: Union[list[str], list[tuple[str, str]]] + ) -> list[str]: + if isinstance(y[0], tuple): + for rem, _ in y: + x.remove(rem) + y = [t[1] for t in y] + return sorted(operator.add(x, y)) + + class State(TypedDict, total=False): + query: str + answer: str + docs: Annotated[list[str], sorted_add] + + async def rewrite_query(data: State) -> State: + return {"query": f'query: {data["query"]}'} + + async def analyzer_one(data: State) -> State: + return {"query": f'analyzed: {data["query"]}'} + + async def retriever_one(data: State) -> State: + return {"docs": ["doc1", "doc2"]} + + async def retriever_two(data: State) -> State: + return {"docs": ["doc3", "doc4"]} + + async def qa(data: State) -> State: + return {"answer": ",".join(data["docs"])} + + async def decider(data: State) -> None: + return None + + async def decider_cond(data: State) -> str: + if data["query"].count("analyzed") > 1: + return "qa" + else: + return "rewrite_query" + + workflow = StateGraph(State) + + workflow.add_node("rewrite_query", rewrite_query) + workflow.add_node("analyzer_one", analyzer_one) + workflow.add_node("retriever_one", retriever_one) + workflow.add_node("retriever_two", retriever_two) + workflow.add_node("decider", decider) + workflow.add_node("qa", qa) + + workflow.set_entry_point("rewrite_query") + workflow.add_edge("rewrite_query", "analyzer_one") + workflow.add_edge("analyzer_one", "retriever_one") + workflow.add_edge("rewrite_query", "retriever_two") + workflow.add_waiting_edge(["retriever_one", "retriever_two"], "decider") + workflow.add_conditional_edges("decider", decider_cond) + workflow.set_finish_point("qa") + + app = workflow.compile() + + assert await app.ainvoke({"query": "what is weather in sf"}) == { + "query": "analyzed: query: analyzed: query: what is weather in sf", + "answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4", + "docs": ["doc1", "doc1", "doc2", "doc2", "doc3", "doc3", "doc4", "doc4"], + } + + assert [c async for c in app.astream({"query": "what is weather in sf"})] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + { + "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, + "retriever_two": {"docs": ["doc3", "doc4"]}, + }, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"decider": None}, + {"rewrite_query": {"query": "query: analyzed: query: what is weather in sf"}}, + { + "analyzer_one": { + "query": "analyzed: query: analyzed: query: what is weather in sf" + }, + "retriever_two": {"docs": ["doc3", "doc4"]}, + }, + { + "retriever_one": {"docs": ["doc1", "doc2"]}, + }, + {"decider": None}, + {"qa": {"answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4"}}, + { + "__end__": { + "query": "analyzed: query: analyzed: query: what is weather in sf", + "answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4", + "docs": [ + "doc1", + "doc1", + "doc2", + "doc2", + "doc3", + "doc3", + "doc4", + "doc4", + ], + } + }, + ] From ccea39500299e98c6a893ef8a83cc35696d3fc7a Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 21 Mar 2024 16:06:23 -0700 Subject: [PATCH 4/7] Fox --- tests/test_pregel_async.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index b4d8cb1cf..4cca085e1 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -3069,7 +3069,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None: async def decider(data: State) -> None: return None - async def decider_cond(data: State) -> str: + def decider_cond(data: State) -> str: if data["query"].count("analyzed") > 1: return "qa" else: From e11153cf0959427634a75483273619bd314b0fb3 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 21 Mar 2024 16:53:59 -0700 Subject: [PATCH 5/7] Fix graph repr --- langgraph/graph/graph.py | 9 ++++++--- langgraph/graph/state.py | 13 +++++++------ tests/test_pregel.py | 36 ++++++++++++++++++++++++++++++++++++ 3 files changed, 49 insertions(+), 9 deletions(-) diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py index c7b6a2074..097e75af0 100644 --- a/langgraph/graph/graph.py +++ b/langgraph/graph/graph.py @@ -50,6 +50,10 @@ class Graph: self.entry_point: Optional[str] = None self.entry_point_branch: Optional[Branch] = None + @property + def _all_edges(self) -> set[tuple[str, str]]: + return self.edges + def add_node(self, key: str, action: RunnableLike) -> None: if self.compiled: logger.warning( @@ -149,9 +153,8 @@ class Graph: def validate( self, interrupt: Optional[Sequence[str]] = None, - additional_edges: Optional[set[tuple[str, str]]] = None, ) -> None: - edges = self.edges.union(additional_edges or []) + edges = self._all_edges all_starts = {src for src, _ in edges} | {src for src in self.branches} for node in self.nodes: if node not in all_starts: @@ -280,7 +283,7 @@ class CompiledGraph(Pregel): n = graph.add_node(node, key) start_nodes[key] = n end_nodes[key] = n - for start, end in self.graph.edges: + for start, end in self.graph._all_edges: graph.add_edge(start_nodes[start], end_nodes[end]) for start, branches in self.graph.branches.items(): for i, branch in enumerate(branches): diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py index c12f2f9d7..2c766c410 100644 --- a/langgraph/graph/state.py +++ b/langgraph/graph/state.py @@ -31,6 +31,12 @@ class StateGraph(Graph): self.support_multiple_edges = True self.w_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 + } + def add_node(self, key: str, action: RunnableLike) -> None: if key in self.channels: raise ValueError( @@ -66,12 +72,7 @@ class StateGraph(Graph): ) -> CompiledGraph: interrupt_before = interrupt_before or [] interrupt_after = interrupt_after or [] - self.validate( - interrupt=interrupt_before + interrupt_after, - additional_edges={ - (start, end) for starts, end in self.w_edges for start in starts - }, - ) + self.validate(interrupt=interrupt_before + interrupt_after) state_keys = list(self.channels) state_keys_read = state_keys[0] if state_keys == ["__root__"] else state_keys diff --git a/tests/test_pregel.py b/tests/test_pregel.py index ab190751f..0cd238c66 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -2832,6 +2832,42 @@ def test_in_one_fan_out_state_graph_waiting_edge() -> None: app = workflow.compile() + assert app.get_graph().draw_ascii() == ( + """ +-----------+ + | __start__ | + +-----------+ + * + * + * + +---------------+ + | rewrite_query | + +---------------+ + *** *** + * * + ** *** ++--------------+ * +| analyzer_one | * ++--------------+ * + * * + * * + * * ++---------------+ +---------------+ +| retriever_one | | retriever_two | ++---------------+ +---------------+ + *** *** + * * + ** ** + +----+ + | qa | + +----+ + * + * + * + +---------+ + | __end__ | + +---------+ """ + ) + assert app.invoke({"query": "what is weather in sf"}) == { "query": "analyzed: query: what is weather in sf", "docs": ["doc1", "doc2", "doc3", "doc4"], From db2ec5b2de78bddaa32a623f9badb9fae4cf10b1 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 21 Mar 2024 17:56:35 -0700 Subject: [PATCH 6/7] Rename to `add_edge(string[], string)` --- langgraph/graph/state.py | 17 ++++++++++------- tests/test_pregel.py | 6 +++--- tests/test_pregel_async.py | 6 +++--- 3 files changed, 16 insertions(+), 13 deletions(-) diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py index 2c766c410..3bb0adf5c 100644 --- a/langgraph/graph/state.py +++ b/langgraph/graph/state.py @@ -2,7 +2,7 @@ import logging from collections import defaultdict from functools import partial from inspect import signature -from typing import Any, Optional, Sequence, Type +from typing import Any, Optional, Sequence, Type, Union from langchain_core.runnables import RunnableLambda from langchain_core.runnables.base import RunnableLike @@ -45,23 +45,26 @@ class StateGraph(Graph): ) return super().add_node(key, action) - def add_waiting_edge(self, starts: Sequence[str], end: str) -> None: + def add_edge(self, start_key: Union[str, list[str]], end_key: str) -> None: + if isinstance(start_key, str): + return super().add_edge(start_key, end_key) + if self.compiled: logger.warning( "Adding an edge to a graph that has already been compiled. This will " "not be reflected in the compiled graph." ) - for start in starts: + for start in start_key: if start == END: raise ValueError("END cannot be a start node") if start not in self.nodes: raise ValueError(f"Need to add_node `{start}` first") - if end == END: + if end_key == END: raise ValueError("END cannot be an end node") - if end not in self.nodes: - raise ValueError(f"Need to add_node `{end}` first") + if end_key not in self.nodes: + raise ValueError(f"Need to add_node `{end_key}` first") - self.w_edges.add((tuple(starts), end)) + self.w_edges.add((tuple(start_key), end_key)) def compile( self, diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 0cd238c66..03e46379e 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -2827,7 +2827,7 @@ def test_in_one_fan_out_state_graph_waiting_edge() -> None: workflow.add_edge("rewrite_query", "analyzer_one") workflow.add_edge("analyzer_one", "retriever_one") workflow.add_edge("rewrite_query", "retriever_two") - workflow.add_waiting_edge(["retriever_one", "retriever_two"], "qa") + workflow.add_edge(["retriever_one", "retriever_two"], "qa") workflow.set_finish_point("qa") app = workflow.compile() @@ -2961,7 +2961,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> None: workflow.add_edge("rewrite_query", "analyzer_one") workflow.add_edge("analyzer_one", "retriever_one") workflow.add_edge("rewrite_query", "retriever_two") - workflow.add_waiting_edge(["retriever_one", "retriever_two"], "qa") + workflow.add_edge(["retriever_one", "retriever_two"], "qa") workflow.set_finish_point("qa") # silly edge, to make sure having been triggered before doesn't break @@ -3089,7 +3089,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None: workflow.add_edge("rewrite_query", "analyzer_one") workflow.add_edge("analyzer_one", "retriever_one") workflow.add_edge("rewrite_query", "retriever_two") - workflow.add_waiting_edge(["retriever_one", "retriever_two"], "decider") + workflow.add_edge(["retriever_one", "retriever_two"], "decider") workflow.add_conditional_edges("decider", decider_cond) workflow.set_finish_point("qa") diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 4cca085e1..5ada8b3c3 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -2856,7 +2856,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge() -> None: workflow.add_edge("rewrite_query", "analyzer_one") workflow.add_edge("analyzer_one", "retriever_one") workflow.add_edge("rewrite_query", "retriever_two") - workflow.add_waiting_edge(["retriever_one", "retriever_two"], "qa") + workflow.add_edge(["retriever_one", "retriever_two"], "qa") workflow.set_finish_point("qa") app = workflow.compile() @@ -2957,7 +2957,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> None: workflow.add_edge("rewrite_query", "analyzer_one") workflow.add_edge("analyzer_one", "retriever_one") workflow.add_edge("rewrite_query", "retriever_two") - workflow.add_waiting_edge(["retriever_one", "retriever_two"], "qa") + workflow.add_edge(["retriever_one", "retriever_two"], "qa") workflow.set_finish_point("qa") # silly edge, to make sure having been triggered before doesn't break @@ -3088,7 +3088,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None: workflow.add_edge("rewrite_query", "analyzer_one") workflow.add_edge("analyzer_one", "retriever_one") workflow.add_edge("rewrite_query", "retriever_two") - workflow.add_waiting_edge(["retriever_one", "retriever_two"], "decider") + workflow.add_edge(["retriever_one", "retriever_two"], "decider") workflow.add_conditional_edges("decider", decider_cond) workflow.set_finish_point("qa") From 62dfa48ce57c46665eaa0aee9047a6e3db827016 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 21 Mar 2024 17:58:23 -0700 Subject: [PATCH 7/7] Lint --- langgraph/graph/graph.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py index 097e75af0..e3f15798b 100644 --- a/langgraph/graph/graph.py +++ b/langgraph/graph/graph.py @@ -150,12 +150,10 @@ class Graph: def set_finish_point(self, key: str) -> None: return self.add_edge(key, END) - def validate( - self, - interrupt: Optional[Sequence[str]] = None, - ) -> None: - edges = self._all_edges - all_starts = {src for src, _ in edges} | {src for src in self.branches} + def validate(self, interrupt: Optional[Sequence[str]] = None) -> None: + all_starts = {src for src, _ in self._all_edges} | { + src for src in self.branches + } for node in self.nodes: if node not in all_starts: raise ValueError(f"Node `{node}` is a dead-end") @@ -166,7 +164,7 @@ class Graph: if self.entry_point_branch is not None: branches.append(self.entry_point_branch) - all_hard_ends = {end for _, end in edges} + all_hard_ends = {end for _, end in self._all_edges} if self.entry_point is not None: all_hard_ends.add(self.entry_point)