[build] Serialize boards as subgraphs when passed as node inputs (#2185)

Input ports with the "board" behavior can now receive board objects declared using the build API. This will be serialized as an embedded graph.
breadboard-ai · Jun 18, 2024 · 6ada218 · 6ada218
1 parent e7bdeaa
commit 6ada218
Show file tree

Hide file tree

Showing 11 changed files with 345 additions and 23 deletions.
diff --git a/.changeset/silly-baboons-smell.md b/.changeset/silly-baboons-smell.md
@@ -0,0 +1,5 @@
+---
+"@breadboard-ai/build": minor
+---
+
+Input ports with the "board" behavior can now receive board objects declared using the build API. This will be serialized as an embedded graph.
diff --git a/packages/build/src/internal/board/board.ts b/packages/build/src/internal/board/board.ts
@@ -71,6 +71,7 @@ export function board<
     description,
     version,
     metadata,
+    isBoard: true,
   });
 }
 
@@ -264,3 +265,9 @@ type ExtractPortTypes<PORTS extends BoardInputPorts | BoardOutputPorts> = {
     ? TYPE
     : never;
 };
+
+export function isBoard(value: unknown): value is GenericBoardDefinition {
+  return (
+    typeof value === "function" && "isBoard" in value && value.isBoard === true
+  );
+}
diff --git a/packages/build/src/internal/board/serialize.ts b/packages/build/src/internal/board/serialize.ts
@@ -23,8 +23,9 @@ import { ConstantVersionOf, isConstant } from "./constant.js";
 import { isConvergence } from "./converge.js";
 import type { GenericSpecialInput, Input, InputWithDefault } from "./input.js";
 import { isLoopback } from "./loopback.js";
-import { isOptional, OptionalVersionOf } from "./optional.js";
+import { OptionalVersionOf, isOptional } from "./optional.js";
 import type { Output } from "./output.js";
+import { isBoard, type GenericBoardDefinition } from "./board.js";
 
 /**
  * Serialize a Breadboard board to Breadboard Graph Language (BGL) so that it
@@ -33,8 +34,10 @@ import type { Output } from "./output.js";
 export function serialize(board: SerializableBoard): GraphDescriptor {
   const nodes = new Map<object, NodeDescriptor>();
   const edges: Edge[] = [];
+  const graphs = new Map<string, GraphDescriptor>();
   const errors: string[] = [];
   const typeCounts = new Map<string, number>();
+  let nextEmbeddedGraphId = 0;
 
   // Prepare our main input and output nodes. They represent the overall
   // signature of the board.
@@ -275,7 +278,7 @@ export function serialize(board: SerializableBoard): GraphDescriptor {
     );
   }
 
-  return {
+  const bgl: GraphDescriptor = {
     ...(board.title ? { title: board.title } : {}),
     ...(board.description ? { description: board.description } : {}),
     ...(board.version ? { version: board.version } : {}),
@@ -302,6 +305,10 @@ export function serialize(board: SerializableBoard): GraphDescriptor {
       ...[...nodes.values()].sort((a, b) => a.id.localeCompare(b.id)),
     ],
   };
+  if (graphs.size > 0) {
+    bgl.graphs = Object.fromEntries([...graphs]);
+  }
+  return bgl;
 
   function visitNodeAndReturnItsId(node: SerializableNode): string {
     let descriptor = nodes.get(node);
@@ -409,6 +416,14 @@ export function serialize(board: SerializableBoard): GraphDescriptor {
           throw new Error(
             `Internal error: value was a symbol (${String(value)}) for a ${inputPort.node.type}:${inputPort.name} port.`
           );
+        } else if (isBoard(value)) {
+          configurationEntries.push([
+            portName,
+            {
+              kind: "board",
+              path: `#${embedBoardAndReturnItsId(value)}`,
+            },
+          ]);
         } else {
           configurationEntries.push([
             portName,
@@ -450,6 +465,13 @@ export function serialize(board: SerializableBoard): GraphDescriptor {
     typeCounts.set(type, count + 1);
     return `${type}-${count}`;
   }
+
+  function embedBoardAndReturnItsId(board: GenericBoardDefinition): string {
+    const id = `subgraph-${nextEmbeddedGraphId}`;
+    nextEmbeddedGraphId++;
+    graphs.set(id, serialize(board));
+    return id;
+  }
 }
 
 function isSpecialInput(value: unknown): value is GenericSpecialInput {

diff --git a/packages/build/src/internal/common/serializable.ts b/packages/build/src/internal/common/serializable.ts
@@ -4,17 +4,18 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import type { GraphMetadata } from "@google-labs/breadboard-schema/graph.js";
+import type { GenericBoardDefinition } from "../board/board.js";
 import type { Convergence } from "../board/converge.js";
 import type {
   GenericSpecialInput,
   Input,
   InputWithDefault,
 } from "../board/input.js";
-import type { Output } from "../board/output.js";
 import type { Loopback } from "../board/loopback.js";
+import type { Output } from "../board/output.js";
 import type { BreadboardType, JsonSerializable } from "../type-system/type.js";
 import type { DefaultValue, OutputPortGetter } from "./port.js";
-import type { GraphMetadata } from "@google-labs/breadboard-schema/graph.js";
 
 export interface SerializableBoard {
   inputs: Record<
@@ -71,6 +72,7 @@ export interface SerializableInputPort {
     | GenericSpecialInput
     | Loopback<JsonSerializable>
     | Convergence<JsonSerializable>
+    | GenericBoardDefinition
     | typeof DefaultValue;
 }
 

diff --git a/packages/build/src/internal/define/define.ts b/packages/build/src/internal/define/define.ts
@@ -156,7 +156,8 @@ export function defineNodeType<
   GetOptionalInputs<I> & keyof Expand<GetStaticTypes<I>>,
   GetReflective<O>,
   Expand<GetPrimary<I>>,
-  Expand<GetPrimary<O>>
+  Expand<GetPrimary<O>>,
+  Expand<ExtractInputMetadata<I>>
 > {
   if (!params.name) {
     throw new Error("params.name is required");
@@ -178,7 +179,8 @@ export function defineNodeType<
     GetOptionalInputs<I> & keyof Expand<GetStaticTypes<I>>,
     GetReflective<O>,
     Expand<GetPrimary<I>>,
-    Expand<GetPrimary<O>>
+    Expand<GetPrimary<O>>,
+    Expand<ExtractInputMetadata<I>>
   >(
     params.name,
     omitDynamic(params.inputs),
@@ -194,9 +196,30 @@ export function defineNodeType<
     invoke: impl.invoke.bind(impl),
     describe: impl.describe.bind(impl),
     metadata: params.metadata || {},
-  });
+    // TODO(aomarks) Should not need cast.
+  }) as Definition<
+    Expand<GetStaticTypes<I>>,
+    Expand<GetStaticTypes<O>>,
+    GetDynamicTypes<I>,
+    GetDynamicTypes<O>,
+    GetOptionalInputs<I> & keyof Expand<GetStaticTypes<I>>,
+    GetReflective<O>,
+    Expand<GetPrimary<I>>,
+    Expand<GetPrimary<O>>,
+    Expand<ExtractInputMetadata<I>>
+  >;
 }
 
+type ExtractInputMetadata<I extends Record<string, InputPortConfig>> = {
+  [K in keyof I as K extends "*" ? never : K]: {
+    board: I[K]["behavior"] extends Array<unknown>
+      ? "board" extends I[K]["behavior"][number]
+        ? true
+        : false
+      : false;
+  };
+};
+
 function omitDynamic(configs: PortConfigs): PortConfigs {
   return Object.fromEntries(
     Object.entries(configs).filter(([name]) => name !== "*")

diff --git a/packages/build/src/internal/define/definition.ts b/packages/build/src/internal/define/definition.ts
@@ -51,6 +51,7 @@ import { array } from "../type-system/array.js";
 import { object } from "../type-system/object.js";
 import { normalizeBreadboardError } from "../common/error.js";
 import type { Convergence } from "../board/converge.js";
+import type { BoardDefinition } from "../board/board.js";
 
 export interface Definition<
   /* Static Inputs   */ SI extends { [K: string]: JsonSerializable },
@@ -61,9 +62,10 @@ export interface Definition<
   /* Reflective?     */ R extends boolean,
   /* Primary Input   */ PI extends string | false,
   /* Primary Output  */ PO extends string | false,
+  /* Input Metadata  */ IM extends { [K: string]: InputMetadata },
 > extends StrictNodeHandler {
   <A extends LooseInstantiateArgs>(
-    args: A & StrictInstantiateArgs<SI, OI, DI, A>
+    args: A & StrictInstantiateArgs<SI, OI, DI, A, IM>
   ): Instance<
     InstanceInputs<SI, DI, A>,
     InstanceOutputs<SI, SO, DO, R, A>,
@@ -83,6 +85,7 @@ export class DefinitionImpl<
   /* Reflective?     */ R extends boolean,
   /* Primary Input   */ PI extends string | false,
   /* Primary Output  */ PO extends string | false,
+  /* Input Metadata  */ IM extends { [K: string]: InputMetadata },
 > implements StrictNodeHandler
 {
   readonly #name: string;
@@ -135,7 +138,7 @@ export class DefinitionImpl<
   }
 
   instantiate<A extends LooseInstantiateArgs>(
-    args: A & StrictInstantiateArgs<SI, OI, DI, A>
+    args: A & StrictInstantiateArgs<SI, OI, DI, A, IM>
   ): Instance<
     InstanceInputs<SI, DI, A>,
     InstanceOutputs<SI, SO, DO, R, A>,
@@ -396,6 +399,14 @@ export class DefinitionImpl<
   }
 }
 
+/**
+ * Extra data about inputs. This is here to stop growing the number of
+ * parameters on Definition.
+ */
+export type InputMetadata = {
+  board: boolean;
+};
+
 function parseDynamicPorts(
   ports: Exclude<CustomDescribePortManifest, UnsafeSchema>,
   base: DynamicInputPortConfig | DynamicOutputPortConfig
@@ -426,14 +437,20 @@ type StrictInstantiateArgs<
   OI extends keyof SI,
   DI extends JsonSerializable | undefined,
   A extends LooseInstantiateArgs,
+  IM extends { [K: string]: InputMetadata },
 > = {
   $id?: string;
   $metadata?: {
     title?: string;
     description?: string;
   };
 } & {
-  [K in keyof Omit<SI, OI | "$id" | "$metadata">]: InstantiateArg<SI[K]>;
+  [K in keyof Omit<SI, OI | "$id" | "$metadata">]: IM[K extends keyof IM
+    ? K
+    : never]["board"] extends true
+    ? // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      InstantiateArg<SI[K]> | BoardDefinition<any, any>
+    : InstantiateArg<SI[K]>;
 } & {
   [K in OI]?:
     | InstantiateArg<SI[K]>

diff --git a/packages/build/src/internal/define/node-factory.ts b/packages/build/src/internal/define/node-factory.ts
@@ -18,7 +18,7 @@ import type { Expand } from "../common/type-util.js";
  * for use with {@link KitBuilder}.
  */
 export type NodeFactoryFromDefinition<
-  D extends Definition<any, any, any, any, any, any, any, any>,
+  D extends Definition<any, any, any, any, any, any, any, any, any>,
 > =
   D extends Definition<
     infer SI,
@@ -28,6 +28,7 @@ export type NodeFactoryFromDefinition<
     infer OI,
     any,
     any,
+    any,
     any
   >
     ? NewNodeFactory<

diff --git a/packages/build/src/test/compatibility_test.ts b/packages/build/src/test/compatibility_test.ts
@@ -26,7 +26,7 @@ function setupKits<
     // TODO(aomarks) See TODO about `any` at {@link NodeFactoryFromDefinition}.
     //
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    Definition<any, any, any, any, any, any, any, any>
+    Definition<any, any, any, any, any, any, any, any, any>
   >,
 >(definitions: DEFS) {
   const ctr = new KitBuilder({ url: "N/A" }).build(definitions);

diff --git a/packages/build/src/test/define_test.ts b/packages/build/src/test/define_test.ts
@@ -19,7 +19,7 @@ import type { BreadboardError } from "../internal/common/error.js";
 test("mono/mono", async () => {
   const values = { si1: "foo", si2: 123 };
 
-  // $ExpectType Definition<{ si1: string; si2: number; }, { so1: boolean; so2: null; }, undefined, undefined, never, false, false, false>
+  // $ExpectType Definition<{ si1: string; si2: number; }, { so1: boolean; so2: null; }, undefined, undefined, never, false, false, false, { si1: { board: false; }; si2: { board: false; }; }>
   const d = defineNodeType({
     name: "foo",
     inputs: {
@@ -139,7 +139,7 @@ test("mono/mono", async () => {
 test("poly/mono", async () => {
   const values = { si1: "si1", di1: 1, di2: 2 };
 
-  // $ExpectType Definition<{ si1: string; }, { so1: boolean; }, number, undefined, never, false, false, false>
+  // $ExpectType Definition<{ si1: string; }, { so1: boolean; }, number, undefined, never, false, false, false, { si1: { board: false; }; }>
   const d = defineNodeType({
     name: "foo",
     inputs: {
@@ -287,7 +287,7 @@ test("poly/mono", async () => {
 test("mono/poly", async () => {
   const values = { si1: "si1" };
 
-  // $ExpectType Definition<{ si1: string; }, { so1: boolean; }, undefined, number, never, false, false, false>
+  // $ExpectType Definition<{ si1: string; }, { so1: boolean; }, undefined, number, never, false, false, false, { si1: { board: false; }; }>
   const d = defineNodeType({
     name: "foo",
     inputs: {
@@ -390,7 +390,7 @@ test("mono/poly", async () => {
 test("poly/poly", async () => {
   const values = { si1: "si1", di1: 1, di2: 2 };
 
-  // $ExpectType Definition<{ si1: string; }, { so1: boolean; }, number, number, never, false, false, false>
+  // $ExpectType Definition<{ si1: string; }, { so1: boolean; }, number, number, never, false, false, false, { si1: { board: false; }; }>
   const d = defineNodeType({
     name: "foo",
     inputs: {
@@ -573,7 +573,7 @@ test("async invoke function", async () => {
 test("reflective", async () => {
   const values = { si1: "si1", di1: 1, di2: 2 };
 
-  // $ExpectType Definition<{ si1: string; }, { so1: boolean; }, number, string, never, true, false, false>
+  // $ExpectType Definition<{ si1: string; }, { so1: boolean; }, number, string, never, true, false, false, { si1: { board: false; }; }>
   const d = defineNodeType({
     name: "foo",
     inputs: {
@@ -708,7 +708,7 @@ test("reflective", async () => {
 test("primary input with no other inputs", () => {
   const values = { si1: 123 };
 
-  // $ExpectType Definition<{ si1: number; }, { so1: boolean; }, undefined, undefined, never, false, "si1", false>
+  // $ExpectType Definition<{ si1: number; }, { so1: boolean; }, undefined, undefined, never, false, "si1", false, { si1: { board: false; }; }>
   const d = defineNodeType({
     name: "foo",
     inputs: {
@@ -772,7 +772,7 @@ test("primary input with no other inputs", () => {
 test("primary input with another input", () => {
   const values = { si1: 123, si2: true };
 
-  // $ExpectType Definition<{ si1: number; si2: boolean; }, { so1: boolean; }, undefined, undefined, never, false, "si1", false>
+  // $ExpectType Definition<{ si1: number; si2: boolean; }, { so1: boolean; }, undefined, undefined, never, false, "si1", false, { si1: { board: false; }; si2: { board: false; }; }>
   const d = defineNodeType({
     name: "foo",
     inputs: {
@@ -835,7 +835,7 @@ test("primary input with another input", () => {
 });
 
 test("primary output with no other outputs", () => {
-  // $ExpectType Definition<{ si1: number; }, { so1: boolean; }, undefined, undefined, never, false, false, "so1">
+  // $ExpectType Definition<{ si1: number; }, { so1: boolean; }, undefined, undefined, never, false, false, "so1", { si1: { board: false; }; }>
   const d = defineNodeType({
     name: "foo",
     inputs: {
@@ -897,7 +897,7 @@ test("primary output with no other outputs", () => {
 });
 
 test("primary output with other outputs", () => {
-  // $ExpectType Definition<{ si1: number; }, { so1: boolean; so2: number; }, undefined, undefined, never, false, false, "so1">
+  // $ExpectType Definition<{ si1: number; }, { so1: boolean; so2: number; }, undefined, undefined, never, false, false, "so1", { si1: { board: false; }; }>
   const d = defineNodeType({
     name: "foo",
     inputs: {
@@ -961,7 +961,7 @@ test("primary output with other outputs", () => {
 });
 
 test("primary input + output", () => {
-  // $ExpectType Definition<{ si1: number; }, { so1: boolean; }, undefined, undefined, never, false, "si1", "so1">
+  // $ExpectType Definition<{ si1: number; }, { so1: boolean; }, undefined, undefined, never, false, "si1", "so1", { si1: { board: false; }; }>
   const d = defineNodeType({
     name: "foo",
     inputs: {