From 62a7dc77d6071a375a23b49cafd0dd8c56a53589 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Thu, 29 Feb 2024 23:39:36 -0300 Subject: [PATCH] Add state management methods to CustomComponent class --- .../custom_component/custom_component.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) 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 c3b06f90a..b8ec433f3 100644 --- a/src/backend/langflow/interface/custom/custom_component/custom_component.py +++ b/src/backend/langflow/interface/custom/custom_component/custom_component.py @@ -74,6 +74,24 @@ class CustomComponent(Component): status: Optional[Any] = None """The status of the component. This is displayed on the frontend. Defaults to None.""" + def update_state(self, name: str, value: Any): + try: + self.vertex.graph.state_manager.update_state(key=name, new_state=value) + 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) + except Exception as e: + raise ValueError(f"Error appending state: {e}") + + def get_state(self, name: str): + try: + return self.vertex.graph.state_manager.get_state(key=name) + except Exception as e: + raise ValueError(f"Error getting state: {e}") + _tree: Optional[dict] = None def __init__(self, **data):