alfred/alfred/agent/registry.py

"""Tool registry — defines and registers all available tools for the agent."""

import inspect
import logging
from collections.abc import Callable
from dataclasses import dataclass
from typing import Any

from .tools.spec import ToolSpec, ToolSpecError
from .tools.spec_loader import load_tool_specs

logger = logging.getLogger(__name__)


@dataclass
class Tool:
    """Represents a tool that can be used by the agent."""

    name: str
    description: str
    func: Callable[..., dict[str, Any]]
    parameters: dict[str, Any]
    cache_key: str | None = None  # Parameter name to use as STM cache key.


_PY_TYPE_TO_JSON = {
    str: "string",
    int: "integer",
    float: "number",
    bool: "boolean",
    list: "array",
    dict: "object",
}


def _json_type_for(annotation) -> str:
    """Map a Python type annotation to a JSON Schema 'type' string."""
    if annotation is inspect.Parameter.empty:
        return "string"
    # Strip Optional[X] / X | None to X.
    args = getattr(annotation, "__args__", None)
    if args:
        non_none = [a for a in args if a is not type(None)]
        if len(non_none) == 1:
            annotation = non_none[0]
    return _PY_TYPE_TO_JSON.get(annotation, "string")


def _create_tool_from_function(func: Callable, spec: ToolSpec | None = None) -> Tool:
    """
    Create a Tool object from a function, optionally enriched with a spec.

    Types and required-ness always come from the Python signature (source of
    truth for the API contract). When a spec is provided, the description
    and per-parameter docs come from the YAML spec instead of the docstring.
    """
    sig = inspect.signature(func)
    sig_params = {name: p for name, p in sig.parameters.items() if name != "self"}

    if spec is not None:
        _validate_spec_matches_signature(func.__name__, sig_params, spec)
        description = spec.compile_description()
        param_descriptions = {
            name: spec.compile_parameter_description(name) for name in sig_params
        }
    else:
        doc = inspect.getdoc(func)
        description = doc.strip().split("\n")[0] if doc else func.__name__
        param_descriptions = {name: f"Parameter {name}" for name in sig_params}

    properties: dict[str, dict[str, Any]] = {}
    required: list[str] = []

    for param_name, param in sig_params.items():
        properties[param_name] = {
            "type": _json_type_for(param.annotation),
            "description": param_descriptions[param_name],
        }
        if param.default is inspect.Parameter.empty:
            required.append(param_name)

    parameters = {
        "type": "object",
        "properties": properties,
        "required": required,
    }

    cache_key = spec.cache.key if spec is not None and spec.cache is not None else None

    return Tool(
        name=func.__name__,
        description=description,
        func=func,
        parameters=parameters,
        cache_key=cache_key,
    )


def _validate_spec_matches_signature(
    func_name: str,
    sig_params: dict[str, inspect.Parameter],
    spec: ToolSpec,
) -> None:
    """Ensure every signature param has a spec entry and vice versa."""
    sig_names = set(sig_params.keys())
    spec_names = set(spec.parameters.keys())

    missing_in_spec = sig_names - spec_names
    if missing_in_spec:
        raise ToolSpecError(
            f"tool '{func_name}': spec is missing entries for parameter(s) "
            f"{sorted(missing_in_spec)}"
        )

    extra_in_spec = spec_names - sig_names
    if extra_in_spec:
        raise ToolSpecError(
            f"tool '{func_name}': spec has entries for unknown parameter(s) "
            f"{sorted(extra_in_spec)} (not in function signature)"
        )


def make_tools(settings) -> dict[str, Tool]:
    """
    Create and register all available tools.

    Args:
        settings: Application settings instance.

    Returns:
        Dictionary mapping tool names to Tool objects.
    """
    from .tools import api as api_tools  # noqa: PLC0415
    from .tools import filesystem as fs_tools  # noqa: PLC0415
    from .tools import language as lang_tools  # noqa: PLC0415
    from .tools import workflow as wf_tools  # noqa: PLC0415

    tool_functions = [
        fs_tools.set_path_for_folder,
        fs_tools.list_folder,
        fs_tools.read_release_metadata,
        fs_tools.query_library,
        fs_tools.analyze_release,
        fs_tools.probe_media,
        fs_tools.resolve_season_destination,
        fs_tools.resolve_episode_destination,
        fs_tools.resolve_movie_destination,
        fs_tools.resolve_series_destination,
        fs_tools.move_media,
        fs_tools.move_to_destination,
        fs_tools.manage_subtitles,
        fs_tools.create_seed_links,
        fs_tools.learn,
        api_tools.find_media_imdb_id,
        api_tools.find_torrent,
        api_tools.add_torrent_by_index,
        api_tools.add_torrent_to_qbittorrent,
        api_tools.get_torrent_by_index,
        lang_tools.set_language,
        wf_tools.start_workflow,
        wf_tools.end_workflow,
    ]

    specs = load_tool_specs()

    tools: dict[str, Tool] = {}
    for func in tool_functions:
        spec = specs.get(func.__name__)
        tool = _create_tool_from_function(func, spec=spec)
        tools[tool.name] = tool

    with_spec = sum(1 for fn in tool_functions if fn.__name__ in specs)
    logger.info(
        f"Registered {len(tools)} tools "
        f"({with_spec} with YAML spec, {len(tools) - with_spec} doc-only): "
        f"{list(tools.keys())}"
    )
    return tools