From 4d6c731a9d4070f7c95859bb3cf666004685eef2 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Fri, 6 Oct 2023 16:19:42 -0300 Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20fix(base.py):=20change=20assignm?= =?UTF-8?q?ent=20of=20source=5Fhandle=20and=20target=5Fhandle=20to=20use?= =?UTF-8?q?=20get()=20method=20to=20handle=20missing=20keys=20gracefully?= =?UTF-8?q?=20=F0=9F=90=9B=20fix(base.py):=20change=20assignment=20of=20pa?= =?UTF-8?q?rent=5Fnode=5Fid=20and=20parent=5Fis=5Ftop=5Flevel=20to=20use?= =?UTF-8?q?=20state=20dictionary=20instead=20of=20=5Fdata=20attribute=20?= =?UTF-8?q?=F0=9F=90=9B=20fix(test=5Fgraph.py):=20change=20arguments=20pas?= =?UTF-8?q?sed=20to=20update=5Fsource=5Fhandle()=20function=20to=20match?= =?UTF-8?q?=20its=20signature?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/backend/langflow/graph/edge/base.py | 4 ++-- src/backend/langflow/graph/vertex/base.py | 6 ++++-- tests/test_graph.py | 2 +- 3 files changed, 7 insertions(+), 5 deletions(-) diff --git a/src/backend/langflow/graph/edge/base.py b/src/backend/langflow/graph/edge/base.py index bd585c168..82714e395 100644 --- a/src/backend/langflow/graph/edge/base.py +++ b/src/backend/langflow/graph/edge/base.py @@ -72,8 +72,8 @@ class Edge: self.source = state["source"] self.target = state["target"] self.target_param = state["target_param"] - self.source_handle = state["source_handle"] - self.target_handle = state["target_handle"] + self.source_handle = state.get("source_handle") + self.target_handle = state.get("target_handle") def reset(self) -> None: self.source._build_params() diff --git a/src/backend/langflow/graph/vertex/base.py b/src/backend/langflow/graph/vertex/base.py index c31cca3ec..8ae0bb5f8 100644 --- a/src/backend/langflow/graph/vertex/base.py +++ b/src/backend/langflow/graph/vertex/base.py @@ -38,6 +38,8 @@ class Vertex: self.task_id: Optional[str] = None self.is_task = is_task self.params = params or {} + self.parent_node_id: Optional[str] = self._data.get("parent_node_id") + self.parent_is_top_level = False def reset_params(self): for edge in self.edges: @@ -88,8 +90,8 @@ class Vertex: self._built = False self.artifacts: Dict[str, Any] = {} self.task_id: Optional[str] = None - self.parent_node_id: Optional[str] = self._data.get("parent_node_id") - self.parent_is_top_level = False + self.parent_node_id = state["parent_node_id"] + self.parent_is_top_level = state["parent_is_top_level"] def set_top_level(self, top_level_nodes: List[str]) -> None: self.parent_is_top_level = self.parent_node_id in top_level_nodes diff --git a/tests/test_graph.py b/tests/test_graph.py index 0dc9fc372..8d442ce76 100644 --- a/tests/test_graph.py +++ b/tests/test_graph.py @@ -497,7 +497,7 @@ def test_update_source_handle(): "edges": [{"source": "some_node"}], } updated_edge = update_source_handle( - new_edge, {"nodes": flow_data["nodes"], "edges": flow_data["edges"]} + new_edge, flow_data["nodes"], flow_data["edges"] ) assert updated_edge["source"] == "last_node" assert updated_edge["data"]["sourceHandle"]["id"] == "last_node"