Skip to content

Commit

Permalink
[bbrt] Support running breadboard components directly (#4118)
Browse files Browse the repository at this point in the history
Regular breadboard components (e.g. GeminiKit, GoogleDriveKit) can now
be executed directly by the main conversation. (Boards from the board
server could already be executed).
  • Loading branch information
aomarks authored Jan 13, 2025
1 parent 8001cd3 commit 3420dc6
Show file tree
Hide file tree
Showing 6 changed files with 200 additions and 21 deletions.
1 change: 1 addition & 0 deletions package-lock.json

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

3 changes: 3 additions & 0 deletions packages/bbrt/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,7 @@
"../breadboard:build",
"../core-kit:build",
"../gemini-kit:build",
"../google-drive-kit:build",
"../json-kit:build",
"../shared-ui:build",
"../template-kit:build"
Expand All @@ -70,6 +71,7 @@
"command": "tsc --pretty",
"dependencies": [
"../agent-kit:build",
"../google-drive-kit:build",
"../breadboard:build:tsc",
"../core-kit:build:tsc",
"../gemini-kit:build:tsc",
Expand Down Expand Up @@ -125,6 +127,7 @@
}
},
"dependencies": {
"@breadboard-ai/google-drive-kit": "0.4.0",
"@breadboard-ai/shared-ui": "1.21.0",
"@google-labs/agent-kit": "^0.14.0",
"@google-labs/breadboard": "^0.30.0",
Expand Down
146 changes: 146 additions & 0 deletions packages/bbrt/src/breadboard/breadboard-component-tool.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,146 @@
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/

import {
createGraphStore,
createLoader,
type GraphDescriptor,
type Kit,
type NodeDescriberFunction,
type NodeDescriptor,
type NodeHandler,
} from "@google-labs/breadboard";
import type { ArtifactHandle } from "../artifacts/artifact-interface.js";
import type { ArtifactStore } from "../artifacts/artifact-store.js";
import type { SecretsProvider } from "../secrets/secrets-provider.js";
import type {
BBRTTool,
BBRTToolAPI,
BBRTToolMetadata,
} from "../tools/tool-types.js";
import type { JsonSerializableObject } from "../util/json-serializable.js";
import type { Result } from "../util/result.js";
import { BreadboardToolInvocation } from "./breadboard-tool.js";
import { makeToolSafeName } from "./make-tool-safe-name.js";
import { standardizeBreadboardSchema } from "./standardize-breadboard-schema.js";

export class BreadboardComponentTool implements BBRTTool {
readonly #kit: Kit;
readonly #id: string;
readonly #describe?: NodeDescriberFunction;
readonly #secrets: SecretsProvider;
readonly #artifacts: ArtifactStore;
readonly #kits: Kit[];

constructor(
kit: Kit,
id: string,
handler: NodeHandler,
secrets: SecretsProvider,
artifacts: ArtifactStore,
kits: Kit[]
) {
this.#kit = kit;
this.#id = id;
this.#secrets = secrets;
this.#artifacts = artifacts;
this.#kits = kits;
if ("describe" in handler) {
this.#describe = handler.describe;
}
}

get metadata(): BBRTToolMetadata {
return {
id: makeToolSafeName(`${this.#kit.url}_${this.#id}`),
title: `${this.#kit.title}: ${this.#id}`,
description: "TODO",
icon: "/bbrt/images/tool.svg",
};
}

async api(): Promise<Result<BBRTToolAPI>> {
const graphStore = createGraphStore({
kits: this.#kits,
loader: createLoader(),
sandbox: {
runModule: () => {
throw new Error("TODO: runModule not implemented");
},
},
});

if (this.#describe) {
const { inputSchema, outputSchema } = await this.#describe(
undefined,
undefined,
undefined,
{
graphStore,
outerGraph: { nodes: [], edges: [] },
wires: { incoming: {}, outgoing: {} },
}
);
return {
ok: true,
value: {
inputSchema: standardizeBreadboardSchema(inputSchema),
outputSchema: standardizeBreadboardSchema(outputSchema),
},
};
}
return {
ok: true,
value: {
inputSchema: { type: "object", properties: {}, required: [] },
outputSchema: { type: "object", properties: {}, required: [] },
},
};
}

execute(args: JsonSerializableObject) {
return { result: this.#execute(args) };
}

async #execute(
args: JsonSerializableObject
): Promise<
Result<{ data: JsonSerializableObject; artifacts: ArtifactHandle[] }>
> {
const component: NodeDescriptor = {
id: "component",
type: this.#id,
configuration: args,
};
const bgl: GraphDescriptor = {
nodes: [component],
edges: [],
};
const invocation = new BreadboardToolInvocation(
args,
async () => ({ ok: true, value: bgl }),
this.#secrets,
this.#artifacts,
this.#kits
);
await invocation.start();
const state = invocation.state.get();
if (state.status === "success") {
return {
ok: true,
value: {
data: state.value.output as JsonSerializableObject,
artifacts: state.value.artifacts,
},
};
} else if (state.status === "error") {
return { ok: false, error: state.error };
} else {
state.status satisfies "running" | "unstarted";
return { ok: false, error: `Internal error: state was ${state.status}` };
}
}
}
30 changes: 11 additions & 19 deletions packages/bbrt/src/breadboard/breadboard-tool.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,9 +4,7 @@
* SPDX-License-Identifier: Apache-2.0
*/

import AgentKit from "@google-labs/agent-kit/agent.kit.json" assert { type: "json" };
import {
asRuntimeKit,
createDefaultDataStore,
createLoader,
type DataStore,
Expand All @@ -16,11 +14,6 @@ import {
type OutputValues,
} from "@google-labs/breadboard";
import { createRunner, type RunConfig } from "@google-labs/breadboard/harness";
import { kitFromGraphDescriptor } from "@google-labs/breadboard/kits";
import CoreKit from "@google-labs/core-kit";
import GeminiKit from "@google-labs/gemini-kit";
import JSONKit from "@google-labs/json-kit";
import TemplateKit from "@google-labs/template-kit";
import { Signal } from "signal-polyfill";
import type { ArtifactHandle } from "../artifacts/artifact-interface.js";
import type { ArtifactStore } from "../artifacts/artifact-store.js";
Expand Down Expand Up @@ -52,17 +45,20 @@ export class BreadboardTool implements BBRTTool {
readonly #server: BreadboardServiceClient;
readonly #secrets: SecretsProvider;
readonly #artifacts: ArtifactStore;
readonly #kits: Kit[];

constructor(
listing: BreadboardBoardListing,
server: BreadboardServiceClient,
secrets: SecretsProvider,
artifacts: ArtifactStore
artifacts: ArtifactStore,
kits: Kit[]
) {
this.#listing = listing;
this.#server = server;
this.#secrets = secrets;
this.#artifacts = artifacts;
this.#kits = kits;
}

get metadata(): BBRTToolMetadata {
Expand Down Expand Up @@ -111,7 +107,8 @@ export class BreadboardTool implements BBRTTool {
args,
() => this.bgl(),
this.#secrets,
this.#artifacts
this.#artifacts,
this.#kits
);
// TODO(aomarks) Broadly rethink this. Executions should be a stream of
// events, and those should be saved in state too.
Expand Down Expand Up @@ -153,6 +150,7 @@ export class BreadboardToolInvocation {
readonly #secrets: SecretsProvider;
readonly #getBgl: () => Promise<Result<GraphDescriptor>>;
readonly #artifacts: ArtifactStore;
readonly #kits: Kit[];

readonly state = new Signal.State<InvocationState<unknown>>({
status: "unstarted",
Expand All @@ -162,12 +160,14 @@ export class BreadboardToolInvocation {
args: unknown,
getBgl: () => Promise<Result<GraphDescriptor>>,
secrets: SecretsProvider,
artifacts: ArtifactStore
artifacts: ArtifactStore,
kits: Kit[]
) {
this.#args = args;
this.#getBgl = getBgl;
this.#secrets = secrets;
this.#artifacts = artifacts;
this.#kits = kits;
}

async start(): Promise<void> {
Expand All @@ -185,14 +185,6 @@ export class BreadboardToolInvocation {
}

const loader = createLoader();
const kits: Kit[] = [
asRuntimeKit(CoreKit),
asRuntimeKit(TemplateKit),
asRuntimeKit(JSONKit),
asRuntimeKit(GeminiKit),
kitFromGraphDescriptor(AgentKit as GraphDescriptor)!,
];

const store = createDefaultDataStore();
const storeGroupId = crypto.randomUUID();
store.createGroup(storeGroupId);
Expand All @@ -201,7 +193,7 @@ export class BreadboardToolInvocation {
// TODO(aomarks) What should this be, it matters for relative imports,
// right?
url: `https://example.com/fake`,
kits,
kits: this.#kits,
runner: bgl.value,
loader,
store,
Expand Down
2 changes: 1 addition & 1 deletion packages/bbrt/src/breadboard/make-tool-safe-name.ts
Original file line number Diff line number Diff line change
Expand Up @@ -9,5 +9,5 @@
* OpenAI requires [a-zA-Z0-9_\-]{1,??}
*/
export function makeToolSafeName(path: string) {
return path.replace(/[^a-zA-Z0-9_\\-]/g, "").slice(0, 63);
return path.replace(/[^a-zA-Z0-9_\\-]/g, "_").slice(0, 63);
}
39 changes: 38 additions & 1 deletion packages/bbrt/src/components/main.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,18 @@
* SPDX-License-Identifier: Apache-2.0
*/

import GoogleDriveKit from "@breadboard-ai/google-drive-kit/google-drive.kit.json" with { type: "json" };
import AgentKit from "@google-labs/agent-kit/agent.kit.json" assert { type: "json" };
import {
asRuntimeKit,
type GraphDescriptor,
type Kit,
} from "@google-labs/breadboard";
import { kitFromGraphDescriptor } from "@google-labs/breadboard/kits";
import CoreKit from "@google-labs/core-kit";
import GeminiKit from "@google-labs/gemini-kit";
import JSONKit from "@google-labs/json-kit";
import TemplateKit from "@google-labs/template-kit";
import { SignalWatcher } from "@lit-labs/signals";
import { provide } from "@lit/context";
import { LitElement, css, html } from "lit";
Expand All @@ -15,6 +27,7 @@ import {
artifactStoreContext,
} from "../artifacts/artifact-store.js";
import { IdbArtifactReaderWriter } from "../artifacts/idb-artifact-reader-writer.js";
import { BreadboardComponentTool } from "../breadboard/breadboard-component-tool.js";
import { BreadboardTool } from "../breadboard/breadboard-tool.js";
import { readBoardServersFromIndexedDB } from "../breadboard/indexed-db-servers.js";
import type { Config } from "../config.js";
Expand Down Expand Up @@ -71,6 +84,15 @@ export class BBRTMain extends SignalWatcher(LitElement) {
readonly #leftBar = createRef();
readonly #rightBar = createRef();

readonly #breadboardKits: Kit[] = [
asRuntimeKit(CoreKit),
asRuntimeKit(TemplateKit),
asRuntimeKit(JSONKit),
asRuntimeKit(GeminiKit),
kitFromGraphDescriptor(AgentKit as GraphDescriptor)!,
kitFromGraphDescriptor(GoogleDriveKit as GraphDescriptor)!,
];

static override styles = css`
:host {
display: grid;
Expand Down Expand Up @@ -317,6 +339,20 @@ export class BBRTMain extends SignalWatcher(LitElement) {

async #loadBreadboardTools(): Promise<BBRTTool[]> {
const tools: BBRTTool[] = [];
for (const kit of this.#breadboardKits) {
for (const [id, handler] of Object.entries(kit.handlers)) {
tools.push(
new BreadboardComponentTool(
kit,
id,
handler,
this.#secrets,
this.#artifacts,
this.#breadboardKits
)
);
}
}
const servers = await readBoardServersFromIndexedDB();
if (!servers.ok) {
console.error(
Expand All @@ -338,7 +374,8 @@ export class BBRTMain extends SignalWatcher(LitElement) {
board,
server,
this.#secrets,
this.#artifacts
this.#artifacts,
this.#breadboardKits
);
tools.push(tool);
}
Expand Down

0 comments on commit 3420dc6

Please sign in to comment.