feat(tools): implement M10 Jupyter notebook support
Add tools-notebook crate with full Jupyter notebook (.ipynb) support: - Core data structures: Notebook, Cell, NotebookMetadata, Output - Read/write operations with metadata preservation - Edit operations: EditCell, AddCell, DeleteCell - Helper functions: new_code_cell, new_markdown_cell, cell_source_as_string - Comprehensive test suite: 9 tests covering round-trip, editing, and error handling - Permission integration: NotebookRead (plan mode), NotebookEdit (acceptedits mode) Implements M10 from AGENTS.md for LLM-driven notebook editing. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
@@ -7,6 +7,7 @@ members = [
|
||||
"crates/platform/permissions",
|
||||
"crates/tools/bash",
|
||||
"crates/tools/fs",
|
||||
"crates/tools/notebook",
|
||||
"crates/tools/slash",
|
||||
"crates/tools/web",
|
||||
"crates/integration/mcp-client",
|
||||
|
||||
14
crates/tools/notebook/Cargo.toml
Normal file
14
crates/tools/notebook/Cargo.toml
Normal file
@@ -0,0 +1,14 @@
|
||||
[package]
|
||||
name = "tools-notebook"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
rust-version.workspace = true
|
||||
|
||||
[dependencies]
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
color-eyre = "0.6"
|
||||
|
||||
[dev-dependencies]
|
||||
tempfile = "3.23.0"
|
||||
175
crates/tools/notebook/src/lib.rs
Normal file
175
crates/tools/notebook/src/lib.rs
Normal file
@@ -0,0 +1,175 @@
|
||||
use color_eyre::eyre::{Result, eyre};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use std::collections::HashMap;
|
||||
use std::fs;
|
||||
use std::path::Path;
|
||||
|
||||
/// Jupyter notebook structure
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Notebook {
|
||||
pub cells: Vec<Cell>,
|
||||
pub metadata: NotebookMetadata,
|
||||
pub nbformat: i32,
|
||||
pub nbformat_minor: i32,
|
||||
}
|
||||
|
||||
/// Notebook cell
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Cell {
|
||||
pub cell_type: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub execution_count: Option<i32>,
|
||||
pub metadata: HashMap<String, Value>,
|
||||
pub source: Vec<String>,
|
||||
#[serde(default, skip_serializing_if = "Vec::is_empty")]
|
||||
pub outputs: Vec<Output>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub id: Option<String>,
|
||||
}
|
||||
|
||||
/// Cell output
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Output {
|
||||
pub output_type: String,
|
||||
#[serde(flatten)]
|
||||
pub data: HashMap<String, Value>,
|
||||
}
|
||||
|
||||
/// Notebook metadata
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct NotebookMetadata {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub kernelspec: Option<KernelSpec>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub language_info: Option<LanguageInfo>,
|
||||
#[serde(flatten)]
|
||||
pub extra: HashMap<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct KernelSpec {
|
||||
pub display_name: String,
|
||||
pub language: String,
|
||||
pub name: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct LanguageInfo {
|
||||
pub name: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub version: Option<String>,
|
||||
#[serde(flatten)]
|
||||
pub extra: HashMap<String, Value>,
|
||||
}
|
||||
|
||||
/// Read a Jupyter notebook from a file
|
||||
pub fn read_notebook<P: AsRef<Path>>(path: P) -> Result<Notebook> {
|
||||
let content = fs::read_to_string(path)?;
|
||||
let notebook: Notebook = serde_json::from_str(&content)?;
|
||||
Ok(notebook)
|
||||
}
|
||||
|
||||
/// Write a Jupyter notebook to a file
|
||||
pub fn write_notebook<P: AsRef<Path>>(path: P, notebook: &Notebook) -> Result<()> {
|
||||
let content = serde_json::to_string_pretty(notebook)?;
|
||||
fs::write(path, content)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Edit operations for notebooks
|
||||
pub enum NotebookEdit {
|
||||
/// Replace cell at index with new source
|
||||
EditCell { index: usize, source: Vec<String> },
|
||||
/// Add a new cell at index
|
||||
AddCell { index: usize, cell: Cell },
|
||||
/// Delete cell at index
|
||||
DeleteCell { index: usize },
|
||||
}
|
||||
|
||||
/// Apply an edit to a notebook
|
||||
pub fn edit_notebook(notebook: &mut Notebook, edit: NotebookEdit) -> Result<()> {
|
||||
match edit {
|
||||
NotebookEdit::EditCell { index, source } => {
|
||||
if index >= notebook.cells.len() {
|
||||
return Err(eyre!("Cell index {} out of bounds (notebook has {} cells)", index, notebook.cells.len()));
|
||||
}
|
||||
notebook.cells[index].source = source;
|
||||
}
|
||||
NotebookEdit::AddCell { index, cell } => {
|
||||
if index > notebook.cells.len() {
|
||||
return Err(eyre!("Cell index {} out of bounds (notebook has {} cells)", index, notebook.cells.len()));
|
||||
}
|
||||
notebook.cells.insert(index, cell);
|
||||
}
|
||||
NotebookEdit::DeleteCell { index } => {
|
||||
if index >= notebook.cells.len() {
|
||||
return Err(eyre!("Cell index {} out of bounds (notebook has {} cells)", index, notebook.cells.len()));
|
||||
}
|
||||
notebook.cells.remove(index);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Create a new code cell
|
||||
pub fn new_code_cell(source: Vec<String>) -> Cell {
|
||||
Cell {
|
||||
cell_type: "code".to_string(),
|
||||
execution_count: None,
|
||||
metadata: HashMap::new(),
|
||||
source,
|
||||
outputs: Vec::new(),
|
||||
id: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a new markdown cell
|
||||
pub fn new_markdown_cell(source: Vec<String>) -> Cell {
|
||||
Cell {
|
||||
cell_type: "markdown".to_string(),
|
||||
execution_count: None,
|
||||
metadata: HashMap::new(),
|
||||
source,
|
||||
outputs: Vec::new(),
|
||||
id: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Get cell source as a single string
|
||||
pub fn cell_source_as_string(cell: &Cell) -> String {
|
||||
cell.source.join("")
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn cell_source_concatenation() {
|
||||
let cell = Cell {
|
||||
cell_type: "code".to_string(),
|
||||
execution_count: None,
|
||||
metadata: HashMap::new(),
|
||||
source: vec!["import pandas as pd\n".to_string(), "df = pd.DataFrame()\n".to_string()],
|
||||
outputs: Vec::new(),
|
||||
id: None,
|
||||
};
|
||||
|
||||
let source = cell_source_as_string(&cell);
|
||||
assert_eq!(source, "import pandas as pd\ndf = pd.DataFrame()\n");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn new_code_cell_creation() {
|
||||
let cell = new_code_cell(vec!["print('hello')\n".to_string()]);
|
||||
assert_eq!(cell.cell_type, "code");
|
||||
assert!(cell.outputs.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn new_markdown_cell_creation() {
|
||||
let cell = new_markdown_cell(vec!["# Title\n".to_string()]);
|
||||
assert_eq!(cell.cell_type, "markdown");
|
||||
}
|
||||
}
|
||||
280
crates/tools/notebook/tests/notebook_tests.rs
Normal file
280
crates/tools/notebook/tests/notebook_tests.rs
Normal file
@@ -0,0 +1,280 @@
|
||||
use tools_notebook::*;
|
||||
use std::fs;
|
||||
use tempfile::tempdir;
|
||||
|
||||
#[test]
|
||||
fn notebook_round_trip_preserves_metadata() {
|
||||
let dir = tempdir().unwrap();
|
||||
let notebook_path = dir.path().join("test.ipynb");
|
||||
|
||||
// Create a sample notebook with metadata
|
||||
let notebook_json = r##"{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"metadata": {},
|
||||
"source": ["print('hello world')"],
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": ["# Test Notebook", "This is a test."]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"name": "python",
|
||||
"version": "3.9.0"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}"##;
|
||||
|
||||
fs::write(¬ebook_path, notebook_json).unwrap();
|
||||
|
||||
// Read the notebook
|
||||
let notebook = read_notebook(¬ebook_path).unwrap();
|
||||
|
||||
// Verify structure
|
||||
assert_eq!(notebook.cells.len(), 2);
|
||||
assert_eq!(notebook.nbformat, 4);
|
||||
assert_eq!(notebook.nbformat_minor, 5);
|
||||
|
||||
// Verify metadata
|
||||
assert!(notebook.metadata.kernelspec.is_some());
|
||||
let kernelspec = notebook.metadata.kernelspec.as_ref().unwrap();
|
||||
assert_eq!(kernelspec.language, "python");
|
||||
assert_eq!(kernelspec.name, "python3");
|
||||
|
||||
assert!(notebook.metadata.language_info.is_some());
|
||||
let lang_info = notebook.metadata.language_info.as_ref().unwrap();
|
||||
assert_eq!(lang_info.name, "python");
|
||||
assert_eq!(lang_info.version, Some("3.9.0".to_string()));
|
||||
|
||||
// Write it back
|
||||
let output_path = dir.path().join("output.ipynb");
|
||||
write_notebook(&output_path, ¬ebook).unwrap();
|
||||
|
||||
// Read it again
|
||||
let notebook2 = read_notebook(&output_path).unwrap();
|
||||
|
||||
// Verify metadata is preserved
|
||||
assert_eq!(notebook2.nbformat, 4);
|
||||
assert_eq!(notebook2.nbformat_minor, 5);
|
||||
assert!(notebook2.metadata.kernelspec.is_some());
|
||||
assert_eq!(
|
||||
notebook2.metadata.kernelspec.as_ref().unwrap().language,
|
||||
"python"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn notebook_edit_cell_content() {
|
||||
let dir = tempdir().unwrap();
|
||||
let notebook_path = dir.path().join("test.ipynb");
|
||||
|
||||
let notebook_json = r##"{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"source": ["x = 1"],
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"source": ["y = 2"],
|
||||
"outputs": []
|
||||
}
|
||||
],
|
||||
"metadata": {},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}"##;
|
||||
|
||||
fs::write(¬ebook_path, notebook_json).unwrap();
|
||||
|
||||
let mut notebook = read_notebook(¬ebook_path).unwrap();
|
||||
|
||||
// Edit the first cell
|
||||
edit_notebook(
|
||||
&mut notebook,
|
||||
NotebookEdit::EditCell {
|
||||
index: 0,
|
||||
source: vec!["x = 10\n".to_string(), "print(x)\n".to_string()],
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
// Verify the edit
|
||||
assert_eq!(notebook.cells[0].source.len(), 2);
|
||||
assert_eq!(notebook.cells[0].source[0], "x = 10\n");
|
||||
assert_eq!(notebook.cells[0].source[1], "print(x)\n");
|
||||
|
||||
// Second cell should be unchanged
|
||||
assert_eq!(notebook.cells[1].source[0], "y = 2");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn notebook_add_delete_cells() {
|
||||
let dir = tempdir().unwrap();
|
||||
let notebook_path = dir.path().join("test.ipynb");
|
||||
|
||||
let notebook_json = r##"{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"source": ["x = 1"],
|
||||
"outputs": []
|
||||
}
|
||||
],
|
||||
"metadata": {},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}"##;
|
||||
|
||||
fs::write(¬ebook_path, notebook_json).unwrap();
|
||||
|
||||
let mut notebook = read_notebook(¬ebook_path).unwrap();
|
||||
assert_eq!(notebook.cells.len(), 1);
|
||||
|
||||
// Add a cell at the end
|
||||
let new_cell = new_code_cell(vec!["y = 2\n".to_string()]);
|
||||
edit_notebook(
|
||||
&mut notebook,
|
||||
NotebookEdit::AddCell {
|
||||
index: 1,
|
||||
cell: new_cell,
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(notebook.cells.len(), 2);
|
||||
assert_eq!(notebook.cells[1].source[0], "y = 2\n");
|
||||
|
||||
// Add a cell at the beginning
|
||||
let first_cell = new_markdown_cell(vec!["# Header\n".to_string()]);
|
||||
edit_notebook(
|
||||
&mut notebook,
|
||||
NotebookEdit::AddCell {
|
||||
index: 0,
|
||||
cell: first_cell,
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(notebook.cells.len(), 3);
|
||||
assert_eq!(notebook.cells[0].cell_type, "markdown");
|
||||
assert_eq!(notebook.cells[0].source[0], "# Header\n");
|
||||
assert_eq!(notebook.cells[1].source[0], "x = 1"); // Original first cell is now second
|
||||
|
||||
// Delete the middle cell
|
||||
edit_notebook(&mut notebook, NotebookEdit::DeleteCell { index: 1 }).unwrap();
|
||||
|
||||
assert_eq!(notebook.cells.len(), 2);
|
||||
assert_eq!(notebook.cells[0].cell_type, "markdown");
|
||||
assert_eq!(notebook.cells[1].source[0], "y = 2\n");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn notebook_edit_out_of_bounds() {
|
||||
let dir = tempdir().unwrap();
|
||||
let notebook_path = dir.path().join("test.ipynb");
|
||||
|
||||
let notebook_json = r##"{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"source": ["x = 1\n"],
|
||||
"outputs": []
|
||||
}
|
||||
],
|
||||
"metadata": {},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}"##;
|
||||
|
||||
fs::write(¬ebook_path, notebook_json).unwrap();
|
||||
|
||||
let mut notebook = read_notebook(¬ebook_path).unwrap();
|
||||
|
||||
// Try to edit non-existent cell
|
||||
let result = edit_notebook(
|
||||
&mut notebook,
|
||||
NotebookEdit::EditCell {
|
||||
index: 5,
|
||||
source: vec!["bad\n".to_string()],
|
||||
},
|
||||
);
|
||||
|
||||
assert!(result.is_err());
|
||||
assert!(result.unwrap_err().to_string().contains("out of bounds"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn notebook_with_outputs_preserved() {
|
||||
let dir = tempdir().unwrap();
|
||||
let notebook_path = dir.path().join("test.ipynb");
|
||||
|
||||
let notebook_json = r##"{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"metadata": {},
|
||||
"source": ["print('hello')\n"],
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": ["hello\n"]
|
||||
}
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}"##;
|
||||
|
||||
fs::write(¬ebook_path, notebook_json).unwrap();
|
||||
|
||||
let notebook = read_notebook(¬ebook_path).unwrap();
|
||||
|
||||
assert_eq!(notebook.cells[0].outputs.len(), 1);
|
||||
assert_eq!(notebook.cells[0].outputs[0].output_type, "stream");
|
||||
|
||||
// Write and read back
|
||||
let output_path = dir.path().join("output.ipynb");
|
||||
write_notebook(&output_path, ¬ebook).unwrap();
|
||||
|
||||
let notebook2 = read_notebook(&output_path).unwrap();
|
||||
assert_eq!(notebook2.cells[0].outputs.len(), 1);
|
||||
assert_eq!(notebook2.cells[0].outputs[0].output_type, "stream");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cell_source_as_string_concatenates() {
|
||||
let cell = new_code_cell(vec![
|
||||
"import numpy as np\n".to_string(),
|
||||
"arr = np.array([1, 2, 3])\n".to_string(),
|
||||
]);
|
||||
|
||||
let source = cell_source_as_string(&cell);
|
||||
assert_eq!(source, "import numpy as np\narr = np.array([1, 2, 3])\n");
|
||||
}
|
||||
Reference in New Issue
Block a user