From a909d90b0819b78f3b26921a2c33b85fa35ad30a Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 14 May 2024 15:36:38 -0700 Subject: [PATCH] More small fixes following JS impl --- langgraph/pregel/__init__.py | 39 +++++++++++++----------------------- langgraph/pregel/io.py | 10 ++++++++- tests/test_io.py | 18 +++++++++++++++++ 3 files changed, 41 insertions(+), 26 deletions(-) create mode 100644 tests/test_io.py diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 63a78833f..e6fffd3ce 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -88,6 +88,7 @@ from langgraph.pregel.io import ( map_output_values, read_channel, read_channels, + single, ) from langgraph.pregel.log import logger from langgraph.pregel.read import PregelNode @@ -751,9 +752,7 @@ class Pregel( ) checkpoint_config = { "configurable": { - "thread_id": checkpoint_config["configurable"][ - "thread_id" - ], + **checkpoint_config["configurable"], "thread_ts": checkpoint["ts"], } } @@ -786,7 +785,7 @@ class Pregel( # if no more tasks, we're done if not next_tasks: - if step == 0: + if step == start: raise ValueError("No tasks to run in graph.") else: break @@ -884,25 +883,21 @@ class Pregel( { "source": "loop", "step": step, - "writes": next( - map_output_updates(output_keys, next_tasks), - None, + "writes": single( + map_output_updates(output_keys, next_tasks) ) if self.stream_mode == "updates" - else next( + else single( map_output_values( output_keys, pending_writes, channels ), - None, ), }, ) ) checkpoint_config = { "configurable": { - "thread_id": checkpoint_config["configurable"][ - "thread_id" - ], + **checkpoint_config["configurable"], "thread_ts": checkpoint["ts"], } } @@ -1048,9 +1043,7 @@ class Pregel( ) checkpoint_config = { "configurable": { - "thread_id": checkpoint_config["configurable"][ - "thread_id" - ], + **checkpoint_config["configurable"], "thread_ts": checkpoint["ts"], } } @@ -1083,7 +1076,7 @@ class Pregel( # if no more tasks, we're done if not next_tasks: - if step == 0: + if step == start: raise ValueError("No tasks to run in graph.") else: break @@ -1191,16 +1184,14 @@ class Pregel( { "source": "loop", "step": step, - "writes": next( - map_output_updates(output_keys, next_tasks), - None, + "writes": single( + map_output_updates(output_keys, next_tasks) ) if self.stream_mode == "updates" - else next( + else single( map_output_values( output_keys, pending_writes, channels - ), - None, + ) ), }, ) @@ -1208,9 +1199,7 @@ class Pregel( ) checkpoint_config = { "configurable": { - "thread_id": checkpoint_config["configurable"][ - "thread_id" - ], + **checkpoint_config["configurable"], "thread_ts": checkpoint["ts"], } } diff --git a/langgraph/pregel/io.py b/langgraph/pregel/io.py index 67382fc77..6228167cc 100644 --- a/langgraph/pregel/io.py +++ b/langgraph/pregel/io.py @@ -1,4 +1,4 @@ -from typing import Any, Iterator, Mapping, Optional, Sequence, Union +from typing import Any, Iterator, Mapping, Optional, Sequence, TypeVar, Union from langchain_core.runnables.utils import AddableDict @@ -120,3 +120,11 @@ def map_output_updates( } ): yield updated + + +T = TypeVar("T") + + +def single(iter: Iterator[T]) -> Optional[T]: + for item in iter: + return item diff --git a/tests/test_io.py b/tests/test_io.py new file mode 100644 index 000000000..cbc928a19 --- /dev/null +++ b/tests/test_io.py @@ -0,0 +1,18 @@ +from typing import Iterator + +from langgraph.pregel.io import single + + +def test_single() -> None: + closed = False + + def myiter() -> Iterator[int]: + try: + yield 1 + yield 2 + finally: + nonlocal closed + closed = True + + assert single(myiter()) == 1 + assert closed