Skip to content

Commit

Permalink
fix: makes outputs be correctly retrieved from edge (#3392)
Browse files Browse the repository at this point in the history
* feat: Add optional target handle name in get_result method.

* fix: Improve logic to consider target handle name in ComponentVertex.

Fixes #3380
  • Loading branch information
ogabrielluiz authored Aug 16, 2024
1 parent c00e687 commit 16afd44
Show file tree
Hide file tree
Showing 2 changed files with 12 additions and 8 deletions.
12 changes: 6 additions & 6 deletions src/backend/base/langflow/graph/vertex/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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.
Expand All @@ -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.
Expand All @@ -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)
Expand All @@ -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):
Expand Down
8 changes: 6 additions & 2 deletions src/backend/base/langflow/graph/vertex/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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)
Expand Down

0 comments on commit 16afd44

Please sign in to comment.