From 64c42a9280189852967105ea5d6569fa110294b9 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Thu, 16 May 2024 15:48:49 -0700 Subject: [PATCH] feat: Update selected_output_type handling in Vertex class This commit updates the handling of the `selected_output_type` attribute in the `Vertex` class. Previously, the attribute was assigned directly from the `data` dictionary, which could result in unexpected behavior. Now, the attribute is properly stripped and converted to a string before assignment, ensuring consistent behavior. This change improves the reliability and accuracy of the `selected_output_type` attribute in the `Vertex` class. --- src/backend/base/langflow/graph/vertex/base.py | 4 +++- src/backend/base/langflow/interface/initialize/loading.py | 1 - src/backend/base/langflow/utils/validate.py | 3 +++ 3 files changed, 6 insertions(+), 2 deletions(-) diff --git a/src/backend/base/langflow/graph/vertex/base.py b/src/backend/base/langflow/graph/vertex/base.py index e250e9419..27110a510 100644 --- a/src/backend/base/langflow/graph/vertex/base.py +++ b/src/backend/base/langflow/graph/vertex/base.py @@ -201,7 +201,9 @@ class Vertex: self.description = self.data["node"].get("description", "") self.frozen = self.data["node"].get("frozen", False) - self.selected_output_type = self.data["node"].get("selected_output_type") + self.selected_output_type = ( + str(self.data.get("selected_output_type")).strip() if self.data.get("selected_output_type") else None + ) self.is_input = self.data["node"].get("is_input") or self.is_input self.is_output = self.data["node"].get("is_output") or self.is_output template_dicts = {key: value for key, value in self.data["node"]["template"].items() if isinstance(value, dict)} diff --git a/src/backend/base/langflow/interface/initialize/loading.py b/src/backend/base/langflow/interface/initialize/loading.py index c1014ef90..96fe8832f 100644 --- a/src/backend/base/langflow/interface/initialize/loading.py +++ b/src/backend/base/langflow/interface/initialize/loading.py @@ -2,7 +2,6 @@ import inspect import json from typing import TYPE_CHECKING, Any, Callable, Dict, Sequence, Type - import orjson from langchain.agents import agent as agent_module from langchain.agents.agent import AgentExecutor diff --git a/src/backend/base/langflow/utils/validate.py b/src/backend/base/langflow/utils/validate.py index 0871dbd82..bd7827199 100644 --- a/src/backend/base/langflow/utils/validate.py +++ b/src/backend/base/langflow/utils/validate.py @@ -253,6 +253,9 @@ def build_class_constructor(compiled_class, exec_globals, class_name): globals()[module_name] = module instance = exec_globals[class_name](*args, **kwargs) + # Get selected type from global scope + if instance.selected_output_type in exec_globals: + instance.selected_output_type = exec_globals[instance.selected_output_type] return instance build_custom_class.__globals__.update(exec_globals)