ai-agents/agents/pirate/tools/sonarr.py

"""Sonarr tools — TV show queries (read-only)."""

from ._common import arr_get


def _compact_series(s):
    return {
        "id": s.get("id"),
        "title": s.get("title"),
        "year": s.get("year"),
        "status": s.get("status"),
        "monitored": s.get("monitored"),
        "network": s.get("network"),
        "episode_count": s.get("statistics", {}).get("episodeCount"),
        "episode_file_count": s.get("statistics", {}).get("episodeFileCount"),
        "size_on_disk_gb": round((s.get("statistics", {}).get("sizeOnDisk", 0) or 0) / 1e9, 1),
    }


def sonarr_queue(limit=20):
    """What TV episodes are currently downloading / pending import."""
    data = arr_get("sonarr", "/api/v3/queue", {"pageSize": limit, "includeSeries": "true"})
    items = []
    for r in data.get("records", []):
        items.append({
            "title": r.get("title"),
            "series": r.get("series", {}).get("title") if r.get("series") else None,
            "status": r.get("status"),
            "timeleft": r.get("timeleft"),
            "size_gb": round((r.get("size", 0) or 0) / 1e9, 2),
            "size_left_gb": round((r.get("sizeleft", 0) or 0) / 1e9, 2),
            "protocol": r.get("protocol"),
            "download_client": r.get("downloadClient"),
        })
    return {"total_records": data.get("totalRecords", 0), "items": items}


def sonarr_upcoming(days=14):
    """TV episodes expected to air (or just released) in the next N days."""
    from datetime import datetime, timedelta, timezone
    now = datetime.now(timezone.utc)
    end = now + timedelta(days=days)
    data = arr_get("sonarr", "/api/v3/calendar", {
        "start": now.date().isoformat(),
        "end": end.date().isoformat(),
        "includeSeries": "true",
    })
    out = []
    for ep in data:
        out.append({
            "series": ep.get("series", {}).get("title"),
            "season": ep.get("seasonNumber"),
            "episode": ep.get("episodeNumber"),
            "title": ep.get("title"),
            "air_date": ep.get("airDateUtc") or ep.get("airDate"),
            "has_file": ep.get("hasFile"),
            "monitored": ep.get("monitored"),
        })
    return {"window_days": days, "episodes": out}


def sonarr_series_search(query):
    """Look up a series in the user's Sonarr library by partial title match.
    Note: this does NOT search Prowlarr/indexers — it only searches what's already tracked."""
    all_series = arr_get("sonarr", "/api/v3/series")
    q = query.lower().strip()
    hits = [s for s in all_series if q in (s.get("title") or "").lower()]
    return {"query": query, "count": len(hits), "series": [_compact_series(s) for s in hits[:20]]}


def sonarr_library_stats():
    """Summary of the Sonarr library: how many series, how many episodes on disk, total size."""
    all_series = arr_get("sonarr", "/api/v3/series")
    total_eps = total_on_disk = total_size = 0
    for s in all_series:
        st = s.get("statistics", {}) or {}
        total_eps += st.get("episodeCount", 0) or 0
        total_on_disk += st.get("episodeFileCount", 0) or 0
        total_size += st.get("sizeOnDisk", 0) or 0
    return {
        "series_count": len(all_series),
        "episode_total": total_eps,
        "episodes_on_disk": total_on_disk,
        "total_size_tb": round(total_size / 1e12, 2),
    }


TOOLS = [
    {
        "name": "sonarr_queue",
        "description": "List TV episodes that Sonarr is currently downloading or has just imported. Use when the user asks 'what's downloading', 'what episodes are coming in', 'is the new episode in yet'.",
        "input_schema": {
            "type": "object",
            "properties": {"limit": {"type": "integer", "description": "Max rows to return (default 20)", "default": 20}},
        },
        "read_only": True,
        "fn": sonarr_queue,
    },
    {
        "name": "sonarr_upcoming",
        "description": "List TV episodes expected to air in the next N days (default 14). Use when the user asks 'what's coming up this week', 'when does the next episode drop', 'what's airing soon'.",
        "input_schema": {
            "type": "object",
            "properties": {"days": {"type": "integer", "description": "Look ahead window in days (default 14)", "default": 14}},
        },
        "read_only": True,
        "fn": sonarr_upcoming,
    },
    {
        "name": "sonarr_series_search",
        "description": "Search the user's Sonarr library for a TV series by title. Does NOT search for new series to add — only checks what's already tracked.",
        "input_schema": {
            "type": "object",
            "properties": {"query": {"type": "string", "description": "Partial or full series title"}},
            "required": ["query"],
        },
        "read_only": True,
        "fn": sonarr_series_search,
    },
    {
        "name": "sonarr_library_stats",
        "description": "High-level stats about the Sonarr TV library: number of series, episodes on disk, total size. Use when the user asks 'how big is the TV library' or 'how many shows do we have'.",
        "input_schema": {"type": "object", "properties": {}},
        "read_only": True,
        "fn": sonarr_library_stats,
    },
]