caudate-cli 0.1.0__py3-none-any.whl

execution/tools/notebook_tool.py

@@ -0,0 +1,279 @@
+"""NotebookEdit — surgical edits to Jupyter `.ipynb` notebook cells.
+
+`.ipynb` files are JSON; you *can* edit them with `Edit` if you carefully
+match the surrounding JSON, but cell `source` is stored as a list of
+strings with embedded escapes — it's brittle and easy to corrupt. This
+tool parses the notebook properly, performs the requested cell-level
+operation, and writes valid notebook JSON back out.
+
+Modes (`mode` argument):
+  - `replace` (default) — replace the entire source of one cell.
+  - `insert`            — insert a new cell at `cell_index` (existing
+                           cells shift down). Existing default if no
+                           cell_index → append at the end.
+  - `delete`            — delete the cell at `cell_index`.
+
+Cell types: `code` (default), `markdown`, `raw`. New cells default to
+`code` unless `cell_type` is set.
+
+Identifies the cell by either:
+  - `cell_index` (0-based integer), or
+  - `cell_id` (notebook v4.5+ stable id like `abc123`).
+
+Outputs (cell `outputs` and `execution_count`) are reset on edits to
+`code` cells so the file looks "fresh" after editing — same convention
+Claude Code's NotebookEdit uses.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+from typing import Any
+
+from core.schemas import ToolResult
+from execution.tools.base import BaseTool
+
+logger = logging.getLogger(__name__)
+
+
+_VALID_MODES = {"replace", "insert", "delete"}
+_VALID_CELL_TYPES = {"code", "markdown", "raw"}
+
+
+def _normalize_source(text: str) -> list[str]:
+    """Notebook spec wants `source` as a list of lines (each line ends
+    in `\\n` except possibly the last). Most notebooks store it that
+    way; some store a single string. We always emit the list form."""
+    if not text:
+        return []
+    # Preserve final newline behavior of the original text.
+    lines = text.splitlines(keepends=True)
+    return lines
+
+
+class NotebookEditTool(BaseTool):
+    mutates = True
+    name = "NotebookEdit"
+    description = (
+        "Edit Jupyter `.ipynb` notebook cells in place. Modes: "
+        "`replace` an existing cell's source, `insert` a new cell, or "
+        "`delete` a cell. Identify the target by `cell_index` (0-based) "
+        "or `cell_id` (v4.5+). New cells default to `code`; pass "
+        "`cell_type` for `markdown` / `raw`. Code-cell outputs are "
+        "reset on edit. Use this instead of the generic `Edit` tool — "
+        "the JSON-with-escaped-strings shape is too brittle for "
+        "string replacement."
+    )
+
+    @property
+    def input_schema(self) -> dict[str, Any]:
+        return {
+            "type": "object",
+            "properties": {
+                "path": {
+                    "type": "string",
+                    "description": "Absolute or relative path to the .ipynb file.",
+                },
+                "mode": {
+                    "type": "string",
+                    "enum": list(_VALID_MODES),
+                    "description": "replace (default) | insert | delete",
+                    "default": "replace",
+                },
+                "cell_index": {
+                    "type": "integer",
+                    "description": (
+                        "0-based cell index. For `insert`, this is "
+                        "where the new cell goes; omit to append at the "
+                        "end. For `replace`/`delete`, this identifies "
+                        "the target cell (or use `cell_id`)."
+                    ),
+                },
+                "cell_id": {
+                    "type": "string",
+                    "description": (
+                        "Stable cell ID (notebook v4.5+). Alternative "
+                        "to `cell_index` for `replace`/`delete`."
+                    ),
+                },
+                "source": {
+                    "type": "string",
+                    "description": (
+                        "New cell content (required for `replace` and "
+                        "`insert`). Pass the raw text — newlines are "
+                        "handled automatically."
+                    ),
+                },
+                "cell_type": {
+                    "type": "string",
+                    "enum": list(_VALID_CELL_TYPES),
+                    "description": "code (default) | markdown | raw. Used for `insert`.",
+                    "default": "code",
+                },
+            },
+            "required": ["path"],
+        }
+
+    async def execute(self, **kwargs: Any) -> ToolResult:
+        path_str = (kwargs.get("path") or "").strip()
+        mode = (kwargs.get("mode") or "replace").lower()
+        cell_index = kwargs.get("cell_index")
+        cell_id = kwargs.get("cell_id")
+        source = kwargs.get("source")
+        cell_type = (kwargs.get("cell_type") or "code").lower()
+
+        # ---- validation ----
+        if not path_str:
+            return ToolResult(
+                tool_name=self.name, status="error",
+                error="`path` is required",
+            )
+        if mode not in _VALID_MODES:
+            return ToolResult(
+                tool_name=self.name, status="error",
+                error=f"unknown mode {mode!r}; use one of {sorted(_VALID_MODES)}",
+            )
+        if cell_type not in _VALID_CELL_TYPES:
+            return ToolResult(
+                tool_name=self.name, status="error",
+                error=f"unknown cell_type {cell_type!r}; use one of {sorted(_VALID_CELL_TYPES)}",
+            )
+        if mode in ("replace", "insert") and source is None:
+            return ToolResult(
+                tool_name=self.name, status="error",
+                error=f"`source` is required for mode={mode!r}",
+            )
+
+        path = Path(path_str).expanduser().resolve()
+        if not path.exists() or not path.is_file():
+            return ToolResult(
+                tool_name=self.name, status="error",
+                error=f"notebook not found: {path}",
+            )
+        if path.suffix != ".ipynb":
+            return ToolResult(
+                tool_name=self.name, status="error",
+                error=f"not a .ipynb file: {path.suffix}",
+            )
+
+        # ---- load ----
+        try:
+            nb = json.loads(path.read_text(encoding="utf-8"))
+        except (OSError, json.JSONDecodeError) as e:
+            return ToolResult(
+                tool_name=self.name, status="error",
+                error=f"failed to parse notebook: {e}",
+            )
+
+        cells = nb.get("cells")
+        if not isinstance(cells, list):
+            return ToolResult(
+                tool_name=self.name, status="error",
+                error="notebook has no `cells` array — corrupted file?",
+            )
+
+        # ---- resolve target index ----
+        target_idx: int | None = None
+        if cell_id is not None and cell_id != "":
+            for i, c in enumerate(cells):
+                if isinstance(c, dict) and c.get("id") == cell_id:
+                    target_idx = i
+                    break
+            if target_idx is None and mode != "insert":
+                return ToolResult(
+                    tool_name=self.name, status="error",
+                    error=f"no cell with id {cell_id!r} (notebook has {len(cells)} cells)",
+                )
+        elif cell_index is not None:
+            try:
+                target_idx = int(cell_index)
+            except (TypeError, ValueError):
+                return ToolResult(
+                    tool_name=self.name, status="error",
+                    error=f"cell_index must be an integer, got {cell_index!r}",
+                )
+            if mode != "insert":
+                if not (0 <= target_idx < len(cells)):
+                    return ToolResult(
+                        tool_name=self.name, status="error",
+                        error=(f"cell_index {target_idx} out of range "
+                               f"(notebook has {len(cells)} cells)"),
+                    )
+
+        # ---- apply ----
+        action_summary = ""
+        if mode == "replace":
+            if target_idx is None:
+                return ToolResult(
+                    tool_name=self.name, status="error",
+                    error="replace mode needs `cell_index` or `cell_id`",
+                )
+            cell = cells[target_idx]
+            if not isinstance(cell, dict):
+                return ToolResult(
+                    tool_name=self.name, status="error",
+                    error=f"cell at index {target_idx} is not an object",
+                )
+            cell["source"] = _normalize_source(source)
+            # Reset transient state on code cells so the file looks
+            # fresh after edit. Mirrors Claude Code's NotebookEdit.
+            if cell.get("cell_type") == "code":
+                cell["outputs"] = []
+                cell["execution_count"] = None
+            action_summary = (
+                f"replaced cell {target_idx} "
+                f"({cell.get('cell_type', '?')}, {len(source)} chars)"
+            )
+
+        elif mode == "delete":
+            if target_idx is None:
+                return ToolResult(
+                    tool_name=self.name, status="error",
+                    error="delete mode needs `cell_index` or `cell_id`",
+                )
+            removed = cells.pop(target_idx)
+            action_summary = (
+                f"deleted cell {target_idx} "
+                f"({removed.get('cell_type', '?')})"
+            )
+
+        elif mode == "insert":
+            new_cell: dict[str, Any] = {
+                "cell_type": cell_type,
+                "metadata": {},
+                "source": _normalize_source(source),
+            }
+            if cell_type == "code":
+                new_cell["outputs"] = []
+                new_cell["execution_count"] = None
+            insert_at = target_idx if target_idx is not None else len(cells)
+            insert_at = max(0, min(insert_at, len(cells)))
+            cells.insert(insert_at, new_cell)
+            action_summary = (
+                f"inserted {cell_type} cell at index {insert_at} "
+                f"({len(source)} chars)"
+            )
+
+        # ---- write back ----
+        try:
+            path.write_text(
+                json.dumps(nb, indent=1, ensure_ascii=False) + "\n",
+                encoding="utf-8",
+            )
+        except OSError as e:
+            return ToolResult(
+                tool_name=self.name, status="error",
+                error=f"failed to write notebook: {e}",
+            )
+
+        return ToolResult(
+            tool_name=self.name, status="success",
+            output=f"{path}: {action_summary} (notebook now has {len(cells)} cells)",
+            metadata={
+                "path": str(path),
+                "mode": mode,
+                "cell_count": len(cells),
+            },
+        )