From 88b2ab328ff87f1ffde993aa2bbdd0a5f30ae3a8 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Fri, 1 Mar 2024 23:57:28 -0300 Subject: [PATCH] Update state and append state methods in SharedState and CustomComponent --- src/backend/langflow/components/utilities/SharedState.py | 2 +- .../interface/custom/custom_component/custom_component.py | 8 ++++++-- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/src/backend/langflow/components/utilities/SharedState.py b/src/backend/langflow/components/utilities/SharedState.py index 7d29da9bb..845252dab 100644 --- a/src/backend/langflow/components/utilities/SharedState.py +++ b/src/backend/langflow/components/utilities/SharedState.py @@ -28,7 +28,7 @@ class SharedState(CustomComponent): self.update_state(name, record) state = self.get_state(name) - if not isinstance(state, Record): + if state and not isinstance(state, Record): if isinstance(state, str): state = Record(text=state) elif isinstance(state, dict): diff --git a/src/backend/langflow/interface/custom/custom_component/custom_component.py b/src/backend/langflow/interface/custom/custom_component/custom_component.py index b8ec433f3..fd8cd7cd2 100644 --- a/src/backend/langflow/interface/custom/custom_component/custom_component.py +++ b/src/backend/langflow/interface/custom/custom_component/custom_component.py @@ -76,13 +76,17 @@ class CustomComponent(Component): def update_state(self, name: str, value: Any): try: - self.vertex.graph.state_manager.update_state(key=name, new_state=value) + self.vertex.graph.update_state( + name=name, record=value, caller=self.vertex.id + ) except Exception as e: raise ValueError(f"Error updating state: {e}") def append_state(self, name: str, value: Any): try: - self.vertex.graph.state_manager.append_state(key=name, new_state=value) + self.vertex.graph.append_state( + key=name, record=value, caller=self.vertex.id + ) except Exception as e: raise ValueError(f"Error appending state: {e}")