From 1b0964603b2e4e80a1204bb687097ac41f0a734b Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Sat, 23 Mar 2024 00:21:40 -0300
Subject: [PATCH] Remove unused import and function

---
 .../langflow/interface/importing/utils.py       |  7 -------
 .../langflow/interface/initialize/loading.py    | 17 +++++++----------
 2 files changed, 7 insertions(+), 17 deletions(-)

diff --git a/src/backend/langflow/interface/importing/utils.py b/src/backend/langflow/interface/importing/utils.py
index d9d9198b6..24612df3c 100644
--- a/src/backend/langflow/interface/importing/utils.py
+++ b/src/backend/langflow/interface/importing/utils.py
@@ -10,7 +10,6 @@ from langchain.prompts import PromptTemplate
 from langchain.tools import BaseTool
 from langchain_core.language_models.chat_models import BaseChatModel
 
-from langflow.interface.custom.custom_component import CustomComponent
 from langflow.interface.wrappers.base import wrapper_creator
 
 
@@ -48,7 +47,6 @@ def import_by_type(_type: str, name: str) -> Any:
         "utilities": import_utility,
         "output_parsers": import_output_parser,
         "retrievers": import_retriever,
-        "custom_components": import_custom_component,
     }
     if _type == "models":
         key = "chat" if "chat" in name.lower() else "llm"
@@ -59,11 +57,6 @@ def import_by_type(_type: str, name: str) -> Any:
     return loaded_func(name)
 
 
-def import_custom_component(custom_component: str) -> CustomComponent:
-    """Import custom component from custom component name"""
-    return import_class("langflow.interface.custom.custom_component.CustomComponent")
-
-
 def import_output_parser(output_parser: str) -> Any:
     """Import output parser from output parser name"""
     return import_module(f"from langchain.output_parsers import {output_parser}")
diff --git a/src/backend/langflow/interface/initialize/loading.py b/src/backend/langflow/interface/initialize/loading.py
index 581f853d5..5cb5ba39a 100644
--- a/src/backend/langflow/interface/initialize/loading.py
+++ b/src/backend/langflow/interface/initialize/loading.py
@@ -14,16 +14,11 @@ from langchain_core.documents import Document
 from loguru import logger
 from pydantic import ValidationError
 
+from langflow.custom.customs import CUSTOM_NODES
 from langflow.interface.custom.eval import eval_custom_component_code
-from langflow.interface.custom.utils import get_function
-from langflow.interface.custom_lists import CUSTOM_NODES
 from langflow.interface.importing.utils import import_by_type
 from langflow.interface.initialize.llm import initialize_vertexai
-from langflow.interface.initialize.utils import (
-    handle_format_kwargs,
-    handle_node_type,
-    handle_partial_variables,
-)
+from langflow.interface.initialize.utils import handle_format_kwargs, handle_node_type, handle_partial_variables
 from langflow.interface.initialize.vector_store import vecstore_initializer
 from langflow.interface.output_parsers.base import output_parser_creator
 from langflow.interface.retrievers.base import retriever_creator
@@ -58,6 +53,8 @@ async def instantiate_class(
     logger.debug(f"Instantiating {vertex_type} of type {base_type}")
     if not base_type:
         raise ValueError("No base type provided for vertex")
+    if base_type == "custom_components":
+        return await instantiate_custom_component(params, user_id, vertex)
     class_object = import_by_type(_type=base_type, name=vertex_type)
     return await instantiate_based_on_type(
         class_object=class_object,
@@ -135,8 +132,6 @@ async def instantiate_based_on_type(
         return instantiate_memory(node_type, class_object, params)
     elif base_type == "custom_components":
         return await instantiate_custom_component(
-            node_type,
-            class_object,
             params,
             user_id,
             vertex,
@@ -147,7 +142,7 @@ async def instantiate_based_on_type(
         return class_object(**params)
 
 
-async def instantiate_custom_component(node_type, class_object, params, user_id, vertex):
+async def instantiate_custom_component(params, user_id, vertex):
     params_copy = params.copy()
     class_object: Type["CustomComponent"] = eval_custom_component_code(params_copy.pop("code"))
     custom_component: "CustomComponent" = class_object(
@@ -289,6 +284,8 @@ def instantiate_tool(node_type, class_object: Type[BaseTool], params: Dict):
             raise ValueError("Invalid file")
         return class_object(**params)
     elif node_type == "PythonFunctionTool":
+        from langflow.interface.custom.utils import get_function
+
         params["func"] = get_function(params.get("code"))
         return class_object(**params)
     elif node_type == "PythonFunction":