From ce1077da40bc10c95d3bd6b80f8a052ef0ba5913 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 17 Mar 2025 21:04:05 -0700 Subject: [PATCH] Speed up task triggers check - Using a sentinel value is faster than raising-catching an exception --- libs/langgraph/langgraph/channels/base.py | 9 +++++++++ .../langgraph/channels/ephemeral_value.py | 16 ++++++++++------ libs/langgraph/langgraph/channels/last_value.py | 15 +++++++++++---- .../langgraph/channels/untracked_value.py | 10 +++++++--- libs/langgraph/langgraph/pregel/algo.py | 5 ++--- libs/langgraph/langgraph/pregel/io.py | 7 ++----- 6 files changed, 41 insertions(+), 21 deletions(-) diff --git a/libs/langgraph/langgraph/channels/base.py b/libs/langgraph/langgraph/channels/base.py index 4aaeb5681..c4b49c650 100644 --- a/libs/langgraph/langgraph/channels/base.py +++ b/libs/langgraph/langgraph/channels/base.py @@ -3,6 +3,7 @@ from typing import Any, Generic, Optional, Sequence, TypeVar from typing_extensions import Self +from langgraph.constants import MISSING from langgraph.errors import EmptyChannelError, InvalidUpdateError Value = TypeVar("Value") @@ -64,6 +65,14 @@ class BaseChannel(Generic[Value, Update, C], ABC): """ return False + def get_catch(self) -> Value: + """Return the current value of the channel, or MISSING if the channel + is empty. Subclasses can override to skip the EmptyChannelError check.""" + try: + return self.get() + except EmptyChannelError: + return MISSING + __all__ = [ "BaseChannel", diff --git a/libs/langgraph/langgraph/channels/ephemeral_value.py b/libs/langgraph/langgraph/channels/ephemeral_value.py index 537a8763c..4a64c6b32 100644 --- a/libs/langgraph/langgraph/channels/ephemeral_value.py +++ b/libs/langgraph/langgraph/channels/ephemeral_value.py @@ -3,6 +3,7 @@ from typing import Any, Generic, Optional, Sequence, Type from typing_extensions import Self from langgraph.channels.base import BaseChannel, Value +from langgraph.constants import MISSING from langgraph.errors import EmptyChannelError, InvalidUpdateError @@ -14,6 +15,7 @@ class EphemeralValue(Generic[Value], BaseChannel[Value, Value, Value]): def __init__(self, typ: Any, guard: bool = True) -> None: super().__init__(typ) self.guard = guard + self.value = MISSING def __eq__(self, value: object) -> bool: return isinstance(value, EphemeralValue) and value.guard == self.guard @@ -37,10 +39,10 @@ class EphemeralValue(Generic[Value], BaseChannel[Value, Value, Value]): def update(self, values: Sequence[Value]) -> bool: if len(values) == 0: - try: - del self.value + if self.value is not MISSING: + self.value = MISSING return True - except AttributeError: + else: return False if len(values) != 1 and self.guard: raise InvalidUpdateError( @@ -51,7 +53,9 @@ class EphemeralValue(Generic[Value], BaseChannel[Value, Value, Value]): return True def get(self) -> Value: - try: - return self.value - except AttributeError: + if self.value is MISSING: raise EmptyChannelError() + return self.value + + def get_catch(self) -> Value: + return self.value diff --git a/libs/langgraph/langgraph/channels/last_value.py b/libs/langgraph/langgraph/channels/last_value.py index 5065f4fc5..13fac983e 100644 --- a/libs/langgraph/langgraph/channels/last_value.py +++ b/libs/langgraph/langgraph/channels/last_value.py @@ -1,8 +1,9 @@ -from typing import Generic, Optional, Sequence, Type +from typing import Any, Generic, Optional, Sequence, Type from typing_extensions import Self from langgraph.channels.base import BaseChannel, Value +from langgraph.constants import MISSING from langgraph.errors import ( EmptyChannelError, ErrorCode, @@ -16,6 +17,10 @@ class LastValue(Generic[Value], BaseChannel[Value, Value, Value]): __slots__ = ("value",) + def __init__(self, typ: Any, key: str = "") -> None: + super().__init__(typ, key) + self.value = MISSING + def __eq__(self, value: object) -> bool: return isinstance(value, LastValue) @@ -50,7 +55,9 @@ class LastValue(Generic[Value], BaseChannel[Value, Value, Value]): return True def get(self) -> Value: - try: - return self.value - except AttributeError: + if self.value is MISSING: raise EmptyChannelError() + return self.value + + def get_catch(self) -> Value: + return self.value diff --git a/libs/langgraph/langgraph/channels/untracked_value.py b/libs/langgraph/langgraph/channels/untracked_value.py index 9b1020710..2560d61a3 100644 --- a/libs/langgraph/langgraph/channels/untracked_value.py +++ b/libs/langgraph/langgraph/channels/untracked_value.py @@ -3,6 +3,7 @@ from typing import Generic, Optional, Sequence, Type from typing_extensions import Self from langgraph.channels.base import BaseChannel, Value +from langgraph.constants import MISSING from langgraph.errors import EmptyChannelError, InvalidUpdateError @@ -14,6 +15,7 @@ class UntrackedValue(Generic[Value], BaseChannel[Value, Value, Value]): def __init__(self, typ: Type[Value], guard: bool = True) -> None: super().__init__(typ) self.guard = guard + self.value = MISSING def __eq__(self, value: object) -> bool: return isinstance(value, UntrackedValue) and value.guard == self.guard @@ -48,7 +50,9 @@ class UntrackedValue(Generic[Value], BaseChannel[Value, Value, Value]): return True def get(self) -> Value: - try: - return self.value - except AttributeError: + if self.value is MISSING: raise EmptyChannelError() + return self.value + + def get_catch(self) -> Value: + return self.value diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 03b2af6f4..b2bcdcedf 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -48,6 +48,7 @@ from langgraph.constants import ( EMPTY_SEQ, ERROR, INTERRUPT, + MISSING, NO_WRITES, NS_END, NS_SEP, @@ -649,9 +650,7 @@ def prepare_single_task( if triggers := sorted( chan for chan in proc.triggers - if not isinstance( - read_channel(channels, chan, return_exception=True), EmptyChannelError - ) + if channels[chan].get_catch() is not MISSING and checkpoint["channel_versions"].get(chan, null_version) # type: ignore[operator] > seen.get(chan, null_version) ): diff --git a/libs/langgraph/langgraph/pregel/io.py b/libs/langgraph/langgraph/pregel/io.py index e9963f5c8..30e976d99 100644 --- a/libs/langgraph/langgraph/pregel/io.py +++ b/libs/langgraph/langgraph/pregel/io.py @@ -38,14 +38,11 @@ def read_channel( chan: str, *, catch: bool = True, - return_exception: bool = False, ) -> Any: try: return channels[chan].get() - except EmptyChannelError as exc: - if return_exception: - return exc - elif catch: + except EmptyChannelError: + if catch: return None else: raise