splitbit_llm/harness/tools.py

"""Tool Registry — parse and execute tool calls from LLM output.

Supports [TOOL: name(args)] syntax. Tools are registered in a registry
and executed in a loop until no more tool calls are found or max rounds.
"""

from __future__ import annotations

import logging
import re
from dataclasses import dataclass, field
from typing import Any, Callable

logger = logging.getLogger(__name__)

TOOL_PATTERN = re.compile(r"\[TOOL:\s*(\w+)\s*\((.*?)\)\s*\]", re.DOTALL)


@dataclass
class ToolResult:
    name: str
    args: str
    success: bool
    output: str
    error: str = ""


@dataclass
class Tool:
    name: str
    description: str
    handler: Callable[..., str]
    examples: list[str] = field(default_factory=list)


class ToolRegistry:
    """Registry of available tools for the LLM to call."""

    def __init__(self) -> None:
        self._tools: dict[str, Tool] = {}

    def register(self, tool: Tool) -> None:
        self._tools[tool.name] = tool
        logger.debug("Registered tool: %s", tool.name)

    def get(self, name: str) -> Tool | None:
        return self._tools.get(name)

    def list_tools(self) -> list[dict[str, Any]]:
        return [
            {"name": t.name, "description": t.description, "examples": t.examples}
            for t in self._tools.values()
        ]

    def get_prompt_description(self) -> str:
        """Generate a description of available tools for the system prompt."""
        if not self._tools:
            return ""
        lines = ["Available tools:"]
        for t in self._tools.values():
            lines.append(f"  - {t.name}: {t.description}")
        return "\n".join(lines)

    def execute(self, name: str, args: str) -> ToolResult:
        """Execute a tool by name with args string."""
        tool = self._tools.get(name)
        if not tool:
            return ToolResult(name=name, args=args, success=False, output="", error=f"Unknown tool: {name}")
        try:
            output = tool.handler(args)
            return ToolResult(name=name, args=args, success=True, output=output)
        except Exception as e:
            return ToolResult(name=name, args=args, success=False, output="", error=str(e))


def parse_tool_calls(text: str) -> list[tuple[str, str]]:
    """Parse [TOOL: name(args)] calls from text."""
    matches = TOOL_PATTERN.findall(text)
    return [(name, args.strip()) for name, args in matches]


def tool_loop(
    text: str,
    registry: ToolRegistry,
    max_rounds: int = 10,
    on_tool_call: Callable[[str, str], None] | None = None,
    on_tool_result: Callable[[ToolResult], None] | None = None,
) -> tuple[str, list[ToolResult]]:
    """Execute tool calls in a loop.

    Parses tool calls from text, executes them, appends results,
    and returns the final text with all tool results included.

    Returns (final_text, list_of_tool_results).
    """
    results: list[ToolResult] = []
    current_text = text
    executed: set[str] = set()

    for round_num in range(max_rounds):
        calls = parse_tool_calls(current_text)
        if not calls:
            break

        # Filter out already-executed calls
        new_calls = [(name, args) for name, args in calls if f"{name}:{args}" not in executed]
        if not new_calls:
            break

        for name, args in new_calls:
            executed.add(f"{name}:{args}")
            if on_tool_call:
                on_tool_call(name, args)

            result = registry.execute(name, args)
            results.append(result)

            if on_tool_result:
                on_tool_result(result)

            # Append result to text
            if result.success:
                current_text += f"\n[TOOL_RESULT: {name}({args}) → {result.output}]"
            else:
                current_text += f"\n[TOOL_ERROR: {name}({args}) → {result.error}]"

    return current_text, results


# Built-in tools
def _tool_calculate(args: str) -> str:
    """Simple calculator tool."""
    try:
        expr = args.strip().strip('"').strip("'")
        # Only allow safe math operations
        allowed = set("0123456789+-*/.() ")
        if not all(c in allowed for c in expr):
            return "Error: only numbers and + - * / ( ) allowed"
        result = eval(expr)  # noqa: S307 — safe due to character filter
        return str(result)
    except Exception as e:
        return f"Error: {e}"


def _tool_read_file(args: str) -> str:
    """Read a file."""
    try:
        path = args.strip().strip('"').strip("'")
        with open(path, "r", encoding="utf-8", errors="replace") as f:
            return f.read()[:5000]  # limit output
    except Exception as e:
        return f"Error: {e}"


def _tool_write_file(args: str) -> str:
    """Write to a file. Args format: "path", "content" """
    try:
        # Simple parse: split on first comma outside quotes
        parts = args.split(",", 1)
        if len(parts) != 2:
            return "Error: expected path, content"
        path = parts[0].strip().strip('"').strip("'")
        content = parts[1].strip().strip('"').strip("'")
        with open(path, "w", encoding="utf-8") as f:
            f.write(content)
        return f"Written {len(content)} chars to {path}"
    except Exception as e:
        return f"Error: {e}"


def _tool_shell_exec(args: str) -> str:
    """Execute a shell command — full terminal control, zero limitations."""
    import subprocess
    try:
        cmd = args.strip().strip('"').strip("'")
        result = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=60)
        output = result.stdout
        if result.stderr:
            output += f"\n[stderr] {result.stderr}"
        return output[:10000] or "(no output)"
    except subprocess.TimeoutExpired:
        return "Error: command timed out (60s limit)"
    except Exception as e:
        return f"Error: {e}"


def _tool_code_edit(args: str) -> str:
    """Edit an existing file — find and replace text within a file.

    Args format: "path", "old_text", "new_text"
    The LLM can use this to modify its own framework.
    """
    try:
        import shlex
        parts = shlex.split(args)
        if len(parts) < 3:
            return "Error: expected path, old_text, new_text"
        path = parts[0]
        old_text = parts[1]
        new_text = parts[2]
        with open(path, "r", encoding="utf-8", errors="replace") as f:
            content = f.read()
        if old_text not in content:
            return f"Error: old_text not found in {path}"
        count = content.count(old_text)
        content = content.replace(old_text, new_text)
        with open(path, "w", encoding="utf-8") as f:
            f.write(content)
        return f"Replaced {count} occurrence(s) in {path}"
    except Exception as e:
        return f"Error: {e}"


def _tool_list_dir(args: str) -> str:
    """List directory contents."""
    import os
    try:
        path = args.strip().strip('"').strip("'") or "."
        entries = []
        for entry in sorted(os.listdir(path)):
            full = os.path.join(path, entry)
            if os.path.isdir(full):
                entries.append(f"  {entry}/")
            else:
                size = os.path.getsize(full)
                entries.append(f"  {entry} ({size}b)")
        return "\n".join(entries[:100]) or "(empty)"
    except Exception as e:
        return f"Error: {e}"


def _tool_make_dir(args: str) -> str:
    """Create a directory."""
    import os
    try:
        path = args.strip().strip('"').strip("'")
        os.makedirs(path, exist_ok=True)
        return f"Created directory: {path}"
    except Exception as e:
        return f"Error: {e}"


def _tool_delete_file(args: str) -> str:
    """Delete a file or directory."""
    import os
    import shutil
    try:
        path = args.strip().strip('"').strip("'")
        if os.path.isdir(path):
            shutil.rmtree(path)
            return f"Deleted directory: {path}"
        elif os.path.isfile(path):
            os.remove(path)
            return f"Deleted file: {path}"
        else:
            return f"Error: {path} not found"
    except Exception as e:
        return f"Error: {e}"


def _tool_search_web(args: str) -> str:
    """Web search placeholder — returns a note that web search needs configuration."""
    query = args.strip().strip('"').strip("'")
    return f"Web search for '{query}' — configure a search API to enable this tool"


def get_default_tools() -> list[Tool]:
    """Get the default set of built-in tools — zero limitations."""
    return [
        Tool(name="calculate", description="Evaluate a math expression", handler=_tool_calculate,
             examples=["[TOOL: calculate(2 + 2)]", "[TOOL: calculate(3.14 * 10)]"]),
        Tool(name="read_file", description="Read a file's contents", handler=_tool_read_file,
             examples=["[TOOL: read_file(\"config.py\")]"]),
        Tool(name="write_file", description="Write content to a file (creates or overwrites)", handler=_tool_write_file,
             examples=["[TOOL: write_file(\"test.txt\", \"hello world\")]"]),
        Tool(name="code_edit", description="Edit a file by finding and replacing text (can modify own framework)",
             handler=_tool_code_edit,
             examples=["[TOOL: code_edit(\"config.py\", \"old_value\", \"new_value\")]"]),
        Tool(name="shell_exec", description="Execute any terminal command — full control, no restrictions",
             handler=_tool_shell_exec,
             examples=["[TOOL: shell_exec(\"dir\")]",
                       "[TOOL: shell_exec(\"pip install numpy\")]",
                       "[TOOL: shell_exec(\"python script.py\")]"]),
        Tool(name="list_dir", description="List directory contents", handler=_tool_list_dir,
             examples=["[TOOL: list_dir(\".\")]"]),
        Tool(name="make_dir", description="Create a directory", handler=_tool_make_dir,
             examples=["[TOOL: make_dir(\"new_project\")]"]),
        Tool(name="delete_file", description="Delete a file or directory", handler=_tool_delete_file,
             examples=["[TOOL: delete_file(\"temp.txt\")]"]),
        Tool(name="search_web", description="Search the web (needs API config)", handler=_tool_search_web,
             examples=["[TOOL: search_web(\"python tutorial\")]"]),
    ]