@ghchinoy/litflow 0.2.8 → 0.3.0

package/src/breadboard/engine/inspector/graph/module.ts

@@ -0,0 +1,69 @@
+// @ts-nocheck
+/**
+ * @license
+ * Copyright 2024 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import type {
+  GraphDescriptor,
+  InspectableModule,
+  InspectableModuleCache,
+  InspectableModules,
+  ModuleCode,
+  ModuleIdentifier,
+  ModuleMetadata,
+  Module as ModuleType,
+} from "@breadboard-ai/types";
+
+class Module implements InspectableModule {
+  #code: ModuleCode | undefined;
+  #metadata: ModuleMetadata | undefined;
+
+  constructor({ code, metadata }: ModuleType) {
+    this.#code = code;
+    this.#metadata = metadata;
+  }
+
+  code(): ModuleCode {
+    return this.#code ?? "";
+  }
+
+  metadata(): ModuleMetadata {
+    return this.#metadata ?? {};
+  }
+}
+
+export class ModuleCache implements InspectableModuleCache {
+  #modules: Record<ModuleIdentifier, InspectableModule> = {};
+
+  get(id: ModuleIdentifier): InspectableModule | undefined {
+    return this.#modules[id];
+  }
+
+  add(id: ModuleIdentifier, module: ModuleType): void {
+    this.#modules[id] = new Module(module);
+  }
+
+  remove(id: ModuleIdentifier): void {
+    if (!this.#modules[id]) {
+      return;
+    }
+
+    delete this.#modules[id];
+  }
+
+  modules(): InspectableModules {
+    return this.#modules ?? {};
+  }
+
+  rebuild(graph: GraphDescriptor) {
+    if (!graph.modules) {
+      return;
+    }
+
+    for (const [id, module] of Object.entries(graph.modules)) {
+      this.add(id, module);
+    }
+  }
+}

package/src/breadboard/engine/inspector/graph/mutable-graph.ts

@@ -0,0 +1,150 @@
+// @ts-nocheck
+/**
+ * @license
+ * Copyright 2024 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { GraphRepresentationImpl } from "../../runtime/legacy.js";
+import type {
+  AffectedNode,
+  GraphDescriptor,
+  GraphIdentifier,
+  GraphRepresentation,
+  InspectableDescriberResultCache,
+  InspectableEdgeCache,
+  InspectableGraphCache,
+  InspectableKitCache,
+  InspectableModuleCache,
+  InspectableNodeCache,
+  InspectablePortCache,
+  KitDescriptor,
+  MainGraphIdentifier,
+  ModuleIdentifier,
+  MutableGraph,
+  MutableGraphStore,
+} from "@breadboard-ai/types";
+import { isImperativeGraph, toDeclarativeGraph } from "@breadboard-ai/utils";
+import { DescribeResultCache } from "./describe-cache.js";
+import { EdgeCache } from "./edge-cache.js";
+import { Edge } from "./edge.js";
+import { UpdateEvent } from "./event.js";
+import { GraphCache } from "./graph-cache.js";
+import { Graph } from "./graph.js";
+import { KitCache } from "./kits.js";
+import { ModuleCache } from "./module.js";
+import { NodeCache } from "./node-cache.js";
+import { NodeDescriberManager } from "./node-describer-manager.js";
+import { Node } from "./node.js";
+import { PortCache } from "./port-cache.js";
+
+export { MutableGraphImpl };
+
+class MutableGraphImpl implements MutableGraph {
+  readonly store: MutableGraphStore;
+  readonly id: MainGraphIdentifier;
+
+  legacyKitMetadata: KitDescriptor | null = null;
+
+  graph!: GraphDescriptor;
+  graphs!: InspectableGraphCache;
+  nodes!: InspectableNodeCache;
+  edges!: InspectableEdgeCache;
+  modules!: InspectableModuleCache;
+  describe!: InspectableDescriberResultCache;
+  kits!: InspectableKitCache;
+  ports!: InspectablePortCache;
+  representation!: GraphRepresentation;
+
+  constructor(graph: GraphDescriptor, store: MutableGraphStore) {
+    this.store = store;
+    this.id = crypto.randomUUID();
+    this.rebuild(graph);
+  }
+
+  update(
+    graph: GraphDescriptor,
+    visualOnly: boolean,
+    affectedNodes: AffectedNode[],
+    affectedModules: ModuleIdentifier[],
+    topologyChange: boolean
+  ): void {
+    // TODO: Handle this a better way?
+    for (const id of affectedModules) {
+      this.modules.remove(id);
+      if (!graph.modules || !graph.modules[id]) {
+        continue;
+      }
+
+      this.modules.add(id, graph.modules[id]);
+
+      // Find any nodes configured to use this module and clear its describer.
+      const runModulesNodes = this.nodes.byType("runModule", "");
+      for (const node of runModulesNodes) {
+        if (
+          node.configuration().$module &&
+          node.configuration().$module === id &&
+          !affectedNodes.find((n) => n.id === node.descriptor.id)
+        ) {
+          affectedNodes.push({
+            id: node.descriptor.id,
+            graphId: "",
+          });
+          visualOnly = false;
+        }
+      }
+    }
+
+    // TODO: Handle removals, etc.
+    if (!visualOnly) {
+      this.describe.update(affectedNodes);
+      this.store.dispatchEvent(
+        new UpdateEvent(this.id, "", "", [], topologyChange)
+      );
+    }
+    this.representation = new GraphRepresentationImpl(graph);
+    this.graph = graph;
+  }
+
+  addSubgraph(subgraph: GraphDescriptor, graphId: GraphIdentifier): void {
+    this.graphs.add(graphId);
+    this.nodes.addSubgraphNodes(subgraph, graphId);
+    this.edges.addSubgraphEdges(subgraph, graphId);
+  }
+
+  removeSubgraph(graphId: GraphIdentifier): void {
+    this.graphs.remove(graphId);
+    this.nodes.removeSubgraphNodes(graphId);
+    this.edges.removeSubgraphEdges(graphId);
+  }
+
+  rebuild(graph: GraphDescriptor) {
+    if (isImperativeGraph(graph)) {
+      graph = toDeclarativeGraph(graph);
+    }
+    this.representation = new GraphRepresentationImpl(graph);
+    this.graph = graph;
+    this.nodes = new NodeCache((descriptor, graphId) => {
+      const graph = graphId ? this.graphs.get(graphId) : this;
+      if (!graph) {
+        throw new Error(
+          `Inspect API Integrity error: unable to find subgraph "${graphId}"`
+        );
+      }
+      return new Node(descriptor, this, graphId);
+    });
+    this.edges = new EdgeCache(
+      (edge, graphId) => new Edge(this, edge, graphId)
+    );
+    this.modules = new ModuleCache();
+    this.describe = new DescribeResultCache(new NodeDescriberManager(this));
+    this.kits = new KitCache(this);
+    this.graphs = new GraphCache((id) => new Graph(id, this));
+    this.ports = new PortCache();
+    this.graphs.rebuild(graph);
+    this.nodes.rebuild(graph);
+    this.edges.rebuild(graph);
+    this.modules.rebuild(graph);
+    this.kits.rebuild(graph);
+  }
+}

package/src/breadboard/engine/inspector/graph/node-cache.ts

@@ -0,0 +1,123 @@
+// @ts-nocheck
+/**
+ * @license
+ * Copyright 2024 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import type {
+  GraphDescriptor,
+  GraphIdentifier,
+  InspectableNode,
+  InspectableNodeCache,
+  NodeDescriptor,
+  NodeIdentifier,
+  NodeTypeIdentifier,
+} from "@breadboard-ai/types";
+import { Node } from "./node.js";
+
+type NodeFactory = (
+  node: NodeDescriptor,
+  graphId: GraphIdentifier
+) => InspectableNode;
+
+export class NodeCache implements InspectableNodeCache {
+  #factory: NodeFactory;
+  #map: Map<GraphIdentifier, Map<NodeIdentifier, InspectableNode>> = new Map();
+  #typeMap: Map<GraphIdentifier, Map<NodeTypeIdentifier, InspectableNode[]>> =
+    new Map();
+
+  constructor(factory: NodeFactory) {
+    this.#factory = factory;
+  }
+
+  rebuild(graph: GraphDescriptor) {
+    graph.nodes.forEach((node) => this.#addNodeInternal(node, ""));
+    Object.entries(graph.graphs || {}).forEach(([graphId, graph]) => {
+      graph.nodes.forEach((node) => this.#addNodeInternal(node, graphId));
+    });
+  }
+
+  addSubgraphNodes(subgraph: GraphDescriptor, graphId: GraphIdentifier): void {
+    subgraph.nodes.forEach((node) => this.#addNodeInternal(node, graphId));
+  }
+
+  removeSubgraphNodes(graphId: GraphIdentifier): void {
+    const subgraph = this.#map.get(graphId);
+    subgraph?.forEach((node) => {
+      (node as Node).setDeleted();
+    });
+    this.#map.delete(graphId);
+  }
+
+  #addNodeInternal(node: NodeDescriptor, graphId: GraphIdentifier) {
+    const graphTypes = getOrCreate(this.#typeMap, graphId, () => new Map());
+    const inspectableNode = this.#factory(node, graphId);
+    const type = node.type;
+    let list = graphTypes.get(type);
+    if (!list) {
+      list = [];
+      graphTypes.set(type, list);
+    }
+    list.push(inspectableNode);
+    const graphNodes = getOrCreate(this.#map, graphId, () => new Map());
+    graphNodes.set(node.id, inspectableNode);
+    return inspectableNode;
+
+    function getOrCreate<K, V>(map: Map<K, V>, key: K, factory: () => V): V {
+      let v = map.get(key);
+      if (v) return v;
+      v = factory();
+      map.set(key, v);
+      return v;
+    }
+  }
+
+  byType(
+    type: NodeTypeIdentifier,
+    graphId: GraphIdentifier
+  ): InspectableNode[] {
+    return this.#typeMap.get(graphId)?.get(type) || [];
+  }
+
+  get(
+    id: NodeIdentifier,
+    graphId: GraphIdentifier
+  ): InspectableNode | undefined {
+    return this.#map.get(graphId)?.get(id);
+  }
+
+  add(node: NodeDescriptor, graphId: GraphIdentifier) {
+    if (!this.#map) {
+      return;
+    }
+    this.#addNodeInternal(node, graphId);
+  }
+
+  remove(id: NodeIdentifier, graphId: GraphIdentifier) {
+    if (!this.#map) {
+      return;
+    }
+    const nodeMap = this.#map.get(graphId);
+    if (!nodeMap) {
+      console.error(
+        `Can't remove node "${id}": graph "${graphId}" was not found`
+      );
+      return;
+    }
+    const node = nodeMap.get(id) as Node;
+    console.assert(node, "Node does not exist in cache.");
+    const type = node!.descriptor.type;
+    const list = this.#typeMap?.get(graphId)?.get(type);
+    if (list) {
+      const index = list.indexOf(node!);
+      list.splice(index, 1);
+    }
+    nodeMap.delete(id);
+    node.setDeleted();
+  }
+
+  nodes(graphId: GraphIdentifier): InspectableNode[] {
+    return Array.from(this.#map.get(graphId)?.values() || []);
+  }
+}

package/src/breadboard/engine/inspector/graph/node-describer-manager.ts

@@ -0,0 +1,279 @@
+// @ts-nocheck
+/**
+ * @license
+ * Copyright 2024 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { createLoader } from "../../loader/index.js";
+import { getHandler, invokeMainDescriber } from "../../runtime/legacy.js";
+import type {
+  DescribeResultCacheArgs,
+  GraphIdentifier,
+  InspectableEdge,
+  MutableGraph,
+  NodeDescriberContext,
+  NodeDescriberFunction,
+  NodeDescriberResult,
+  NodeHandler,
+  NodeIdentifier,
+  NodeTypeDescriberOptions,
+  NodeTypeIdentifier,
+} from "@breadboard-ai/types";
+import { SchemaDiffer } from "@breadboard-ai/utils";
+import { contextFromMutableGraph } from "../graph-store.js";
+import { UpdateEvent } from "./event.js";
+import { GraphDescriptorHandle } from "./graph-descriptor-handle.js";
+import {
+  describeInput,
+  describeOutput,
+  edgesToSchema,
+  EdgeType,
+} from "./schemas.js";
+import {
+  assetsFromGraphDescriptor,
+  envFromGraphDescriptor,
+} from "../../../data/file-system.js";
+
+export { emptyResult, NodeDescriberManager };
+
+function emptyResult(): NodeDescriberResult {
+  return {
+    inputSchema: { type: "object" },
+    outputSchema: { type: "object" },
+  };
+}
+
+class NodeDescriberManager implements DescribeResultCacheArgs {
+  public constructor(public readonly mutable: MutableGraph) {}
+
+  initial(
+    graphId: GraphIdentifier,
+    nodeId: NodeIdentifier
+  ): NodeDescriberResult {
+    const node = this.mutable.nodes.get(nodeId, graphId);
+    if (!node) {
+      return emptyResult();
+    }
+    return NodeDescriberManager.asWired(node.incoming(), node.outgoing());
+  }
+
+  async latest(
+    graphId: GraphIdentifier,
+    nodeId: NodeIdentifier
+  ): Promise<NodeDescriberResult> {
+    const node = this.mutable.nodes.get(nodeId, graphId);
+    if (!node) {
+      return emptyResult();
+    }
+    const result = await this.getLatestDescription(
+      node.descriptor.type,
+      graphId,
+      {
+        incoming: node.incoming(),
+        outgoing: node.outgoing(),
+        inputs: { ...node.configuration() },
+      }
+    );
+    return result;
+  }
+
+  willUpdate(
+    previous: NodeDescriberResult,
+    current: NodeDescriberResult
+  ): void {
+    const inputsDiffer = new SchemaDiffer(
+      previous.inputSchema,
+      current.inputSchema
+    );
+    inputsDiffer.computeDiff();
+
+    const outputsDiffer = new SchemaDiffer(
+      previous.outputSchema,
+      current.outputSchema
+    );
+    outputsDiffer.computeDiff();
+
+    if (
+      inputsDiffer.same() &&
+      outputsDiffer.same() &&
+      sameMetadata(previous, current)
+    ) {
+      return;
+    }
+  }
+
+  updated(graphId: GraphIdentifier, nodeId: NodeIdentifier): void {
+    this.mutable.store.dispatchEvent(
+      new UpdateEvent(this.mutable.id, graphId, nodeId, [], false)
+    );
+  }
+
+  async #getDescriber(
+    type: NodeTypeIdentifier
+  ): Promise<NodeDescriberFunction | undefined> {
+    let handler: NodeHandler | undefined;
+    try {
+      handler = await getHandler(type, contextFromMutableGraph(this.mutable));
+    } catch (e) {
+      console.warn(`Error getting describer for node type ${type}`, e);
+    }
+    if (!handler || !("describe" in handler) || !handler.describe) {
+      return undefined;
+    }
+    return handler.describe;
+  }
+
+  async getLatestDescription(
+    type: NodeTypeIdentifier,
+    graphId: GraphIdentifier,
+    options: NodeTypeDescriberOptions = {}
+  ) {
+    const gettingHandle = GraphDescriptorHandle.create(
+      this.mutable.graph,
+      graphId
+    );
+    if (!gettingHandle.success) {
+      throw new Error(gettingHandle.error);
+    }
+    const handle = gettingHandle.result;
+    // The schema of an input or an output is defined by their
+    // configuration schema or their incoming/outgoing edges.
+    if (type === "input") {
+      if (handle.main()) {
+        if (!this.mutable.store.sandbox) {
+          throw new Error(
+            "Sandbox not supplied, won't be able to describe this graph correctly"
+          );
+        }
+        const result = await invokeMainDescriber(
+          {},
+          this.mutable,
+          handle.graph(),
+          options.inputs!,
+          {},
+          {}
+        );
+        if (result)
+          return describeInput({
+            inputs: {
+              schema: result.inputSchema,
+            },
+            incoming: options?.incoming,
+            outgoing: options?.outgoing,
+          });
+        return describeInput(options);
+      }
+      return describeInput(options);
+    }
+    if (type === "output") {
+      if (handle.main()) {
+        if (!this.mutable.store.sandbox) {
+          throw new Error(
+            "Sandbox not supplied, won't be able to describe this graph correctly"
+          );
+        }
+        const result = await invokeMainDescriber(
+          {},
+          this.mutable,
+          handle.graph(),
+          options.inputs!,
+          {},
+          {}
+        );
+        if (result)
+          return describeOutput({
+            inputs: {
+              schema: result.outputSchema,
+            },
+            incoming: options?.incoming,
+            outgoing: options?.outgoing,
+          });
+        return describeInput(options);
+      }
+      return describeOutput(options);
+    }
+
+    const kits = [...this.mutable.store.kits];
+    const describer = await this.#getDescriber(type);
+    const asWired = NodeDescriberManager.asWired(
+      options.incoming,
+      options.outgoing
+    );
+    if (!describer) {
+      return asWired;
+    }
+    const loader = this.mutable.store.loader || createLoader();
+    const context: NodeDescriberContext = {
+      outerGraph: handle.outerGraph(),
+      loader,
+      kits,
+      sandbox: this.mutable.store.sandbox,
+      graphStore: this.mutable.store,
+      fileSystem: this.mutable.store.fileSystem.createRunFileSystem({
+        graphUrl: handle.outerGraph().url!,
+        env: envFromGraphDescriptor(
+          this.mutable.store.fileSystem.env(),
+          handle.outerGraph()
+        ),
+        assets: assetsFromGraphDescriptor(handle.outerGraph()),
+      }),
+      flags: this.mutable.store.flags,
+      wires: {
+        incoming: Object.fromEntries(
+          (options?.incoming ?? []).map((edge) => [
+            edge.in,
+            {
+              outputPort: {
+                describe: async () => (await edge.outPort()).type.schema,
+              },
+            },
+          ])
+        ),
+        outgoing: Object.fromEntries(
+          (options?.outgoing ?? []).map((edge) => [
+            edge.out,
+            {
+              inputPort: {
+                describe: async () => (await edge.inPort()).type.schema,
+              },
+            },
+          ])
+        ),
+      },
+      asType: !!options?.asType,
+    };
+    if (handle.url()) {
+      context.base = handle.url();
+    }
+    try {
+      return describer(
+        options?.inputs || undefined,
+        asWired.inputSchema,
+        asWired.outputSchema,
+        context
+      );
+    } catch (e) {
+      console.warn(`Error describing node type ${type}`, e);
+      return asWired;
+    }
+  }
+
+  static asWired(
+    incoming: InspectableEdge[] = [],
+    outgoing: InspectableEdge[] = []
+  ) {
+    return {
+      inputSchema: edgesToSchema(EdgeType.In, incoming),
+      outputSchema: edgesToSchema(EdgeType.Out, outgoing),
+    } satisfies NodeDescriberResult;
+  }
+}
+
+function sameMetadata(a: NodeDescriberResult, b: NodeDescriberResult) {
+  if (a.title !== b.title) return false;
+  if (a.description !== b.description) return false;
+  if (a.metadata?.icon !== b.metadata?.icon) return false;
+  // TODO: Compare tags and help.
+  return true;
+}