From 98d1eecd40b0258e05d05712692ba9b69b92aae1 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Thu, 30 Nov 2023 15:42:41 -0300
Subject: [PATCH] Refactor code and fix custom_fields bug

---
 src/backend/langflow/api/v1/validate.py       | 19 +++++++++----------
 .../langflow/template/frontend_node/base.py   |  7 ++++---
 2 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/src/backend/langflow/api/v1/validate.py b/src/backend/langflow/api/v1/validate.py
index c1ef04e54..e17a95fb8 100644
--- a/src/backend/langflow/api/v1/validate.py
+++ b/src/backend/langflow/api/v1/validate.py
@@ -1,15 +1,10 @@
 from fastapi import APIRouter, HTTPException
-
-from langflow.api.v1.base import (
-    Code,
-    CodeValidationResponse,
-    ValidatePromptRequest,
-    PromptValidationResponse,
-    validate_prompt,
-)
+from langflow.api.v1.base import (Code, CodeValidationResponse,
+                                  PromptValidationResponse,
+                                  ValidatePromptRequest, validate_prompt)
 from langflow.template.field.base import TemplateField
-from loguru import logger
 from langflow.utils.validate import validate_code
+from loguru import logger
 
 # build router
 router = APIRouter(prefix="/validate", tags=["Validate"])
@@ -61,7 +56,11 @@ def get_old_custom_fields(prompt_request):
             # then we are dealing with the first prompt request after the node was created
             prompt_request.name = list(prompt_request.frontend_node.custom_fields.keys())[0]
 
-        old_custom_fields = prompt_request.frontend_node.custom_fields[prompt_request.name].copy()
+        old_custom_fields = prompt_request.frontend_node.custom_fields[prompt_request.name]
+        if old_custom_fields is None:
+            old_custom_fields = []
+
+        old_custom_fields = old_custom_fields.copy()
     except KeyError:
         old_custom_fields = []
     prompt_request.frontend_node.custom_fields[prompt_request.name] = []
diff --git a/src/backend/langflow/template/frontend_node/base.py b/src/backend/langflow/template/frontend_node/base.py
index 2946f1e9f..361d644dc 100644
--- a/src/backend/langflow/template/frontend_node/base.py
+++ b/src/backend/langflow/template/frontend_node/base.py
@@ -1,9 +1,10 @@
 import re
 from collections import defaultdict
-from typing import ClassVar, DefaultDict, Dict, List, Optional
+from typing import ClassVar, Dict, List, Optional
 
 from langflow.template.field.base import TemplateField
-from langflow.template.frontend_node.constants import CLASSES_TO_REMOVE, FORCE_SHOW_FIELDS
+from langflow.template.frontend_node.constants import (CLASSES_TO_REMOVE,
+                                                       FORCE_SHOW_FIELDS)
 from langflow.template.frontend_node.formatter import field_formatters
 from langflow.template.template.base import Template
 from langflow.utils import constants
@@ -45,7 +46,7 @@ class FrontendNode(BaseModel):
     name: str = ""
     display_name: str = ""
     documentation: str = ""
-    custom_fields: Optional[DefaultDict[str, List[str]]] = defaultdict(list)
+    custom_fields: Optional[Dict] = defaultdict(list)
     output_types: List[str] = []
     full_path: Optional[str] = None
     field_formatters: FieldFormatters = Field(default_factory=FieldFormatters)