From 11888121af9c17babe4450fc6cdd1d7b0f67cac0 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@langflow.org>
Date: Mon, 11 Nov 2024 13:10:58 -0300
Subject: [PATCH] fix: Simplify exception handling and refactor error messaging
 (#4496)

* Refactor `ExceptionWithMessageError` to remove redundant exception parameter

* Handle nested exceptions in message schema initialization

* fix: Simplify exception handling in LCAgentComponent
---
 src/backend/base/langflow/base/agents/agent.py  | 2 +-
 src/backend/base/langflow/base/agents/events.py | 5 ++---
 src/backend/base/langflow/schema/message.py     | 5 ++++-
 3 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/src/backend/base/langflow/base/agents/agent.py b/src/backend/base/langflow/base/agents/agent.py
index a05edd795..7e41ade26 100644
--- a/src/backend/base/langflow/base/agents/agent.py
+++ b/src/backend/base/langflow/base/agents/agent.py
@@ -160,7 +160,7 @@ class LCAgentComponent(Component):
             msg_id = e.agent_message.id
             await asyncio.to_thread(delete_message, id_=msg_id)
             self._send_message_event(e.agent_message, category="remove_message")
-            raise e.exception  # noqa: B904
+            raise
         except Exception:
             raise
 
diff --git a/src/backend/base/langflow/base/agents/events.py b/src/backend/base/langflow/base/agents/events.py
index 61ba6497d..f95940111 100644
--- a/src/backend/base/langflow/base/agents/events.py
+++ b/src/backend/base/langflow/base/agents/events.py
@@ -14,9 +14,8 @@ from langflow.schema.message import Message
 
 
 class ExceptionWithMessageError(Exception):
-    def __init__(self, e: Exception, agent_message: Message):
+    def __init__(self, agent_message: Message):
         self.agent_message = agent_message
-        self.exception = e
         super().__init__()
 
 
@@ -253,5 +252,5 @@ async def process_agent_events(
                 start_time = start_time or perf_counter()
         agent_message.properties.state = "complete"
     except Exception as e:
-        raise ExceptionWithMessageError(e, agent_message) from e
+        raise ExceptionWithMessageError(agent_message) from e
     return Message(**agent_message.model_dump())
diff --git a/src/backend/base/langflow/schema/message.py b/src/backend/base/langflow/schema/message.py
index ea78f0f51..9b1ace474 100644
--- a/src/backend/base/langflow/schema/message.py
+++ b/src/backend/base/langflow/schema/message.py
@@ -339,12 +339,15 @@ class ErrorMessage(Message):
 
     def __init__(
         self,
-        exception: Exception,
+        exception: BaseException,
         session_id: str,
         source: Source,
         trace_name: str | None = None,
         flow_id: str | None = None,
     ) -> None:
+        # This is done to avoid circular imports
+        if exception.__class__.__name__ == "ExceptionWithMessageError" and exception.__cause__ is not None:
+            exception = exception.__cause__
         # Get the error reason
         reason = f"**{exception.__class__.__name__}**\n"
         if hasattr(exception, "body") and "message" in exception.body: