feat(ui): add deleteNode and getEdges to graph util

invokeai/frontend/web/src/features/nodes/util/graph/Graph.test.ts

@@ -123,6 +123,34 @@ describe('Graph', () => {
     });
   });
 
+  describe('deleteNode', () => {
+    it('should delete the node with the provided id', () => {
+      const g = new Graph();
+      const n1 = g.addNode({
+        id: 'n1',
+        type: 'add',
+      });
+      const n2 = g.addNode({
+        id: 'n2',
+        type: 'add',
+      });
+      const n3 = g.addNode({
+        id: 'n3',
+        type: 'add',
+      });
+      g.addEdge(n1, 'value', n2, 'a');
+      g.addEdge(n2, 'value', n3, 'a');
+      // This edge should not be deleted bc it doesn't touch n2
+      g.addEdge(n1, 'value', n3, 'a');
+      g.deleteNode(n2.id);
+      expect(g.hasNode(n1.id)).toBe(true);
+      expect(g.hasNode(n2.id)).toBe(false);
+      expect(g.hasNode(n3.id)).toBe(true);
+      // Should delete edges to and from the node
+      expect(g.getEdges().length).toBe(1);
+    });
+  });
+
   describe('hasNode', () => {
     const g = new Graph();
     g.addNode({
@@ -160,6 +188,27 @@ describe('Graph', () => {
     });
   });
 
+  describe('getEdges', () => {
+    it('should get all edges in the graph', () => {
+      const g = new Graph();
+      const n1 = g.addNode({
+        id: 'n1',
+        type: 'add',
+      });
+      const n2 = g.addNode({
+        id: 'n2',
+        type: 'add',
+      });
+      const n3 = g.addNode({
+        id: 'n3',
+        type: 'add',
+      });
+      const e1 = g.addEdge(n1, 'value', n2, 'a');
+      const e2 = g.addEdge(n2, 'value', n3, 'a');
+      expect(g.getEdges()).toEqual([e1, e2]);
+    });
+  });
+
   describe('hasEdge', () => {
     const g = new Graph();
     const add: Invocation<'add'> = {

invokeai/frontend/web/src/features/nodes/util/graph/Graph.ts

@@ -68,6 +68,19 @@ export class Graph {
     return node;
   }
 
+  /**
+   * Deletes a node from the graph. All edges to and from the node are also deleted.
+   * @param id The id of the node to delete.
+   */
+  deleteNode(id: string): void {
+    const node = this._graph.nodes[id];
+    if (node) {
+      this.deleteEdgesFrom(node);
+      this.deleteEdgesTo(node);
+      delete this._graph.nodes[id];
+    }
+  }
+
   /**
    * Check if a node exists in the graph.
    * @param id The id of the node to check.
@@ -182,6 +195,14 @@ export class Graph {
     return edge;
   }
 
+  /**
+   * Get all edges in the graph.
+   * @returns The edges.
+   */
+  getEdges(): Edge[] {
+    return this._graph.edges;
+  }
+
   /**
    * Check if a graph has an edge.
    * Provide the from and to node types as generics to get type hints for from and to field names.