@ghchinoy/litflow 0.2.8 → 0.3.0

package/src/breadboard/engine/runtime/bubble.ts

@@ -0,0 +1,269 @@
+// @ts-nocheck
+/**
+ * @license
+ * Copyright 2023 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import type {
+  GraphInlineMetadata,
+  InputValues,
+  NodeDescriptor,
+  NodeHandlerContext,
+  NodeValue,
+  OutputValues,
+  RunArguments,
+  Schema,
+  TraversalResult,
+} from "@breadboard-ai/types";
+import { InputStageResult, OutputStageResult, RunResult } from "./run.js";
+
+export const createErrorMessage = (
+  inputName: string | string[],
+  metadata: GraphInlineMetadata = {},
+  required: boolean
+): string => {
+  const boardTitle = metadata.title ?? metadata?.url;
+  const inputNames = typeof inputName === "string" ? [inputName] : inputName;
+  const requiredText = required ? "required " : "";
+  return `Missing ${requiredText}input ${inputNames.map((name) => `"${name}"`).join(", ")}${
+    boardTitle ? ` for board "${boardTitle}".` : "."
+  }`;
+};
+
+export const bubbleUpInputsIfNeeded = async (
+  metadata: GraphInlineMetadata,
+  context: NodeHandlerContext,
+  descriptor: NodeDescriptor,
+  result: TraversalResult,
+  path: number[]
+): Promise<void> => {
+  // If we have no way to bubble up inputs, we just return and not
+  // enforce required inputs.
+  if (!context.requestInput) return;
+
+  const outputs = result.outputs ?? {};
+  const reader = new InputSchemaReader(outputs, result.inputs, path);
+  result.outputs = await reader.read(
+    createBubbleHandler(metadata, context, descriptor)
+  );
+};
+
+export const createBubbleHandler = (
+  metadata: GraphInlineMetadata,
+  context: NodeHandlerContext,
+  descriptor: NodeDescriptor
+) => {
+  return (async (propertiesSchema, path) => {
+    const entries = Object.entries(propertiesSchema.properties || {});
+    const defaultOutputs: OutputValues = {};
+    const propertiesToRequest: [string, Schema][] = [];
+    // Pre-process the entries and prepare to request input.
+    for (const [name, schema] of entries) {
+      const required = propertiesSchema.required?.includes(name) ?? false;
+
+      if (required) {
+        throw new Error(createErrorMessage(name, metadata, required));
+      }
+      if (schema.default !== undefined) {
+        if ("type" in schema && schema.type !== "string") {
+          defaultOutputs[name] = JSON.parse(schema.default);
+        } else {
+          defaultOutputs[name] = schema.default;
+        }
+        continue;
+      }
+      propertiesToRequest.push([name, schema]);
+    }
+    // Exit early if we already have all properties.
+    if (propertiesToRequest.length === 0) {
+      return defaultOutputs;
+    }
+    // Request properties that did not have default values.
+    const inputRequestSchema: Schema = {
+      properties: Object.fromEntries(propertiesToRequest),
+    };
+    const outputs = await context.requestInput?.(
+      inputRequestSchema,
+      descriptor,
+      path
+    );
+    // If the run was aborted, let's bail
+    if (context?.signal?.aborted) {
+      throw context.signal.throwIfAborted();
+    }
+    // Finally, let's make sure we have all values and if not, throw an error.
+    if (!outputs || Object.keys(outputs).length === 0) {
+      throw new Error(
+        createErrorMessage(
+          propertiesToRequest.map(([name]) => name),
+          metadata,
+          false
+        )
+      );
+    }
+    return { ...defaultOutputs, ...outputs };
+  }) satisfies InputSchemaHandler;
+};
+
+export type InputSchemaHandler = (
+  schema: Schema,
+  path: number[]
+) => Promise<OutputValues>;
+
+export class InputSchemaReader {
+  #currentOutputs: OutputValues;
+  #inputs: InputValues;
+  #path: number[];
+
+  constructor(
+    currentOutputs: OutputValues,
+    inputs: InputValues,
+    path: number[]
+  ) {
+    this.#currentOutputs = currentOutputs;
+    this.#inputs = inputs;
+    this.#path = path;
+  }
+
+  async read(handler: InputSchemaHandler): Promise<OutputValues> {
+    if (!("schema" in this.#inputs)) return this.#currentOutputs;
+
+    const schema = this.#inputs.schema as Schema;
+
+    if (!schema.properties) return this.#currentOutputs;
+
+    const unfulfilled = structuredClone(schema);
+    for (const name of Object.keys(schema.properties)) {
+      if (name in this.#currentOutputs) {
+        delete unfulfilled.properties?.[name];
+      }
+    }
+
+    const willAsk = Object.keys(unfulfilled.properties || {}).length > 0;
+
+    const newOutputs = willAsk ? await handler(unfulfilled, this.#path) : {};
+
+    return { ...this.#currentOutputs, ...newOutputs };
+  }
+}
+
+export class RequestedInputsManager {
+  #context: NodeHandlerContext;
+  #cache: Map<string, NodeValue> = new Map();
+
+  constructor(args: RunArguments) {
+    const { inputs, ...context } = args;
+    this.#context = context;
+    this.#cache = new Map(inputs ? Object.entries(inputs) : []);
+  }
+
+  createHandler(
+    next: (result: RunResult) => Promise<void>,
+    result: TraversalResult
+  ): NodeHandlerContext["requestInput"] {
+    return async (schema: Schema, node: NodeDescriptor, path: number[]) => {
+      // Retrieve all cached values
+      const propertiesToRequest = Object.entries(schema.properties || {});
+      const cachedValues: OutputValues = {};
+      const uncachedProperties = propertiesToRequest.filter(([name]) => {
+        const cachedValue = this.#cache.get(name);
+        if (cachedValue !== undefined) {
+          cachedValues[name] = cachedValue;
+          return false;
+        }
+        return true;
+      });
+      // Early return when all properties are cached
+      if (uncachedProperties.length === 0) return cachedValues;
+
+      const configuration = node.configuration?.schema
+        ? {
+            configuration: { schema: node.configuration.schema },
+          }
+        : {};
+      const descriptor = { id: node.id, type: node.type, ...configuration };
+      const requestInputResult = {
+        ...result,
+        descriptor,
+        inputs: {
+          schema: {
+            type: "object",
+            properties: Object.fromEntries(uncachedProperties),
+          } satisfies Schema,
+        },
+      };
+      await next(new InputStageResult(requestInputResult, -1, path));
+      const outputs = requestInputResult.outputs;
+      const requestedProperties = schema.properties || {};
+      const remainingProperties = outputs
+        ? Object.fromEntries(
+            Object.entries(requestedProperties).filter(([name]) => {
+              return !(name in outputs);
+            })
+          )
+        : requestedProperties;
+      if (Object.keys(remainingProperties).length === 0) {
+        return outputs;
+      }
+      // Bubble up: request outer context to request input
+      const bubbledOutputs: OutputValues =
+        (await this.#context.requestInput?.(
+          { properties: remainingProperties },
+          descriptor,
+          path
+        )) || {};
+      // Cache all non-transient properties.
+      for (const [name, propertySchema] of uncachedProperties) {
+        if (!isTransient(propertySchema)) {
+          const value = bubbledOutputs[name];
+          if (value) {
+            this.#cache.set(name, bubbledOutputs[name]);
+          }
+        }
+      }
+      return { ...outputs, ...bubbledOutputs };
+    };
+  }
+}
+
+function isTransient(schema: Schema): boolean {
+  return schema.behavior?.includes("transient") ?? false;
+}
+
+export const bubbleUpOutputsIfNeeded = async (
+  outputs: OutputValues,
+  descriptor: NodeDescriptor,
+  context: NodeHandlerContext,
+  path: number[]
+): Promise<boolean> => {
+  if (!context.provideOutput) return false;
+  const schema = descriptor.configuration?.schema as Schema;
+  const shouldBubble = schema?.behavior?.includes("bubble");
+  if (!shouldBubble) return false;
+
+  await context.provideOutput(outputs, descriptor, path);
+  return true;
+};
+
+export const createOutputProvider = (
+  next: (result: RunResult) => Promise<void>,
+  result: TraversalResult,
+  context: NodeHandlerContext
+) => {
+  if (context.provideOutput) {
+    return context.provideOutput;
+  }
+  return async (
+    outputs: OutputValues,
+    descriptor: NodeDescriptor,
+    path: number[]
+  ) => {
+    const provideOutputResult = {
+      ...result,
+      descriptor,
+      inputs: outputs,
+    };
+    await next(new OutputStageResult(provideOutputResult, -1, path));
+  };
+};

package/src/breadboard/engine/runtime/graph-based-node-handler.ts

@@ -0,0 +1,174 @@
+// @ts-nocheck
+/**
+ * @license
+ * Copyright 2024 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import {
+  GraphToRun,
+  InputValues,
+  MutableGraph,
+  NodeDescriberContext,
+  NodeDescriberResult,
+  NodeHandlerContext,
+  NodeHandlerMetadata,
+  NodeHandlerObject,
+  NodeTypeIdentifier,
+  Schema,
+} from "@breadboard-ai/types";
+import { emptyDescriberResult, filterEmptyValues } from "@breadboard-ai/utils";
+import { invokeGraph } from "./run/invoke-graph.js";
+import { CapabilitiesManagerImpl } from "./sandbox/capabilities-manager.js";
+import { invokeDescriber } from "./sandbox/invoke-describer.js";
+
+export {
+  describerResultToNodeHandlerMetadata,
+  GraphBasedNodeHandler,
+  toNodeHandlerMetadata,
+};
+
+class GraphBasedNodeHandler implements NodeHandlerObject {
+  #graph: GraphToRun;
+  #type: NodeTypeIdentifier;
+
+  constructor(graph: GraphToRun, type: NodeTypeIdentifier) {
+    this.#graph = graph;
+    this.#type = type;
+    this.invoke = this.invoke.bind(this);
+    this.describe = this.describe.bind(this);
+  }
+
+  async invoke(inputs: InputValues, context: NodeHandlerContext) {
+    const base = context.board?.url && new URL(context.board?.url);
+    const invocationContext = base
+      ? {
+          ...context,
+          base,
+        }
+      : { ...context };
+
+    return await invokeGraph(this.#graph, inputs, invocationContext);
+  }
+
+  async describe(
+    inputs?: InputValues,
+    inputSchema?: Schema,
+    outputSchema?: Schema,
+    context?: NodeDescriberContext
+  ) {
+    if (!context) {
+      return { inputSchema: {}, outputSchema: {} };
+    }
+
+    const graphStore = context?.graphStore;
+    if (!graphStore) {
+      return emptyDescriberResult();
+    }
+    const url = this.#graph.graph.url;
+    let mutable: MutableGraph;
+    if (url) {
+      // In the most common case, just use the URL of the graph.
+      const adding = graphStore.addByURL(url, [], context);
+      mutable = adding.mutable;
+    } else {
+      const adding = graphStore.addByDescriptor(this.#graph.graph);
+      if (!adding.success) {
+        return emptyDescriberResult();
+      }
+      mutable = graphStore.get(adding.result)!;
+    }
+    mutable = await graphStore.getLatest(mutable);
+    if (this.#graph.moduleId) {
+      const moduleId = this.#graph.moduleId;
+      const graph = this.#graph.graph;
+      const sandbox = context.graphStore?.sandbox;
+      if (!sandbox) {
+        console.warn("Sandbox was not supplied to describe node type");
+        return emptyDescriberResult();
+      }
+      const result = await invokeDescriber(
+        context,
+        moduleId,
+        mutable,
+        graph,
+        inputs || {},
+        inputSchema,
+        outputSchema,
+        new CapabilitiesManagerImpl(context),
+        context.asType || false
+      );
+      if (!result) {
+        return emptyDescriberResult();
+      }
+      return result;
+    } else {
+      const inspectableGraph = graphStore.inspect(
+        mutable.id,
+        this.#graph.subGraphId || ""
+      );
+      if (!inspectableGraph) {
+        return emptyDescriberResult();
+      }
+      const result = await inspectableGraph.describe(inputs, context);
+      return result;
+    }
+  }
+
+  get metadata() {
+    return toNodeHandlerMetadata(this.#graph, this.#type, false);
+  }
+}
+
+function describerResultToNodeHandlerMetadata(
+  result: NodeDescriberResult,
+  updating: boolean
+): NodeHandlerMetadata {
+  return filterEmptyValues({
+    title: result.title,
+    description: result.description,
+    url: result.url,
+    icon: result.metadata?.icon,
+    help: result.metadata?.help,
+    updating,
+  });
+}
+
+function toNodeHandlerMetadata(
+  graphToRun: GraphToRun,
+  url: NodeTypeIdentifier,
+  updating: boolean
+): NodeHandlerMetadata | undefined {
+  const graph = graphToRun.graph;
+  if (graphToRun.moduleId) {
+    const module = graph.modules?.[graphToRun.moduleId];
+    if (!module) return undefined;
+    const {
+      title = graphToRun.moduleId,
+      description,
+      icon,
+      help,
+    } = module.metadata || {};
+    return filterEmptyValues({ title, description, url, icon, help, updating });
+  } else if (graphToRun.subGraphId) {
+    const descriptor = graph.graphs?.[graphToRun.subGraphId];
+    if (!descriptor) return undefined;
+    return filterEmptyValues({
+      title: descriptor.title,
+      description: descriptor.description,
+      url,
+      icon: descriptor.metadata?.icon,
+      help: descriptor.metadata?.help,
+      updating,
+    });
+  } else {
+    return filterEmptyValues({
+      title: graph.title,
+      description: graph.description,
+      url: graph.url,
+      icon: graph.metadata?.icon,
+      help: graph.metadata?.help,
+      updating,
+    });
+  }
+}

package/src/breadboard/engine/runtime/handler.ts

@@ -0,0 +1,166 @@
+// @ts-nocheck
+/**
+ * @license
+ * Copyright 2023 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import type {
+  InputValues,
+  Kit,
+  MutableGraph,
+  NodeHandler,
+  NodeHandlerContext,
+  NodeHandlerFunction,
+  NodeHandlerObject,
+  NodeHandlers,
+  NodeTypeIdentifier,
+  OutputValues,
+} from "@breadboard-ai/types";
+import { graphUrlLike } from "@breadboard-ai/utils";
+import { GraphBasedNodeHandler } from "./graph-based-node-handler.js";
+import { getGraphUrl } from "../loader/loader.js";
+
+// TODO: Deduplicate.
+function contextFromMutableGraph(mutable: MutableGraph): NodeHandlerContext {
+  const store = mutable.store;
+  return {
+    kits: [...store.kits],
+    loader: store.loader,
+    sandbox: store.sandbox,
+    graphStore: store,
+    outerGraph: mutable.graph,
+  };
+}
+
+const getHandlerFunction = (handler: NodeHandler): NodeHandlerFunction => {
+  if ("invoke" in handler && handler.invoke) return handler.invoke;
+  if (handler instanceof Function) return handler;
+  throw new Error("Invalid handler");
+};
+
+export const callHandler = async (
+  handler: NodeHandler,
+  inputs: InputValues,
+  context: NodeHandlerContext
+): Promise<OutputValues | void> => {
+  // if (handler instanceof Function) return handler(inputs, context);
+  // if (handler.invoke) return handler.invoke(inputs, context);
+  const handlerFunction = getHandlerFunction(handler);
+  return new Promise((resolve) => {
+    handlerFunction(inputs, context)
+      .then(resolve)
+      .catch((error) => {
+        resolve({ $error: { error } });
+      });
+  });
+};
+
+export const handlersFromKits = (kits: Kit[]): NodeHandlers => {
+  return kits.reduce((handlers, kit) => {
+    // If multiple kits have the same handler, the kit earlier in the list
+    // gets precedence, including upstream kits getting precedence over kits
+    // defined in the graph.
+    //
+    // TODO: This means kits are fallback, consider non-fallback as well.
+    return { ...kit.handlers, ...handlers };
+  }, {} as NodeHandlers);
+};
+
+/**
+ * The single entry point for getting a handler for a node type.
+ * The handler can be one of the two types:
+ * - A graph-based handler, where the `type` is actually a
+ *   URL-like string that points to a graph.
+ * - A kit-based handler, where the `type` is a string that
+ *   corresponds to a node type in a kit.
+ *
+ * The function throws an error if no handler is found for the
+ * given node type.
+ *
+ * @param type    -- The node type to get a handler for.
+ * @param context -- The context in which the handler is
+ *                   being requested.
+ * @returns       -- The handler for the node type.
+ */
+export async function getHandler(
+  type: NodeTypeIdentifier,
+  context: NodeHandlerContext
+): Promise<NodeHandler> {
+  if (graphUrlLike(type)) {
+    const graphHandler = await getGraphHandler(type, context);
+    if (graphHandler) {
+      return graphHandler;
+    }
+  }
+  const handlers = handlersFromKits(context.kits ?? []);
+  const kitHandler = handlers[type];
+  if (kitHandler) {
+    return kitHandler;
+  }
+  throw new Error(`No handler for node type "${type}"`);
+}
+
+export async function getGraphHandlerFromMutableGraph(
+  type: NodeTypeIdentifier,
+  mutable: MutableGraph
+): Promise<NodeHandlerObject | undefined> {
+  const nodeTypeUrl = graphUrlLike(type)
+    ? getGraphUrl(type, contextFromMutableGraph(mutable))
+    : undefined;
+  if (!nodeTypeUrl) {
+    return undefined;
+  }
+  const store = mutable.store;
+  const result = store.addByURL(type, [], {
+    outerGraph: mutable.graph,
+  });
+  const latest = await store.getLatest(result.mutable);
+  return new GraphBasedNodeHandler(
+    {
+      graph: latest.graph,
+      subGraphId: result.graphId,
+      moduleId: result.moduleId,
+    },
+    type
+  );
+}
+
+export async function getGraphHandler(
+  type: NodeTypeIdentifier,
+  context: NodeHandlerContext,
+  allow3PModules = false
+): Promise<NodeHandlerObject | undefined> {
+  if (is3pModule(type) && !allow3PModules) {
+    return undefined;
+  }
+  const nodeTypeUrl = graphUrlLike(type)
+    ? getGraphUrl(type, context)
+    : undefined;
+  if (!nodeTypeUrl) {
+    return undefined;
+  }
+  const { graphStore } = context;
+  if (!graphStore) {
+    throw new Error(
+      `Cannot load graph for type "${type}" without a graph store.`
+    );
+  }
+
+  const loadResult = await graphStore.load(type, context);
+  if (!loadResult.success) {
+    throw new Error(
+      `Loading graph for type "${type}" failed: ${loadResult.error}`
+    );
+  }
+  return new GraphBasedNodeHandler(loadResult, type);
+}
+
+/**
+ * This is a somewhat hacky module invocation filter.
+ * TODO: Instead of hard-coding the type, plumb the invocationFilter from
+ * the configuration here.
+ */
+function is3pModule(type: NodeTypeIdentifier) {
+  return !type.startsWith("embed://");
+}

package/src/breadboard/engine/runtime/harness/diagnostics.ts

@@ -0,0 +1,22 @@
+// @ts-nocheck
+/**
+ * @license
+ * Copyright 2023 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import type { Probe, ProbeMessage } from "@breadboard-ai/types";
+
+export type DiagnosticsCallback = (message: ProbeMessage) => Promise<void>;
+
+export class Diagnostics implements Probe {
+  #callback: DiagnosticsCallback;
+
+  constructor(callback: DiagnosticsCallback) {
+    this.#callback = callback;
+  }
+
+  async report(message: ProbeMessage): Promise<void> {
+    return this.#callback(message);
+  }
+}