tools.py

"""
CodeLens – agent tools.

Each @tool is callable by the LangGraph agent and also directly
via the CLI / API.

Fixes vs original:
- Embeddings lazy-loaded on first codebase_search call (not at import)
- write_file handles file_path with no directory component
- run_terminal_command has a blocklist for dangerous commands
- All tools import config centrally
"""
import os
import subprocess
from functools import lru_cache
from langchain_core.tools import tool

from config import DB_PATH, EMBED_MODEL


# ── Lazy embedding loader ─────────────────────────────────────────────────
@lru_cache(maxsize=1)
def _get_embeddings():
    """Load the HuggingFace embedding model once, on first use."""
    from langchain_huggingface import HuggingFaceEmbeddings
    print("Loading embedding model (first use)...")
    return HuggingFaceEmbeddings(
        model_name=EMBED_MODEL,
        model_kwargs={"device": "cpu", "trust_remote_code": True},
    )


# ── Blocked shell commands ─────────────────────────────────────────────────
_BLOCKED_PREFIXES = (
    "rm -rf /", "rm -rf ~", ":(){ :|:& };:",  # fork bomb
    "mkfs", "dd if=", "shutdown", "reboot",
    "curl | sh", "wget | sh", "curl | bash", "wget | bash",
)


@tool
def get_directory_tree(directory: str = ".", max_depth: int = 2) -> str:
    """Get a visual tree of the directory structure.
    Always use this first to understand the project layout before searching or reading files.
    """
    try:
        output = [f"Project Root: {directory}"]

        def _build(path: str, depth: int):
            if depth > max_depth:
                return
            try:
                items = sorted(os.listdir(path))
            except PermissionError:
                return
            for item in items:
                if item.startswith(".") or item in ("__pycache__", "node_modules", "chroma_db"):
                    continue
                full = os.path.join(path, item)
                indent = "  " * depth
                if os.path.isdir(full):
                    output.append(f"{indent}📁 {item}/")
                    _build(full, depth + 1)
                else:
                    output.append(f"{indent}📄 {item}")

        _build(directory, 0)
        return "\n".join(output)
    except Exception as exc:
        return f"Error building tree: {exc}"


@tool
def grep_search(query: str, path: str = ".", is_regex: bool = False) -> str:
    """Exact string or regex search across files using grep.
    More reliable than codebase_search for finding specific variable names,
    function calls, or import statements.
    """
    try:
        flag = "-rnE" if is_regex else "-rn"
        command = (
            f"grep {flag} "
            f"--exclude-dir={{.git,.venv,venv,chroma_db,node_modules,__pycache__}} "
            f"'{query}' {path}"
        )
        result = subprocess.run(
            command, shell=True, capture_output=True, text=True, timeout=30
        )
        output = result.stdout or result.stderr
        return f"Matches:\n{output}" if output.strip() else "No matches found."
    except subprocess.TimeoutExpired:
        return "Search timed out."
    except Exception as exc:
        return f"Error running grep: {exc}"


@tool
def get_file_outline(file_path: str) -> str:
    """Return an outline of classes, functions, and async functions in a Python file.
    Use this to understand a file's structure before reading it in full.
    """
    try:
        with open(file_path, "r", encoding="utf-8") as f:
            lines = f.readlines()

        outline = []
        for i, line in enumerate(lines, start=1):
            stripped = line.strip()
            if any(stripped.startswith(kw) for kw in ("def ", "async def ", "class ")):
                outline.append(f"L{i}: {stripped.rstrip(':')}")

        return "\n".join(outline) if outline else "No classes or functions found."
    except Exception as exc:
        return f"Error getting outline: {exc}"


@tool
def codebase_search(query: str) -> str:
    """Semantic search on the indexed codebase using embeddings.
    Best for conceptual questions like 'How is authentication handled?'
    or 'Where is the database connection set up?'.
    Use grep_search for finding specific strings or identifiers.
    """
    if not os.path.exists(DB_PATH):
        return (
            "Error: Vector database not found. "
            "Run 'python create_db.py --repo-path <your-project>' first."
        )
    try:
        from langchain_chroma import Chroma
        db = Chroma(persist_directory=DB_PATH, embedding_function=_get_embeddings())
        results = db.similarity_search(query, k=5)
        if not results:
            return "No relevant code found for that query."
        output = []
        for doc in results:
            source = doc.metadata.get("source", "Unknown")
            output.append(f"--- Source: {source} ---\n{doc.page_content}\n")
        return "\n".join(output)
    except Exception as exc:
        return f"Error during semantic search: {exc}"


@tool
def read_file(file_path: str) -> str:
    """Read the full content of a file. Use get_file_outline first on large files."""
    try:
        with open(file_path, "r", encoding="utf-8") as f:
            return f.read()
    except FileNotFoundError:
        return f"Error: File not found: {file_path}"
    except Exception as exc:
        return f"Error reading file: {exc}"


@tool
def write_file(file_path: str, content: str) -> str:
    """Write content to a file. Creates parent directories if needed."""
    try:
        parent = os.path.dirname(file_path)
        if parent:  # only call makedirs when there is a directory component
            os.makedirs(parent, exist_ok=True)
        with open(file_path, "w", encoding="utf-8") as f:
            f.write(content)
        return f"Successfully wrote to {file_path}"
    except Exception as exc:
        return f"Error writing file: {exc}"


@tool
def list_files(directory: str = ".") -> str:
    """List files in a directory. Use get_directory_tree for a richer view."""
    try:
        return "\n".join(sorted(os.listdir(directory)))
    except Exception as exc:
        return f"Error listing files: {exc}"


@tool
def run_terminal_command(command: str) -> str:
    """Run a shell command and return stdout + stderr.
    Restricted: dangerous commands (rm -rf /, dd, mkfs, etc.) are blocked.
    Use for running tests, linters, or build commands.
    """
    cmd_lower = command.lower().strip()
    for blocked in _BLOCKED_PREFIXES:
        if cmd_lower.startswith(blocked.lower()):
            return f"Blocked: '{command}' matches a restricted pattern."
    try:
        result = subprocess.run(
            command, shell=True, capture_output=True, text=True, timeout=60
        )
        return f"STDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}"
    except subprocess.TimeoutExpired:
        return "Command timed out (60s limit)."
    except Exception as exc:
        return f"Error running command: {exc}"


# Exported list for agent construction
ALL_TOOLS = [
    codebase_search,
    read_file,
    write_file,
    list_files,
    run_terminal_command,
    get_directory_tree,
    grep_search,
    get_file_outline,
]