Use __slots__ for Channels

langchain-ai · Sep 13, 2024 · 2cf6661 · 2cf6661
1 parent 911f433
commit 2cf6661
Show file tree

Hide file tree

Showing 9 changed files with 26 additions and 13 deletions.
diff --git a/libs/langgraph/langgraph/channels/any_value.py b/libs/langgraph/langgraph/channels/any_value.py
@@ -10,8 +10,7 @@ class AnyValue(Generic[Value], BaseChannel[Value, Value, Value]):
     """Stores the last value received, assumes that if multiple values are
     received, they are all equal."""
 
-    def __init__(self, typ: Type[Value]) -> None:
-        self.typ = typ
+    __slots__ = ("typ", "value")
 
     def __eq__(self, value: object) -> bool:
         return isinstance(value, AnyValue)

diff --git a/libs/langgraph/langgraph/channels/base.py b/libs/langgraph/langgraph/channels/base.py
@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Any, Generic, Optional, Sequence, TypeVar
+from typing import Any, Generic, Optional, Sequence, Type, TypeVar
 
 from typing_extensions import Self
 
@@ -11,7 +11,11 @@
 
 
 class BaseChannel(Generic[Value, Update, C], ABC):
-    key: str = ""
+    __slots__ = ("key", "typ")
+
+    def __init__(self, typ: Type[Any], key: str = "") -> None:
+        self.typ = typ
+        self.key = key
 
     @property
     @abstractmethod

diff --git a/libs/langgraph/langgraph/channels/binop.py b/libs/langgraph/langgraph/channels/binop.py
@@ -34,10 +34,11 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]):
     ```
     """
 
+    __slots__ = ("value", "operator")
+
     def __init__(self, typ: Type[Value], operator: Callable[[Value, Value], Value]):
+        super().__init__(typ)
         self.operator = operator
-        # keep the type exposed by ValueType/UpdateType as-is
-        self.typ = typ
         # special forms from typing or collections.abc are not instantiable
         # so we need to replace them with their concrete counterparts
         typ = _strip_extras(typ)

diff --git a/libs/langgraph/langgraph/channels/dynamic_barrier_value.py b/libs/langgraph/langgraph/channels/dynamic_barrier_value.py
@@ -22,11 +22,13 @@ class DynamicBarrierValue(
           back to the "priming" state.
     """
 
+    __slots__ = ("names", "seen")
+
     names: Optional[set[Value]]
     seen: set[Value]
 
     def __init__(self, typ: Type[Value]) -> None:
-        self.typ = typ
+        super().__init__(typ)
         self.names = None
         self.seen = set()
 

diff --git a/libs/langgraph/langgraph/channels/ephemeral_value.py b/libs/langgraph/langgraph/channels/ephemeral_value.py
@@ -9,8 +9,10 @@
 class EphemeralValue(Generic[Value], BaseChannel[Value, Value, Value]):
     """Stores the value received in the step immediately preceding, clears after."""
 
+    __slots__ = ("value", "guard")
+
     def __init__(self, typ: Type[Value], guard: bool = True) -> None:
-        self.typ = typ
+        super().__init__(typ)
         self.guard = guard
 
     def __eq__(self, value: object) -> bool:

diff --git a/libs/langgraph/langgraph/channels/last_value.py b/libs/langgraph/langgraph/channels/last_value.py
@@ -9,8 +9,7 @@
 class LastValue(Generic[Value], BaseChannel[Value, Value, Value]):
     """Stores the last value received, can receive at most one value per step."""
 
-    def __init__(self, typ: Type[Value]) -> None:
-        self.typ = typ
+    __slots__ = ("value",)
 
     def __eq__(self, value: object) -> bool:
         return isinstance(value, LastValue)

diff --git a/libs/langgraph/langgraph/channels/named_barrier_value.py b/libs/langgraph/langgraph/channels/named_barrier_value.py
@@ -9,8 +9,10 @@
 class NamedBarrierValue(Generic[Value], BaseChannel[Value, Value, set[Value]]):
     """A channel that waits until all named values are received before making the value available."""
 
+    __slots__ = ("names", "seen")
+
     def __init__(self, typ: Type[Value], names: set[Value]) -> None:
-        self.typ = typ
+        super().__init__(typ)
         self.names = names
         self.seen = set()
 

diff --git a/libs/langgraph/langgraph/channels/topic.py b/libs/langgraph/langgraph/channels/topic.py
@@ -28,9 +28,11 @@ class Topic(
         accumulate: Whether to accumulate values across steps. If False, the channel will be emptied after each step.
     """
 
+    __slots__ = ("values", "accumulate")
+
     def __init__(self, typ: Type[Value], accumulate: bool = False) -> None:
+        super().__init__(typ)
         # attrs
-        self.typ = typ
         self.accumulate = accumulate
         # state
         self.values = list[Value]()

diff --git a/libs/langgraph/langgraph/channels/untracked_value.py b/libs/langgraph/langgraph/channels/untracked_value.py
@@ -9,8 +9,10 @@
 class UntrackedValue(Generic[Value], BaseChannel[Value, Value, Value]):
     """Stores the last value received, never checkpointed."""
 
+    __slots__ = ("value", "guard")
+
     def __init__(self, typ: Type[Value], guard: bool = True) -> None:
-        self.typ = typ
+        super().__init__(typ)
         self.guard = guard
 
     def __eq__(self, value: object) -> bool: