hanzo-mcp 0.1.21__py3-none-any.whl

hanzo_mcp/tools/jupyter/__init__.py

@@ -0,0 +1,8 @@
+"""Jupyter notebook tools for Hanzo Dev MCP.
+
+This module provides tools for reading and editing Jupyter notebooks (.ipynb files).
+"""
+
+from hanzo_mcp.tools.jupyter.notebook_operations import JupyterNotebookTools
+
+__all__ = ["JupyterNotebookTools"]

hanzo_mcp/tools/jupyter/notebook_operations.py

@@ -0,0 +1,554 @@
+"""Jupyter notebook operations tools for Hanzo Dev MCP.
+
+This module provides tools for reading and editing Jupyter notebook (.ipynb) files.
+It supports reading notebook cells with their outputs and modifying notebook contents.
+"""
+
+import json
+import re
+from pathlib import Path
+from typing import Any, final, Literal
+
+from mcp.server.fastmcp import Context as MCPContext
+from mcp.server.fastmcp import FastMCP
+
+from hanzo_mcp.tools.common.context import DocumentContext, create_tool_context
+from hanzo_mcp.tools.common.permissions import PermissionManager
+from hanzo_mcp.tools.common.validation import validate_path_parameter
+
+
+# Pattern to match ANSI escape sequences
+ANSI_ESCAPE_PATTERN = re.compile(r"\x1B\[[0-9;]*[a-zA-Z]")
+
+
+# Function to clean ANSI escape codes from text
+def clean_ansi_escapes(text: str) -> str:
+    """Remove ANSI escape sequences from text.
+
+    Args:
+        text: Text containing ANSI escape sequences
+
+    Returns:
+        Text with ANSI escape sequences removed
+    """
+    return ANSI_ESCAPE_PATTERN.sub("", text)
+
+
+# Type definitions for Jupyter notebooks based on the nbformat spec
+CellType = Literal["code", "markdown"]
+OutputType = Literal["stream", "display_data", "execute_result", "error"]
+EditMode = Literal["replace", "insert", "delete"]
+
+
+# Define a structure for notebook cell outputs
+@final
+class NotebookOutputImage:
+    """Representation of an image output in a notebook cell."""
+
+    def __init__(self, image_data: str, media_type: str):
+        """Initialize a notebook output image.
+
+        Args:
+            image_data: Base64-encoded image data
+            media_type: Media type of the image (e.g., "image/png")
+        """
+        self.image_data = image_data
+        self.media_type = media_type
+
+
+@final
+class NotebookCellOutput:
+    """Representation of an output from a notebook cell."""
+
+    def __init__(
+        self,
+        output_type: OutputType,
+        text: str | None = None,
+        image: NotebookOutputImage | None = None,
+    ):
+        """Initialize a notebook cell output.
+
+        Args:
+            output_type: Type of output
+            text: Text output (if any)
+            image: Image output (if any)
+        """
+        self.output_type = output_type
+        self.text = text
+        self.image = image
+
+
+@final
+class NotebookCellSource:
+    """Representation of a source cell from a notebook."""
+
+    def __init__(
+        self,
+        cell_index: int,
+        cell_type: CellType,
+        source: str,
+        language: str,
+        execution_count: int | None = None,
+        outputs: list[NotebookCellOutput] | None = None,
+    ):
+        """Initialize a notebook cell source.
+
+        Args:
+            cell_index: Index of the cell in the notebook
+            cell_type: Type of cell (code or markdown)
+            source: Source code or text of the cell
+            language: Programming language of the cell
+            execution_count: Execution count of the cell (if any)
+            outputs: Outputs from the cell (if any)
+        """
+        self.cell_index = cell_index
+        self.cell_type = cell_type
+        self.source = source
+        self.language = language
+        self.execution_count = execution_count
+        self.outputs = outputs or []
+
+
+@final
+class JupyterNotebookTools:
+    """Tools for working with Jupyter notebooks."""
+
+    def __init__(
+        self, document_context: DocumentContext, permission_manager: PermissionManager
+    ) -> None:
+        """Initialize notebook tools.
+
+        Args:
+            document_context: Document context for tracking file contents
+            permission_manager: Permission manager for access control
+        """
+        self.document_context: DocumentContext = document_context
+        self.permission_manager: PermissionManager = permission_manager
+
+    def register_tools(self, mcp_server: FastMCP) -> None:
+        """Register jupyter notebook tools with the MCP server.
+
+        Args:
+            mcp_server: The FastMCP server instance
+        """
+
+        @mcp_server.tool()
+        async def read_notebook(path: str, ctx: MCPContext) -> str:
+            """Extract and read source code from all cells in a Jupyter notebook.
+
+            Reads a Jupyter notebook (.ipynb file) and returns all of the cells with
+            their outputs. Jupyter notebooks are interactive documents that combine
+            code, text, and visualizations, commonly used for data analysis and
+            scientific computing.
+
+            Args:
+                path: Absolute path to the notebook file (must be absolute, not relative)
+
+            Returns:
+                Formatted content of all notebook cells with their outputs
+            """
+            tool_ctx = create_tool_context(ctx)
+            tool_ctx.set_tool_info("read_notebook")
+
+            # Validate path parameter
+            path_validation = validate_path_parameter(path)
+            if path_validation.is_error:
+                await tool_ctx.error(path_validation.error_message)
+                return f"Error: {path_validation.error_message}"
+
+            await tool_ctx.info(f"Reading notebook: {path}")
+
+            # Check if path is allowed to be read
+            if not self.permission_manager.is_path_allowed(path):
+                await tool_ctx.error(
+                    f"Access denied - path outside allowed directories: {path}"
+                )
+                return (
+                    f"Error: Access denied - path outside allowed directories: {path}"
+                )
+
+            try:
+                file_path = Path(path)
+
+                if not file_path.exists():
+                    await tool_ctx.error(f"File does not exist: {path}")
+                    return f"Error: File does not exist: {path}"
+
+                if not file_path.is_file():
+                    await tool_ctx.error(f"Path is not a file: {path}")
+                    return f"Error: Path is not a file: {path}"
+
+                # Check file extension
+                if file_path.suffix.lower() != ".ipynb":
+                    await tool_ctx.error(f"File is not a Jupyter notebook: {path}")
+                    return f"Error: File is not a Jupyter notebook: {path}"
+
+                # Read and parse the notebook
+                try:
+                    with open(file_path, "r", encoding="utf-8") as f:
+                        content = f.read()
+                        notebook = json.loads(content)
+                except json.JSONDecodeError:
+                    await tool_ctx.error(f"Invalid notebook format: {path}")
+                    return f"Error: Invalid notebook format: {path}"
+                except UnicodeDecodeError:
+                    await tool_ctx.error(f"Cannot read notebook file: {path}")
+                    return f"Error: Cannot read notebook file: {path}"
+
+                # Add to document context
+                self.document_context.add_document(path, content)
+
+                # Process notebook cells
+                language = (
+                    notebook.get("metadata", {})
+                    .get("language_info", {})
+                    .get("name", "python")
+                )
+                cells = notebook.get("cells", [])
+                processed_cells = []
+
+                for i, cell in enumerate(cells):
+                    cell_type = cell.get("cell_type", "code")
+
+                    # Skip if not code or markdown
+                    if cell_type not in ["code", "markdown"]:
+                        continue
+
+                    # Get source
+                    source = cell.get("source", "")
+                    if isinstance(source, list):
+                        source = "".join(source)
+
+                    # Get execution count for code cells
+                    execution_count = None
+                    if cell_type == "code":
+                        execution_count = cell.get("execution_count")
+
+                    # Process outputs for code cells
+                    outputs = []
+                    if cell_type == "code" and "outputs" in cell:
+                        for output in cell["outputs"]:
+                            output_type = output.get("output_type", "")
+
+                            # Process different output types
+                            if output_type == "stream":
+                                text = output.get("text", "")
+                                if isinstance(text, list):
+                                    text = "".join(text)
+                                outputs.append(
+                                    NotebookCellOutput(output_type="stream", text=text)
+                                )
+
+                            elif output_type in ["execute_result", "display_data"]:
+                                # Process text output
+                                text = None
+                                if "data" in output and "text/plain" in output["data"]:
+                                    text_data = output["data"]["text/plain"]
+                                    if isinstance(text_data, list):
+                                        text = "".join(text_data)
+                                    else:
+                                        text = text_data
+
+                                # Process image output
+                                image = None
+                                if "data" in output:
+                                    if "image/png" in output["data"]:
+                                        image = NotebookOutputImage(
+                                            image_data=output["data"]["image/png"],
+                                            media_type="image/png",
+                                        )
+                                    elif "image/jpeg" in output["data"]:
+                                        image = NotebookOutputImage(
+                                            image_data=output["data"]["image/jpeg"],
+                                            media_type="image/jpeg",
+                                        )
+
+                                outputs.append(
+                                    NotebookCellOutput(
+                                        output_type=output_type, text=text, image=image
+                                    )
+                                )
+
+                            elif output_type == "error":
+                                # Format error traceback
+                                ename = output.get("ename", "")
+                                evalue = output.get("evalue", "")
+                                traceback = output.get("traceback", [])
+
+                                # Handle raw text strings and lists of strings
+                                if isinstance(traceback, list):
+                                    # Clean ANSI escape codes and join the list but preserve the formatting
+                                    clean_traceback = [
+                                        clean_ansi_escapes(line) for line in traceback
+                                    ]
+                                    traceback_text = "\n".join(clean_traceback)
+                                else:
+                                    traceback_text = clean_ansi_escapes(str(traceback))
+
+                                error_text = f"{ename}: {evalue}\n{traceback_text}"
+                                outputs.append(
+                                    NotebookCellOutput(
+                                        output_type="error", text=error_text
+                                    )
+                                )
+
+                    # Create cell object
+                    processed_cell = NotebookCellSource(
+                        cell_index=i,
+                        cell_type=cell_type,
+                        source=source,
+                        language=language,
+                        execution_count=execution_count,
+                        outputs=outputs,
+                    )
+
+                    processed_cells.append(processed_cell)
+
+                # Format the notebook content as a readable string
+                result = []
+                for cell in processed_cells:
+                    # Format the cell header
+                    cell_header = f"Cell [{cell.cell_index}] {cell.cell_type}"
+                    if cell.execution_count is not None:
+                        cell_header += f" (execution_count: {cell.execution_count})"
+                    if cell.cell_type == "code" and cell.language != "python":
+                        cell_header += f" [{cell.language}]"
+
+                    # Add cell to result
+                    result.append(f"{cell_header}:")
+                    result.append(
+                        f"```{cell.language if cell.cell_type == 'code' else ''}"
+                    )
+                    result.append(cell.source)
+                    result.append("```")
+
+                    # Add outputs if any
+                    if cell.outputs:
+                        result.append("Outputs:")
+                        for output in cell.outputs:
+                            if output.output_type == "error":
+                                result.append("Error:")
+                                result.append("```")
+                                result.append(output.text)
+                                result.append("```")
+                            elif output.text:
+                                result.append("Output:")
+                                result.append("```")
+                                result.append(output.text)
+                                result.append("```")
+                            if output.image:
+                                result.append(
+                                    f"[Image output: {output.image.media_type}]"
+                                )
+
+                    result.append("")  # Empty line between cells
+
+                await tool_ctx.info(
+                    f"Successfully read notebook: {path} ({len(processed_cells)} cells)"
+                )
+                return "\n".join(result)
+            except Exception as e:
+                await tool_ctx.error(f"Error reading notebook: {str(e)}")
+                return f"Error reading notebook: {str(e)}"
+
+        @mcp_server.tool()
+        async def edit_notebook(
+            path: str,
+            cell_number: int,
+            new_source: str,
+            ctx: MCPContext,
+            cell_type: CellType | None = None,
+            edit_mode: EditMode = "replace",
+        ) -> str:
+            """Edit a specific cell in a Jupyter notebook.
+
+            Enables editing, inserting, or deleting cells in a Jupyter notebook (.ipynb file).
+            In replace mode, the specified cell's source is updated with the new content.
+            In insert mode, a new cell is added at the specified index.
+            In delete mode, the specified cell is removed.
+
+            Args:
+                path: Absolute path to the notebook file (must be absolute, not relative)
+                cell_number: Zero-based index of the cell to edit
+                new_source: New source code or text for the cell (ignored in delete mode)
+                cell_type: Type of cell (code or markdown), default is to keep existing type
+                edit_mode: Type of edit operation (replace, insert, delete), default is replace
+
+            Returns:
+                Result of the edit operation with details on changes made
+            """
+            tool_ctx = create_tool_context(ctx)
+            tool_ctx.set_tool_info("edit_notebook")
+
+            # Validate path parameter
+            path_validation = validate_path_parameter(path)
+            if path_validation.is_error:
+                await tool_ctx.error(path_validation.error_message)
+                return f"Error: {path_validation.error_message}"
+
+            # Validate cell_number
+            if cell_number < 0:
+                await tool_ctx.error("Cell number must be non-negative")
+                return "Error: Cell number must be non-negative"
+
+            # Validate edit_mode
+            if edit_mode not in ["replace", "insert", "delete"]:
+                await tool_ctx.error("Edit mode must be replace, insert, or delete")
+                return "Error: Edit mode must be replace, insert, or delete"
+
+            # In insert mode, cell_type is required
+            if edit_mode == "insert" and cell_type is None:
+                await tool_ctx.error("Cell type is required when using insert mode")
+                return "Error: Cell type is required when using insert mode"
+
+            # Don't validate new_source for delete mode
+            if edit_mode != "delete" and not new_source:
+                await tool_ctx.error(
+                    "New source is required for replace or insert operations"
+                )
+                return "Error: New source is required for replace or insert operations"
+
+            await tool_ctx.info(
+                f"Editing notebook: {path} (cell: {cell_number}, mode: {edit_mode})"
+            )
+
+            # Check if path is allowed
+            if not self.permission_manager.is_path_allowed(path):
+                await tool_ctx.error(
+                    f"Access denied - path outside allowed directories: {path}"
+                )
+                return (
+                    f"Error: Access denied - path outside allowed directories: {path}"
+                )
+
+            try:
+                file_path = Path(path)
+
+                if not file_path.exists():
+                    await tool_ctx.error(f"File does not exist: {path}")
+                    return f"Error: File does not exist: {path}"
+
+                if not file_path.is_file():
+                    await tool_ctx.error(f"Path is not a file: {path}")
+                    return f"Error: Path is not a file: {path}"
+
+                # Check file extension
+                if file_path.suffix.lower() != ".ipynb":
+                    await tool_ctx.error(f"File is not a Jupyter notebook: {path}")
+                    return f"Error: File is not a Jupyter notebook: {path}"
+
+                # Read and parse the notebook
+                try:
+                    with open(file_path, "r", encoding="utf-8") as f:
+                        content = f.read()
+                        notebook = json.loads(content)
+                except json.JSONDecodeError:
+                    await tool_ctx.error(f"Invalid notebook format: {path}")
+                    return f"Error: Invalid notebook format: {path}"
+                except UnicodeDecodeError:
+                    await tool_ctx.error(f"Cannot read notebook file: {path}")
+                    return f"Error: Cannot read notebook file: {path}"
+
+                # Check cell_number is valid
+                cells = notebook.get("cells", [])
+
+                if edit_mode == "insert":
+                    if cell_number > len(cells):
+                        await tool_ctx.error(
+                            f"Cell number {cell_number} is out of bounds for insert (max: {len(cells)})"
+                        )
+                        return f"Error: Cell number {cell_number} is out of bounds for insert (max: {len(cells)})"
+                else:  # replace or delete
+                    if cell_number >= len(cells):
+                        await tool_ctx.error(
+                            f"Cell number {cell_number} is out of bounds (max: {len(cells) - 1})"
+                        )
+                        return f"Error: Cell number {cell_number} is out of bounds (max: {len(cells) - 1})"
+
+                language = (
+                    notebook.get("metadata", {})
+                    .get("language_info", {})
+                    .get("name", "python")
+                )
+
+                # Perform the requested operation
+                if edit_mode == "replace":
+                    # Get the target cell
+                    target_cell = cells[cell_number]
+
+                    # Store previous contents for reporting
+                    old_type = target_cell.get("cell_type", "code")
+                    old_source = target_cell.get("source", "")
+                    if isinstance(old_source, list):
+                        old_source = "".join(old_source)
+
+                    # Update source
+                    target_cell["source"] = new_source
+
+                    # Update type if specified
+                    if cell_type is not None:
+                        target_cell["cell_type"] = cell_type
+
+                    # If changing to markdown, remove code-specific fields
+                    if cell_type == "markdown":
+                        if "outputs" in target_cell:
+                            del target_cell["outputs"]
+                        if "execution_count" in target_cell:
+                            del target_cell["execution_count"]
+
+                    # If code cell, reset execution
+                    if target_cell["cell_type"] == "code":
+                        target_cell["outputs"] = []
+                        target_cell["execution_count"] = None
+
+                    change_description = f"Replaced cell {cell_number}"
+                    if cell_type is not None and cell_type != old_type:
+                        change_description += (
+                            f" (changed type from {old_type} to {cell_type})"
+                        )
+
+                elif edit_mode == "insert":
+                    # Create new cell
+                    new_cell: dict[str, Any] = {
+                        "cell_type": cell_type,
+                        "source": new_source,
+                        "metadata": {},
+                    }
+
+                    # Add code-specific fields
+                    if cell_type == "code":
+                        new_cell["outputs"] = []
+                        new_cell["execution_count"] = None
+
+                    # Insert the cell
+                    cells.insert(cell_number, new_cell)
+                    change_description = (
+                        f"Inserted new {cell_type} cell at position {cell_number}"
+                    )
+
+                else:  # delete
+                    # Store deleted cell info for reporting
+                    deleted_cell = cells[cell_number]
+                    deleted_type = deleted_cell.get("cell_type", "code")
+
+                    # Remove the cell
+                    del cells[cell_number]
+                    change_description = (
+                        f"Deleted {deleted_type} cell at position {cell_number}"
+                    )
+
+                # Write the updated notebook back to file
+                with open(file_path, "w", encoding="utf-8") as f:
+                    json.dump(notebook, f, indent=1)
+
+                # Update document context
+                updated_content = json.dumps(notebook, indent=1)
+                self.document_context.update_document(path, updated_content)
+
+                await tool_ctx.info(
+                    f"Successfully edited notebook: {path} - {change_description}"
+                )
+                return f"Successfully edited notebook: {path} - {change_description}"
+            except Exception as e:
+                await tool_ctx.error(f"Error editing notebook: {str(e)}")
+                return f"Error editing notebook: {str(e)}"

hanzo_mcp/tools/project/__init__.py

@@ -0,0 +1 @@
+"""Project analysis and management tools for Hanzo Dev MCP."""