ripperdoc 0.1.0__py3-none-any.whl

ripperdoc/tools/notebook_edit_tool.py

@@ -0,0 +1,325 @@
+"""Notebook edit tool.
+
+Allows performing insert/replace/delete operations on Jupyter notebook cells.
+"""
+
+import json
+import random
+import string
+from pathlib import Path
+from textwrap import dedent
+from typing import AsyncGenerator, List, Optional
+from pydantic import BaseModel, Field
+
+from ripperdoc.core.tool import (
+    Tool,
+    ToolUseContext,
+    ToolResult,
+    ToolOutput,
+    ToolUseExample,
+    ValidationResult,
+)
+from ripperdoc.utils.log import get_logger
+
+
+logger = get_logger()
+
+
+def _resolve_path(path_str: str) -> Path:
+    """Return an absolute Path, interpreting relative paths from CWD."""
+    path = Path(path_str).expanduser()
+    return path if path.is_absolute() else Path.cwd() / path
+
+
+def _generate_cell_id() -> str:
+    """Generate a short random cell id."""
+    return "".join(random.choices(string.ascii_lowercase + string.digits, k=12))
+
+
+NOTEBOOK_EDIT_DESCRIPTION = dedent(
+    """\
+    Replace, insert, or delete a specific cell in a Jupyter notebook (.ipynb file).
+    notebook_path must be an absolute path. cell_id may be a 0-based index or a cell id.
+    Use edit_mode=insert to add a new cell after the referenced cell (or at the start if omitted).
+    Use edit_mode=delete to delete the referenced cell. Defaults to edit_mode=replace.
+
+    Usage:
+    - cell_type: 'code' or 'markdown'. Required for insert; defaults to existing type for replace.
+    - new_source: New content for the cell.
+    - Edits are applied atomically; failures leave the file unchanged.
+    - Code cell replacements clear execution_count and outputs.
+    - Only use emojis if explicitly requested; avoid adding emojis otherwise.
+    """
+)
+
+
+class NotebookEditInput(BaseModel):
+    """Input schema for NotebookEditTool."""
+
+    notebook_path: str = Field(description="Absolute path to the Jupyter notebook file to edit")
+    cell_id: Optional[str] = Field(
+        default=None,
+        description="Cell ID or 0-based index. For insert, omitted means insert at start.",
+    )
+    new_source: str = Field(description="New source content for the target cell")
+    cell_type: Optional[str] = Field(
+        default=None,
+        description="Cell type: 'code' or 'markdown'. Required for insert.",
+    )
+    edit_mode: Optional[str] = Field(
+        default="replace",
+        description="Edit mode: 'replace' (default), 'insert', or 'delete'.",
+    )
+
+
+class NotebookEditOutput(BaseModel):
+    """Output from notebook editing."""
+
+    new_source: str
+    cell_id: Optional[str] = None
+    cell_type: str
+    language: str
+    edit_mode: str
+    error: Optional[str] = None
+
+
+class NotebookEditTool(Tool[NotebookEditInput, NotebookEditOutput]):
+    """Tool for editing Jupyter notebooks."""
+
+    @property
+    def name(self) -> str:
+        return "NotebookEdit"
+
+    async def description(self) -> str:
+        return NOTEBOOK_EDIT_DESCRIPTION
+
+    @property
+    def input_schema(self) -> type[NotebookEditInput]:
+        return NotebookEditInput
+
+    def input_examples(self) -> List[ToolUseExample]:
+        return [
+            ToolUseExample(
+                description="Replace a markdown cell by id",
+                input={
+                    "notebook_path": "/repo/notebooks/analysis.ipynb",
+                    "cell_id": "abc123",
+                    "new_source": "# Updated overview\\nThis notebook analyzes revenue.",
+                    "cell_type": "markdown",
+                    "edit_mode": "replace",
+                },
+            ),
+            ToolUseExample(
+                description="Insert a new code cell at the beginning",
+                input={
+                    "notebook_path": "/repo/notebooks/analysis.ipynb",
+                    "cell_type": "code",
+                    "edit_mode": "insert",
+                    "new_source": "import pandas as pd\\nimport numpy as np",
+                },
+            ),
+        ]
+
+    async def prompt(self, safe_mode: bool = False) -> str:
+        return NOTEBOOK_EDIT_DESCRIPTION
+
+    def is_read_only(self) -> bool:
+        return False
+
+    def is_concurrency_safe(self) -> bool:
+        return False
+
+    def needs_permissions(self, input_data: Optional[NotebookEditInput] = None) -> bool:
+        return True
+
+    async def validate_input(
+        self, input_data: NotebookEditInput, context: Optional[ToolUseContext] = None
+    ) -> ValidationResult:
+        path = _resolve_path(input_data.notebook_path)
+
+        if not path.exists():
+            return ValidationResult(result=False, message="Notebook file does not exist.")
+        if path.suffix != ".ipynb":
+            return ValidationResult(
+                result=False,
+                message="File must be a Jupyter notebook (.ipynb file). Use Edit for other file types.",
+            )
+
+        mode = (input_data.edit_mode or "replace").lower()
+        if mode not in {"replace", "insert", "delete"}:
+            return ValidationResult(
+                result=False, message="edit_mode must be replace, insert, or delete."
+            )
+        if mode == "insert" and not input_data.cell_type:
+            return ValidationResult(
+                result=False,
+                message="cell_type is required when using edit_mode=insert.",
+            )
+        if mode != "insert" and not input_data.cell_id:
+            return ValidationResult(
+                result=False,
+                message="cell_id must be specified when using edit_mode=replace or delete.",
+            )
+
+        # Validate notebook structure and target cell.
+        try:
+            raw = path.read_text(encoding="utf-8")
+            nb_json = json.loads(raw)
+        except Exception as exc:
+            logger.error(f"Failed to parse notebook {path}: {exc}")
+            return ValidationResult(
+                result=False, message="Notebook is not valid JSON.", error_code=6
+            )
+
+        cells = nb_json.get("cells", [])
+        target_index, _ = self._resolve_cell_index(cells, input_data.cell_id, mode)
+        if target_index is None:
+            if mode == "insert" and input_data.cell_id is None:
+                return ValidationResult(result=True)
+            return ValidationResult(
+                result=False,
+                message=f"Cell '{input_data.cell_id}' not found in notebook.",
+                error_code=7,
+            )
+
+        return ValidationResult(result=True)
+
+    def render_result_for_assistant(self, output: NotebookEditOutput) -> str:
+        if output.error:
+            return output.error
+        action = output.edit_mode or "replace"
+        cell_label = output.cell_id or "(new cell)"
+        if action == "delete":
+            return f"Deleted cell {cell_label}"
+        if action == "insert":
+            return f"Inserted cell {cell_label}"
+        return f"Updated cell {cell_label}"
+
+    def render_tool_use_message(self, input_data: NotebookEditInput, verbose: bool = False) -> str:
+        parts = [f"path: {input_data.notebook_path}"]
+        if input_data.cell_id:
+            parts.append(f"cell_id: {input_data.cell_id}")
+        if verbose:
+            parts.append(f"mode: {input_data.edit_mode or 'replace'}")
+        return ", ".join(parts)
+
+    async def call(
+        self, input_data: NotebookEditInput, context: ToolUseContext
+    ) -> AsyncGenerator[ToolOutput, None]:
+        path = _resolve_path(input_data.notebook_path)
+        mode = (input_data.edit_mode or "replace").lower()
+        cell_type = (input_data.cell_type or "").lower() or None
+        new_source = input_data.new_source
+        cell_id = input_data.cell_id
+
+        try:
+            raw = path.read_text(encoding="utf-8")
+            nb_json = json.loads(raw)
+            cells = nb_json.get("cells", [])
+
+            target_index, matched_id = self._resolve_cell_index(cells, cell_id, mode)
+
+            final_mode = mode
+            if final_mode == "replace" and target_index is not None and target_index == len(cells):
+                final_mode = "insert"
+                if not cell_type:
+                    cell_type = "code"
+
+            if final_mode == "delete":
+                if target_index is None:
+                    raise ValueError("Target cell not found for delete.")
+                cells.pop(target_index)
+            elif final_mode == "insert":
+                insert_at = target_index if target_index is not None else 0
+                new_id = matched_id or _generate_cell_id()
+                new_cell_type = cell_type or "code"
+                new_cell = (
+                    {
+                        "cell_type": "markdown",
+                        "id": new_id,
+                        "source": new_source,
+                        "metadata": {},
+                    }
+                    if new_cell_type == "markdown"
+                    else {
+                        "cell_type": "code",
+                        "id": new_id,
+                        "source": new_source,
+                        "metadata": {},
+                        "execution_count": 0,
+                        "outputs": [],
+                    }
+                )
+                cells.insert(insert_at, new_cell)
+                matched_id = new_id
+                cell_type = new_cell_type
+            else:  # replace
+                if target_index is None:
+                    raise ValueError("Target cell not found for replace.")
+                target_cell = cells[target_index]
+                target_cell["source"] = new_source
+                if target_cell.get("cell_type") == "code":
+                    target_cell["execution_count"] = None
+                    target_cell["outputs"] = []
+                if cell_type and cell_type != target_cell.get("cell_type"):
+                    target_cell["cell_type"] = cell_type
+                matched_id = target_cell.get("id") or matched_id
+                cell_type = target_cell.get("cell_type", cell_type or "code")
+
+            nb_json["cells"] = cells
+            notebook_language = (
+                nb_json.get("metadata", {}).get("language_info", {}).get("name", "python")
+            )
+
+            path.write_text(json.dumps(nb_json, indent=1), encoding="utf-8")
+
+            output = NotebookEditOutput(
+                new_source=new_source,
+                cell_type=cell_type or "code",
+                language=notebook_language,
+                edit_mode=final_mode,
+                cell_id=matched_id,
+                error=None,
+            )
+            yield ToolResult(
+                data=output, result_for_assistant=self.render_result_for_assistant(output)
+            )
+        except Exception as exc:  # pragma: no cover - error path
+            logger.error(f"Error editing notebook {input_data.notebook_path}: {exc}")
+            output = NotebookEditOutput(
+                new_source=new_source,
+                cell_type=cell_type or "code",
+                language="python",
+                edit_mode=mode,
+                cell_id=cell_id,
+                error=str(exc),
+            )
+            yield ToolResult(
+                data=output, result_for_assistant=self.render_result_for_assistant(output)
+            )
+
+    def _resolve_cell_index(
+        self, cells: list, cell_id: Optional[str], mode: str
+    ) -> tuple[Optional[int], Optional[str]]:
+        """Return target index and resolved id."""
+        if cell_id is None:
+            return (0 if mode == "insert" else None, None)
+
+        # Try numeric index first.
+        try:
+            idx = int(cell_id)
+            if mode == "insert":
+                if idx < 0 or idx > len(cells):
+                    return None, None
+                return min(idx + 1, len(cells)), None
+            else:
+                if idx < 0 or idx >= len(cells):
+                    return None, None
+                return idx, None
+        except (ValueError, TypeError):
+            pass
+
+        for i, cell in enumerate(cells):
+            if cell.get("id") == cell_id:
+                return (i if mode != "insert" else i + 1, cell.get("id"))
+        return None, None

ripperdoc/tools/task_tool.py

@@ -0,0 +1,282 @@
+"""Task tool that delegates work to configured subagents."""
+
+from __future__ import annotations
+
+import time
+from typing import Any, AsyncGenerator, Callable, Dict, Iterable, List, Optional
+
+from pydantic import BaseModel, Field
+
+from ripperdoc.core.agents import (
+    AgentDefinition,
+    AgentLoadResult,
+    clear_agent_cache,
+    load_agent_definitions,
+    resolve_agent_tools,
+    summarize_agent,
+)
+from ripperdoc.core.query import QueryContext, query
+from ripperdoc.core.system_prompt import build_environment_prompt
+from ripperdoc.core.tool import Tool, ToolOutput, ToolProgress, ToolResult, ToolUseContext
+from ripperdoc.utils.messages import AssistantMessage, create_user_message
+
+
+class TaskToolInput(BaseModel):
+    """Input schema for delegating to a subagent."""
+
+    prompt: str = Field(description="Detailed task description for the subagent to perform")
+    subagent_type: str = Field(description="Agent type to run (matches agent frontmatter name)")
+
+
+class TaskToolOutput(BaseModel):
+    """Summary of a completed subagent run."""
+
+    agent_type: str
+    result_text: str
+    duration_ms: float
+    tool_use_count: int
+    missing_tools: List[str] = Field(default_factory=list)
+    model_used: Optional[str] = None
+
+
+class TaskTool(Tool[TaskToolInput, TaskToolOutput]):
+    """Launches a configured agent in a fresh context."""
+
+    def __init__(self, available_tools_provider: Callable[[], Iterable[Tool[Any, Any]]]) -> None:
+        super().__init__()
+        self._available_tools_provider = available_tools_provider
+
+    @property
+    def name(self) -> str:
+        return "Task"
+
+    async def description(self) -> str:
+        clear_agent_cache()
+        agents = load_agent_definitions()
+        agent_lines = "\n".join(summarize_agent(agent) for agent in agents.active_agents)
+        return (
+            "Launch a specialized subagent in its own context window to handle a task.\n"
+            f"Available agents:\n{agent_lines or '- general-purpose (built-in)'}"
+        )
+
+    @property
+    def input_schema(self) -> type[TaskToolInput]:
+        return TaskToolInput
+
+    async def prompt(self, safe_mode: bool = False) -> str:
+        del safe_mode
+        clear_agent_cache()
+        agents: AgentLoadResult = load_agent_definitions()
+        agent_lines = "\n".join(summarize_agent(agent) for agent in agents.active_agents)
+        return (
+            "Use this tool to delegate a well-scoped task to a subagent. "
+            "Always set subagent_type to one of the available agent types below. "
+            "Provide a detailed prompt so the agent can work autonomously and return a single, concise report.\n\n"
+            f"Available agents:\n{agent_lines or '- general-purpose (built-in)'}"
+        )
+
+    def is_read_only(self) -> bool:
+        return True
+
+    def is_concurrency_safe(self) -> bool:
+        return True
+
+    def render_result_for_assistant(self, output: TaskToolOutput) -> str:
+        details: List[str] = []
+        if output.tool_use_count:
+            details.append(f"{output.tool_use_count} tool uses")
+        details.append(f"{output.duration_ms/1000:.1f}s")
+        if output.missing_tools:
+            details.append(f"missing tools: {', '.join(output.missing_tools)}")
+
+        suffix = f" ({'; '.join(details)})" if details else ""
+        return f"[subagent:{output.agent_type}] {output.result_text}{suffix}"
+
+    def render_tool_use_message(self, input_data: TaskToolInput, verbose: bool = False) -> str:
+        del verbose
+        return f"Task via {input_data.subagent_type}: {input_data.prompt}"
+
+    async def call(
+        self,
+        input_data: TaskToolInput,
+        context: ToolUseContext,
+    ) -> AsyncGenerator[ToolOutput, None]:
+        clear_agent_cache()
+        agents = load_agent_definitions()
+        target_agent = next(
+            (
+                agent
+                for agent in agents.active_agents
+                if agent.agent_type == input_data.subagent_type
+            ),
+            None,
+        )
+        if not target_agent:
+            raise ValueError(
+                f"Agent type '{input_data.subagent_type}' not found. "
+                f"Available agents: {', '.join(agent.agent_type for agent in agents.active_agents)}"
+            )
+
+        available_tools = list(self._available_tools_provider())
+        agent_tools, missing_tools = resolve_agent_tools(target_agent, available_tools, self.name)
+        if not agent_tools:
+            raise ValueError(
+                f"Agent '{target_agent.agent_type}' has no usable tools. "
+                f"Missing or unknown tools: {', '.join(missing_tools) if missing_tools else 'none'}"
+            )
+
+        agent_system_prompt = self._build_agent_prompt(target_agent, agent_tools)
+        subagent_messages = [create_user_message(input_data.prompt)]
+
+        subagent_context = QueryContext(
+            tools=agent_tools,
+            safe_mode=context.safe_mode,
+            verbose=context.verbose,
+            model=target_agent.model or "task",
+        )
+
+        start = time.time()
+        assistant_messages: List[AssistantMessage] = []
+        tool_use_count = 0
+
+        yield ToolProgress(content=f"Launching subagent '{target_agent.agent_type}'")
+
+        async for message in query(
+            subagent_messages,
+            agent_system_prompt,
+            {},
+            subagent_context,
+            context.permission_checker,
+        ):
+            if getattr(message, "type", "") == "assistant":
+                # Surface subagent tool requests as progress so the user sees activity.
+                msg_content = getattr(message, "message", None)
+                blocks = getattr(msg_content, "content", []) if msg_content else []
+                if isinstance(blocks, list):
+                    for block in blocks:
+                        block_type = getattr(block, "type", None) or (
+                            block.get("type") if isinstance(block, Dict) else None
+                        )
+                        if block_type == "tool_use":
+                            tool_name = getattr(block, "name", None) or (
+                                block.get("name") if isinstance(block, Dict) else "unknown tool"
+                            )
+                            block_input = (
+                                getattr(block, "input", None)
+                                if hasattr(block, "input")
+                                else (block.get("input") if isinstance(block, Dict) else None)
+                            )
+                            summary = self._summarize_tool_input(block_input)
+                            label = f"Subagent requesting {tool_name}"
+                            if summary:
+                                label += f" — {summary}"
+                            yield ToolProgress(content=label)
+                        if block_type == "text":
+                            text_val = getattr(block, "text", None) or (
+                                block.get("text") if isinstance(block, Dict) else ""
+                            )
+                            if text_val:
+                                snippet = str(text_val).strip()
+                                if snippet:
+                                    short = (
+                                        snippet if len(snippet) <= 200 else snippet[:197] + "..."
+                                    )
+                                    yield ToolProgress(content=f"Subagent: {short}")
+                assistant_messages.append(message)  # type: ignore[arg-type]
+                tool_use_count += self._count_tool_uses(message)
+
+        duration_ms = (time.time() - start) * 1000
+        result_text = (
+            self._extract_text(assistant_messages[-1])
+            if assistant_messages
+            else "Agent returned no response."
+        )
+
+        output = TaskToolOutput(
+            agent_type=target_agent.agent_type,
+            result_text=result_text.strip(),
+            duration_ms=duration_ms,
+            tool_use_count=tool_use_count,
+            missing_tools=missing_tools,
+            model_used=target_agent.model or "task",
+        )
+
+        yield ToolResult(data=output, result_for_assistant=self.render_result_for_assistant(output))
+
+    def _build_agent_prompt(self, agent: AgentDefinition, tools: List[Tool[Any, Any]]) -> str:
+        tool_names = ", ".join(tool.name for tool in tools if getattr(tool, "name", None))
+        guidance = (
+            "You are a specialized Ripperdoc subagent working autonomously. "
+            "Execute the task completely using the allowed tools. "
+            "Return a single, concise summary for the parent agent that includes what you did, "
+            "important findings, and any follow-ups. Do not ask the user questions."
+        )
+        sections = [
+            guidance,
+            f"Agent type: {agent.agent_type}",
+            f"When to use: {agent.when_to_use}",
+            f"Allowed tools: {tool_names}",
+            "Agent system prompt:",
+            agent.system_prompt or "(no additional prompt)",
+            build_environment_prompt(),
+        ]
+        return "\n\n".join(sections)
+
+    def _extract_text(self, message: AssistantMessage) -> str:
+        content = message.message.content
+        if isinstance(content, str):
+            return content
+        if isinstance(content, list):
+            parts: List[str] = []
+            for block in content:
+                text = getattr(block, "text", None) or (
+                    block.get("text") if isinstance(block, Dict) else None
+                )
+                if text:
+                    parts.append(str(text))
+            return "\n".join(parts)
+        return ""
+
+    def _count_tool_uses(self, message: AssistantMessage) -> int:
+        content = message.message.content
+        if not isinstance(content, list):
+            return 0
+        count = 0
+        for block in content:
+            block_type = getattr(block, "type", None) or (
+                block.get("type") if isinstance(block, Dict) else None
+            )
+            if block_type == "tool_use":
+                count += 1
+        return count
+
+    def _summarize_tool_input(self, inp: Any) -> str:
+        """Generate a short human-readable summary of a tool_use input."""
+        if not inp or not isinstance(inp, (dict, Dict)):
+            return ""
+
+        pieces: List[str] = []
+        # Prioritize common keys
+        for key in ("command", "file_path", "path", "glob", "pattern", "description", "prompt"):
+            if key in inp and inp[key]:
+                val = str(inp[key])
+                short = val if len(val) <= 80 else val[:77] + "..."
+                pieces.append(f"{key}={short}")
+
+        # Include range info if present
+        start = inp.get("start_line") or inp.get("offset")
+        end = inp.get("end_line") or inp.get("limit")
+        if start is not None or end is not None:
+            pieces.append(f"range={start or 0}-{end or '…'}")
+
+        if not pieces:
+            # Fallback to truncated dict representation
+            import json
+
+            try:
+                serialized = json.dumps(inp, ensure_ascii=False)
+            except Exception:
+                serialized = str(inp)
+            return serialized if len(serialized) <= 120 else serialized[:117] + "..."
+
+        return ", ".join(pieces)