Skip to content

Commit 1fe8539

Browse files
Light2DarkLight2Dark
authored
simplify edit notebook tool and add google tool schema cleanup (#7228)
## 📝 Summary The motivation to clean up the edit notebook tool is that discriminated unions are more verbose in context size: <details> <summary>tool schema parameters with discriminated union</summary> ```json "properties": { "edit": { "anyOf": [ { "type": "object", "properties": { "type": { "type": "string", "const": "update_cell", }, "cellId": {"type": "string"}, "code": {"type": "string"}, }, "required": ["type", "cellId", "code"], "additionalProperties": False, }, { "type": "object", "properties": { "type": { "type": "string", "const": "add_cell", }, "position": { "anyOf": [ { "type": "object", "properties": { "type": { "type": "string", "const": "relative", }, "cellId": {"type": "string"}, "before": {"type": "boolean"}, }, "required": [ "type", "cellId", "before", ], "additionalProperties": False, }, { "type": "object", "properties": { "type": { "type": "string", "const": "column_end", }, "columnIndex": { "type": "number" }, }, "required": [ "type", "columnIndex", ], "additionalProperties": False, }, { "type": "object", "properties": { "type": { "type": "string", "const": "notebook_end", } }, "required": ["type"], "additionalProperties": False, }, ] }, "code": {"type": "string"}, }, "required": ["type", "position", "code"], "additionalProperties": False, }, { "type": "object", "properties": { "type": { "type": "string", "const": "delete_cell", }, "cellId": {"type": "string"}, }, "required": ["type", "cellId"], "additionalProperties": False, }, ] } }, ``` </details> <details> <summary>after simplifying</summary> ```json "properties": { "edit": { "type": "object", "properties": { "type": { "type": "string", "enum": ["update_cell", "add_cell", "delete_cell"], }, "code": {"type": "string"}, "position": { "type": "object", "properties": { "type": { "type": "string", "enum": [ "relative", "column_end", "notebook_end", ], }, "cellId": {"type": "string"}, "columnIndex": {"type": "number"}, "before": {"type": "boolean"}, }, "required": ["type"], "additionalProperties": False, }, }, "required": ["type", "position"], "additionalProperties": False, } }, ``` </details> However, the google tools spec is still quite restrictive, even with the cleaned up edit notebook tool. So I implemented a thin wrapper inspired by #7211 on the backend to make sure schemas are compliant. <!-- Provide a concise summary of what this pull request is addressing. If this PR fixes any issues, list them here by number (e.g., Fixes #123). --> ## 🔍 Description of Changes <!-- Detail the specific changes made in this pull request. Explain the problem addressed and how it was resolved. If applicable, provide before and after comparisons, screenshots, or any relevant details to help reviewers understand the changes easily. --> ## 📋 Checklist - [x] I have read the [contributor guidelines](https://github.com/marimo-team/marimo/blob/main/CONTRIBUTING.md). - [ ] For large changes, or changes that affect the public API: this change was discussed or approved through an issue, on [Discord](https://marimo.io/discord?ref=pr), or the community [discussions](https://github.com/marimo-team/marimo/discussions) (Please provide a link if applicable). - [x] I have added tests for the changes made. - [x] I have run the code and verified that it works as expected.
1 parent df841e2 commit 1fe8539

File tree

6 files changed

+349
-66
lines changed

6 files changed

+349
-66
lines changed

‎frontend/src/core/ai/tools/__tests__/edit-notebook-tool.test.ts‎

Lines changed: 156 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -128,7 +128,7 @@ describe("EditNotebookTool", () => {
128128
{
129129
edit: {
130130
type: "update_cell",
131-
cellId: cellId1,
131+
position: { type: "relative", cellId: cellId1 },
132132
code: newCode,
133133
},
134134
},
@@ -154,7 +154,7 @@ describe("EditNotebookTool", () => {
154154
{
155155
edit: {
156156
type: "update_cell",
157-
cellId: cellId1,
157+
position: { type: "relative", cellId: cellId1 },
158158
code: "x = 3",
159159
},
160160
},
@@ -188,7 +188,7 @@ describe("EditNotebookTool", () => {
188188
{
189189
edit: {
190190
type: "update_cell",
191-
cellId: "nonexistent" as CellId,
191+
position: { cellId: "nonexistent" as CellId },
192192
code: "x = 2",
193193
},
194194
},
@@ -201,7 +201,7 @@ describe("EditNotebookTool", () => {
201201
{
202202
edit: {
203203
type: "update_cell",
204-
cellId: "nonexistent" as CellId,
204+
position: { cellId: "nonexistent" as CellId },
205205
code: "x = 2",
206206
},
207207
},
@@ -225,14 +225,59 @@ describe("EditNotebookTool", () => {
225225
{
226226
edit: {
227227
type: "update_cell",
228-
cellId: cellId1,
228+
position: { type: "relative", cellId: cellId1 },
229229
code: "x = 2",
230230
},
231231
},
232232
toolContext as never,
233233
),
234234
).rejects.toThrow("Cell editor not found");
235235
});
236+
237+
it("should throw error when cellId is missing for update_cell", async () => {
238+
const notebook = MockNotebook.notebookState({
239+
cellData: {
240+
[cellId1]: { code: "x = 1" },
241+
},
242+
});
243+
store.set(notebookAtom, notebook);
244+
245+
await expect(
246+
tool.handler(
247+
{
248+
edit: {
249+
type: "update_cell",
250+
position: {},
251+
code: "x = 2",
252+
},
253+
},
254+
toolContext as never,
255+
),
256+
).rejects.toThrow("Cell ID is required for update_cell");
257+
});
258+
259+
it("should throw error when code is missing for update_cell", async () => {
260+
const editorView = createMockEditorView("x = 1");
261+
const notebook = MockNotebook.notebookState({
262+
cellData: {
263+
[cellId1]: { code: "x = 1" },
264+
},
265+
});
266+
notebook.cellHandles[cellId1] = { current: { editorView } } as never;
267+
store.set(notebookAtom, notebook);
268+
269+
await expect(
270+
tool.handler(
271+
{
272+
edit: {
273+
type: "update_cell",
274+
position: { cellId: cellId1 },
275+
},
276+
},
277+
toolContext as never,
278+
),
279+
).rejects.toThrow("Code is required for update_cell");
280+
});
236281
});
237282

238283
describe("add_cell operation", () => {
@@ -404,6 +449,86 @@ describe("EditNotebookTool", () => {
404449
),
405450
).rejects.toThrow("Column index is out of range");
406451
});
452+
453+
it("should throw error when cellId is missing for add_cell with relative position", async () => {
454+
const notebook = MockNotebook.notebookState({
455+
cellData: {
456+
[cellId1]: { code: "x = 1" },
457+
},
458+
});
459+
store.set(notebookAtom, notebook);
460+
461+
await expect(
462+
tool.handler(
463+
{
464+
edit: {
465+
type: "add_cell",
466+
position: {
467+
type: "relative",
468+
before: true,
469+
},
470+
code: "y = 2",
471+
},
472+
},
473+
toolContext as never,
474+
),
475+
).rejects.toThrow(
476+
"Cell ID is required for add_cell with relative position",
477+
);
478+
});
479+
480+
it("should throw error when before is missing for add_cell with relative position", async () => {
481+
const notebook = MockNotebook.notebookState({
482+
cellData: {
483+
[cellId1]: { code: "x = 1" },
484+
},
485+
});
486+
store.set(notebookAtom, notebook);
487+
488+
await expect(
489+
tool.handler(
490+
{
491+
edit: {
492+
type: "add_cell",
493+
position: {
494+
type: "relative",
495+
cellId: cellId1,
496+
},
497+
code: "y = 2",
498+
},
499+
},
500+
toolContext as never,
501+
),
502+
).rejects.toThrow(
503+
"Before is required for add_cell with relative position",
504+
);
505+
});
506+
507+
it("should throw error when columnIndex is missing for add_cell with column_end position", async () => {
508+
const notebook = MockNotebook.notebookState({
509+
cellData: {
510+
[cellId1]: { code: "x = 1" },
511+
},
512+
});
513+
store.set(notebookAtom, notebook);
514+
515+
await expect(
516+
tool.handler(
517+
{
518+
edit: {
519+
type: "add_cell",
520+
position: {
521+
type: "column_end",
522+
},
523+
code: "y = 2",
524+
},
525+
},
526+
toolContext as never,
527+
),
528+
).rejects.toThrow(
529+
"Column index is required for add_cell with column_end position",
530+
);
531+
});
407532
});
408533

409534
describe("delete_cell operation", () => {
@@ -423,7 +548,7 @@ describe("EditNotebookTool", () => {
423548
{
424549
edit: {
425550
type: "delete_cell",
426-
cellId: cellId1,
551+
position: { cellId: cellId1 },
427552
},
428553
},
429554
toolContext as never,
@@ -453,7 +578,7 @@ describe("EditNotebookTool", () => {
453578
{
454579
edit: {
455580
type: "delete_cell",
456-
cellId: "nonexistent" as CellId,
581+
position: { cellId: "nonexistent" as CellId },
457582
},
458583
},
459584
toolContext as never,
@@ -476,13 +601,34 @@ describe("EditNotebookTool", () => {
476601
{
477602
edit: {
478603
type: "delete_cell",
479-
cellId: cellId1,
604+
position: { cellId: cellId1 },
480605
},
481606
},
482607
toolContext as never,
483608
),
484609
).rejects.toThrow("Cell editor not found");
485610
});
611+
612+
it("should throw error when cellId is missing for delete_cell", async () => {
613+
const notebook = MockNotebook.notebookState({
614+
cellData: {
615+
[cellId1]: { code: "x = 1" },
616+
},
617+
});
618+
store.set(notebookAtom, notebook);
619+
620+
await expect(
621+
tool.handler(
622+
{
623+
edit: {
624+
type: "delete_cell",
625+
position: {},
626+
},
627+
},
628+
toolContext as never,
629+
),
630+
).rejects.toThrow("Cell ID is required for delete_cell");
631+
});
486632
});
487633

488634
describe("validation", () => {
@@ -508,7 +654,7 @@ describe("EditNotebookTool", () => {
508654
{
509655
edit: {
510656
type: "update_cell",
511-
cellId: cellId1,
657+
position: { cellId: cellId1 },
512658
code: "y = 1",
513659
},
514660
},
@@ -521,7 +667,7 @@ describe("EditNotebookTool", () => {
521667
{
522668
edit: {
523669
type: "update_cell",
524-
cellId: cellId3,
670+
position: { cellId: cellId3 },
525671
code: "y = 3",
526672
},
527673
},

‎frontend/src/core/ai/tools/__tests__/registry.test.ts‎

Lines changed: 35 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,14 @@
11
/* Copyright 2024 Marimo. All rights reserved. */
22

33
import { describe, expect, it } from "vitest";
4-
import { FrontendToolRegistry } from "../registry";
4+
import { COPILOT_MODES } from "@/core/config/config-schema";
5+
import { FRONTEND_TOOL_REGISTRY, FrontendToolRegistry } from "../registry";
56
import { TestFrontendTool } from "../sample-tool";
67

8+
// Tools that have these keys in their parameters are likely not supported by Google
9+
// So we should be careful to add these types of schemas
10+
const INVALID_KEYS_IN_SCHEMA_PARAMS = new Set(["anyOf", "oneOf"]);
11+
712
describe("FrontendToolRegistry", () => {
813
it("registers tools via constructor and supports has()", () => {
914
const registry = new FrontendToolRegistry([new TestFrontendTool()]);
@@ -87,4 +92,33 @@ describe("FrontendToolRegistry", () => {
8792
const schemas3 = registry.getToolSchemas("agent");
8893
expect(schemas3.length).toBe(0);
8994
});
95+
96+
it("tool schemas should not contain invalid keys", () => {
97+
const registry = FRONTEND_TOOL_REGISTRY;
98+
99+
function hasInvalidKeys(obj: unknown): boolean {
100+
if (obj && typeof obj === "object") {
101+
for (const [key, value] of Object.entries(obj)) {
102+
if (INVALID_KEYS_IN_SCHEMA_PARAMS.has(key)) {
103+
return true;
104+
}
105+
if (
106+
typeof value === "object" &&
107+
value !== null &&
108+
hasInvalidKeys(value)
109+
) {
110+
return true;
111+
}
112+
}
113+
}
114+
return false;
115+
}
116+
117+
for (const mode of COPILOT_MODES) {
118+
const schemas = registry.getToolSchemas(mode);
119+
for (const schema of schemas) {
120+
expect(hasInvalidKeys(schema.parameters)).toBe(false);
121+
}
122+
}
123+
});
90124
});

0 commit comments

Comments
 (0)