Lyra/worker/lyra_worker/discovery.py

from dataclasses import dataclass

import psycopg

from lyra_worker.browser import MbBrowser
from lyra_worker.release_filter import is_core_release
from lyra_worker.similarity.base import SimilaritySource

_TRUE = {"1", "true", "yes", "on"}


@dataclass(frozen=True)
class DiscoveryConfig:
    enabled: bool = False
    interval_hours: int = 168
    chunk_size: int = 5
    similar_per_seed: int = 20
    albums_per_artist: int = 1
    min_score: float = 0.0

    @classmethod
    def from_config(cls, config: dict) -> "DiscoveryConfig":
        def _int(key: str, default: int) -> int:
            try:
                return int(config[key])
            except (KeyError, TypeError, ValueError):
                return default

        def _float(key: str, default: float) -> float:
            try:
                return float(config[key])
            except (KeyError, TypeError, ValueError):
                return default

        return cls(
            enabled=str(config.get("discover.enabled", "")).strip().lower() in _TRUE,
            interval_hours=_int("discover.intervalHours", 168),
            chunk_size=_int("discover.chunkSize", 5),
            similar_per_seed=_int("discover.similarPerSeed", 20),
            albums_per_artist=_int("discover.albumsPerArtist", 1),
            min_score=_float("discover.minScore", 0.0),
        )


@dataclass(frozen=True)
class DiscoveryResult:
    artists: int
    albums: int
    seeds: int = 0


def _healthy_sources(sources):
    """Return sources whose health() is truthy; a source whose health() raises
    is logged and skipped, never aborting the sweep."""
    live = []
    for s in sources:
        try:
            if s.health():
                live.append(s)
        except Exception as e:  # a bad health check must not abort the sweep
            print(f"worker: discovery source {getattr(s, 'name', '?')} health check failed: {e}", flush=True)
    return live


def _seed_artists(conn: psycopg.Connection, cfg: DiscoveryConfig):
    with conn.cursor() as cur:
        cur.execute(
            'SELECT mbid, name FROM "WatchedArtist" '
            'WHERE "lastDiscoveredAt" IS NULL '
            '   OR "lastDiscoveredAt" < now() - make_interval(hours => %s) '
            'ORDER BY "lastDiscoveredAt" ASC NULLS FIRST LIMIT %s',
            (cfg.interval_hours, cfg.chunk_size),
        )
        return cur.fetchall()


def _followed_mbids(conn: psycopg.Connection) -> set[str]:
    with conn.cursor() as cur:
        cur.execute('SELECT mbid FROM "WatchedArtist"')
        return {r[0] for r in cur.fetchall()}


def _upsert_artist(conn: psycopg.Connection, cand: dict) -> int:
    dedupe = f"artist:{cand['mbid']}:"
    with conn.cursor() as cur:
        cur.execute(
            'INSERT INTO "DiscoverySuggestion" (id, kind, "artistMbid", "artistName", '
            'score, "seedCount", sources, status, "dedupeKey", "createdAt", "updatedAt") '
            "VALUES (gen_random_uuid()::text, 'artist', %s, %s, %s, %s, %s, 'pending', %s, now(), now()) "
            'ON CONFLICT ("dedupeKey") DO UPDATE SET '
            'score = EXCLUDED.score, "seedCount" = EXCLUDED."seedCount", '
            'sources = EXCLUDED.sources, "updatedAt" = now() '
            "WHERE \"DiscoverySuggestion\".status = 'pending'",
            (cand["mbid"], cand["name"], cand["score"], cand["seed_count"],
             sorted(cand["sources"]), dedupe),
        )
        return cur.rowcount


def _existing_rg_mbids(conn: psycopg.Connection) -> set[str]:
    with conn.cursor() as cur:
        cur.execute('SELECT "rgMbid" FROM "MonitoredRelease"')
        return {r[0] for r in cur.fetchall()}


def _upsert_album(conn: psycopg.Connection, artist: dict, rg) -> int:
    dedupe = f"album:{artist['mbid']}:{rg.rg_mbid}"
    with conn.cursor() as cur:
        cur.execute(
            'INSERT INTO "DiscoverySuggestion" (id, kind, "artistMbid", "artistName", '
            '"rgMbid", album, "primaryType", "secondaryTypes", "firstReleaseDate", '
            'score, "seedCount", sources, status, "dedupeKey", "createdAt", "updatedAt") '
            "VALUES (gen_random_uuid()::text, 'album', %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, "
            "'pending', %s, now(), now()) "
            'ON CONFLICT ("dedupeKey") DO UPDATE SET score = EXCLUDED.score, '
            '"updatedAt" = now() '
            "WHERE \"DiscoverySuggestion\".status = 'pending'",
            (artist["mbid"], artist["name"], rg.rg_mbid, rg.title,
             rg.primary_type or None, list(rg.secondary_types), rg.first_release_date or None,
             artist["score"], artist["seed_count"], sorted(artist["sources"]), dedupe),
        )
        return cur.rowcount


def _derive_albums(conn: psycopg.Connection, browser: MbBrowser, surfaced: list[dict],
                   cfg: DiscoveryConfig) -> int:
    if cfg.albums_per_artist <= 0:
        return 0
    have = _existing_rg_mbids(conn)
    albums = 0
    for artist in surfaced:
        try:
            groups = browser.browse_release_groups(artist["mbid"])
        except Exception as e:  # one artist's browse failure must not abort the sweep
            print(f"worker: discovery album browse failed for {artist['mbid']}: {e}", flush=True)
            continue
        core = [g for g in groups
                if is_core_release(g.primary_type, g.secondary_types) and g.rg_mbid not in have]
        core.sort(key=lambda g: g.first_release_date or "", reverse=True)
        for rg in core[: cfg.albums_per_artist]:
            albums += _upsert_album(conn, artist, rg)
            have.add(rg.rg_mbid)
    return albums


def run_discovery(conn: psycopg.Connection, sources: list[SimilaritySource],
                  browser: MbBrowser, cfg: DiscoveryConfig) -> DiscoveryResult:
    """Aggregate similar artists across the library's seed artists; upsert suggestions."""
    seeds = _seed_artists(conn, cfg)
    followed = _followed_mbids(conn)
    live = _healthy_sources(sources)

    agg: dict[str, dict] = {}
    for seed_mbid, _seed_name in seeds:
        for src in live:
            try:
                similar = src.similar_artists(seed_mbid)
            except Exception as e:  # a bad source/seed must not abort the sweep
                print(f"worker: discovery source {src.name} failed for {seed_mbid}: {e}", flush=True)
                continue
            for sa in sorted(similar, key=lambda a: a.score, reverse=True)[: cfg.similar_per_seed]:
                if sa.mbid in followed:
                    continue
                c = agg.setdefault(sa.mbid, {"mbid": sa.mbid, "name": sa.name,
                                             "score": 0.0, "seeds": set(), "sources": set()})
                c["score"] += sa.score
                c["seeds"].add(seed_mbid)
                c["sources"].add(src.name)
                if sa.name and not c["name"]:
                    c["name"] = sa.name

    artists = 0
    surfaced: list[dict] = []
    for c in sorted(agg.values(), key=lambda c: c["score"], reverse=True):
        if c["score"] < cfg.min_score:
            continue
        cand = {"mbid": c["mbid"], "name": c["name"], "score": c["score"],
                "seed_count": len(c["seeds"]), "sources": c["sources"]}
        inserted = _upsert_artist(conn, cand)
        artists += inserted
        # Only derive albums for candidates that produced a live pending suggestion.
        # A no-op upsert (rowcount 0) means the artist was already dismissed/wanted —
        # don't surface their albums against a dismissal.
        if inserted:
            surfaced.append(cand)

    albums = _derive_albums(conn, browser, surfaced, cfg)

    with conn.cursor() as cur:
        for seed_mbid, _seed_name in seeds:
            cur.execute('UPDATE "WatchedArtist" SET "lastDiscoveredAt" = now() WHERE mbid = %s',
                        (seed_mbid,))
    conn.commit()
    return DiscoveryResult(artists=artists, albums=albums, seeds=len(seeds))