From 11cece92b84806ff5ff6f20d4e7567d40c42bcf6 Mon Sep 17 00:00:00 2001
From: gustavoschaedler <gustavopoa@gmail.com>
Date: Thu, 6 Jul 2023 01:02:34 +0100
Subject: [PATCH] Refactor code for CustomComponent class

- Refactored code to replace the deprecated is_valid property with the is_check_valid method.
- Added validation for the entrypoint function name and return type.
- Modified the error messages to provide more specific details.

This commit implements the necessary changes to refactor the CustomComponent class in the endpoints.py file. The is_valid property has been replaced with the is_check_valid method to check the validity of the custom component. Additionally, validation has been added to ensure the presence of a valid entrypoint function and a valid return type. If any of the validation checks fail, appropriate error messages are raised to provide detailed traceback information.
---
 src/backend/langflow/api/v1/endpoints.py      |  5 +---
 .../langflow/interface/tools/custom.py        | 28 +++++++++++++------
 2 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/src/backend/langflow/api/v1/endpoints.py b/src/backend/langflow/api/v1/endpoints.py
index 6909a299f..c9f671bcf 100644
--- a/src/backend/langflow/api/v1/endpoints.py
+++ b/src/backend/langflow/api/v1/endpoints.py
@@ -100,9 +100,6 @@ async def custom_component(
     raw_code: CustomComponentCode,
 ):
     extractor = CustomComponent(code=raw_code.code)
-
-    if not extractor.is_valid:
-        print("ERROR")
-        # TODO: Raise error
+    extractor.is_check_valid()
 
     return build_langchain_template_custom_component(extractor)
diff --git a/src/backend/langflow/interface/tools/custom.py b/src/backend/langflow/interface/tools/custom.py
index 60541960e..d18577599 100644
--- a/src/backend/langflow/interface/tools/custom.py
+++ b/src/backend/langflow/interface/tools/custom.py
@@ -197,21 +197,32 @@ class CustomComponent(BaseModel):
             )
 
         functions = code.get("functions", [])
-        if build_function := next(
+        build_function = next(
             (f for f in functions if f["name"] == self.function_entrypoint_name),
             None,
-        ):
-            # Check if the return type of the build function is valid
-            return build_function.get("return_type") in self.return_type_valid_list
-        else:
+        )
+
+        if not build_function:
             raise HTTPException(
                 status_code=400,
                 detail={
-                    "error": f"The class return [{str(build_function.get('return_type'))}] needs to be an item from this list. [{str(self.return_type_valid_list)}]",
-                    "traceback": "",
+                    "error": "Invalid entrypoint function name",
+                    "traceback": f"There needs to be at least one entrypoint function named '{self.function_entrypoint_name}' and it needs to return one of the types from this list {str(self.return_type_valid_list)}.",
                 },
             )
 
+        return_type = build_function.get("return_type")
+        if return_type not in self.return_type_valid_list:
+            raise HTTPException(
+                status_code=400,
+                detail={
+                    "error": "Invalid entrypoint function return",
+                    "traceback": f"The entrypoint function return '{return_type}' needs to be an item from this list {str(self.return_type_valid_list)}.",
+                },
+            )
+
+        return True
+
     def get_function(self):
         return validate.create_function(self.code, self.function_entrypoint_name)
 
@@ -219,8 +230,7 @@ class CustomComponent(BaseModel):
     def data(self):
         return self.extract_class_info()
 
-    @property
-    def is_valid(self):
+    def is_check_valid(self):
         return self._class_template_validation(self.data)
 
     @property