Files
langgraph/libs/langgraph/tests/test_channels.py
T
Nuno CamposandGitHub 66fc7c96aa More performance improvements in checkpointing and channels (#1685)
* Performance improvements in checkpointer libs

- Use sha1 instead of md5 for hashing (faster in python 3.x)
- Use orjson instead of json for json dumping (sadly can't use for json loading)

* Update tests

* Update

* Use random number instead of hash for get_version_number

* Avoid saving writes for the last task to complete in each step

- only when possible, exceptions for ERROR, INTERRUPT, SEND

* Make Channel.from_checkpoint a regular function

- context manager no longer needed since Context became a managed value

* Use __slots__ for Channels

* Fix for kafka
2024-09-12 21:08:30 -07:00

89 lines
2.9 KiB
Python

import operator
from typing import Sequence, Union
import pytest
from langgraph.channels.binop import BinaryOperatorAggregate
from langgraph.channels.last_value import LastValue
from langgraph.channels.topic import Topic
from langgraph.errors import EmptyChannelError, InvalidUpdateError
pytestmark = pytest.mark.anyio
def test_last_value() -> None:
channel = LastValue(int).from_checkpoint(None)
assert channel.ValueType is int
assert channel.UpdateType is int
with pytest.raises(EmptyChannelError):
channel.get()
with pytest.raises(InvalidUpdateError):
channel.update([5, 6])
channel.update([3])
assert channel.get() == 3
channel.update([4])
assert channel.get() == 4
checkpoint = channel.checkpoint()
channel = LastValue(int).from_checkpoint(checkpoint)
assert channel.get() == 4
def test_topic() -> None:
channel = Topic(str).from_checkpoint(None)
assert channel.ValueType is Sequence[str]
assert channel.UpdateType is Union[str, list[str]]
assert channel.update(["a", "b"])
assert channel.get() == ["a", "b"]
assert channel.update([["c", "d"], "d"])
assert channel.get() == ["c", "d", "d"]
assert channel.update([])
with pytest.raises(EmptyChannelError):
channel.get()
assert not channel.update([]), "channel already empty"
assert channel.update(["e"])
assert channel.get() == ["e"]
checkpoint = channel.checkpoint()
channel = Topic(str).from_checkpoint(checkpoint)
assert channel.get() == ["e"]
channel_copy = Topic(str).from_checkpoint(checkpoint)
channel_copy.update(["f"])
assert channel_copy.get() == ["f"]
assert channel.get() == ["e"]
def test_topic_accumulate() -> None:
channel = Topic(str, accumulate=True).from_checkpoint(None)
assert channel.ValueType is Sequence[str]
assert channel.UpdateType is Union[str, list[str]]
assert channel.update(["a", "b"])
assert channel.get() == ["a", "b"]
assert channel.update(["b", ["c", "d"], "d"])
assert channel.get() == ["a", "b", "b", "c", "d", "d"]
assert not channel.update([])
assert channel.get() == ["a", "b", "b", "c", "d", "d"]
checkpoint = channel.checkpoint()
channel = Topic(str, accumulate=True).from_checkpoint(checkpoint)
assert channel.get() == ["a", "b", "b", "c", "d", "d"]
assert channel.update(["e"])
assert channel.get() == ["a", "b", "b", "c", "d", "d", "e"]
def test_binop() -> None:
channel = BinaryOperatorAggregate(int, operator.add).from_checkpoint(None)
assert channel.ValueType is int
assert channel.UpdateType is int
assert channel.get() == 0
channel.update([1, 2, 3])
assert channel.get() == 6
channel.update([4])
assert channel.get() == 10
checkpoint = channel.checkpoint()
channel = BinaryOperatorAggregate(int, operator.add).from_checkpoint(checkpoint)
assert channel.get() == 10