hanzo-mcp 0.9.0__py3-none-any.whl → 0.9.1__py3-none-any.whl

hanzo_mcp/tools/jupyter/notebook_read.py

@@ -0,0 +1,148 @@
+"""Read notebook tool implementation.
+
+This module provides the NotebookReadTool for reading Jupyter notebook files.
+"""
+
+import json
+from typing import Unpack, Annotated, TypedDict, final, override
+from pathlib import Path
+
+from pydantic import Field
+from mcp.server import FastMCP
+
+from hanzo_mcp.tools.common.auto_timeout import auto_timeout
+from mcp.server.fastmcp import Context as MCPContext
+
+from hanzo_mcp.tools.jupyter.base import JupyterBaseTool
+
+NotebookPath = Annotated[
+    str,
+    Field(
+        description="The absolute path to the Jupyter notebook file to read (must be absolute, not relative)",
+    ),
+]
+
+
+class NotebookReadToolParams(TypedDict):
+    """Parameters for the NotebookReadTool.
+
+    Attributes:
+        notebook_path: The absolute path to the Jupyter notebook file to read (must be absolute, not relative)
+    """
+
+    notebook_path: NotebookPath
+
+
+@final
+class NotebookReadTool(JupyterBaseTool):
+    """Tool for reading Jupyter notebook files."""
+
+    @property
+    @override
+    def name(self) -> str:
+        """Get the tool name.
+
+        Returns:
+            Tool name
+        """
+        return "notebook_read"
+
+    @property
+    @override
+    def description(self) -> str:
+        """Get the tool description.
+
+        Returns:
+            Tool description
+        """
+        return "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. The notebook_path parameter must be an absolute path, not a relative path."
+
+    @override
+    @auto_timeout("notebook_read")
+
+
+    async def call(
+        self,
+        ctx: MCPContext,
+        **params: Unpack[NotebookReadToolParams],
+    ) -> str:
+        """Execute the tool with the given parameters.
+
+        Args:
+            ctx: MCP context
+            **params: Tool parameters
+
+        Returns:
+            Tool result
+        """
+        tool_ctx = self.create_tool_context(ctx)
+        self.set_tool_context_info(tool_ctx)
+
+        # Extract parameters
+        notebook_path: NotebookPath = params["notebook_path"]
+
+        # Validate path parameter
+        path_validation = self.validate_path(notebook_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: {notebook_path}")
+
+        # Check if path is allowed
+        if not self.is_path_allowed(notebook_path):
+            await tool_ctx.error(f"Access denied - path outside allowed directories: {notebook_path}")
+            return f"Error: Access denied - path outside allowed directories: {notebook_path}"
+
+        try:
+            file_path = Path(notebook_path)
+
+            if not file_path.exists():
+                await tool_ctx.error(f"File does not exist: {notebook_path}")
+                return f"Error: File does not exist: {notebook_path}"
+
+            if not file_path.is_file():
+                await tool_ctx.error(f"Path is not a file: {notebook_path}")
+                return f"Error: Path is not a file: {notebook_path}"
+
+            # Check file extension
+            if file_path.suffix.lower() != ".ipynb":
+                await tool_ctx.error(f"File is not a Jupyter notebook: {notebook_path}")
+                return f"Error: File is not a Jupyter notebook: {notebook_path}"
+
+            # Read and parse the notebook
+            try:
+                # This will read the file, so we don't need to read it separately
+                _, processed_cells = await self.parse_notebook(file_path)
+
+                # Format the notebook content as a readable string
+                result = self.format_notebook_cells(processed_cells)
+
+                await tool_ctx.info(f"Successfully read notebook: {notebook_path} ({len(processed_cells)} cells)")
+                return result
+            except json.JSONDecodeError:
+                await tool_ctx.error(f"Invalid notebook format: {notebook_path}")
+                return f"Error: Invalid notebook format: {notebook_path}"
+            except UnicodeDecodeError:
+                await tool_ctx.error(f"Cannot read notebook file: {notebook_path}")
+                return f"Error: Cannot read notebook file: {notebook_path}"
+        except Exception as e:
+            await tool_ctx.error(f"Error reading notebook: {str(e)}")
+            return f"Error reading notebook: {str(e)}"
+
+    @override
+    def register(self, mcp_server: FastMCP) -> None:
+        """Register this read notebook tool with the MCP server.
+
+        Creates a wrapper function with explicitly defined parameters that match
+        the tool's parameter schema and registers it with the MCP server.
+
+        Args:
+            mcp_server: The FastMCP server instance
+        """
+
+        tool_self = self  # Create a reference to self for use in the closure
+
+        @mcp_server.tool(name=self.name, description=self.description)
+        async def notebook_read(notebook_path: NotebookPath, ctx: MCPContext) -> str:
+            return await tool_self.call(ctx, notebook_path=notebook_path)

hanzo_mcp/tools/llm/consensus_tool.py

@@ -15,9 +15,10 @@ from typing import (
 from pydantic import Field
 from mcp.server.fastmcp import Context as MCPContext
 
+from hanzo_mcp.tools.common.auto_timeout import auto_timeout
 from hanzo_mcp.tools.common.base import BaseTool
 from hanzo_mcp.tools.llm.llm_tool import LLMTool
-from hanzo_mcp.tools.common.context import create_tool_context
+from hanzo_mcp.tools.common.context import create_tool_context, ToolContext
 
 Prompt = Annotated[
     str,
@@ -155,6 +156,7 @@ The tool will:
 """
 
     @override
+    @auto_timeout("consensus")
     async def call(
         self,
         ctx: MCPContext,
@@ -269,10 +271,10 @@ The tool will:
                 if max_tokens:
                     params["max_tokens"] = max_tokens
 
-                # Create a mock context for the LLM tool
-                mock_ctx = type("MockContext", (), {"client": None})()
+                # Create a proper context for the LLM tool
+                tool_ctx = create_tool_context("consensus", self.server)
 
-                result = await asyncio.wait_for(self.llm_tool.call(mock_ctx, **params), timeout=timeout)
+                result = await asyncio.wait_for(self.llm_tool.call(tool_ctx, **params), timeout=timeout)
                 return (model, result)
             except asyncio.TimeoutError:
                 return (model, f"Error: Timeout after {timeout} seconds")
@@ -317,7 +319,7 @@ Be concise but thorough. Focus on providing actionable insights."""
 
         try:
             # Use the LLM tool to get the aggregation
-            mock_ctx = type("MockContext", (), {"client": None})()
+            tool_ctx = create_tool_context("consensus_aggregation", self.server)
 
             aggregation_params = {
                 "model": aggregation_model,
@@ -326,7 +328,7 @@ Be concise but thorough. Focus on providing actionable insights."""
                 "system_prompt": "You are an expert at analyzing and synthesizing multiple AI responses to provide balanced, insightful consensus.",
             }
 
-            result = await self.llm_tool.call(mock_ctx, **aggregation_params)
+            result = await self.llm_tool.call(tool_ctx, **aggregation_params)
             return result
 
         except Exception:

hanzo_mcp/tools/llm/llm_manage.py

@@ -7,6 +7,8 @@ from pathlib import Path
 from pydantic import Field
 from mcp.server.fastmcp import Context as MCPContext
 
+from hanzo_mcp.tools.common.auto_timeout import auto_timeout
+
 from hanzo_mcp.tools.common.base import BaseTool
 from hanzo_mcp.tools.llm.llm_tool import LLMTool
 from hanzo_mcp.tools.common.context import create_tool_context
@@ -110,6 +112,9 @@ Providers are automatically detected based on environment variables:
 """
 
     @override
+    @auto_timeout("llm_manage")
+
+
     async def call(
         self,
         ctx: MCPContext,

hanzo_mcp/tools/llm/llm_tool.py

@@ -21,6 +21,7 @@ from mcp.server.fastmcp import Context as MCPContext
 
 from hanzo_mcp.tools.common.base import BaseTool
 from hanzo_mcp.tools.common.context import create_tool_context
+from hanzo_mcp.tools.common.auto_timeout import auto_timeout
 
 # Check if litellm is available
 try:
@@ -263,6 +264,7 @@ llm --action models --provider openai
 Available: {", ".join(available) if available else "None"}"""
 
     @override
+    @auto_timeout("llm")
     async def call(
         self,
         ctx: MCPContext,

hanzo_mcp/tools/llm/llm_unified.py

@@ -19,6 +19,8 @@ from pathlib import Path
 from pydantic import Field
 from mcp.server.fastmcp import Context as MCPContext
 
+from hanzo_mcp.tools.common.auto_timeout import auto_timeout
+
 from hanzo_mcp.tools.common.base import BaseTool
 from hanzo_mcp.tools.common.context import create_tool_context
 
@@ -263,6 +265,9 @@ llm --action models --provider openai
 Available: {", ".join(available) if available else "None"}"""
 
     @override
+    @auto_timeout("llm_unified")
+
+
     async def call(
         self,
         ctx: MCPContext,

hanzo_mcp/tools/llm/provider_tools.py

@@ -5,6 +5,8 @@ from typing import Dict, Unpack, Optional, Annotated, TypedDict, final, override
 from pydantic import Field
 from mcp.server.fastmcp import Context as MCPContext
 
+from hanzo_mcp.tools.common.auto_timeout import auto_timeout
+
 from hanzo_mcp.tools.common.base import BaseTool
 from hanzo_mcp.tools.llm.llm_tool import LLMTool
 
@@ -98,6 +100,9 @@ class BaseProviderTool(BaseTool):
         return model
 
     @override
+    @auto_timeout("provider_tools")
+
+
     async def call(
         self,
         ctx: MCPContext,