diff --git a/frontend/src/core/ai/model-registry.ts b/frontend/src/core/ai/model-registry.ts index ea327c3118f..a9982f0a717 100644 --- a/frontend/src/core/ai/model-registry.ts +++ b/frontend/src/core/ai/model-registry.ts @@ -13,19 +13,6 @@ import { once } from "@/utils/once"; import type { ProviderId } from "./ids/ids"; import { AiModelId, type QualifiedModelId, type ShortModelId } from "./ids/ids"; -export const PROVIDER_SORT_ORDER: ProviderId[] = [ - // Sort by popular ones - "anthropic", - "openai", - "google", - "github", - "openrouter", - "deepseek", - "azure", - "bedrock", - "ollama", -]; - export interface AiModel extends AiModelType { roles: Role[]; model: ShortModelId; diff --git a/frontend/src/core/ai/tools/__tests__/edit-notebook-tool.test.ts b/frontend/src/core/ai/tools/__tests__/edit-notebook-tool.test.ts index b99db80b4ba..bef5bb7ab4b 100644 --- a/frontend/src/core/ai/tools/__tests__/edit-notebook-tool.test.ts +++ b/frontend/src/core/ai/tools/__tests__/edit-notebook-tool.test.ts @@ -249,7 +249,7 @@ describe("EditNotebookTool", () => { { edit: { type: "add_cell", - position: "end", + position: { type: "notebook_end" }, code: newCode, }, }, @@ -280,7 +280,7 @@ describe("EditNotebookTool", () => { { edit: { type: "add_cell", - position: { cellId: cellId2, before: true }, + position: { type: "relative", cellId: cellId2, before: true }, code: newCode, }, }, @@ -310,7 +310,7 @@ describe("EditNotebookTool", () => { { edit: { type: "add_cell", - position: { cellId: cellId2, before: false }, + position: { type: "relative", cellId: cellId2, before: false }, code: newCode, }, }, @@ -340,7 +340,7 @@ describe("EditNotebookTool", () => { { edit: { type: "add_cell", - position: { type: "end", columnIndex: 1 }, + position: { type: "column_end", columnIndex: 1 }, code: newCode, }, }, @@ -367,7 +367,11 @@ describe("EditNotebookTool", () => { { edit: { type: "add_cell", - position: { cellId: "nonexistent" as CellId, before: true }, + position: { + type: "relative", + cellId: "nonexistent" as CellId, + before: true, + }, code: "y = 2", }, }, @@ -390,7 +394,7 @@ describe("EditNotebookTool", () => { edit: { type: "add_cell", position: { - type: "end", + type: "column_end", columnIndex: -1, }, code: "y = 2", @@ -540,7 +544,7 @@ describe("EditNotebookTool", () => { { edit: { type: "add_cell", - position: "end", + position: { type: "notebook_end" }, code: "y = 2", }, }, diff --git a/frontend/src/core/ai/tools/__tests__/registry.test.ts b/frontend/src/core/ai/tools/__tests__/registry.test.ts index 32c74ed735b..fefbc73d653 100644 --- a/frontend/src/core/ai/tools/__tests__/registry.test.ts +++ b/frontend/src/core/ai/tools/__tests__/registry.test.ts @@ -55,9 +55,9 @@ describe("FrontendToolRegistry", () => { expect(typeof response.error).toBe("string"); // Verify error message contains expected prefix - expect(response.error).toContain("Error invoking tool ToolExecutionError:"); - expect(response.error).toContain('"code":"TOOL_ERROR"'); - expect(response.error).toContain('"is_retryable":false'); + expect(response.error).toMatchInlineSnapshot( + `"Error invoking tool ToolExecutionError: {"message":"Tool test_frontend_tool returned invalid input: ✖ Invalid input: expected string, received undefined\\n → at name","code":"INVALID_ARGUMENTS","is_retryable":true,"suggested_fix":"Please check the arguments and try again."}"`, + ); }); it("returns tool schemas with expected shape and memoizes the result", () => { diff --git a/frontend/src/core/ai/tools/base.ts b/frontend/src/core/ai/tools/base.ts index 7de066dbd10..163aca85fa1 100644 --- a/frontend/src/core/ai/tools/base.ts +++ b/frontend/src/core/ai/tools/base.ts @@ -51,8 +51,8 @@ export class ToolExecutionError extends Error { message: this.message, code: this.code, is_retryable: this.isRetryable, - suggested_fix: this.suggestedFix, - meta: this.meta ?? {}, + ...(this.suggestedFix && { suggested_fix: this.suggestedFix }), + ...(this.meta && { meta: this.meta }), }); return `Error invoking tool ${this.name}: ${stringError}`; } diff --git a/frontend/src/core/ai/tools/edit-notebook-tool.ts b/frontend/src/core/ai/tools/edit-notebook-tool.ts index b390ffd3916..52fdc829905 100644 --- a/frontend/src/core/ai/tools/edit-notebook-tool.ts +++ b/frontend/src/core/ai/tools/edit-notebook-tool.ts @@ -31,12 +31,12 @@ const description: ToolDescription = { additionalInfo: ` Args: edit (object): The editing operation to perform. Must be one of: - - update_cell: Update the code of an existing cell, pass CellId and the new code. - - add_cell: Add a new cell to the notebook. The position of the new cell is specified by the position argument. - Pass "end" to add the new cell at the end of the notebook. - Pass { cellId: cellId, before: true } to add the new cell before the specified cell. And before: false if after the specified cell. - Pass { type: "end", columnIndex: number } to add the new cell at the end of a specified column index. The column index is 0-based. - - delete_cell: Delete an existing cell, pass CellId. For deleting cells, the user needs to accept the deletion to actually delete the cell, so you may still see the cell in the notebook on subsequent edits which is fine. + - update_cell: Update the code of an existing cell, pass cellId and the new code. + - add_cell: Add a new cell to the notebook. The position is specified by the position object with a "type" field: + { type: "notebook_end" } - Add at the end of the notebook + { type: "relative", cellId: "...", before: true } - Add before the specified cell (before: false for after) + { type: "column_end", columnIndex: 0 } - Add at the end of a specific column (0-based index) + - delete_cell: Delete an existing cell, pass cellId. For deleting cells, the user needs to accept the deletion to actually delete the cell, so you may still see the cell in the notebook on subsequent edits which is fine. For adding code, use the following guidelines: - Markdown cells: use mo.md(f"""{content}""") function to insert content. @@ -47,9 +47,9 @@ const description: ToolDescription = { }; type CellPosition = - | { cellId: CellId; before: boolean } - | { type: "end"; columnIndex: number } - | "end"; + | { type: "relative"; cellId: CellId; before: boolean } + | { type: "column_end"; columnIndex: number } + | { type: "notebook_end" }; const editNotebookSchema = z.object({ edit: z.discriminatedUnion("type", [ @@ -60,16 +60,19 @@ const editNotebookSchema = z.object({ }), z.object({ type: z.literal("add_cell"), - position: z.union([ + position: z.discriminatedUnion("type", [ z.object({ + type: z.literal("relative"), cellId: z.string() as unknown as z.ZodType, before: z.boolean(), }), z.object({ - type: z.literal("end"), + type: z.literal("column_end"), columnIndex: z.number(), }), - z.literal("end"), + z.object({ + type: z.literal("notebook_end"), + }), ]) satisfies z.ZodType, code: z.string(), }), @@ -135,21 +138,25 @@ export class EditNotebookTool case "add_cell": { const { position, code } = edit; - // By default, add the new cell to the end of the notebook let notebookPosition: NotebookCellPosition = "__end__"; let before = false; const newCellId = CellId.create(); + const notebook = store.get(notebookAtom); - if (typeof position === "object") { - const notebook = store.get(notebookAtom); - if ("cellId" in position) { + switch (position.type) { + case "relative": this.validateCellIdExists(position.cellId, notebook); notebookPosition = position.cellId; before = position.before; - } else if ("columnIndex" in position) { + break; + case "column_end": { const columnId = this.getColumnId(position.columnIndex, notebook); notebookPosition = { type: "__end__", columnId }; + break; } + case "notebook_end": + // Use default: notebookPosition = "__end__" + break; } createNewCell({ diff --git a/frontend/src/core/ai/tools/registry.ts b/frontend/src/core/ai/tools/registry.ts index a3df19345b9..9dcfd981d55 100644 --- a/frontend/src/core/ai/tools/registry.ts +++ b/frontend/src/core/ai/tools/registry.ts @@ -67,7 +67,12 @@ export class FrontendToolRegistry { const inputResponse = await inputSchema.safeParseAsync(rawArgs); if (inputResponse.error) { const strError = z.prettifyError(inputResponse.error); - throw new Error(`Tool ${toolName} returned invalid input: ${strError}`); + throw new ToolExecutionError( + `Tool ${toolName} returned invalid input: ${strError}`, + "INVALID_ARGUMENTS", + true, + "Please check the arguments and try again.", + ); } const args = inputResponse.data;