From 16afd442951347f70da763e9cab34b898401b97f Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Fri, 16 Aug 2024 16:11:59 -0300 Subject: [PATCH] fix: makes outputs be correctly retrieved from edge (#3392) * feat: Add optional target handle name in get_result method. * fix: Improve logic to consider target handle name in ComponentVertex. Fixes #3380 --- src/backend/base/langflow/graph/vertex/base.py | 12 ++++++------ src/backend/base/langflow/graph/vertex/types.py | 8 ++++++-- 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/src/backend/base/langflow/graph/vertex/base.py b/src/backend/base/langflow/graph/vertex/base.py index b0df08ee923..c2135600ed1 100644 --- a/src/backend/base/langflow/graph/vertex/base.py +++ b/src/backend/base/langflow/graph/vertex/base.py @@ -568,7 +568,7 @@ async def _build_dict_and_update_params( if not self._is_vertex(value): self.params[key][sub_key] = value else: - result = await value.get_result(self) + result = await value.get_result(self, target_handle_name=key) self.params[key][sub_key] = result def _is_vertex(self, value): @@ -583,7 +583,7 @@ def _is_list_of_vertices(self, value): """ return all(self._is_vertex(vertex) for vertex in value) - async def get_result(self, requester: "Vertex") -> Any: + async def get_result(self, requester: "Vertex", target_handle_name: Optional[str] = None) -> Any: """ Retrieves the result of the vertex. @@ -593,9 +593,9 @@ async def get_result(self, requester: "Vertex") -> Any: The result of the vertex. """ async with self._lock: - return await self._get_result(requester) + return await self._get_result(requester, target_handle_name) - async def _get_result(self, requester: "Vertex") -> Any: + async def _get_result(self, requester: "Vertex", target_handle_name: Optional[str] = None) -> Any: """ Retrieves the result of the built component. @@ -620,7 +620,7 @@ async def _build_vertex_and_update_params(self, key, vertex: "Vertex"): Builds a given vertex and updates the params dictionary accordingly. """ - result = await vertex.get_result(self) + result = await vertex.get_result(self, target_handle_name=key) self._handle_func(key, result) if isinstance(result, list): self._extend_params_list_with_result(key, result) @@ -636,7 +636,7 @@ async def _build_list_of_vertices_and_update_params( """ self.params[key] = [] for vertex in vertices: - result = await vertex.get_result(self) + result = await vertex.get_result(self, target_handle_name=key) # Weird check to see if the params[key] is a list # because sometimes it is a Data and breaks the code if not isinstance(self.params[key], list): diff --git a/src/backend/base/langflow/graph/vertex/types.py b/src/backend/base/langflow/graph/vertex/types.py index 0551d653217..1c60b54618d 100644 --- a/src/backend/base/langflow/graph/vertex/types.py +++ b/src/backend/base/langflow/graph/vertex/types.py @@ -84,7 +84,7 @@ def get_edge_with_target(self, target_id: str) -> Generator["CycleEdge", None, N if edge.target_id == target_id: yield edge - async def _get_result(self, requester: "Vertex") -> Any: + async def _get_result(self, requester: "Vertex", target_handle_name: str | None = None) -> Any: """ Retrieves the result of the built component. @@ -114,7 +114,11 @@ async def _get_result(self, requester: "Vertex") -> Any: edges = self.get_edge_with_target(requester.id) result = UNDEFINED for edge in edges: - if edge is not None and edge.source_handle.name in self.results: + if ( + edge is not None + and edge.source_handle.name in self.results + and edge.target_handle.field_name == target_handle_name + ): # Get the result from the output instead of the results dict try: output = self.get_output(edge.source_handle.name)