From f8ef220a6e8c0e7d92d22a92fde00b8275917f0f Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@langflow.org>
Date: Thu, 20 Jun 2024 16:29:27 -0300
Subject: [PATCH] refactor: Improve error handling in ComponentBuildException

---
 src/backend/base/langflow/api/v1/chat.py       | 18 ++++++++----------
 .../base/langflow/exceptions/__init__.py       |  0
 .../base/langflow/exceptions/component.py      |  6 ++++++
 src/backend/base/langflow/graph/graph/base.py  |  6 ++++--
 src/backend/base/langflow/graph/vertex/base.py |  5 ++++-
 5 files changed, 22 insertions(+), 13 deletions(-)
 create mode 100644 src/backend/base/langflow/exceptions/__init__.py
 create mode 100644 src/backend/base/langflow/exceptions/component.py

diff --git a/src/backend/base/langflow/api/v1/chat.py b/src/backend/base/langflow/api/v1/chat.py
index f57aaa0b8..8ceae1d22 100644
--- a/src/backend/base/langflow/api/v1/chat.py
+++ b/src/backend/base/langflow/api/v1/chat.py
@@ -24,13 +24,7 @@ from langflow.api.v1.schemas import (
     VertexBuildResponse,
     VerticesOrderResponse,
 )
-from langflow.schema.schema import Log
-from langflow.services.auth.utils import get_current_active_user
-from langflow.services.chat.service import ChatService
-from langflow.services.deps import get_chat_service, get_session, get_session_service
-from langflow.services.monitor.utils import log_vertex_build
-
-if TYPE_CHECKING:
+from langflow.exceptions.component import ComponentBuildException
     from langflow.graph.vertex.types import InterfaceVertex
     from langflow.services.session.service import SessionService
 
@@ -190,9 +184,13 @@ async def build_vertex(
 
             result_data_response = ResultDataResponse.model_validate(result_dict, from_attributes=True)
         except Exception as exc:
-            tb = traceback.format_exc()
-            logger.exception(f"Error building Component: {exc}")
-            params = format_exception_message(exc)
+            if isinstance(exc, ComponentBuildException):
+                params = exc.message
+                tb = exc.formatted_traceback
+            else:
+                tb = traceback.format_exc()
+                logger.exception(f"Error building Component: {exc}")
+                params = format_exception_message(exc)
             message = {"errorMessage": params, "stackTrace": tb}
             valid = False
             output_label = vertex.outputs[0]["name"] if vertex.outputs else "output"
diff --git a/src/backend/base/langflow/exceptions/__init__.py b/src/backend/base/langflow/exceptions/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/src/backend/base/langflow/exceptions/component.py b/src/backend/base/langflow/exceptions/component.py
new file mode 100644
index 000000000..a4a557565
--- /dev/null
+++ b/src/backend/base/langflow/exceptions/component.py
@@ -0,0 +1,6 @@
+# Create an exception class that receives the message and the formatted traceback
+class ComponentBuildException(Exception):
+    def __init__(self, message: str, formatted_traceback: str):
+        self.message = message
+        self.formatted_traceback = formatted_traceback
+        super().__init__(message)
diff --git a/src/backend/base/langflow/graph/graph/base.py b/src/backend/base/langflow/graph/graph/base.py
index a6c1e830a..038791a2f 100644
--- a/src/backend/base/langflow/graph/graph/base.py
+++ b/src/backend/base/langflow/graph/graph/base.py
@@ -4,9 +4,10 @@ from collections import defaultdict, deque
 from functools import partial
 from itertools import chain
 from typing import TYPE_CHECKING, Dict, Generator, List, Optional, Tuple, Type, Union
-
 from loguru import logger
 
+from langflow.exceptions.component import ComponentBuildException
+from langflow.exceptions.component import ComponentBuildException
 from langflow.graph.edge.base import ContractEdge
 from langflow.graph.graph.constants import lazy_load_vertex_dict
 from langflow.graph.graph.runnable_vertices_manager import RunnableVerticesManager
@@ -830,7 +831,8 @@ class Graph:
             log_transaction(flow_id, vertex, status="success")
             return result_dict, params, valid, artifacts, vertex
         except Exception as exc:
-            logger.exception(f"Error building Component:\n\n{exc}")
+            if not isinstance(exc, ComponentBuildException):
+                logger.exception(f"Error building Component:\n\n{exc}")
             flow_id = self.flow_id
             log_transaction(flow_id, vertex, status="failure", error=str(exc))
             raise exc
diff --git a/src/backend/base/langflow/graph/vertex/base.py b/src/backend/base/langflow/graph/vertex/base.py
index cdaae72c3..e7e12e86c 100644
--- a/src/backend/base/langflow/graph/vertex/base.py
+++ b/src/backend/base/langflow/graph/vertex/base.py
@@ -2,12 +2,14 @@ import ast
 import asyncio
 import inspect
 import os
+import traceback
 import types
 from enum import Enum
 from typing import TYPE_CHECKING, Any, AsyncIterator, Callable, Dict, Iterator, List, Mapping, Optional
 
 from loguru import logger
 
+from langflow.exceptions.component import ComponentBuildException
 from langflow.graph.schema import INPUT_COMPONENTS, OUTPUT_COMPONENTS, InterfaceComponentTypes, ResultData
 from langflow.graph.utils import UnbuiltObject, UnbuiltResult
 from langflow.interface.initialize import loading
@@ -621,8 +623,9 @@ class Vertex:
             self.logs = build_logs(self, result)
             self._update_built_object_and_artifacts(result)
         except Exception as exc:
+            tb = traceback.format_exc()
             logger.exception(exc)
-            raise ValueError(f"Error building Component {self.display_name}:\n\n{exc}") from exc
+            raise ComponentBuildException(f"Error building Component {self.display_name}:\n\n{exc}", tb) from exc
 
     def _update_built_object_and_artifacts(self, result):
         """