From c2411d4c22b6855d1e51dd3c97879e5769a84940 Mon Sep 17 00:00:00 2001 From: Cristhian Zanforlin Lousa Date: Fri, 31 Jan 2025 16:11:46 -0300 Subject: [PATCH] refactor: Simplify _process_raw_result method in custom component processing (#6040) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 🐛 (component.py): fix logic in _process_raw_result method to correctly extract data from result object based on conditions and return it * [autofix.ci] apply automated fixes * ♻️ (component.py): refactor extract_data method to improve readability and maintainability by using more descriptive variable names and simplifying the logic. * [autofix.ci] apply automated fixes * 🐛 (component.py): update isinstance check to use union type for better type handling --------- Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> Co-authored-by: Gabriel Luiz Freitas Almeida --- .../custom/custom_component/component.py | 23 +++++++++++-------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/src/backend/base/langflow/custom/custom_component/component.py b/src/backend/base/langflow/custom/custom_component/component.py index 82673419a2e4..44984c86baed 100644 --- a/src/backend/base/langflow/custom/custom_component/component.py +++ b/src/backend/base/langflow/custom/custom_component/component.py @@ -978,17 +978,20 @@ def _build_artifact(self, result): return {"repr": custom_repr, "raw": raw, "type": artifact_type} def _process_raw_result(self, result): + if len(self.outputs) == 1: + return self.status or self.extract_data(result) + return self.extract_data(result) + + def extract_data(self, result): + if hasattr(result, "data"): + return result.data + if hasattr(result, "model_dump"): + return result.model_dump() + if isinstance(result, Data | dict | str): + return result.data if isinstance(result, Data) else result if self.status: - raw = self.status - elif hasattr(result, "data"): - raw = result.data - elif hasattr(result, "model_dump"): - raw = result.model_dump() - elif isinstance(result, dict | Data | str): - raw = result.data if isinstance(result, Data) else result - else: - raw = result - return raw + return self.status + return result def _log_output(self, output): self._output_logs[output.name] = self._logs