From 1b961f68b9d4f09abb207d27a02f5b0d0dab4052 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 28 May 2025 14:09:02 -0700 Subject: [PATCH] Remove Channel node builder - Replaced by NodeBuilder introduced in earlier PR --- docs/docs/reference/pregel.md | 3 +- libs/langgraph/langgraph/pregel/__init__.py | 75 +------------------ libs/langgraph/tests/test_pregel.py | 79 +++++---------------- 3 files changed, 21 insertions(+), 136 deletions(-) diff --git a/docs/docs/reference/pregel.md b/docs/docs/reference/pregel.md index a115b1fde..3bc7db105 100644 --- a/docs/docs/reference/pregel.md +++ b/docs/docs/reference/pregel.md @@ -6,8 +6,9 @@ show_root_heading: true show_root_full_path: false members: - - subscribe_single + - subscribe_only - subscribe_to + - read_from - do - write_to - meta diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 1cd6a225d..5ae746f89 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -8,7 +8,7 @@ import weakref from collections import defaultdict, deque from collections.abc import AsyncIterator, Iterator, Mapping, Sequence from functools import partial -from typing import Any, Callable, Union, cast, get_type_hints, overload +from typing import Any, Callable, Union, cast, get_type_hints from uuid import UUID, uuid5 from langchain_core.globals import get_debug @@ -297,79 +297,6 @@ class NodeBuilder: ) -# Deprecated, remove in 2.0 -class Channel: - @overload - @classmethod - def subscribe_to( - cls, - channels: str, - *, - key: str | None = None, - tags: list[str] | None = None, - ) -> PregelNode: ... - - @overload - @classmethod - def subscribe_to( - cls, - channels: Sequence[str], - *, - key: None = None, - tags: list[str] | None = None, - ) -> PregelNode: ... - - @classmethod - def subscribe_to( - cls, - channels: str | Sequence[str], - *, - key: str | None = None, - tags: list[str] | None = None, - ) -> PregelNode: - """Runs process.invoke() each time channels are updated, - with a dict of the channel values as input.""" - if not isinstance(channels, str) and key is not None: - raise ValueError( - "Can't specify a key when subscribing to multiple channels" - ) - return PregelNode( - channels=cast( - Union[list[str], Mapping[str, str]], - ( - {key: channels} - if isinstance(channels, str) and key is not None - else ( - [channels] - if isinstance(channels, str) - else {chan: chan for chan in channels} - ) - ), - ), - triggers=[channels] if isinstance(channels, str) else channels, - tags=tags, - ) - - @classmethod - def write_to( - cls, - *channels: str | ChannelWriteEntry, - **kwargs: WriteValue, - ) -> ChannelWrite: - """Writes to channels the result of the lambda, or None to skip writing.""" - return ChannelWrite( - [ChannelWriteEntry(c) if isinstance(c, str) else c for c in channels] - + [ - ( - ChannelWriteEntry(k, mapper=v) - if callable(v) - else ChannelWriteEntry(k, value=v) - ) - for k, v in kwargs.items() - ] - ) - - class Pregel(PregelProtocol): """Pregel manages the runtime behavior for LangGraph applications. diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 60bb05385..6ba5e8b43 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -49,14 +49,12 @@ from langgraph.graph import END, Graph, StateGraph from langgraph.graph.message import MessageGraph, MessagesState, add_messages from langgraph.prebuilt.tool_node import ToolNode from langgraph.pregel import ( - Channel, GraphRecursionError, NodeBuilder, Pregel, StateSnapshot, ) from langgraph.pregel.loop import SyncPregelLoop -from langgraph.pregel.read import PregelNode from langgraph.pregel.retry import RetryPolicy from langgraph.pregel.runner import PregelRunner from langgraph.store.base import BaseStore @@ -266,15 +264,9 @@ def test_checkpoint_errors() -> None: graph.invoke("", {"configurable": {"thread_id": "thread-1"}}) -@pytest.mark.parametrize("use_node_builder", [True, False]) -def test_config_json_schema(use_node_builder: bool) -> None: +def test_config_json_schema() -> None: """Test that config json schema is generated properly.""" - if use_node_builder: - chain: Union[NodeBuilder, PregelNode] = ( - NodeBuilder().subscribe_only("input").write_to("output") - ) - else: - chain = Channel.subscribe_to("input") | Channel.write_to("output") + chain = NodeBuilder().subscribe_only("input").write_to("output") @dataclass class Foo: @@ -481,15 +473,9 @@ def test_reducer_before_first_node() -> None: } -@pytest.mark.parametrize("use_node_builder", [True, False]) -def test_invoke_single_process_in_out( - mocker: MockerFixture, use_node_builder: bool -) -> None: +def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - if use_node_builder: - chain = NodeBuilder().subscribe_only("input").do(add_one).write_to("output") - else: - chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") + chain = NodeBuilder().subscribe_only("input").do(add_one).write_to("output") app = Pregel( nodes={ @@ -544,24 +530,14 @@ def test_invoke_single_process_in_out_falsy_values(falsy_value: Any) -> None: assert gapp.invoke(1) == falsy_value -@pytest.mark.parametrize("use_node_builder", [True, False]) -def test_invoke_single_process_in_write_kwargs( - mocker: MockerFixture, use_node_builder: bool -) -> None: +def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - if use_node_builder: - chain = ( - NodeBuilder() - .subscribe_only("input") - .do(add_one) - .write_to("output", fixed=5, output_plus_one=lambda x: x + 1) - ) - else: - chain = ( - Channel.subscribe_to("input") - | add_one - | Channel.write_to("output", fixed=5, output_plus_one=lambda x: x + 1) - ) + chain = ( + NodeBuilder() + .subscribe_only("input") + .do(add_one) + .write_to("output", fixed=5, output_plus_one=lambda x: x + 1) + ) app = Pregel( nodes={"one": chain}, @@ -595,15 +571,9 @@ def test_invoke_single_process_in_write_kwargs( assert app.invoke(2) == {"output": 3, "fixed": 5, "output_plus_one": 4} -@pytest.mark.parametrize("use_node_builder", [True, False]) -def test_invoke_single_process_in_out_dict( - mocker: MockerFixture, use_node_builder: bool -) -> None: +def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - if use_node_builder: - chain = NodeBuilder().subscribe_only("input").do(add_one).write_to("output") - else: - chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") + chain = NodeBuilder().subscribe_only("input").do(add_one).write_to("output") app = Pregel( nodes={"one": chain}, @@ -626,15 +596,9 @@ def test_invoke_single_process_in_out_dict( assert app.invoke(2) == {"output": 3} -@pytest.mark.parametrize("use_node_builder", [True, False]) -def test_invoke_single_process_in_dict_out_dict( - mocker: MockerFixture, use_node_builder: bool -) -> None: +def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - if use_node_builder: - chain = NodeBuilder().subscribe_only("input").do(add_one).write_to("output") - else: - chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") + chain = NodeBuilder().subscribe_only("input").do(add_one).write_to("output") app = Pregel( nodes={"one": chain}, @@ -657,17 +621,10 @@ def test_invoke_single_process_in_dict_out_dict( assert app.invoke({"input": 2}) == {"output": 3} -@pytest.mark.parametrize("use_node_builder", [True, False]) -def test_invoke_two_processes_in_out( - mocker: MockerFixture, use_node_builder: bool -) -> None: +def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - if use_node_builder: - one = NodeBuilder().subscribe_only("input").do(add_one).write_to("inbox") - two = NodeBuilder().subscribe_only("inbox").do(add_one).write_to("output") - else: - one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") - two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output") + one = NodeBuilder().subscribe_only("input").do(add_one).write_to("inbox") + two = NodeBuilder().subscribe_only("inbox").do(add_one).write_to("output") app = Pregel( nodes={"one": one, "two": two},