mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-08 18:57:52 +02:00
Add name of subgraph to streaming output
This commit is contained in:
Nuno Campos
@@ -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(
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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 (
|
||||
|
||||
@@ -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 (
|
||||
|
||||
Reference in New Issue
Block a user