agent_code_lib/tools/
notebook_edit.rs

1//! NotebookEdit tool: edit Jupyter notebook cells.
2//!
3//! Supports editing code cells, markdown cells, and outputs
4//! in `.ipynb` files.
5
6use async_trait::async_trait;
7use serde_json::json;
8use std::path::PathBuf;
9
10use super::{Tool, ToolContext, ToolResult};
11use crate::error::ToolError;
12
13pub struct NotebookEditTool;
14
15#[async_trait]
16impl Tool for NotebookEditTool {
17    fn name(&self) -> &'static str {
18        "NotebookEdit"
19    }
20
21    fn description(&self) -> &'static str {
22        "Edit cells in Jupyter notebooks (.ipynb files). Can replace cell \
23         content, insert new cells, or delete cells."
24    }
25
26    fn input_schema(&self) -> serde_json::Value {
27        json!({
28            "type": "object",
29            "required": ["file_path", "edit"],
30            "properties": {
31                "file_path": {
32                    "type": "string",
33                    "description": "Path to the .ipynb file"
34                },
35                "edit": {
36                    "type": "object",
37                    "description": "The edit to apply",
38                    "properties": {
39                        "action": {
40                            "type": "string",
41                            "enum": ["replace", "insert", "delete"],
42                            "description": "Edit action"
43                        },
44                        "cell_index": {
45                            "type": "integer",
46                            "description": "Index of the cell to edit (0-based)"
47                        },
48                        "cell_type": {
49                            "type": "string",
50                            "enum": ["code", "markdown"],
51                            "description": "Cell type (for insert)"
52                        },
53                        "content": {
54                            "type": "string",
55                            "description": "New cell content"
56                        }
57                    }
58                }
59            }
60        })
61    }
62
63    fn is_read_only(&self) -> bool {
64        false
65    }
66
67    fn get_path(&self, input: &serde_json::Value) -> Option<PathBuf> {
68        input
69            .get("file_path")
70            .and_then(|v| v.as_str())
71            .map(PathBuf::from)
72    }
73
74    async fn call(
75        &self,
76        input: serde_json::Value,
77        _ctx: &ToolContext,
78    ) -> Result<ToolResult, ToolError> {
79        let file_path = input
80            .get("file_path")
81            .and_then(|v| v.as_str())
82            .ok_or_else(|| ToolError::InvalidInput("'file_path' is required".into()))?;
83
84        let edit = input
85            .get("edit")
86            .ok_or_else(|| ToolError::InvalidInput("'edit' is required".into()))?;
87
88        let action = edit
89            .get("action")
90            .and_then(|v| v.as_str())
91            .ok_or_else(|| ToolError::InvalidInput("'edit.action' is required".into()))?;
92
93        // Read and parse the notebook.
94        let content = tokio::fs::read_to_string(file_path)
95            .await
96            .map_err(|e| ToolError::ExecutionFailed(format!("Failed to read {file_path}: {e}")))?;
97
98        let mut notebook: serde_json::Value = serde_json::from_str(&content)
99            .map_err(|e| ToolError::ExecutionFailed(format!("Invalid notebook JSON: {e}")))?;
100
101        let cells = notebook
102            .get_mut("cells")
103            .and_then(|v| v.as_array_mut())
104            .ok_or_else(|| ToolError::ExecutionFailed("Notebook has no 'cells' array".into()))?;
105
106        match action {
107            "replace" => {
108                let idx = edit
109                    .get("cell_index")
110                    .and_then(|v| v.as_u64())
111                    .ok_or_else(|| {
112                        ToolError::InvalidInput("'cell_index' required for replace".into())
113                    })? as usize;
114
115                let new_content =
116                    edit.get("content")
117                        .and_then(|v| v.as_str())
118                        .ok_or_else(|| {
119                            ToolError::InvalidInput("'content' required for replace".into())
120                        })?;
121
122                if idx >= cells.len() {
123                    return Err(ToolError::InvalidInput(format!(
124                        "Cell index {idx} out of range (notebook has {} cells)",
125                        cells.len()
126                    )));
127                }
128
129                // Replace the cell source.
130                let source_lines: Vec<serde_json::Value> = new_content
131                    .lines()
132                    .map(|l| serde_json::Value::String(format!("{l}\n")))
133                    .collect();
134                cells[idx]["source"] = serde_json::Value::Array(source_lines);
135
136                // Clear outputs for code cells.
137                if cells[idx].get("cell_type").and_then(|v| v.as_str()) == Some("code") {
138                    cells[idx]["outputs"] = serde_json::Value::Array(vec![]);
139                    cells[idx]["execution_count"] = serde_json::Value::Null;
140                }
141            }
142            "insert" => {
143                let idx = edit
144                    .get("cell_index")
145                    .and_then(|v| v.as_u64())
146                    .unwrap_or(cells.len() as u64) as usize;
147
148                let cell_type = edit
149                    .get("cell_type")
150                    .and_then(|v| v.as_str())
151                    .unwrap_or("code");
152
153                let new_content = edit.get("content").and_then(|v| v.as_str()).unwrap_or("");
154
155                let source_lines: Vec<serde_json::Value> = new_content
156                    .lines()
157                    .map(|l| serde_json::Value::String(format!("{l}\n")))
158                    .collect();
159
160                let new_cell = if cell_type == "code" {
161                    json!({
162                        "cell_type": "code",
163                        "source": source_lines,
164                        "metadata": {},
165                        "outputs": [],
166                        "execution_count": null
167                    })
168                } else {
169                    json!({
170                        "cell_type": "markdown",
171                        "source": source_lines,
172                        "metadata": {}
173                    })
174                };
175
176                let idx = idx.min(cells.len());
177                cells.insert(idx, new_cell);
178            }
179            "delete" => {
180                let idx = edit
181                    .get("cell_index")
182                    .and_then(|v| v.as_u64())
183                    .ok_or_else(|| {
184                        ToolError::InvalidInput("'cell_index' required for delete".into())
185                    })? as usize;
186
187                if idx >= cells.len() {
188                    return Err(ToolError::InvalidInput(format!(
189                        "Cell index {idx} out of range (notebook has {} cells)",
190                        cells.len()
191                    )));
192                }
193
194                cells.remove(idx);
195            }
196            other => {
197                return Err(ToolError::InvalidInput(format!(
198                    "Unknown action '{other}'. Use 'replace', 'insert', or 'delete'."
199                )));
200            }
201        }
202
203        // Write the modified notebook.
204        let output = serde_json::to_string_pretty(&notebook)
205            .map_err(|e| ToolError::ExecutionFailed(format!("Failed to serialize: {e}")))?;
206
207        tokio::fs::write(file_path, &output)
208            .await
209            .map_err(|e| ToolError::ExecutionFailed(format!("Failed to write {file_path}: {e}")))?;
210
211        Ok(ToolResult::success(format!(
212            "Notebook {file_path}: {action} applied ({} cells total)",
213            notebook["cells"].as_array().map(|c| c.len()).unwrap_or(0)
214        )))
215    }
216}
agent_code_lib/tools/notebook_edit.rs

agent_code_lib/tools/
notebook_edit.rs
