WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-31 12:19:58 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #451 from langchain-ai/nc/14may/small-fixes-ii

Browse Source 
More small fixes following JS impl
This commit is contained in:
Nuno Campos
2024-05-14 16:29:09 -07:00
committed by GitHub
parent a694aaa2bb a909d90b08
commit 9a43234ae5
3 changed files with 41 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
langgraph/pregel/__init__.py
+14 -25
View File
@@ -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"],
}
}
langgraph/pregel/io.py
+9 -1
View File
@@ -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
tests/test_io.py
+18
View File
@@ -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