adam/external_tools.py

from __future__ import annotations

import ast
import json
import re
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any


ENTRY_NAMES = (
    "train.py", "trainer.py", "finetune.py", "fine_tune.py",
    "main.py", "run.py", "app.py",
)
CONFIG_SUFFIXES = {".json", ".yaml", ".yml", ".toml", ".ini"}
DATASET_EXTENSIONS = {
    ".png", ".jpg", ".jpeg", ".webp", ".bmp", ".gif",
    ".txt", ".csv", ".json", ".jsonl", ".parquet", ".wav", ".mp3",
}


@dataclass(slots=True)
class ToolAnalysis:
    folder: str
    python_files: list[str] = field(default_factory=list)
    entry_candidates: list[str] = field(default_factory=list)
    config_files: list[str] = field(default_factory=list)
    selected_entry: str = ""
    selected_configs: list[str] = field(default_factory=list)
    arguments: list[str] = field(default_factory=list)
    required_arguments: list[str] = field(default_factory=list)
    dataset_format: str = "Not detected"
    output_behavior: str = "Not detected"
    checkpoint_behavior: str = "Not detected"
    progress_behavior: str = "Not detected"
    resume_behavior: str = "Not detected"
    score: int = 1
    reasons: list[str] = field(default_factory=list)
    warnings: list[str] = field(default_factory=list)


def _relative(path: Path, folder: Path) -> str:
    return path.relative_to(folder).as_posix()


def scan_folder(folder_value: str) -> ToolAnalysis:
    folder = Path(folder_value).expanduser().resolve()
    if not folder.is_dir():
        return ToolAnalysis(str(folder), reasons=["The selected folder does not exist."])
    python_paths = sorted(folder.rglob("*.py"))
    python_paths = [
        path for path in python_paths
        if not any(part.casefold() in {".git", ".venv", "venv", "__pycache__", "site-packages"}
                   for part in path.parts)
    ][:500]
    python_files = [_relative(path, folder) for path in python_paths]
    ranked = sorted(
        python_files,
        key=lambda value: (
            0 if Path(value).name.casefold() in ENTRY_NAMES else 1,
            value.count("/"),
            len(value),
        ),
    )
    configs = sorted(
        _relative(path, folder)
        for path in folder.rglob("*")
        if path.is_file()
        and path.suffix.casefold() in CONFIG_SUFFIXES
        and not any(part.casefold() in {".git", ".venv", "venv", "__pycache__", "site-packages"}
                    for part in path.parts)
    )[:100]
    result = ToolAnalysis(
        folder=str(folder),
        python_files=python_files,
        entry_candidates=ranked[:30],
        config_files=configs,
        selected_entry=ranked[0] if ranked else "",
    )
    return analyze_selection(result, result.selected_entry, [])


def _string_literals(tree: ast.AST) -> list[str]:
    return [
        node.value
        for node in ast.walk(tree)
        if isinstance(node, ast.Constant) and isinstance(node.value, str)
    ]


def _call_name(call: ast.Call) -> str:
    parts: list[str] = []
    value: ast.AST = call.func
    while isinstance(value, ast.Attribute):
        parts.append(value.attr)
        value = value.value
    if isinstance(value, ast.Name):
        parts.append(value.id)
    return ".".join(reversed(parts))


def analyze_selection(
    base: ToolAnalysis,
    entry_relative: str,
    selected_configs: list[str],
) -> ToolAnalysis:
    result = ToolAnalysis(
        folder=base.folder,
        python_files=list(base.python_files),
        entry_candidates=list(base.entry_candidates),
        config_files=list(base.config_files),
        selected_entry=entry_relative,
        selected_configs=list(selected_configs),
    )
    if not entry_relative:
        result.reasons = ["No Python entry script was selected."]
        return result
    folder = Path(result.folder)
    entry = (folder / entry_relative).resolve()
    try:
        entry.relative_to(folder)
    except ValueError:
        result.reasons = ["The entry script is outside the selected tool folder."]
        return result
    try:
        source = entry.read_text(encoding="utf-8", errors="replace")
        tree = ast.parse(source, filename=str(entry))
    except (OSError, SyntaxError) as exc:
        result.reasons = [f"The entry script could not be parsed: {exc}"]
        return result

    score = 2
    reasons: list[str] = []
    warnings: list[str] = []
    literals = _string_literals(tree)
    lowered_source = source.casefold()
    calls = [node for node in ast.walk(tree) if isinstance(node, ast.Call)]
    call_names = [_call_name(call).casefold() for call in calls]

    arguments: dict[str, bool] = {}
    for call in calls:
        name = _call_name(call).casefold()
        if not name.endswith("add_argument"):
            continue
        flags = [
            value.value
            for value in call.args
            if isinstance(value, ast.Constant)
            and isinstance(value.value, str)
            and value.value.startswith("-")
        ]
        if not flags:
            continue
        flag = max(flags, key=len).lstrip("-").replace("-", "_")
        required = any(
            keyword.arg == "required"
            and isinstance(keyword.value, ast.Constant)
            and keyword.value.value is True
            for keyword in call.keywords
        )
        arguments[flag] = required
    result.arguments = sorted(arguments)
    result.required_arguments = sorted(key for key, required in arguments.items() if required)
    if arguments:
        score += 2
        reasons.append(f"Detected {len(arguments)} command-line option(s).")
    else:
        reasons.append("No argparse-style command-line options were detected.")

    has_main_guard = "__name__" in lowered_source and "__main__" in lowered_source
    if has_main_guard:
        score += 1
        reasons.append("Has a standard Python main entry point.")
    else:
        warnings.append("No standard __main__ entry point was detected.")

    training_terms = ("train", "epoch", "optimizer", "loss", "backward", "fit(")
    training_hits = sum(term in lowered_source for term in training_terms)
    if training_hits >= 2:
        score += 1
        reasons.append("The script contains recognizable training logic.")
    else:
        warnings.append("Little recognizable training logic was found in the selected file.")

    data_exts = sorted(
        extension for extension in DATASET_EXTENSIONS
        if extension in lowered_source
    )
    dataset_args = [key for key in arguments if any(term in key for term in ("data", "dataset", "input"))]
    if data_exts or dataset_args:
        score += 1
        parts = []
        if data_exts:
            parts.append(", ".join(data_exts))
        if dataset_args:
            parts.append("arguments: " + ", ".join(dataset_args))
        result.dataset_format = "; ".join(parts)
        reasons.append("Dataset input is discoverable.")

    output_args = [key for key in arguments if any(term in key for term in ("output", "save", "model_dir"))]
    output_literals = [
        value for value in literals
        if any(term in value.casefold() for term in ("output", "checkpoint", ".pt", ".pth", ".ckpt", ".safetensors"))
    ][:5]
    if output_args or output_literals:
        score += 1
        result.output_behavior = (
            ("arguments: " + ", ".join(output_args)) if output_args
            else "Paths referenced by the script: " + ", ".join(output_literals)
        )
        reasons.append("Output or model-save behavior is visible.")

    checkpoint_terms = [
        value for value in ("checkpoint", "state_dict", ".ckpt", ".pt", ".pth", ".safetensors")
        if value in lowered_source
    ]
    result.checkpoint_behavior = (
        "Detected: " + ", ".join(checkpoint_terms)
        if checkpoint_terms else "No checkpoint-writing pattern detected"
    )
    if checkpoint_terms:
        score += 1

    resume_args = [key for key in arguments if any(term in key for term in ("resume", "checkpoint", "load"))]
    if resume_args:
        result.resume_behavior = "Supported through: " + ", ".join(resume_args)
        score += 1
    elif "load_state_dict" in lowered_source or "resume" in lowered_source:
        result.resume_behavior = "Possible in code, but no clear command-line option was detected"
    else:
        result.resume_behavior = "No safe resume interface detected"

    progress_patterns = []
    if "tqdm" in lowered_source:
        progress_patterns.append("tqdm")
    if "print(" in lowered_source:
        progress_patterns.append("console output")
    if any(term in lowered_source for term in ("tensorboard", "wandb", "mlflow")):
        progress_patterns.append("experiment tracker")
    result.progress_behavior = (
        ", ".join(progress_patterns) if progress_patterns else "No progress reporting detected"
    )
    if progress_patterns:
        score += 1

    risky = {
        "os.system": "uses os.system",
        "eval": "uses eval",
        "exec": "uses exec",
        "pickle.loads": "loads unrestricted pickle data",
        "shutil.rmtree": "can recursively delete folders",
        "os.remove": "can delete files",
        "pathlib.path.unlink": "can delete files",
    }
    for call_name in call_names:
        for pattern, message in risky.items():
            if call_name == pattern or call_name.endswith("." + pattern):
                warnings.append(message)
                score -= 2
    for call in calls:
        if _call_name(call).casefold().endswith(("subprocess.run", "subprocess.popen", "subprocess.call")):
            shell_true = any(
                keyword.arg == "shell"
                and isinstance(keyword.value, ast.Constant)
                and keyword.value.value is True
                for keyword in call.keywords
            )
            warnings.append(
                "launches another process with a shell" if shell_true
                else "launches another process"
            )
            score -= 2 if shell_true else 1

    if selected_configs:
        valid_configs = [
            name for name in selected_configs
            if name in result.config_files and (folder / name).is_file()
        ]
        result.selected_configs = valid_configs
        if valid_configs:
            reasons.append("Selected configuration files are present.")

    result.score = max(1, min(10, score))
    result.reasons = reasons or ["Only basic Python compatibility was detected."]
    result.warnings = sorted(set(warnings))
    return result


class ExternalToolStore:
    def __init__(self, root: Path) -> None:
        self.path = root.resolve() / "config" / "external_tools.json"

    def load(self) -> list[dict[str, Any]]:
        try:
            payload = json.loads(self.path.read_text(encoding="utf-8"))
            tools = payload.get("tools", [])
            return [dict(item) for item in tools if isinstance(item, dict)]
        except (OSError, ValueError, TypeError, json.JSONDecodeError):
            return []

    def save_connector(
        self,
        *,
        name: str,
        description: str,
        analysis: ToolAnalysis,
        arguments: list[str],
        required_arguments: list[str],
    ) -> dict[str, Any]:
        safe_id = re.sub(r"[^a-z0-9]+", "_", name.casefold()).strip("_") or "tool"
        tool_id = "external_" + safe_id
        entry = (Path(analysis.folder) / analysis.selected_entry).resolve()
        entry.relative_to(Path(analysis.folder).resolve())
        record = {
            "id": tool_id,
            "name": name.strip(),
            "description": description.strip() or "User-connected external Python tool.",
            "category": "External",
            "entry_function": entry.stem,
            "arguments": arguments,
            "required_arguments": required_arguments,
            "capabilities": ["external_script", "progress", "cancel"],
            "requires_confirmation": True,
            "enabled": True,
            "demo": False,
            "backend": {
                "type": "script",
                "path": str(entry),
                "root": analysis.folder,
            },
            "analysis": asdict(analysis),
        }
        tools = self.load()
        replaced = False
        for index, existing in enumerate(tools):
            if existing.get("id") == tool_id:
                tools[index] = record
                replaced = True
                break
        if not replaced:
            tools.append(record)
        self.path.parent.mkdir(parents=True, exist_ok=True)
        temporary = self.path.with_suffix(".tmp")
        temporary.write_text(json.dumps({"tools": tools}, indent=2), encoding="utf-8")
        temporary.replace(self.path)
        return record