Jonathan
139 lines
5.3 KiB
Python
139 lines
5.3 KiB
Python
from dataclasses import dataclass
|
|
|
|
import psycopg
|
|
|
|
from lyra_worker.browser import MbBrowser
|
|
from lyra_worker.similarity.base import SimilaritySource
|
|
|
|
_TRUE = {"1", "true", "yes", "on"}
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DiscoveryConfig:
|
|
enabled: bool = False
|
|
interval_hours: int = 168
|
|
max_seeds: int = 50
|
|
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),
|
|
max_seeds=_int("discover.maxSeeds", 50),
|
|
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
|
|
|
|
|
|
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.max_seeds),
|
|
)
|
|
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 _derive_albums(conn: psycopg.Connection, browser: MbBrowser, surfaced: list[dict],
|
|
cfg: DiscoveryConfig) -> int:
|
|
return 0 # replaced in the album-derivation task
|
|
|
|
|
|
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 = [s for s in sources if s.health()]
|
|
|
|
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)
|