diff --git a/worker/lyra_worker/confidence.py b/worker/lyra_worker/confidence.py new file mode 100644 index 0000000..574576f --- /dev/null +++ b/worker/lyra_worker/confidence.py @@ -0,0 +1,21 @@ +from difflib import SequenceMatcher + +from lyra_worker.types import Candidate, MBTarget + + +def _sim(a: str, b: str) -> float: + return SequenceMatcher(None, a.strip().casefold(), b.strip().casefold()).ratio() + + +def score_confidence(target: MBTarget, c: Candidate) -> float: + """0..1 confidence that candidate c is the requested target.""" + artist_sim = _sim(target.artist, c.matched_artist) + album_sim = _sim(target.album, c.matched_album) + + if target.track_count is None or target.track_count == c.track_count: + track_factor = 1.0 + else: + diff = abs(target.track_count - c.track_count) + track_factor = max(0.0, 1.0 - diff / max(target.track_count, 1)) + + return 0.5 * artist_sim + 0.4 * album_sim + 0.1 * track_factor diff --git a/worker/tests/test_confidence.py b/worker/tests/test_confidence.py new file mode 100644 index 0000000..9806d0b --- /dev/null +++ b/worker/tests/test_confidence.py @@ -0,0 +1,37 @@ +from lyra_worker.confidence import score_confidence +from lyra_worker.types import Candidate, MBTarget, Quality + +Q = Quality(fmt="FLAC", lossless=True, bit_depth=16, sample_rate=44100) + + +def _cand(artist, album, tracks=10): + return Candidate(source="qobuz", source_ref="x", matched_artist=artist, + matched_album=album, quality=Q, track_count=tracks) + + +def test_exact_match_scores_high(): + target = MBTarget(artist="Radiohead", album="In Rainbows", track_count=10) + assert score_confidence(target, _cand("Radiohead", "In Rainbows", 10)) > 0.95 + + +def test_case_and_whitespace_insensitive(): + target = MBTarget(artist="Radiohead", album="In Rainbows") + assert score_confidence(target, _cand(" radiohead ", "IN RAINBOWS")) > 0.95 + + +def test_wrong_album_scores_low(): + target = MBTarget(artist="Radiohead", album="In Rainbows", track_count=10) + assert score_confidence(target, _cand("Radiohead", "OK Computer", 12)) < 0.7 + + +def test_unknown_track_count_does_not_penalize(): + target = MBTarget(artist="Radiohead", album="In Rainbows", track_count=None) + # track_count unknown -> track_factor neutral; name match still drives a high score + assert score_confidence(target, _cand("Radiohead", "In Rainbows", 99)) > 0.95 + + +def test_track_count_mismatch_lowers_score(): + target = MBTarget(artist="Radiohead", album="In Rainbows", track_count=10) + good = score_confidence(target, _cand("Radiohead", "In Rainbows", 10)) + bad = score_confidence(target, _cand("Radiohead", "In Rainbows", 4)) + assert bad < good