adam/registry.py

from __future__ import annotations

import json
import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from adam.model_plugins import ModelPluginRegistry


class RegistryError(RuntimeError):
    pass


@dataclass(frozen=True, slots=True)
class ToolSpec:
    id: str
    name: str
    description: str
    category: str
    entry_function: str
    arguments: tuple[str, ...] = ()
    required_arguments: tuple[str, ...] = ()
    capabilities: tuple[str, ...] = ()
    model_trainers: tuple[str, ...] = ()
    generation_options: dict[str, Any] = field(default_factory=dict)
    requires_confirmation: bool = False
    enabled: bool = True
    demo: bool = False
    backend: dict[str, Any] = field(default_factory=dict)

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "ToolSpec":
        required = {"id", "name", "description", "category", "backend"}
        missing = required - data.keys()
        if missing:
            raise RegistryError(f"Tool entry is missing: {', '.join(sorted(missing))}")
        return cls(
            id=str(data["id"]),
            name=str(data["name"]),
            description=str(data["description"]),
            category=str(data["category"]),
            entry_function=str(data.get("entry_function", "")),
            arguments=tuple(str(item) for item in data.get("arguments", [])),
            required_arguments=tuple(
                str(item) for item in data.get("required_arguments", [])
            ),
            capabilities=tuple(str(item) for item in data.get("capabilities", [])),
            model_trainers=tuple(
                str(item) for item in data.get("model_trainers", [])
            ),
            generation_options=dict(data.get("generation_options", {})),
            requires_confirmation=bool(data.get("requires_confirmation", False)),
            enabled=bool(data.get("enabled", True)),
            demo=bool(data.get("demo", False)),
            backend=dict(data["backend"]),
        )


class ToolRegistry:
    def __init__(self, root: Path) -> None:
        self.root = root.resolve()
        self.path = self.root / "config" / "tools.json"
        self._tools: dict[str, ToolSpec] = {}
        self.model_plugins = ModelPluginRegistry(self.root, logging.getLogger(__name__))
        self.load()

    def load(self) -> None:
        try:
            payload = json.loads(self.path.read_text(encoding="utf-8"))
        except FileNotFoundError as exc:
            raise RegistryError(f"Tool registry not found: {self.path}") from exc
        except json.JSONDecodeError as exc:
            raise RegistryError(f"Tool registry is invalid JSON: {exc}") from exc

        entries = payload.get("tools")
        if not isinstance(entries, list):
            raise RegistryError("Tool registry must contain a 'tools' list.")

        loaded: dict[str, ToolSpec] = {}
        for entry in entries:
            spec = ToolSpec.from_dict(entry)
            if spec.id in loaded:
                raise RegistryError(f"Duplicate tool id: {spec.id}")
            loaded[spec.id] = spec
        external_path = self.root / "config" / "external_tools.json"
        try:
            external_payload = json.loads(external_path.read_text(encoding="utf-8"))
            external_entries = external_payload.get("tools", [])
        except FileNotFoundError:
            external_entries = []
        except json.JSONDecodeError as exc:
            raise RegistryError(f"External tool registry is invalid JSON: {exc}") from exc
        if not isinstance(external_entries, list):
            raise RegistryError("External tool registry must contain a 'tools' list.")
        for entry in external_entries:
            if not isinstance(entry, dict):
                raise RegistryError("External tool entry must be an object.")
            safe_entry = dict(entry)
            tool_id = str(safe_entry.get("id", ""))
            backend = dict(safe_entry.get("backend", {}))
            path = Path(str(backend.get("path", ""))).expanduser()
            root = Path(str(backend.get("root", ""))).expanduser()
            if not tool_id.startswith("external_"):
                raise RegistryError("External tool ids must start with 'external_'.")
            if tool_id in loaded:
                raise RegistryError(f"External tool cannot replace registered tool: {tool_id}")
            if backend.get("type") != "script" or not path.is_absolute() or not root.is_absolute():
                raise RegistryError(f"External tool {tool_id} has an invalid script backend.")
            try:
                path.resolve().relative_to(root.resolve())
            except ValueError as exc:
                raise RegistryError(f"External tool {tool_id} script is outside its folder.") from exc
            safe_entry["requires_confirmation"] = True
            safe_entry["demo"] = False
            spec = ToolSpec.from_dict(safe_entry)
            loaded[spec.id] = spec
        for entry in self.model_plugins.training_tool_specs() + self.model_plugins.generation_tool_specs():
            if not isinstance(entry, dict):
                continue
            safe_entry = dict(entry)
            backend = dict(safe_entry.get("backend", {}))
            if not backend:
                safe_entry["backend"] = {
                    "type": "python",
                    "module": "adam.model_plugin_backend",
                    "function": "train" if safe_entry.get("category") == "Training" else "generate",
                }
            try:
                spec = ToolSpec.from_dict(safe_entry)
            except RegistryError as exc:
                self.model_plugins.errors.append(f"{safe_entry.get('id', 'unknown')}: {exc}")
                continue
            if spec.id in loaded:
                loaded[spec.id] = _merge_tool_specs(loaded[spec.id], spec)
            else:
                loaded[spec.id] = spec
        self._tools = loaded

    def get(self, tool_id: str, *, require_enabled: bool = True) -> ToolSpec:
        try:
            tool = self._tools[tool_id]
        except KeyError as exc:
            raise RegistryError(f"Unregistered tool: {tool_id}") from exc
        if require_enabled and not tool.enabled:
            raise RegistryError(f"Tool is not configured: {tool.name}")
        return tool

    def all(self) -> list[ToolSpec]:
        return list(self._tools.values())

    def enabled(self) -> list[ToolSpec]:
        return [tool for tool in self._tools.values() if tool.enabled]

    def safe_llm_catalog(self) -> list[dict[str, Any]]:
        return [
            {
                "id": tool.id,
                "name": tool.name,
                "description": tool.description,
                "arguments": list(tool.arguments),
                "required_arguments": list(tool.required_arguments),
                "capabilities": list(tool.capabilities),
                "model_trainers": list(tool.model_trainers),
                "generation_options": dict(tool.generation_options),
                "requires_confirmation": tool.requires_confirmation,
            }
            for tool in self.enabled()
        ]


def _merge_tool_specs(existing: ToolSpec, plugin: ToolSpec) -> ToolSpec:
    """Keep the existing backend while accepting plugin-declared schema arguments."""
    arguments = tuple(dict.fromkeys([*existing.arguments, *plugin.arguments]))
    required_arguments = existing.required_arguments or plugin.required_arguments
    capabilities = tuple(dict.fromkeys([*existing.capabilities, *plugin.capabilities]))
    model_trainers = tuple(
        dict.fromkeys([*existing.model_trainers, *plugin.model_trainers])
    )
    generation_options = dict(existing.generation_options)
    generation_options.update(plugin.generation_options)
    return ToolSpec(
        id=existing.id,
        name=existing.name,
        description=existing.description,
        category=existing.category,
        entry_function=existing.entry_function,
        arguments=arguments,
        required_arguments=required_arguments,
        capabilities=capabilities,
        model_trainers=model_trainers,
        generation_options=generation_options,
        requires_confirmation=existing.requires_confirmation,
        enabled=existing.enabled,
        demo=existing.demo,
        backend=existing.backend,
    )