🐛 fix(base.py): change assignment of source_handle and target_handle to use get() method to handle missing keys gracefully

🐛 fix(base.py): change assignment of parent_node_id and parent_is_top_level to use state dictionary instead of _data attribute
🐛 fix(test_graph.py): change arguments passed to update_source_handle() function to match its signature

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()

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

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"