Data Engineer mode (base 3.11.17)

Author: complyue

src/core/Cline.ts

@@ -81,6 +81,9 @@ 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"
 
 export type ToolResponse = string | Array<Anthropic.TextBlockParam | Anthropic.ImageBlockParam>
 type UserContent = Array<Anthropic.Messages.ContentBlockParam>
@@ -1430,6 +1433,12 @@ 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}']`
 					}
 				}
 
@@ -1674,6 +1683,22 @@ 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
 				}
 
 				break

src/core/assistant-message/index.ts

@@ -26,6 +26,9 @@ export const toolUseNames = [
 	"switch_mode",
 	"new_task",
 	"fetch_instructions",
+	"notebook_read",
+	"notebook_edit",
+	"notebook_execute",
 ] as const
 
 // Converts array of tool call names into a union type ("execute_command" | "read_file" | ...)
@@ -61,6 +64,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 +156,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" | "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">>
+}

src/core/prompts/tools/index.ts

@@ -14,6 +14,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"
@@ -38,6 +41,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 }) : "",
 }
@@ -112,4 +118,7 @@ export {
 	getSwitchModeDescription,
 	getInsertContentDescription,
 	getSearchAndReplaceDescription,
+	getNotebookReadToolDescription,
+	getNotebookEditToolDescription,
+	getNotebookExecuteToolDescription,
 }

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

@@ -0,0 +1,112 @@
+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, 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) A JSON array containing cell definitions. The structure depends on the action:
+  - For insert_cells: Array of objects with properties:
+    - content: (required) The cell content
+    - cell_type: (optional) The type of cell, either "code" or "markdown"
+    - language_id: (optional) The language of the cell (e.g., "python", "javascript")
+  - For modify_cell_content: Array of objects with properties:
+    - index: (required) The index of the cell to modify (0-based)
+    - content: (required) The new content for the cell
+  - For replace_cells: A single object with properties:
+    - start_index: (required) The starting index of the range to replace (0-based, inclusive)
+    - end_index: (required) The ending index of the range to replace (0-based, exclusive)
+    - cells: (required) Array of objects with properties:
+      - content: (required) The cell content
+      - cell_type: (optional) The type of cell, either "code" or "markdown"
+      - language_id: (optional) The language of the cell (e.g., "python", "javascript")
+  - For delete_cells: A single object with properties:
+    - start_index: (required) The starting index of the range to delete (0-based, inclusive)
+    - end_index: (required) The ending index of the range to delete (0-based, exclusive)
+- insert_at_index: (optional, only for insert_cells action) The position to insert cells at (0-based, defaults to end of notebook)
+- noexec: (optional) Boolean value to prevent automatic execution of new/modified code cells.
+Usage:
+<notebook_edit>
+<action>action name here</action>
+<cells>[JSON array of cell definitions]</cells>
+<insert_at_index>index value here (if required)</insert_at_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>[
+  {
+    "cell_type": "markdown",
+    "content": "# Data Analysis\nThis notebook contains the analysis of our dataset."
+  },
+  {
+    "cell_type": "code",
+    "language_id": "python",
+    "content": "import pandas as pd\nimport numpy as np\ndf = pd.read_csv('data.csv')\ndf.head()"
+  }
+]</cells>
+</notebook_edit>
+
+Example 2: Append a new markdown cell (no insert_at_index specified)
+<notebook_edit>
+<action>insert_cells</action>
+<cells>[
+  {
+    "cell_type": "markdown",
+    "content": "# Data Analysis\nThis notebook contains the analysis of our dataset with the following steps:\n1. Data loading and cleaning\n2. Exploratory data analysis\n3. Statistical testing\n4. Visualization"
+  }
+]</cells>
+</notebook_edit>
+
+Example 3: Modify multiple existing cells without execution
+<notebook_edit>
+<action>modify_cell_content</action>
+<cells>[
+  {
+    "index": 2,
+    "content": "import matplotlib.pyplot as plt\nplt.figure(figsize=(10, 6))\nplt.plot(df['x'], df['y'])\nplt.title('Data Visualization')\nplt.xlabel('X Axis')\nplt.ylabel('Y Axis')\nplt.show()"
+  },
+  {
+    "index": 3,
+    "content": "# Statistical Analysis\nLet'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>
+<cells>{
+  "start_index": 2,
+  "end_index": 4,
+  "cells": [
+    {
+      "cell_type": "code",
+      "language_id": "python",
+      "content": "import matplotlib.pyplot as plt\nimport seaborn as sns\n\nplt.figure(figsize=(10, 6))\nsns.scatterplot(x='x', y='y', data=df)\nplt.title('Scatter Plot')\nplt.show()"
+    },
+    {
+      "cell_type": "markdown",
+      "content": "## Observations\nThe 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>
+<cells>{
+  "start_index": 2,
+  "end_index": 4
+}</cells>
+</notebook_edit>`
+}

src/core/prompts/tools/notebook-execute.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-read.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`
+}