From d33e4981e028693384d5b8666e303fcdb3731319 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Thu, 25 Jul 2024 16:33:32 -0300 Subject: [PATCH] refactor: update StateVertex to function like a normal component unless specified not to (#2950) * refactor: update StateVertex class to inherit from ComponentVertex The StateVertex class in types.py has been updated to inherit from the ComponentVertex class instead of the Vertex class. This change ensures that the StateVertex class has access to the methods and attributes of the ComponentVertex class, improving code organization and maintainability. * feat: add _successors_ids attribute to Vertex class The Vertex class in base.py has been updated to include a new attribute `_successors_ids`. This attribute will store the IDs of the successors of the vertex. This change improves the functionality and flexibility of the Vertex class. * feat: add _set_successors_ids method to NotifyComponent This commit adds a new method `_set_successors_ids` to the `NotifyComponent` class in `Notify.py`. This method retrieves the successors of the vertex and sets the `successors_ids` attribute accordingly. This change improves the functionality and flexibility of the `NotifyComponent` class. * refactor: add _set_successors_ids method to ListenComponent This commit adds a new method `_set_successors_ids` to the `ListenComponent` class in `Listen.py`. This method retrieves the successors of the vertex and sets the `successors_ids` attribute accordingly. This change improves the functionality and flexibility of the `ListenComponent` class. * refactor: update StateVertex class to not be a state This commit updates the StateVertex class in types.py to no longer be considered a state. The `is_state` attribute is set to False, ensuring that the vertex is not treated as a state. This change improves code clarity and aligns with the intended functionality of the class. * refactor: set is_state attribute to True in ListenComponent This commit sets the `is_state` attribute to True in the `ListenComponent` class in `Listen.py`. This change ensures that the `ListenComponent` is treated as a state and improves the functionality of the component. * refactor: set is_state attribute to True in NotifyComponent * refactor: remove base_type argument from StateVertex constructor This commit removes the `base_type` argument from the constructor of the `StateVertex` class in `types.py`. The `base_type` argument was not being used and was unnecessary for the functionality of the class. This change improves code clarity and removes unnecessary code. --------- Co-authored-by: Lucas Oliveira <62335616+lucaseduoli@users.noreply.github.com> --- .../base/langflow/components/prototypes/Listen.py | 6 ++++++ .../base/langflow/components/prototypes/Notify.py | 6 ++++++ src/backend/base/langflow/graph/vertex/base.py | 1 + src/backend/base/langflow/graph/vertex/types.py | 12 +++++++----- 4 files changed, 20 insertions(+), 5 deletions(-) diff --git a/src/backend/base/langflow/components/prototypes/Listen.py b/src/backend/base/langflow/components/prototypes/Listen.py index 23b9c3ee1..6e5de723a 100644 --- a/src/backend/base/langflow/components/prototypes/Listen.py +++ b/src/backend/base/langflow/components/prototypes/Listen.py @@ -18,5 +18,11 @@ class ListenComponent(CustomComponent): def build(self, name: str) -> Data: state = self.get_state(name) + self._set_successors_ids() self.status = state return state + + def _set_successors_ids(self): + self.vertex.is_state = True + successors = self.vertex.graph.successor_map.get(self.vertex.id, []) + return successors + self.vertex.graph.activated_vertices diff --git a/src/backend/base/langflow/components/prototypes/Notify.py b/src/backend/base/langflow/components/prototypes/Notify.py index c07baae34..b83331e3a 100644 --- a/src/backend/base/langflow/components/prototypes/Notify.py +++ b/src/backend/base/langflow/components/prototypes/Notify.py @@ -39,4 +39,10 @@ class NotifyComponent(CustomComponent): else: self.status = "No record provided." self.status = data + self._set_successors_ids() return data + + def _set_successors_ids(self): + self.vertex.is_state = True + successors = self.vertex.graph.successor_map.get(self.vertex.id, []) + return successors + self.vertex.graph.activated_vertices diff --git a/src/backend/base/langflow/graph/vertex/base.py b/src/backend/base/langflow/graph/vertex/base.py index 6d1d27f4e..1fc6a08cc 100644 --- a/src/backend/base/langflow/graph/vertex/base.py +++ b/src/backend/base/langflow/graph/vertex/base.py @@ -71,6 +71,7 @@ class Vertex: self._built_object = UnbuiltObject() self._built_result = None self._built = False + self._successors_ids: Optional[List[str]] = None self.artifacts: Dict[str, Any] = {} self.artifacts_raw: Dict[str, Any] = {} self.artifacts_type: Dict[str, str] = {} diff --git a/src/backend/base/langflow/graph/vertex/types.py b/src/backend/base/langflow/graph/vertex/types.py index cd330aa0a..78c541f65 100644 --- a/src/backend/base/langflow/graph/vertex/types.py +++ b/src/backend/base/langflow/graph/vertex/types.py @@ -404,16 +404,18 @@ class InterfaceVertex(ComponentVertex): return self.vertex_type == InterfaceComponentTypes.ChatInput and self.is_input -class StateVertex(Vertex): +class StateVertex(ComponentVertex): def __init__(self, data: Dict, graph): - super().__init__(data, graph=graph, base_type="custom_components") + super().__init__(data, graph=graph) self.steps = [self._build] - self.is_state = True + self.is_state = False @property def successors_ids(self) -> List[str]: - successors = self.graph.successor_map.get(self.id, []) - return successors + self.graph.activated_vertices + if self._successors_ids is None: + self.is_state = False + return super().successors_ids + return self._successors_ids def _built_object_repr(self): if self.artifacts and "repr" in self.artifacts: