From 24344ee0fa379b761fda4e114e38a28933257902 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Thu, 17 Aug 2023 10:15:38 -0300
Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20fix(base.py):=20remove=20unneces?=
 =?UTF-8?q?sary=20import=20statement=20for=20celery.result.AsyncResult=20?=
 =?UTF-8?q?=F0=9F=94=80=20refactor(base.py):=20refactor=20the=20build=20me?=
 =?UTF-8?q?thod=20in=20Vertex=20class=20to=20improve=20readability=20and?=
 =?UTF-8?q?=20remove=20redundant=20code=20=E2=9C=A8=20feat(base.py):=20add?=
 =?UTF-8?q?=20get=5Ftask=20method=20to=20Vertex=20class=20to=20retrieve=20?=
 =?UTF-8?q?the=20task=20from=20celery=20using=20task=5Fid?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/backend/langflow/graph/vertex/base.py | 21 +++++++++++----------
 1 file changed, 11 insertions(+), 10 deletions(-)

diff --git a/src/backend/langflow/graph/vertex/base.py b/src/backend/langflow/graph/vertex/base.py
index bbd2bbb65..50adc7764 100644
--- a/src/backend/langflow/graph/vertex/base.py
+++ b/src/backend/langflow/graph/vertex/base.py
@@ -73,6 +73,13 @@ class Vertex:
                     self.base_type = base_type
                     break
 
+    def get_task(self):
+        # using the task_id, get the task from celery
+        # and return it
+        from celery.result import AsyncResult
+
+        return AsyncResult(self.task_id)
+
     def _build_params(self):
         # sourcery skip: merge-list-append, remove-redundant-if
         # Some params are required, some are optional
@@ -179,16 +186,9 @@ class Vertex:
         if self._built:
             return self._built_object
 
-        # Check if there's a task_id, which means it was sent to a Celery worker
-        try:
-            from celery.result import AsyncResult
-        except ImportError:
-            # If Celery is not installed, just build the vertex locally
-            return self.build()
         if self.is_task and self.task_id is not None:
-            result = AsyncResult(self.task_id).get(
-                timeout=timeout
-            )  # Blocking until result is ready or timeout
+            task = self.get_task()
+            result = task.get(timeout=timeout)
             if result is not None:  # If result is ready
                 self._update_built_object_and_artifacts(result)
                 return self._built_object
@@ -197,7 +197,8 @@ class Vertex:
                 pass
 
         # If there's no task_id, build the vertex locally
-        return self.build()
+        self.build()
+        return self._built_object
 
     def _build_node_and_update_params(self, key, node):
         """