From b479b88c7e2912bebbbc78484f48af1e03a48218 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 29 Aug 2024 13:20:28 -0700 Subject: [PATCH] Add name of subgraph to streaming output --- libs/langgraph/langgraph/pregel/__init__.py | 91 +++++++++++---------- libs/langgraph/langgraph/pregel/loop.py | 21 +++-- libs/langgraph/tests/test_pregel.py | 12 +-- libs/langgraph/tests/test_pregel_async.py | 12 +-- 4 files changed, 73 insertions(+), 63 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 3d676a7bc..515a33aa4 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -1113,6 +1113,25 @@ class Pregel( {'type': 'task_result', 'timestamp': '2024-06-23T...+00:00', 'step': 2, 'payload': {'id': '...', 'name': 'b', 'result': [('alist', ['there'])]}} ``` """ + + def output() -> Iterator: + while loop.stream: + ns, mode, payload = loop.stream.popleft() + ns = ( + NS_SEP.join(p.split(NS_END)[0] for p in ns.split(NS_SEP)) + if ns + else "" + ) + if mode in stream_modes: + if subgraphs and isinstance(stream_mode, list): + yield (ns, mode, payload) + elif isinstance(stream_mode, list): + yield (mode, payload) + elif subgraphs: + yield (ns, payload) + else: + yield payload + config = ensure_config(merge_configs(self.config, config)) callback_manager = get_callback_manager_for_config(config) run_manager = callback_manager.on_chain_start( @@ -1176,13 +1195,7 @@ class Pregel( self.stream_channels_list, ) # emit output - while loop.stream: - mode, payload = loop.stream.popleft() - if mode in stream_modes: - if isinstance(stream_mode, list): - yield (mode, payload) - else: - yield payload + yield from output() # debug flag if debug: print_step_tasks(loop.step, loop.tasks) @@ -1237,13 +1250,8 @@ class Pregel( # remove references to loop vars del fut, task # emit output - while loop.stream: - mode, payload = loop.stream.popleft() - if mode in stream_modes: - if isinstance(stream_mode, list): - yield (mode, payload) - else: - yield payload + yield from output() + # maybe stop other tasks if _should_stop_others(done): break @@ -1259,13 +1267,7 @@ class Pregel( self.stream_channels_list, ) # emit output - while loop.stream: - mode, payload = loop.stream.popleft() - if mode in stream_modes: - if isinstance(stream_mode, list): - yield (mode, payload) - else: - yield payload + yield from output() # handle exit if loop.status == "out_of_steps": raise GraphRecursionError( @@ -1358,6 +1360,25 @@ class Pregel( {'type': 'task_result', 'timestamp': '2024-06-23T...+00:00', 'step': 2, 'payload': {'id': '...', 'name': 'b', 'result': [('alist', ['there'])]}} ``` """ + + def output() -> Iterator: + while loop.stream: + ns, mode, payload = loop.stream.popleft() + ns = ( + NS_SEP.join(p.split(NS_END)[0] for p in ns.split(NS_SEP)) + if ns + else "" + ) + if mode in stream_modes: + if subgraphs and isinstance(stream_mode, list): + yield (ns, mode, payload) + elif isinstance(stream_mode, list): + yield (mode, payload) + elif subgraphs: + yield (ns, payload) + else: + yield payload + config = ensure_config(merge_configs(self.config, config)) callback_manager = get_async_callback_manager_for_config(config) run_manager = await callback_manager.on_chain_start( @@ -1430,13 +1451,8 @@ class Pregel( self.stream_channels_list, ) # emit output - while loop.stream: - mode, payload = loop.stream.popleft() - if mode in stream_modes: - if isinstance(stream_mode, list): - yield (mode, payload) - else: - yield payload + for o in output(): + yield o # debug flag if debug: print_step_tasks(loop.step, loop.tasks) @@ -1493,13 +1509,9 @@ class Pregel( # remove references to loop vars del fut, task # emit output - while loop.stream: - mode, payload = loop.stream.popleft() - if mode in stream_modes: - if isinstance(stream_mode, list): - yield (mode, payload) - else: - yield payload + for o in output(): + yield o + # maybe stop other tasks if _should_stop_others(done): break @@ -1515,13 +1527,8 @@ class Pregel( self.stream_channels_list, ) # emit output - while loop.stream: - mode, payload = loop.stream.popleft() - if mode in stream_modes: - if isinstance(stream_mode, list): - yield (mode, payload) - else: - yield payload + for o in output(): + yield o # handle exit if loop.status == "out_of_steps": raise GraphRecursionError( diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 7ff95cc38..87551b83f 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -91,8 +91,8 @@ EMPTY_SEQ = () class StreamProtocol(Protocol): - def extend(self, values: Iterable[Tuple[str, Any]]) -> None: ... - def popleft(self) -> Tuple[str, Any]: ... + def extend(self, values: Iterable[Tuple[str, str, Any]]) -> None: ... + def popleft(self) -> Tuple[str, str, Any]: ... def __bool__(self) -> bool: ... @@ -100,11 +100,11 @@ class DuplexStream(StreamProtocol): def __init__(self, *streams: StreamProtocol) -> None: self.streams = streams - def extend(self, values: Iterable[Tuple[str, Any]]) -> None: + def extend(self, values: Iterable[Tuple[str, str, Any]]) -> None: for stream, vv in zip(self.streams, tee(values, len(self.streams))): stream.extend(vv) - def popleft(self) -> Tuple[str, Any]: + def popleft(self) -> Tuple[str, str, Any]: return self.streams[0].popleft() def __bool__(self) -> bool: @@ -207,11 +207,11 @@ class PregelLoop: ) if task := next((t for t in self.tasks if t.id == task_id), None): self.stream.extend( - ("updates", v) + (self.config["configurable"].get("checkpoint_ns", ""), "updates", v) for v in map_output_updates(self.output_keys, [(task, writes)]) ) self.stream.extend( - ("debug", v) + (self.config["configurable"].get("checkpoint_ns", ""), "debug", v) for v in map_debug_task_results( self.step, [(task, writes)], self.stream_keys ) @@ -247,7 +247,7 @@ class PregelLoop: self._update_mv(key, values) # produce values output self.stream.extend( - ("values", v) + (self.config["configurable"].get("checkpoint_ns", ""), "values", v) for v in map_output_values(self.output_keys, writes, self.channels) ) # clear pending writes @@ -295,7 +295,7 @@ class PregelLoop: # produce debug output if self._checkpointer_put_after_previous is not None: self.stream.extend( - ("debug", v) + (self.config["configurable"].get("checkpoint_ns", ""), "debug", v) for v in map_debug_checkpoint( self.step - 1, # printing checkpoint for previous step self.checkpoint_config, @@ -339,7 +339,10 @@ class PregelLoop: return False # produce debug output - self.stream.extend(("debug", v) for v in map_debug_tasks(self.step, self.tasks)) + self.stream.extend( + (self.config["configurable"].get("checkpoint_ns", ""), "debug", v) + for v in map_debug_tasks(self.step, self.tasks) + ) return True diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 645761881..a6a338a75 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -10898,8 +10898,8 @@ def test_doubly_nested_graph_state( # test invoke w/ nested interrupt config = {"configurable": {"thread_id": "1"}} assert [c for c in app.stream({"my_key": "my value"}, config, subgraphs=True)] == [ - {"parent_1": {"my_key": "hi my value"}}, - {"grandchild_1": {"my_key": "hi my value here"}}, + ("", {"parent_1": {"my_key": "hi my value"}}), + ("child|child_1", {"grandchild_1": {"my_key": "hi my value here"}}), ] # get state without subgraphs outer_state = app.get_state(config) @@ -11121,10 +11121,10 @@ def test_doubly_nested_graph_state( ) # resume assert [c for c in app.stream(None, config, subgraphs=True)] == [ - {"grandchild_2": {"my_key": "hi my value here and there"}}, - {"child_1": {"my_key": "hi my value here and there"}}, - {"child": {"my_key": "hi my value here and there"}}, - {"parent_2": {"my_key": "hi my value here and there and back again"}}, + ("child|child_1", {"grandchild_2": {"my_key": "hi my value here and there"}}), + ("child", {"child_1": {"my_key": "hi my value here and there"}}), + ("", {"child": {"my_key": "hi my value here and there"}}), + ("", {"parent_2": {"my_key": "hi my value here and there and back again"}}), ] # get state with and without subgraphs assert ( diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index f59003ac4..830d69238 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -9341,8 +9341,8 @@ async def test_doubly_nested_graph_state( assert [ c async for c in app.astream({"my_key": "my value"}, config, subgraphs=True) ] == [ - {"parent_1": {"my_key": "hi my value"}}, - {"grandchild_1": {"my_key": "hi my value here"}}, + ("", {"parent_1": {"my_key": "hi my value"}}), + ("child|child_1", {"grandchild_1": {"my_key": "hi my value here"}}), ] # get state without subgraphs outer_state = await app.aget_state(config) @@ -9564,10 +9564,10 @@ async def test_doubly_nested_graph_state( ) # resume assert [c async for c in app.astream(None, config, subgraphs=True)] == [ - {"grandchild_2": {"my_key": "hi my value here and there"}}, - {"child_1": {"my_key": "hi my value here and there"}}, - {"child": {"my_key": "hi my value here and there"}}, - {"parent_2": {"my_key": "hi my value here and there and back again"}}, + ("child|child_1", {"grandchild_2": {"my_key": "hi my value here and there"}}), + ("child", {"child_1": {"my_key": "hi my value here and there"}}), + ("", {"child": {"my_key": "hi my value here and there"}}), + ("", {"parent_2": {"my_key": "hi my value here and there and back again"}}), ] # get state with and without subgraphs assert (