From f43fd895c69db1b3e70e48416874a6cb36d80829 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@langflow.org>
Date: Mon, 11 Nov 2024 10:43:59 -0300
Subject: [PATCH] fix: Enhance extract_class_name function to identify
 Component subclasses (#4492)

---
 src/backend/base/langflow/utils/validate.py | 36 +++++++++++++++++----
 1 file changed, 29 insertions(+), 7 deletions(-)

diff --git a/src/backend/base/langflow/utils/validate.py b/src/backend/base/langflow/utils/validate.py
index ab9672a55..ea9715612 100644
--- a/src/backend/base/langflow/utils/validate.py
+++ b/src/backend/base/langflow/utils/validate.py
@@ -326,10 +326,32 @@ def extract_function_name(code):
     raise ValueError(msg)
 
 
-def extract_class_name(code):
-    module = ast.parse(code)
-    for node in module.body:
-        if isinstance(node, ast.ClassDef):
-            return node.name
-    msg = f"No class definition found in the code string. Code snippet: {code[:100]}"
-    raise ValueError(msg)
+def extract_class_name(code: str) -> str:
+    """Extract the name of the first Component subclass found in the code.
+
+    Args:
+        code (str): The source code to parse
+
+    Returns:
+        str: Name of the first Component subclass found
+
+    Raises:
+        ValueError: If no Component subclass is found in the code
+    """
+    try:
+        module = ast.parse(code)
+        for node in module.body:
+            if not isinstance(node, ast.ClassDef):
+                continue
+
+            # Check bases for Component inheritance
+            # TODO: Build a more robust check for Component inheritance
+            for base in node.bases:
+                if isinstance(base, ast.Name) and any(pattern in base.id for pattern in ["Component", "LC"]):
+                    return node.name
+
+        msg = f"No Component subclass found in the code string. Code snippet: {code[:100]}"
+        raise TypeError(msg)
+    except SyntaxError as e:
+        msg = f"Invalid Python code: {e!s}"
+        raise ValueError(msg) from e