@ghchinoy/litflow 0.2.7 → 0.3.0

package/src/breadboard/engine/runtime/static/create-plan.ts

@@ -0,0 +1,134 @@
+// @ts-nocheck
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import {
+  Edge,
+  GraphDescriptor,
+  NodeDescriptor,
+  OrchestrationPlan,
+  PlanNodeInfo,
+} from "@breadboard-ai/types";
+
+export { createPlan };
+
+/**
+ * Creates an execution plan from the provided GraphDescription.
+ * The graph is guaranteed to be condensed (no cycles), and each
+ * strongly connected component is represented with a single node
+ * that has a "folded" tag.
+ */
+function createPlan(graph: GraphDescriptor): OrchestrationPlan {
+  const { nodes, edges } = graph;
+
+  if (!nodes || nodes.length === 0) {
+    return { stages: [] };
+  }
+
+  const nodeMap = new Map(nodes.map((node) => [node.id, node]));
+  const inDegree = new Map<string, number>();
+  const outEdges = new Map<string, Edge[]>();
+  const inEdges = new Map<string, Edge[]>();
+
+  nodes.forEach((node) => {
+    inDegree.set(node.id, 0);
+    outEdges.set(node.id, []);
+    inEdges.set(node.id, []);
+  });
+
+  if (edges) {
+    edges.forEach((edge) => {
+      const currentDegree = inDegree.get(edge.to) || 0;
+      inDegree.set(edge.to, currentDegree + 1);
+
+      const fromEdges = outEdges.get(edge.from) || [];
+      fromEdges.push({ ...edge });
+      outEdges.set(edge.from, fromEdges);
+
+      const toEdges = inEdges.get(edge.to) || [];
+      toEdges.push({ ...edge });
+      inEdges.set(edge.to, toEdges);
+    });
+  }
+
+  const stages: PlanNodeInfo[][] = [];
+  const entries = nodes.filter((node) => inDegree.get(node.id) === 0);
+
+  // Now, let's separate out all standalone steps and see if maybe we only
+  // have standalone nodes.
+  const standalone: NodeDescriptor[] = [];
+  const connected: NodeDescriptor[] = [];
+  let onlyStandalone = true;
+  entries.forEach((node) => {
+    const nodeTails = outEdges.get(node.id) || [];
+    if (nodeTails.length == 0) {
+      standalone.push(node);
+    } else {
+      onlyStandalone = false;
+      connected.push(node);
+    }
+  });
+  let queue: NodeDescriptor[];
+
+  // If there are no standalone nodes, return all entries as usual.
+  if (standalone.length === 0) {
+    queue = entries;
+  } else if (onlyStandalone) {
+    // This is the situation when we have a bunch of random nodes in graph
+    // and they are not connected, and there's no designated start node.
+
+    // Just return the first standalone node.
+    queue = [standalone[0]];
+  } else {
+    // If there are both standalone and connected nodes, we just ignore
+    // all standalone nodes.
+    queue = connected;
+  }
+
+  const processed = new Set<string>();
+
+  while (queue.length > 0) {
+    const stageNodes: PlanNodeInfo[] = [];
+    const nextQueue: typeof queue = [];
+
+    while (queue.length > 0) {
+      const node = queue.shift()!;
+
+      if (processed.has(node.id)) continue;
+      processed.add(node.id);
+
+      const downstream = outEdges.get(node.id) || [];
+
+      const upstream = inEdges.get(node.id) || [];
+
+      const planNodeInfo: PlanNodeInfo = {
+        node: nodeMap.get(node.id)!,
+        downstream,
+        upstream,
+      };
+      stageNodes.push(planNodeInfo);
+
+      (outEdges.get(node.id) || []).forEach((edge) => {
+        const targetDegree = inDegree.get(edge.to) || 0;
+        if (targetDegree > 0) {
+          inDegree.set(edge.to, targetDegree - 1);
+          if (targetDegree === 1) {
+            const targetNode = nodeMap.get(edge.to);
+            if (targetNode) nextQueue.push(targetNode);
+          }
+        }
+      });
+    }
+
+    if (stageNodes.length > 0) {
+      stages.push(stageNodes);
+    }
+
+    queue.push(...nextQueue);
+  }
+
+  return { stages };
+}

package/src/breadboard/engine/runtime/static/nodes-to-subgraph.ts

@@ -0,0 +1,168 @@
+// @ts-nocheck
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import {
+  GraphDescriptor,
+  GraphIdentifier,
+  NodeDescriptor,
+  NodeIdentifier,
+} from "@breadboard-ai/types";
+
+export { nodesToSubgraph };
+
+/**
+ * Moves a group of nodes within a graph into a subgraph, mutating
+ * the supplied graph.
+ *
+ * @param graph - The graph where the subgraph will be created.
+ * @param nodeGroup - The group of nodes to include in the subgraph.
+ * @param subgraphId - The unique identifier for the subgraph.
+ * @param title - An optional title for the subgraph.
+ * @param description - An optional description for the subgraph.
+ */
+function nodesToSubgraph(
+  graph: GraphDescriptor,
+  nodeGroup: NodeIdentifier[],
+  subgraphId: GraphIdentifier,
+  title?: string,
+  description?: string
+) {
+  // Find all edges within the node group
+  const subgraph = createSubgraph(
+    graph,
+    nodeGroup,
+    subgraphId,
+    title,
+    description
+  );
+
+  graph.graphs ??= {};
+  graph.graphs[subgraphId] = subgraph;
+
+  const replacementNodeId = subgraphId;
+  const replacementNode: NodeDescriptor = {
+    id: replacementNodeId,
+    type: `#${subgraphId}`,
+    metadata: {
+      title: `Subgraph "${subgraphId}"`,
+      tags: ["folded"],
+    },
+  };
+
+  // Replace folded nodes with a replacement
+  graph.nodes =
+    graph.nodes?.filter((node) => !nodeGroup.includes(node.id)) || [];
+  graph.nodes.push(replacementNode);
+
+  // Update edges to point to/from condensed node
+  updateEdgeForFoldedNode(graph, nodeGroup, replacementNodeId);
+}
+
+function createSubgraph(
+  graph: GraphDescriptor,
+  nodeGroup: NodeIdentifier[],
+  subgraphId: GraphIdentifier,
+  title?: string,
+  description?: string
+): GraphDescriptor {
+  const internalEdges =
+    graph.edges?.filter(
+      (edge) => nodeGroup.includes(edge.from) && nodeGroup.includes(edge.to)
+    ) || [];
+
+  // Find incoming edges to the node group
+  const incomingEdges =
+    graph.edges?.filter(
+      (edge) => !nodeGroup.includes(edge.from) && nodeGroup.includes(edge.to)
+    ) || [];
+
+  // Find outgoing edges from the node group
+  const outgoingEdges =
+    graph.edges?.filter(
+      (edge) => nodeGroup.includes(edge.from) && !nodeGroup.includes(edge.to)
+    ) || [];
+
+  // Get nodes in the group
+  const groupNodes =
+    graph.nodes?.filter((node) => nodeGroup.includes(node.id)) || [];
+
+  // Create input node for capturing incoming edges
+  const inputNode: NodeDescriptor = {
+    id: `input_${subgraphId}`,
+    type: "input",
+    metadata: {
+      title: "Subgraph Input",
+      description: "Captures incoming edges to the subgraph",
+    },
+  };
+
+  // Create output node for capturing outgoing edges
+  const outputNode: NodeDescriptor = {
+    id: `output_${subgraphId}`,
+    type: "output",
+    metadata: {
+      title: "Subgraph Output",
+      description: "Captures outgoing edges from the subgraph",
+    },
+  };
+
+  // Create subgraph edges
+  const subgraphEdges = [...internalEdges];
+
+  // Add edges from input node to group entry points
+  incomingEdges.forEach((edge) => {
+    subgraphEdges.push({
+      from: `input_${subgraphId}`,
+      to: edge.to,
+      in: edge.in,
+      out: edge.in || "out",
+    });
+  });
+
+  // Add edges from group exit points to output node
+  outgoingEdges.forEach((edge) => {
+    subgraphEdges.push({
+      from: edge.from,
+      to: `output_${subgraphId}`,
+      in: edge.out || "in",
+      out: edge.out,
+    });
+  });
+
+  return {
+    title: title || `Subgraph ${subgraphId}`,
+    description: description || `Subgraph containing ${nodeGroup.length} nodes`,
+    nodes: [inputNode, ...groupNodes, outputNode],
+    edges: subgraphEdges,
+  };
+}
+
+function updateEdgeForFoldedNode(
+  graph: GraphDescriptor,
+  nodeGroup: string[],
+  nodeId: string
+) {
+  if (!graph.edges) return;
+
+  // Update edges that cross fold lines
+  graph.edges = graph.edges
+    .map((edge) => {
+      const from = nodeGroup.includes(edge.from);
+      const into = nodeGroup.includes(edge.to);
+
+      if (from && !into) {
+        return { ...edge, from: nodeId };
+      } else if (!from && into) {
+        return { ...edge, to: nodeId };
+      }
+      return edge;
+    })
+    // Remove internal edges (they're now in the subgraph)
+    .filter((edge) => {
+      return !(nodeGroup.includes(edge.from) && nodeGroup.includes(edge.to));
+    });
+}