alfred/agent/registry.py

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

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

from .tools import api as api_tools
from .tools import filesystem as fs_tools

logger = logging.getLogger(__name__)


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

    Attributes:
        name: Unique identifier for the tool.
        description: Human-readable description for the LLM.
        func: The callable that implements the tool.
        parameters: JSON Schema describing the tool's parameters.
    """

    name: str
    description: str
    func: Callable[..., dict[str, Any]]
    parameters: dict[str, Any]


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

    Tools access memory via get_memory() context.

    Returns:
        Dictionary mapping tool names to Tool instances.
    """
    tools = [
        # Filesystem tools
        Tool(
            name="set_path_for_folder",
            description=(
                "Sets a path in the configuration "
                "(download_folder, tvshow_folder, movie_folder, or torrent_folder)."
            ),
            func=fs_tools.set_path_for_folder,
            parameters={
                "type": "object",
                "properties": {
                    "folder_name": {
                        "type": "string",
                        "description": "Name of folder to set",
                        "enum": ["download", "tvshow", "movie", "torrent"],
                    },
                    "path_value": {
                        "type": "string",
                        "description": "Absolute path to the folder",
                    },
                },
                "required": ["folder_name", "path_value"],
            },
        ),
        Tool(
            name="list_folder",
            description="Lists the contents of a configured folder.",
            func=fs_tools.list_folder,
            parameters={
                "type": "object",
                "properties": {
                    "folder_type": {
                        "type": "string",
                        "description": "Type of folder to list",
                        "enum": ["download", "tvshow", "movie", "torrent"],
                    },
                    "path": {
                        "type": "string",
                        "description": "Relative path within the folder",
                        "default": ".",
                    },
                },
                "required": ["folder_type"],
            },
        ),
        # Media search tools
        Tool(
            name="find_media_imdb_id",
            description=(
                "Finds the IMDb ID for a given media title using TMDB API. "
                "Use this to get information about a movie or TV show."
            ),
            func=api_tools.find_media_imdb_id,
            parameters={
                "type": "object",
                "properties": {
                    "media_title": {
                        "type": "string",
                        "description": "Title of the media to search for",
                    },
                },
                "required": ["media_title"],
            },
        ),
        # Torrent tools
        Tool(
            name="find_torrents",
            description=(
                "Finds torrents for a given media title. "
                "Results are numbered (1, 2, 3...) so the user can select by number."
            ),
            func=api_tools.find_torrent,
            parameters={
                "type": "object",
                "properties": {
                    "media_title": {
                        "type": "string",
                        "description": "Title to search for (include quality if specified)",
                    },
                },
                "required": ["media_title"],
            },
        ),
        Tool(
            name="add_torrent_by_index",
            description=(
                "Adds a torrent from the previous search results by its number. "
                "Use when the user says 'download the 3rd one' or 'take number 2'."
            ),
            func=api_tools.add_torrent_by_index,
            parameters={
                "type": "object",
                "properties": {
                    "index": {
                        "type": "integer",
                        "description": "Number of the torrent in search results (1, 2, 3...)",
                    },
                },
                "required": ["index"],
            },
        ),
        Tool(
            name="add_torrent_to_qbittorrent",
            description=(
                "Adds a torrent to qBittorrent using a magnet link directly. "
                "Use add_torrent_by_index if user selected from search results."
            ),
            func=api_tools.add_torrent_to_qbittorrent,
            parameters={
                "type": "object",
                "properties": {
                    "magnet_link": {
                        "type": "string",
                        "description": "The magnet link of the torrent",
                    },
                },
                "required": ["magnet_link"],
            },
        ),
        Tool(
            name="get_torrent_by_index",
            description=(
                "Gets details of a torrent from search results by its number, "
                "without downloading it."
            ),
            func=api_tools.get_torrent_by_index,
            parameters={
                "type": "object",
                "properties": {
                    "index": {
                        "type": "integer",
                        "description": "Number of the torrent in search results (1, 2, 3...)",
                    },
                },
                "required": ["index"],
            },
        ),
    ]

    logger.info(f"Registered {len(tools)} tools")
    return {t.name: t for t in tools}