Data Engineer mode (rebased on 3.9.2)

Author: complyue

src/core/Cline.ts

@@ -25,6 +25,9 @@ export const toolUseNames = [
 	"attempt_completion",
 	"switch_mode",
 	"new_task",
+	"notebook_read",
+	"notebook_edit",
+	"notebook_execute",
 ] as const
 
 // Converts array of tool call names into a union type ("execute_command" | "read_file" | ...)
@@ -58,6 +61,12 @@ export const toolParamNames = [
 	"message",
 	"cwd",
 	"follow_up",
+	"cell_index",
+	"cell_type",
+	"cell_content",
+	"language_id",
+	"start_index",
+	"end_index",
 ] as const
 
 export type ToolParamName = (typeof toolParamNames)[number]
@@ -140,3 +149,23 @@ 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" | "cell_index" | "start_index" | "end_index" | "cell_content" | "cell_type" | "language_id"
+		>
+	>
+}
+
+export interface NotebookExecuteToolUse extends ToolUse {
+	name: "notebook_execute"
+	params: Partial<Pick<Record<ToolParamName, string>, "action" | "start_index" | "end_index">>
+}

src/core/assistant-message/index.ts

@@ -13,6 +13,9 @@ 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 { DiffStrategy } from "../../diff/DiffStrategy"
 import { McpHub } from "../../../services/mcp/McpHub"
 import { Mode, ModeConfig, getModeConfig, isToolAllowedForMode, getGroupName } from "../../../shared/modes"
@@ -36,6 +39,9 @@ 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),
 	apply_diff: (args) =>
 		args.diffStrategy ? args.diffStrategy.getToolDescription({ cwd: args.cwd, toolOptions: args.toolOptions }) : "",
 }
@@ -109,4 +115,7 @@ export {
 	getSwitchModeDescription,
 	getInsertContentDescription,
 	getSearchAndReplaceDescription,
+	getNotebookReadToolDescription,
+	getNotebookEditToolDescription,
+	getNotebookExecuteToolDescription,
 }

src/core/prompts/tools/index.ts

@@ -0,0 +1,85 @@
+import { ToolArgs } from "./types"
+
+export function getNotebookEditToolDescription(args: ToolArgs): string {
+	return `## notebook_edit
+Description: Edit the active notebook in the editor. This tool allows you to insert new cells or replace existing cells. Note that new/modified code cells will be executed immediately, by default, unless the <noexec/> parameter presents.
+Parameters:
+- action: (required) The action to perform. Valid values are:
+  - "insert_cells": Insert a new cell into the notebook
+  - "modify_cell_content": Modify the content of an existing cell in the notebook (note: this will clear any existing outputs)
+  - "replace_cells": Replace a range of cells with a new cell (note: this will clear any existing outputs)
+- cell_index: (required for modify_cell_content, optional for insert_cells) The index of the cell to modify (0-based), or for insert_cells the position to insert at (defaults to end of notebook)
+- start_index: (required for replace_cells) The starting index of the range to replace (0-based, inclusive)
+- end_index: (required for replace_cells) The ending index of the range to replace (0-based, exclusive)
+- cell_type: (optional for insert_cells and replace_cells) The type of cell, either "code" or "markdown"
+- language_id: (optional for insert_cells and replace_cells) The language of the cell (e.g., "python", "javascript")
+- cell_content: (required for all actions) The content of the cell
+- noexec: (optional) Can be used as a flag <noexec/> or with explicit value <noexec>true</noexec> to prevent automatic execution of new/modified code cells.
+Usage:
+<notebook_edit>
+<action>action name here</action>
+<cell_index>index value here (if required)</cell_index>
+<start_index>start index here (if required)</start_index>
+<end_index>end index here (if required)</end_index>
+<cell_type>cell type here (if required)</cell_type>
+<language_id>language here (if required)</language_id>
+<cell_content>cell content here</cell_content>
+<noexec/>
+</notebook_edit>
+
+Example 1: Insert a new Python code cell
+<notebook_edit>
+<action>insert_cells</action>
+<cell_index>0</cell_index>
+<cell_type>code</cell_type>
+<language_id>python</language_id>
+<cell_content>import pandas as pd
+import numpy as np
+df = pd.read_csv('data.csv')
+df.head()</cell_content>
+</notebook_edit>
+
+Example 2: Insert a new markdown cell
+<notebook_edit>
+<action>insert_cells</action>
+<cell_type>markdown</cell_type>
+<cell_content># 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</cell_content>
+</notebook_edit>
+
+Example 3: Modify an existing cell without execution
+<notebook_edit>
+<action>modify_cell_content</action>
+<cell_index>2</cell_index>
+<noexec/>
+<cell_content>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_content>
+</notebook_edit>
+
+Example 4: Replace a range of cells with a new cell
+<notebook_edit>
+<action>replace_cells</action>
+<start_index>2</start_index>
+<end_index>4</end_index>
+<cell_type>code</cell_type>
+<language_id>python</language_id>
+<cell_content>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()</cell_content>
+</notebook_edit>
+
+IMPORTANT: For insert_cells and replace_cells, each cell_content tag creates a new cell. The cell_type and language_id tags must come BEFORE the cell_content tag they apply to. The last cell_type before a cell_content will be used as that cell's type, and the last language_id before a cell_content will be used as that cell's language.`
+}

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

@@ -0,0 +1,32 @@
+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 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-execute.ts

@@ -0,0 +1,33 @@
+import { ToolArgs } from "./types"
+
+export function getNotebookReadToolDescription(args: ToolArgs): string {
+	return `## notebook_read
+Description: Retrieve information about the active notebook in the editor.
+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 "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-read.ts

@@ -2020,6 +2020,26 @@ export class ClineProvider extends EventEmitter<ClineProviderEvents> implements
 						await this.postStateToWebview()
 						break
 					}
+					case "notebookOutputSizeLimit":
+						await this.updateGlobalState("notebookOutputSizeLimit", message.value)
+						await this.postStateToWebview()
+						break
+					case "notebookExecutionTimeoutSeconds":
+						await this.updateGlobalState("notebookExecutionTimeoutSeconds", message.value)
+						await this.postStateToWebview()
+						break
+					case "alwaysAllowReadNotebook":
+						await this.updateGlobalState("alwaysAllowReadNotebook", message.bool ?? undefined)
+						await this.postStateToWebview()
+						break
+					case "alwaysAllowEditNotebook":
+						await this.updateGlobalState("alwaysAllowEditNotebook", message.bool ?? undefined)
+						await this.postStateToWebview()
+						break
+					case "alwaysAllowExecuteNotebook":
+						await this.updateGlobalState("alwaysAllowExecuteNotebook", message.bool ?? undefined)
+						await this.postStateToWebview()
+						break
 				}
 			},
 			null,
@@ -2474,6 +2494,11 @@ export class ClineProvider extends EventEmitter<ClineProviderEvents> implements
 			showRooIgnoredFiles,
 			language,
 			maxReadFileLine,
+			notebookOutputSizeLimit,
+			notebookExecutionTimeoutSeconds,
+			alwaysAllowReadNotebook,
+			alwaysAllowEditNotebook,
+			alwaysAllowExecuteNotebook,
 		} = await this.getState()
 
 		const telemetryKey = process.env.POSTHOG_API_KEY
@@ -2544,6 +2569,11 @@ export class ClineProvider extends EventEmitter<ClineProviderEvents> implements
 			language,
 			renderContext: this.renderContext,
 			maxReadFileLine: maxReadFileLine ?? 500,
+			notebookOutputSizeLimit: notebookOutputSizeLimit ?? 2000,
+			notebookExecutionTimeoutSeconds: notebookExecutionTimeoutSeconds ?? 30,
+			alwaysAllowReadNotebook: alwaysAllowReadNotebook ?? false,
+			alwaysAllowEditNotebook: alwaysAllowEditNotebook ?? false,
+			alwaysAllowExecuteNotebook: alwaysAllowExecuteNotebook ?? false,
 		}
 	}
 
@@ -2703,6 +2733,11 @@ export class ClineProvider extends EventEmitter<ClineProviderEvents> implements
 			telemetrySetting: stateValues.telemetrySetting || "unset",
 			showRooIgnoredFiles: stateValues.showRooIgnoredFiles ?? true,
 			maxReadFileLine: stateValues.maxReadFileLine ?? 500,
+			notebookOutputSizeLimit: stateValues.notebookOutputSizeLimit ?? 2000,
+			notebookExecutionTimeoutSeconds: stateValues.notebookExecutionTimeoutSeconds ?? 30,
+			alwaysAllowReadNotebook: stateValues.alwaysAllowReadNotebook ?? false,
+			alwaysAllowEditNotebook: stateValues.alwaysAllowEditNotebook ?? false,
+			alwaysAllowExecuteNotebook: stateValues.alwaysAllowExecuteNotebook ?? false,
 		}
 	}
 

src/core/webview/ClineProvider.ts

@@ -461,6 +461,8 @@ describe("ClineProvider", () => {
 			showRooIgnoredFiles: true,
 			renderContext: "sidebar",
 			maxReadFileLine: 500,
+			notebookOutputSizeLimit: 1000,
+			notebookExecutionTimeoutSeconds: 30,
 		}
 
 		const message: ExtensionMessage = {

src/core/webview/__tests__/ClineProvider.test.ts

@@ -221,6 +221,11 @@ export type GlobalStateKey =
 	| "writeDelayMs"
 	| "terminalOutputLineLimit"
 	| "terminalShellIntegrationTimeout"
+	| "notebookOutputSizeLimit"
+	| "notebookExecutionTimeoutSeconds"
+	| "alwaysAllowReadNotebook"
+	| "alwaysAllowEditNotebook"
+	| "alwaysAllowExecuteNotebook"
 	| "mcpEnabled"
 	| "enableMcpServerCreation"
 	| "alwaysApproveResubmit"