From a4c9ada178cb4db71e6f1f3888e4123bc1bd750d Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Sun, 23 Jun 2024 22:22:45 -0300 Subject: [PATCH] refactor: Handle exceptions and end traces in Graph class This commit adds exception handling and trace ending functionality to the Graph class in the `base.py` file. It imports the `traceback` module and adds a try-except block to catch any exceptions that occur during graph execution. If an exception is caught, it logs the exception and ends all traces with the error message. Finally, it ensures that all traces are ended regardless of whether an exception occurred or not. This change improves the robustness and error handling of the Graph class. --- src/backend/base/langflow/graph/graph/base.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/backend/base/langflow/graph/graph/base.py b/src/backend/base/langflow/graph/graph/base.py index 2cd5af1fa..46f562c95 100644 --- a/src/backend/base/langflow/graph/graph/base.py +++ b/src/backend/base/langflow/graph/graph/base.py @@ -1,4 +1,5 @@ import asyncio +import traceback import uuid from collections import defaultdict, deque from datetime import datetime, timezone @@ -341,7 +342,11 @@ class Graph: self.increment_run_count() except Exception as exc: logger.exception(exc) + tb = traceback.format_exc() + await self.end_all_traces(error=f"{exc.__class__.__name__}: {exc}\n\n{tb}") raise ValueError(f"Error running graph: {exc}") from exc + finally: + await self.end_all_traces() # Get the outputs vertex_outputs = [] for vertex in self.vertices: