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.
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-05-16 15:48:49 -07:00
commit 64c42a9280
3 changed files with 6 additions and 2 deletions

View file

@ -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)}

View file

@ -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

View file

@ -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)