Data Engineer mode (base 3.12)

Author: complyue

src/core/Cline.ts

@@ -76,6 +76,10 @@ import { askFollowupQuestionTool } from "./tools/askFollowupQuestionTool"
 import { switchModeTool } from "./tools/switchModeTool"
 import { attemptCompletionTool } from "./tools/attemptCompletionTool"
 import { newTaskTool } from "./tools/newTaskTool"
+import { notebookReadTool } from "./tools/notebookReadTool"
+import { notebookEditTool } from "./tools/notebookEditTool"
+import { notebookExecuteTool } from "./tools/notebookExecuteTool"
+import { notebookSaveTool } from "./tools/notebookSaveTool"
 
 export type ToolResponse = string | Array<Anthropic.TextBlockParam | Anthropic.ImageBlockParam>
 type UserContent = Array<Anthropic.Messages.ContentBlockParam>
@@ -1422,6 +1426,14 @@ export class Cline extends EventEmitter<ClineEvents> {
 							const modeName = getModeBySlug(mode, customModes)?.name ?? mode
 							return `[${block.name} in ${modeName} mode: '${message}']`
 						}
+						case "notebook_read":
+							return `[${block.name} action '${block.params.action}']`
+						case "notebook_edit":
+							return `[${block.name} action '${block.params.action}']`
+						case "notebook_execute":
+							return `[${block.name} action '${block.params.action}']`
+						case "notebook_save":
+							return `[${block.name}]`
 					}
 				}
 
@@ -1666,6 +1678,25 @@ export class Cline extends EventEmitter<ClineEvents> {
 							askFinishSubTaskApproval,
 						)
 						break
+					case "notebook_read":
+						await notebookReadTool(this, block, askApproval, handleError, pushToolResult, removeClosingTag)
+						break
+					case "notebook_edit":
+						await notebookEditTool(this, block, askApproval, handleError, pushToolResult, removeClosingTag)
+						break
+					case "notebook_execute":
+						await notebookExecuteTool(
+							this,
+							block,
+							askApproval,
+							handleError,
+							pushToolResult,
+							removeClosingTag,
+						)
+						break
+					case "notebook_save":
+						await notebookSaveTool(this, block, askApproval, handleError, pushToolResult, removeClosingTag)
+						break
 				}
 
 				break

src/core/assistant-message/index.ts

@@ -26,6 +26,10 @@ export const toolUseNames = [
 	"switch_mode",
 	"new_task",
 	"fetch_instructions",
+	"notebook_read",
+	"notebook_edit",
+	"notebook_execute",
+	"notebook_save",
 ] as const
 
 // Converts array of tool call names into a union type ("execute_command" | "read_file" | ...)
@@ -61,6 +65,11 @@ export const toolParamNames = [
 	"follow_up",
 	"task",
 	"size",
+	"cells",
+	"insert_at_index",
+	"start_index",
+	"end_index",
+	"noexec",
 ] as const
 
 export type ToolParamName = (typeof toolParamNames)[number]
@@ -148,3 +157,28 @@ export interface NewTaskToolUse extends ToolUse {
 	name: "new_task"
 	params: Partial<Pick<Record<ToolParamName, string>, "mode" | "message">>
 }
+
+export interface NotebookReadToolUse extends ToolUse {
+	name: "notebook_read"
+	params: Partial<Pick<Record<ToolParamName, string>, "action">>
+}
+
+export interface NotebookEditToolUse extends ToolUse {
+	name: "notebook_edit"
+	params: Partial<
+		Pick<
+			Record<ToolParamName, string>,
+			"action" | "cells" | "insert_at_index" | "start_index" | "end_index" | "noexec"
+		>
+	>
+}
+
+export interface NotebookExecuteToolUse extends ToolUse {
+	name: "notebook_execute"
+	params: Partial<Pick<Record<ToolParamName, string>, "action" | "start_index" | "end_index">>
+}
+
+export interface NotebookSaveToolUse extends ToolUse {
+	name: "notebook_save"
+	params: Partial<Record<ToolParamName, string>>
+}

src/core/prompts/tools/index.ts

@@ -14,6 +14,10 @@ import { getUseMcpToolDescription } from "./use-mcp-tool"
 import { getAccessMcpResourceDescription } from "./access-mcp-resource"
 import { getSwitchModeDescription } from "./switch-mode"
 import { getNewTaskDescription } from "./new-task"
+import { getNotebookReadToolDescription } from "./notebook-read"
+import { getNotebookEditToolDescription } from "./notebook-edit"
+import { getNotebookExecuteToolDescription } from "./notebook-execute"
+import { getNotebookSaveToolDescription } from "./notebook-save"
 import { DiffStrategy } from "../../diff/DiffStrategy"
 import { McpHub } from "../../../services/mcp/McpHub"
 import { Mode, ModeConfig, getModeConfig, isToolAllowedForMode, getGroupName } from "../../../shared/modes"
@@ -38,6 +42,10 @@ const toolDescriptionMap: Record<string, (args: ToolArgs) => string | undefined>
 	new_task: (args) => getNewTaskDescription(args),
 	insert_content: (args) => getInsertContentDescription(args),
 	search_and_replace: (args) => getSearchAndReplaceDescription(args),
+	notebook_read: (args) => getNotebookReadToolDescription(args),
+	notebook_edit: (args) => getNotebookEditToolDescription(args),
+	notebook_execute: (args) => getNotebookExecuteToolDescription(args),
+	notebook_save: (args) => getNotebookSaveToolDescription(args),
 	apply_diff: (args) =>
 		args.diffStrategy ? args.diffStrategy.getToolDescription({ cwd: args.cwd, toolOptions: args.toolOptions }) : "",
 }
@@ -112,4 +120,8 @@ export {
 	getSwitchModeDescription,
 	getInsertContentDescription,
 	getSearchAndReplaceDescription,
+	getNotebookReadToolDescription,
+	getNotebookEditToolDescription,
+	getNotebookExecuteToolDescription,
+	getNotebookSaveToolDescription,
 }

src/core/prompts/tools/notebook-edit.ts

@@ -0,0 +1,131 @@
+import { ToolArgs } from "./types"
+
+export function getNotebookEditToolDescription(args: ToolArgs): string {
+	return `## notebook_edit
+Description: Edit the active notebook. This tool allows you to insert new cells, append cells to the end, or replace existing cells. Note that new/modified code cells will be executed immediately by default, unless the noexec parameter is true.
+Parameters:
+- action: (required) The action to perform. Valid values are:
+  - "insert_cells": Insert new cells into the notebook (or append to the end if insert_at_index is omitted)
+  - "modify_cell_content": Modify the content of existing cells in the notebook (note: this will clear any existing outputs)
+  - "replace_cells": Replace a range of cells with new cells (note: this will clear any existing outputs)
+  - "delete_cells": Delete a range of cells from the notebook
+- cells: (required for insert_cells, modify_cell_content, replace_cells) Contains cell definitions using markdown code blocks:
+  - For insert_cells: A series of code blocks representing the cells to insert
+  - For modify_cell_content: Code blocks with @cell#N tag before each block (where N is the 0-based index)
+  - For replace_cells: Series of code blocks to replace the specified range of cells
+- insert_at_index: (optional, only for insert_cells action) The position to insert cells at (0-based, defaults to end of notebook)
+- start_index: (required for replace_cells and delete_cells) The starting index of the range (0-based, inclusive)
+- end_index: (required for replace_cells and delete_cells) The ending index of the range (0-based, exclusive)
+- noexec: (optional) Boolean value to prevent automatic execution of new/modified code cells.
+
+The language of each cell is determined by the code block header (e.g., \`\`\`python).
+Cells with \`\`\`markdown header are treated as markdown cells, all others are code cells.
+
+Usage:
+<notebook_edit>
+<action>action name here</action>
+<cells>
+(markdown code blocks for cell content here)
+</cells>
+<insert_at_index>index value here (if required)</insert_at_index>
+<start_index>start index (for replace/delete actions)</start_index>
+<end_index>end index (for replace/delete actions)</end_index>
+<noexec>true</noexec>
+</notebook_edit>
+
+Example 1: Insert multiple cells at index 0
+<notebook_edit>
+<action>insert_cells</action>
+<insert_at_index>0</insert_at_index>
+<cells>
+\`\`\`markdown
+# Data Analysis
+This notebook contains the analysis of our dataset.
+\`\`\`
+
+\`\`\`python
+import pandas as pd
+import numpy as np
+df = pd.read_csv('data.csv')
+df.head()
+\`\`\`
+</cells>
+</notebook_edit>
+
+Example 2: Append a new markdown cell (no insert_at_index specified)
+<notebook_edit>
+<action>insert_cells</action>
+<cells>
+\`\`\`markdown
+# Data Analysis
+This notebook contains the analysis of our dataset with the following steps:
+1. Data loading and cleaning
+2. Exploratory data analysis
+3. Statistical testing
+4. Visualization
+\`\`\`
+</cells>
+</notebook_edit>
+
+Example 3: Modify multiple existing cells without execution
+<notebook_edit>
+<action>modify_cell_content</action>
+<cells>
+@cell#2
+\`\`\`python
+import matplotlib.pyplot as plt
+plt.figure(figsize=(10, 6))
+plt.plot(df['x'], df['y'])
+plt.title('Data Visualization')
+plt.xlabel('X Axis')
+plt.ylabel('Y Axis')
+plt.show()
+\`\`\`
+
+@cell#3
+\`\`\`markdown
+# Statistical Analysis
+Let's analyze the correlation between variables.
+\`\`\`
+</cells>
+<noexec>true</noexec>
+</notebook_edit>
+
+Example 4: Replace a range of cells with new cells
+<notebook_edit>
+<action>replace_cells</action>
+<start_index>2</start_index>
+<end_index>4</end_index>
+<cells>
+\`\`\`python
+import matplotlib.pyplot as plt
+import seaborn as sns
+
+plt.figure(figsize=(10, 6))
+sns.scatterplot(x='x', y='y', data=df)
+plt.title('Scatter Plot')
+plt.show()
+\`\`\`
+
+\`\`\`markdown
+## Observations
+The scatter plot shows a positive correlation between x and y variables.
+\`\`\`
+</cells>
+</notebook_edit>
+
+Example 5: Delete a range of cells
+<notebook_edit>
+<action>delete_cells</action>
+<start_index>2</start_index>
+<end_index>4</end_index>
+<cells>
+</cells>
+</notebook_edit>
+
+Notes:
+- The user must have opened a notebook file in VSCode to have an active notebook in the workspace
+- If no active notebook is found, an error message will be returned
+- Cell indices are 0-based (the first cell is at index 0)
+- For modify_cell_content, each code block must have a @cell#N tag to identify which cell to modify`
+}

src/core/prompts/tools/notebook-execute.ts

@@ -0,0 +1,33 @@
+import { ToolArgs } from "./types"
+
+export function getNotebookExecuteToolDescription(args: ToolArgs): string {
+	return `## notebook_execute
+Description: Execute cells in the active notebook.
+Parameters:
+- action: (required) The action to perform. Currently only "execute_cells" is supported.
+- start_index: (required) The starting index of the range to execute (0-based, inclusive). Must be a valid cell index in the notebook.
+- end_index: (required) The ending index of the range to execute (0-based, exclusive). Must be a valid cell index in the notebook and >= start_index.
+
+Usage:
+<notebook_execute>
+<action>execute_cells</action>
+<start_index>start index value here</start_index>
+<end_index>end index value here</end_index>
+</notebook_execute>
+
+Example: Execute cells from index 3 through 5 (inclusive)
+<notebook_execute>
+<action>execute_cells</action>
+<start_index>3</start_index>
+<end_index>6</end_index>
+</notebook_execute>
+
+Notes:
+- The user must have opened a notebook file in VSCode to have an active notebook in the workspace
+- The execution is initiated immediately, but completion time depends on the cells' content
+- The tool will wait for execution to complete and return the results
+- Results include both the cell content and any outputs produced by execution
+- Execution will timeout if it takes too long, and a timeout message will be included in the response
+- Cell outputs are truncated if they exceed the configured size limit
+- To execute a single cell, set end_index to start_index + 1`
+}

src/core/prompts/tools/notebook-read.ts

@@ -0,0 +1,34 @@
+import { ToolArgs } from "./types"
+
+export function getNotebookReadToolDescription(args: ToolArgs): string {
+	return `## notebook_read
+Description: Retrieve information about the active notebook.
+Parameters:
+- action: (required) The action to perform. Valid values are:
+  - "get_info": Get comprehensive information about the active notebook, including URI, kernel info, and cell statistics
+  - "get_cells": Get the cells of the active notebook. Returns detailed information about each cell, including content and outputs
+
+Usage:
+<notebook_read>
+<action>action name here</action>
+</notebook_read>
+
+Example 1: Get comprehensive notebook information
+<notebook_read>
+<action>get_info</action>
+</notebook_read>
+
+Example 2: Get the cells of the active notebook
+<notebook_read>
+<action>get_cells</action>
+</notebook_read>
+
+Notes:
+- The user must have opened a notebook file in VSCode to have an active notebook in the workspace
+- The "get_info" action provides a comprehensive overview including URI, kernel info, and statistics about cell types and languages
+- If no active notebook is found, "get_info" will return a message stating "No active notebook found"
+- The "get_cells" action provides a formatted analysis of all cells, including their content and outputs
+- If no active notebook is found, "get_cells" will return an error message
+- Cell content and outputs are truncated if they exceed the configured size limit
+- For code cells, both the source code and execution outputs are included`
+}

src/core/prompts/tools/notebook-save.ts

@@ -0,0 +1,23 @@
+import { ToolArgs } from "./types"
+
+export function getNotebookSaveToolDescription(args: ToolArgs): string {
+	return `## notebook_save
+Description: Save the active notebook to disk.
+Parameters:
+None - this tool simply saves the active notebook.
+
+Usage:
+<notebook_save>
+</notebook_save>
+
+Example: Save the active notebook
+<notebook_save>
+</notebook_save>
+
+Notes:
+- This tool requires notebook_edit permissions
+- The user must have opened a notebook file in VSCode to have an active notebook in the workspace
+- All unsaved changes in the notebook will be saved
+- This is useful after making changes to a notebook to ensure your work is preserved
+- If no active notebook is found, an error message will be returned`
+}