WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-22 07:32:25 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #252 from langchain-ai/nc/31mar/channel-from-checkpoint-copy

Browse Source 
Rename BaseChannel.empty() to BaseChannel.from_checkpoint()
This commit is contained in:
Nuno Campos
2024-04-01 16:15:41 -07:00
committed by GitHub
parent 2d351e8752 4fd90f8ad6
commit 18112b4df8
9 changed files with 119 additions and 88 deletions
Download Patch File Download Diff File Expand all files Collapse all files
langgraph/channels/any_value.py
+9 -7
View File
@@ -23,8 +23,16 @@ class AnyValue(Generic[Value], BaseChannel[Value, Value, Value]):
"""The type of the update received by the channel."""
return self.typ
def checkpoint(self) -> Value:
try:
return self.value
except AttributeError:
raise EmptyChannelError()
@contextmanager
def empty(self, checkpoint: Optional[Value] = None) -> Generator[Self, None, None]:
def from_checkpoint(
self, checkpoint: Optional[Value] = None
) -> Generator[Self, None, None]:
empty = self.__class__(self.typ)
if checkpoint is not None:
empty.value = checkpoint
@@ -51,9 +59,3 @@ class AnyValue(Generic[Value], BaseChannel[Value, Value, Value]):
return self.value
except AttributeError:
raise EmptyChannelError()
def checkpoint(self) -> Value:
try:
return self.value
except AttributeError:
raise EmptyChannelError()
langgraph/channels/base.py
+23 -14
View File
@@ -45,19 +45,33 @@ class BaseChannel(Generic[Value, Update, C], ABC):
def UpdateType(self) -> Any:
"""The type of the update received by the channel."""
# ser/de methods
@abstractmethod
def checkpoint(self) -> Optional[C]:
"""Return a serializable representation of the channel's current state.
Raises EmptyChannelError if the channel is empty (never updated yet),
or doesn't support checkpoints."""
@contextmanager
@abstractmethod
def empty(self, checkpoint: Optional[C] = None) -> Generator[Self, None, None]:
"""Return a new identical channel, optionally initialized from a checkpoint."""
def from_checkpoint(
self, checkpoint: Optional[C] = None
) -> Generator[Self, None, None]:
"""Return a new identical channel, optionally initialized from a checkpoint.
If the checkpoint contains complex data structures, they should be copied."""
@asynccontextmanager
async def aempty(
async def afrom_checkpoint(
self, checkpoint: Optional[C] = None
) -> AsyncGenerator[Self, None]:
"""Return a new identical channel, optionally initialized from a checkpoint."""
with self.empty(checkpoint) as value:
"""Return a new identical channel, optionally initialized from a checkpoint.
If the checkpoint contains complex data structures, they should be copied."""
with self.from_checkpoint(checkpoint) as value:
yield value
# state methods
@abstractmethod
def update(self, values: Sequence[Update]) -> None:
"""Update the channel's value with the given sequence of updates.
@@ -71,13 +85,6 @@ class BaseChannel(Generic[Value, Update, C], ABC):
Raises EmptyChannelError if the channel is empty (never updated yet)."""
@abstractmethod
def checkpoint(self) -> Optional[C]:
"""Return a string representation of the channel's current state.
Raises EmptyChannelError if the channel is empty (never updated yet),
or doesn't supportcheckpoints."""
@contextmanager
def ChannelsManager(
@@ -87,7 +94,8 @@ def ChannelsManager(
"""Manage channels for the lifetime of a Pregel invocation (multiple steps)."""
# TODO use https://docs.python.org/3/library/contextlib.html#contextlib.ExitStack
empty = {
k: v.empty(checkpoint["channel_values"].get(k)) for k, v in channels.items()
k: v.from_checkpoint(checkpoint["channel_values"].get(k))
for k, v in channels.items()
}
try:
yield {k: v.__enter__() for k, v in empty.items()}
@@ -103,7 +111,8 @@ async def AsyncChannelsManager(
) -> AsyncGenerator[Mapping[str, BaseChannel], None]:
"""Manage channels for the lifetime of a Pregel invocation (multiple steps)."""
empty = {
k: v.aempty(checkpoint["channel_values"].get(k)) for k, v in channels.items()
k: v.afrom_checkpoint(checkpoint["channel_values"].get(k))
for k, v in channels.items()
}
try:
yield {k: await v.__aenter__() for k, v in empty.items()}
langgraph/channels/binop.py
+9 -7
View File
@@ -34,8 +34,16 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]):
"""The type of the update received by the channel."""
return self.typ
def checkpoint(self) -> Value:
try:
return self.value
except AttributeError:
raise EmptyChannelError()
@contextmanager
def empty(self, checkpoint: Optional[Value] = None) -> Generator[Self, None, None]:
def from_checkpoint(
self, checkpoint: Optional[Value] = None
) -> Generator[Self, None, None]:
empty = self.__class__(self.typ, self.operator)
if checkpoint is not None:
empty.value = checkpoint
@@ -62,9 +70,3 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]):
return self.value
except AttributeError:
raise EmptyChannelError()
def checkpoint(self) -> Value:
try:
return self.value
except AttributeError:
raise EmptyChannelError()
langgraph/channels/context.py
+6 -6
View File
@@ -65,8 +65,11 @@ class Context(Generic[Value], BaseChannel[Value, None, None]):
"""The type of the update received by the channel."""
raise InvalidUpdateError()
def checkpoint(self) -> None:
raise EmptyChannelError()
@contextmanager
def empty(self, checkpoint: None = None) -> Generator[Self, None, None]:
def from_checkpoint(self, checkpoint: None = None) -> Generator[Self, None, None]:
if self.ctx is None:
raise ValueError("Cannot enter sync context manager.")
@@ -80,7 +83,7 @@ class Context(Generic[Value], BaseChannel[Value, None, None]):
ctx.__exit__(None, None, None)
@asynccontextmanager
async def aempty(
async def afrom_checkpoint(
self, checkpoint: Optional[str] = None
) -> AsyncGenerator[Self, None]:
if self.actx is not None:
@@ -93,7 +96,7 @@ class Context(Generic[Value], BaseChannel[Value, None, None]):
finally:
await actx.__aexit__(None, None, None)
else:
with self.empty() as empty:
with self.from_checkpoint() as empty:
yield empty
def update(self, values: Sequence[None]) -> None:
@@ -105,6 +108,3 @@ class Context(Generic[Value], BaseChannel[Value, None, None]):
return self.value
except AttributeError:
raise EmptyChannelError()
def checkpoint(self) -> None:
raise EmptyChannelError()
langgraph/channels/ephemeral_value.py
+9 -7
View File
@@ -28,8 +28,16 @@ class EphemeralValue(Generic[Value], BaseChannel[Value, Value, Value]):
"""The type of the update received by the channel."""
return self.typ
def checkpoint(self) -> Value:
try:
return self.value
except AttributeError:
raise EmptyChannelError()
@contextmanager
def empty(self, checkpoint: Optional[Value] = None) -> Generator[Self, None, None]:
def from_checkpoint(
self, checkpoint: Optional[Value] = None
) -> Generator[Self, None, None]:
empty = self.__class__(self.typ, self.guard)
if checkpoint is not None:
empty.value = checkpoint
@@ -59,9 +67,3 @@ class EphemeralValue(Generic[Value], BaseChannel[Value, Value, Value]):
return self.value
except AttributeError:
raise EmptyChannelError()
def checkpoint(self) -> Value:
try:
return self.value
except AttributeError:
raise EmptyChannelError()
langgraph/channels/last_value.py
+9 -7
View File
@@ -27,8 +27,16 @@ class LastValue(Generic[Value], BaseChannel[Value, Value, Value]):
"""The type of the update received by the channel."""
return self.typ
def checkpoint(self) -> Value:
try:
return self.value
except AttributeError:
raise EmptyChannelError()
@contextmanager
def empty(self, checkpoint: Optional[Value] = None) -> Generator[Self, None, None]:
def from_checkpoint(
self, checkpoint: Optional[Value] = None
) -> Generator[Self, None, None]:
empty = self.__class__(self.typ)
if checkpoint is not None:
empty.value = checkpoint
@@ -53,9 +61,3 @@ class LastValue(Generic[Value], BaseChannel[Value, Value, Value]):
return self.value
except AttributeError:
raise EmptyChannelError()
def checkpoint(self) -> Value:
try:
return self.value
except AttributeError:
raise EmptyChannelError()
langgraph/channels/named_barrier_value.py
+9 -7
View File
@@ -11,10 +11,10 @@ from langgraph.channels.base import (
)
class NamedBarrierValue(Generic[Value], BaseChannel[Value, Value, Value]):
class NamedBarrierValue(Generic[Value], BaseChannel[Value, Value, set[Value]]):
"""A channel that waits until all named values are received before making the value available."""
def __init__(self, typ: Type[Value], names: set[str]) -> None:
def __init__(self, typ: Type[Value], names: set[Value]) -> None:
self.typ = typ
self.names = names
self.seen = set()
@@ -29,11 +29,16 @@ class NamedBarrierValue(Generic[Value], BaseChannel[Value, Value, Value]):
"""The type of the update received by the channel."""
return self.typ
def checkpoint(self) -> set[Value]:
return self.seen
@contextmanager
def empty(self, checkpoint: Optional[Value] = None) -> Generator[Self, None, None]:
def from_checkpoint(
self, checkpoint: Optional[set[Value]] = None
) -> Generator[Self, None, None]:
empty = self.__class__(self.typ, self.names)
if checkpoint is not None:
empty.seen = checkpoint
empty.seen = checkpoint.copy()
try:
yield empty
@@ -53,6 +58,3 @@ class NamedBarrierValue(Generic[Value], BaseChannel[Value, Value, Value]):
if self.seen != self.names:
raise EmptyChannelError()
return None
def checkpoint(self) -> Value:
return self.seen
langgraph/channels/topic.py
+6 -6
View File
@@ -49,14 +49,17 @@ class Topic(
"""The type of the update received by the channel."""
return Union[self.typ, list[self.typ]] # type: ignore[name-defined]
def checkpoint(self) -> tuple[set[Value], list[Value]]:
return (self.seen, self.values)
@contextmanager
def empty(
def from_checkpoint(
self, checkpoint: Optional[tuple[set[Value], list[Value]]] = None
) -> Generator[Self, None, None]:
empty = self.__class__(self.typ, self.unique, self.accumulate)
if checkpoint is not None:
empty.seen = checkpoint[0]
empty.values = checkpoint[1]
empty.seen = checkpoint[0].copy()
empty.values = checkpoint[1].copy()
try:
yield empty
finally:
@@ -76,6 +79,3 @@ class Topic(
def get(self) -> Sequence[Value]:
return list(self.values)
def checkpoint(self) -> tuple[set[Value], list[Value]]:
return (self.seen, self.values)
tests/test_channels.py
+39 -27
View File
@@ -14,7 +14,7 @@ from langgraph.channels.topic import Topic
def test_last_value() -> None:
with LastValue(int).empty() as channel:
with LastValue(int).from_checkpoint() as channel:
assert channel.ValueType is int
assert channel.UpdateType is int
@@ -28,12 +28,12 @@ def test_last_value() -> None:
channel.update([4])
assert channel.get() == 4
checkpoint = channel.checkpoint()
with LastValue(int).empty(checkpoint) as channel:
with LastValue(int).from_checkpoint(checkpoint) as channel:
assert channel.get() == 4
async def test_last_value_async() -> None:
async with LastValue(int).aempty() as channel:
async with LastValue(int).afrom_checkpoint() as channel:
assert channel.ValueType is int
assert channel.UpdateType is int
@@ -47,12 +47,12 @@ async def test_last_value_async() -> None:
channel.update([4])
assert channel.get() == 4
checkpoint = channel.checkpoint()
async with LastValue(int).aempty(checkpoint) as channel:
async with LastValue(int).afrom_checkpoint(checkpoint) as channel:
assert channel.get() == 4
def test_topic() -> None:
with Topic(str).empty() as channel:
with Topic(str).from_checkpoint() as channel:
assert channel.ValueType is Sequence[str]
assert channel.UpdateType is Union[str, list[str]]
@@ -65,12 +65,16 @@ def test_topic() -> None:
channel.update(["e"])
assert channel.get() == ["e"]
checkpoint = channel.checkpoint()
with Topic(str).empty(checkpoint) as channel:
with Topic(str).from_checkpoint(checkpoint) as channel:
assert channel.get() == ["e"]
with Topic(str).from_checkpoint(checkpoint) as channel_copy:
channel_copy.update(["f"])
assert channel_copy.get() == ["f"]
assert channel.get() == ["e"]
async def test_topic_async() -> None:
async with Topic(str).aempty() as channel:
async with Topic(str).afrom_checkpoint() as channel:
assert channel.ValueType is Sequence[str]
assert channel.UpdateType is Union[str, list[str]]
@@ -83,12 +87,12 @@ async def test_topic_async() -> None:
channel.update(["e"])
assert channel.get() == ["e"]
checkpoint = channel.checkpoint()
async with Topic(str).aempty(checkpoint) as channel:
async with Topic(str).afrom_checkpoint(checkpoint) as channel:
assert channel.get() == ["e"]
def test_topic_unique() -> None:
with Topic(str, unique=True).empty() as channel:
with Topic(str, unique=True).from_checkpoint() as channel:
assert channel.ValueType is Sequence[str]
assert channel.UpdateType is Union[str, list[str]]
@@ -101,14 +105,14 @@ def test_topic_unique() -> None:
channel.update(["e"])
assert channel.get() == ["e"]
checkpoint = channel.checkpoint()
with Topic(str, unique=True).empty(checkpoint) as channel:
with Topic(str, unique=True).from_checkpoint(checkpoint) as channel:
assert channel.get() == ["e"]
channel.update(["d", "f"])
assert channel.get() == ["f"], "de-dupes from checkpoint"
async def test_topic_unique_async() -> None:
async with Topic(str, unique=True).aempty() as channel:
async with Topic(str, unique=True).afrom_checkpoint() as channel:
assert channel.ValueType is Sequence[str]
assert channel.UpdateType is Union[str, list[str]]
@@ -121,14 +125,14 @@ async def test_topic_unique_async() -> None:
channel.update(["e"])
assert channel.get() == ["e"]
checkpoint = channel.checkpoint()
async with Topic(str, unique=True).aempty(checkpoint) as channel:
async with Topic(str, unique=True).afrom_checkpoint(checkpoint) as channel:
assert channel.get() == ["e"]
channel.update(["d", "f"])
assert channel.get() == ["f"], "de-dupes from checkpoint"
def test_topic_accumulate() -> None:
with Topic(str, accumulate=True).empty() as channel:
with Topic(str, accumulate=True).from_checkpoint() as channel:
assert channel.ValueType is Sequence[str]
assert channel.UpdateType is Union[str, list[str]]
@@ -139,14 +143,14 @@ def test_topic_accumulate() -> None:
channel.update([])
assert channel.get() == ["a", "b", "b", "c", "d", "d"]
checkpoint = channel.checkpoint()
with Topic(str, accumulate=True).empty(checkpoint) as channel:
with Topic(str, accumulate=True).from_checkpoint(checkpoint) as channel:
assert channel.get() == ["a", "b", "b", "c", "d", "d"]
channel.update(["e"])
assert channel.get() == ["a", "b", "b", "c", "d", "d", "e"]
async def test_topic_accumulate_async() -> None:
async with Topic(str, accumulate=True).aempty() as channel:
async with Topic(str, accumulate=True).afrom_checkpoint() as channel:
assert channel.ValueType is Sequence[str]
assert channel.UpdateType is Union[str, list[str]]
@@ -157,14 +161,14 @@ async def test_topic_accumulate_async() -> None:
channel.update([])
assert channel.get() == ["a", "b", "b", "c", "d", "d"]
checkpoint = channel.checkpoint()
async with Topic(str, accumulate=True).aempty(checkpoint) as channel:
async with Topic(str, accumulate=True).afrom_checkpoint(checkpoint) as channel:
assert channel.get() == ["a", "b", "b", "c", "d", "d"]
channel.update(["e"])
assert channel.get() == ["a", "b", "b", "c", "d", "d", "e"]
def test_topic_unique_accumulate() -> None:
with Topic(str, unique=True, accumulate=True).empty() as channel:
with Topic(str, unique=True, accumulate=True).from_checkpoint() as channel:
assert channel.ValueType is Sequence[str]
assert channel.UpdateType is Union[str, list[str]]
@@ -175,14 +179,16 @@ def test_topic_unique_accumulate() -> None:
channel.update([])
assert channel.get() == ["a", "b", "c", "d"]
checkpoint = channel.checkpoint()
with Topic(str, unique=True, accumulate=True).empty(checkpoint) as channel:
with Topic(str, unique=True, accumulate=True).from_checkpoint(
checkpoint
) as channel:
assert channel.get() == ["a", "b", "c", "d"]
channel.update(["d", "e"])
assert channel.get() == ["a", "b", "c", "d", "e"]
async def test_topic_unique_accumulate_async() -> None:
async with Topic(str, unique=True, accumulate=True).aempty() as channel:
async with Topic(str, unique=True, accumulate=True).afrom_checkpoint() as channel:
assert channel.ValueType is Sequence[str]
assert channel.UpdateType is Union[str, list[str]]
@@ -193,14 +199,16 @@ async def test_topic_unique_accumulate_async() -> None:
channel.update([])
assert channel.get() == ["a", "b", "c", "d"]
checkpoint = channel.checkpoint()
async with Topic(str, unique=True, accumulate=True).aempty(checkpoint) as channel:
async with Topic(str, unique=True, accumulate=True).afrom_checkpoint(
checkpoint
) as channel:
assert channel.get() == ["a", "b", "c", "d"]
channel.update(["d", "e"])
assert channel.get() == ["a", "b", "c", "d", "e"]
def test_binop() -> None:
with BinaryOperatorAggregate(int, operator.add).empty() as channel:
with BinaryOperatorAggregate(int, operator.add).from_checkpoint() as channel:
assert channel.ValueType is int
assert channel.UpdateType is int
@@ -211,12 +219,14 @@ def test_binop() -> None:
channel.update([4])
assert channel.get() == 10
checkpoint = channel.checkpoint()
with BinaryOperatorAggregate(int, operator.add).empty(checkpoint) as channel:
with BinaryOperatorAggregate(int, operator.add).from_checkpoint(
checkpoint
) as channel:
assert channel.get() == 10
async def test_binop_async() -> None:
async with BinaryOperatorAggregate(int, operator.add).aempty() as channel:
async with BinaryOperatorAggregate(int, operator.add).afrom_checkpoint() as channel:
assert channel.ValueType is int
assert channel.UpdateType is int
@@ -227,7 +237,9 @@ async def test_binop_async() -> None:
channel.update([4])
assert channel.get() == 10
checkpoint = channel.checkpoint()
async with BinaryOperatorAggregate(int, operator.add).aempty(checkpoint) as channel:
async with BinaryOperatorAggregate(int, operator.add).afrom_checkpoint(
checkpoint
) as channel:
assert channel.get() == 10
@@ -243,7 +255,7 @@ def test_ctx_manager(mocker: MockerFixture) -> None:
finally:
cleanup()
with Context(an_int, None, int).empty() as channel:
with Context(an_int, None, int).from_checkpoint() as channel:
assert setup.call_count == 1
assert cleanup.call_count == 0
@@ -261,7 +273,7 @@ def test_ctx_manager(mocker: MockerFixture) -> None:
def test_ctx_manager_ctx(mocker: MockerFixture) -> None:
with Context(httpx.Client).empty() as channel:
with Context(httpx.Client).from_checkpoint() as channel:
assert channel.ValueType is httpx.Client
with pytest.raises(InvalidUpdateError):
assert channel.UpdateType is None
@@ -294,7 +306,7 @@ async def test_ctx_manager_async(mocker: MockerFixture) -> None:
finally:
cleanup()
async with Context(an_int_sync, an_int, int).aempty() as channel:
async with Context(an_int_sync, an_int, int).afrom_checkpoint() as channel:
assert setup.call_count == 1
assert cleanup.call_count == 0