From e3c42d36bae411c7291cb8eec8ad146fac6fe1f5 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Thu, 22 Feb 2024 18:26:55 -0300
Subject: [PATCH] Refactor RoutingVertex to handle successors based on boolean
 value

---
 src/backend/langflow/graph/vertex/types.py | 40 ++++++----------------
 1 file changed, 10 insertions(+), 30 deletions(-)

diff --git a/src/backend/langflow/graph/vertex/types.py b/src/backend/langflow/graph/vertex/types.py
index 5a9576fe5..34d6d8554 100644
--- a/src/backend/langflow/graph/vertex/types.py
+++ b/src/backend/langflow/graph/vertex/types.py
@@ -4,6 +4,7 @@ from typing import Callable, Dict, List, Optional, Union
 
 import yaml
 from langchain_core.messages import AIMessage
+
 from langflow.graph.utils import UnbuiltObject, flatten_list
 from langflow.graph.vertex.base import StatefulVertex, StatelessVertex
 from langflow.interface.utils import extract_input_variables_from_prompt
@@ -366,42 +367,21 @@ class RoutingVertex(StatelessVertex):
     def __init__(self, data: Dict, graph):
         super().__init__(data, graph=graph, base_type="routing")
         self.use_result = True
-        self.steps = [self._build, self._run]
+        self.steps = [self._build]
 
     def _built_object_repr(self):
         if self.artifacts and "repr" in self.artifacts:
             return self.artifacts["repr"] or super()._built_object_repr()
         return super()._built_object_repr()
 
-    def _build(self, *args, **kwargs):
-        super()._build(*args, **kwargs)
-
-        # After building, the _built_object should be a dict with
-        # {"result": Any, "condition": bool}
-        # if true, we need to set should_run attr in the target of true edge
-        # to true and should_run attr in the target of false edge to false
-        # TODO: Add support for multiple conditions
-
-    def _run(self, *args, **kwargs):
-        if self._built_object:
-            condition = self._built_object.get("condition")
-            result = self._built_object.get("result")
-            if condition is not None:
-                for edge in self.edges:
-                    if edge.source_id == self.id:
-                        target_vertex = self.graph.get_vertex(edge.target_id)
-                        # source_handle.channel and condition should be the same
-                        channel_bool = edge.source_handle.channel == "true"
-                        if condition == channel_bool:
-                            target_vertex.should_run = True
-                        else:
-                            target_vertex.should_run = False
-            else:
-                raise ValueError(f"RoutingVertex {self.id} must have a condition in the _built_object")
-
-            self._built_result = result
-        else:
-            raise ValueError(f"RoutingVertex {self.id} must have a _built_object with a condition and a result")
+    @property
+    def successors_ids(self):
+        if isinstance(self._built_object, bool):
+            ids = super().successors_ids
+            if self._built_object:
+                return ids
+            return []
+        raise ValueError("RoutingVertex should return a boolean value.")
 
 
 def dict_to_codeblock(d: dict) -> str: