jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

Refactor RoutingVertex to handle successors based on boolean value

Browse source
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-02-22 18:26:55 -03:00
commit e3c42d36ba
1 changed files with 10 additions and 30 deletions
Download patch file Download diff file Expand all files Collapse all files

40
src/backend/langflow/graph/vertex/types.py
View file

@ -4,6 +4,7 @@ from typing import Callable, Dict, List, Optional, Union
import yaml import yaml
from langchain_core.messages import AIMessage from langchain_core.messages import AIMessage
from langflow.graph.utils import UnbuiltObject, flatten_list from langflow.graph.utils import UnbuiltObject, flatten_list
from langflow.graph.vertex.base import StatefulVertex, StatelessVertex from langflow.graph.vertex.base import StatefulVertex, StatelessVertex
from langflow.interface.utils import extract_input_variables_from_prompt from langflow.interface.utils import extract_input_variables_from_prompt
@ -366,42 +367,21 @@ class RoutingVertex(StatelessVertex):
def __init__(self, data: Dict, graph): def __init__(self, data: Dict, graph):
super().__init__(data, graph=graph, base_type="routing") super().__init__(data, graph=graph, base_type="routing")
self.use_result = True self.use_result = True
self.steps = [self._build, self._run] self.steps = [self._build]
def _built_object_repr(self): def _built_object_repr(self):
if self.artifacts and "repr" in self.artifacts: if self.artifacts and "repr" in self.artifacts:
return self.artifacts["repr"] or super()._built_object_repr() return self.artifacts["repr"] or super()._built_object_repr()
return super()._built_object_repr() return super()._built_object_repr()
def _build(self, *args, **kwargs): @property
super()._build(*args, **kwargs) def successors_ids(self):
if isinstance(self._built_object, bool):
# After building, the _built_object should be a dict with ids = super().successors_ids
# {"result": Any, "condition": bool} if self._built_object:
# if true, we need to set should_run attr in the target of true edge return ids
# to true and should_run attr in the target of false edge to false return []
# TODO: Add support for multiple conditions raise ValueError("RoutingVertex should return a boolean value.")
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")
def dict_to_codeblock(d: dict) -> str: def dict_to_codeblock(d: dict) -> str: