codekingpro/portable-devtools

1from collections.abc import Sequence2from typing import Generic3 4from typing_extensions import Self5 6from langgraph._internal._typing import MISSING7from langgraph.channels.base import BaseChannel, Value8from langgraph.errors import EmptyChannelError, InvalidUpdateError9 10__all__ = ("NamedBarrierValue", "NamedBarrierValueAfterFinish")11 12 13class NamedBarrierValue(Generic[Value], BaseChannel[Value, Value, set[Value]]):14    """A channel that waits until all named values are received before making the value available."""15 16    __slots__ = ("names", "seen")17 18    names: set[Value]19    seen: set[Value]20 21    def __init__(self, typ: type[Value], names: set[Value]) -> None:22        super().__init__(typ)23        self.names = names24        self.seen: set[str] = set()25 26    def __eq__(self, value: object) -> bool:27        return isinstance(value, NamedBarrierValue) and value.names == self.names28 29    @property30    def ValueType(self) -> type[Value]:31        """The type of the value stored in the channel."""32        return self.typ33 34    @property35    def UpdateType(self) -> type[Value]:36        """The type of the update received by the channel."""37        return self.typ38 39    def copy(self) -> Self:40        """Return a copy of the channel."""41        empty = self.__class__(self.typ, self.names)42        empty.key = self.key43        empty.seen = self.seen.copy()44        return empty45 46    def checkpoint(self) -> set[Value]:47        return self.seen48 49    def from_checkpoint(self, checkpoint: set[Value]) -> Self:50        empty = self.__class__(self.typ, self.names)51        empty.key = self.key52        if checkpoint is not MISSING:53            empty.seen = checkpoint54        return empty55 56    def update(self, values: Sequence[Value]) -> bool:57        updated = False58        for value in values:59            if value in self.names:60                if value not in self.seen:61                    self.seen.add(value)62                    updated = True63            else:64                raise InvalidUpdateError(65                    f"At key '{self.key}': Value {value} not in {self.names}"66                )67        return updated68 69    def get(self) -> Value:70        if self.seen != self.names:71            raise EmptyChannelError()72        return None73 74    def is_available(self) -> bool:75        return self.seen == self.names76 77    def consume(self) -> bool:78        if self.seen == self.names:79            self.seen = set()80            return True81        return False82 83 84class NamedBarrierValueAfterFinish(85    Generic[Value], BaseChannel[Value, Value, set[Value]]86):87    """A channel that waits until all named values are received before making the value ready to be made available. It is only made available after finish() is called."""88 89    __slots__ = ("names", "seen", "finished")90 91    names: set[Value]92    seen: set[Value]93 94    def __init__(self, typ: type[Value], names: set[Value]) -> None:95        super().__init__(typ)96        self.names = names97        self.seen: set[str] = set()98        self.finished = False99 100    def __eq__(self, value: object) -> bool:101        return (102            isinstance(value, NamedBarrierValueAfterFinish)103            and value.names == self.names104        )105 106    @property107    def ValueType(self) -> type[Value]:108        """The type of the value stored in the channel."""109        return self.typ110 111    @property112    def UpdateType(self) -> type[Value]:113        """The type of the update received by the channel."""114        return self.typ115 116    def copy(self) -> Self:117        """Return a copy of the channel."""118        empty = self.__class__(self.typ, self.names)119        empty.key = self.key120        empty.seen = self.seen.copy()121        empty.finished = self.finished122        return empty123 124    def checkpoint(self) -> tuple[set[Value], bool]:125        return (self.seen, self.finished)126 127    def from_checkpoint(self, checkpoint: tuple[set[Value], bool]) -> Self:128        empty = self.__class__(self.typ, self.names)129        empty.key = self.key130        if checkpoint is not MISSING:131            empty.seen, empty.finished = checkpoint132        return empty133 134    def update(self, values: Sequence[Value]) -> bool:135        updated = False136        for value in values:137            if value in self.names:138                if value not in self.seen:139                    self.seen.add(value)140                    updated = True141            else:142                raise InvalidUpdateError(143                    f"At key '{self.key}': Value {value} not in {self.names}"144                )145        return updated146 147    def get(self) -> Value:148        if not self.finished or self.seen != self.names:149            raise EmptyChannelError()150        return None151 152    def is_available(self) -> bool:153        return self.finished and self.seen == self.names154 155    def consume(self) -> bool:156        if self.finished and self.seen == self.names:157            self.finished = False158            self.seen = set()159            return True160        return False161 162    def finish(self) -> bool:163        if not self.finished and self.seen == self.names:164            self.finished = True165            return True166        else:167            return False168