From b2205bba8aa64b61dd150849ddc9782d231bd50f Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Wed, 11 Oct 2023 17:33:32 -0300 Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20fix(types.py):=20add=20output=20?= =?UTF-8?q?types=20to=20the=20frontend=20node=20to=20improve=20data=20inte?= =?UTF-8?q?grity=20and=20error=20handling?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ✨ feat(types.py): add support for adding output types to the frontend node to ensure valid return types and improve error handling --- src/backend/langflow/interface/types.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/src/backend/langflow/interface/types.py b/src/backend/langflow/interface/types.py index 6708c11c9..5f29dbf2e 100644 --- a/src/backend/langflow/interface/types.py +++ b/src/backend/langflow/interface/types.py @@ -288,6 +288,24 @@ def add_base_classes(frontend_node, return_types: List[str]): frontend_node.get("base_classes").append(base_class) +def add_output_types(frontend_node, return_types: List[str]): + """Add output types to the frontend node""" + for return_type in return_types: + if return_type not in CUSTOM_COMPONENT_SUPPORTED_TYPES or return_type is None: + raise HTTPException( + status_code=400, + detail={ + "error": ( + "Invalid return type should be one of: " + f"{list(CUSTOM_COMPONENT_SUPPORTED_TYPES.keys())}" + ), + "traceback": traceback.format_exc(), + }, + ) + + frontend_node.get("output_types").append(return_type) + + def build_langchain_template_custom_component(custom_component: CustomComponent): """Build a custom component template for the langchain""" try: @@ -314,6 +332,9 @@ def build_langchain_template_custom_component(custom_component: CustomComponent) add_base_classes( frontend_node, custom_component.get_function_entrypoint_return_type ) + add_output_types( + frontend_node, custom_component.get_function_entrypoint_return_type + ) logger.debug("Added base classes") return frontend_node except Exception as exc: