From 7f6aebb490226764f70fdca450412677caef7132 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Thu, 20 Jun 2024 18:52:15 -0300 Subject: [PATCH] refactor: Add output_name parameter to stop method in CustomComponent --- .../langflow/custom/custom_component/custom_component.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/backend/base/langflow/custom/custom_component/custom_component.py b/src/backend/base/langflow/custom/custom_component/custom_component.py index 92631c140..7f1e184e1 100644 --- a/src/backend/base/langflow/custom/custom_component/custom_component.py +++ b/src/backend/base/langflow/custom/custom_component/custom_component.py @@ -91,7 +91,11 @@ class CustomComponent(BaseComponent): except Exception as e: raise ValueError(f"Error updating state: {e}") - def stop(self, output_name: str): + def stop(self, output_name: str | None): + if not output_name and len(self.vertex.outputs) == 1: + output_name = self.vertex.outputs[0]["name"] + else: + raise ValueError("You must specify an output name to call stop") if not self.vertex: raise ValueError("Vertex is not set") try: