diff --git a/core/download_engine/engine.py b/core/download_engine/engine.py index 8aaa07e8..d3f04a82 100644 --- a/core/download_engine/engine.py +++ b/core/download_engine/engine.py @@ -28,6 +28,7 @@ from __future__ import annotations import threading from typing import Any, Dict, Iterator, List, Optional, Tuple +from core.quality.selection import load_profile_targets, rank_with_targets from utils.logging_config import get_logger logger = get_logger("download_engine") @@ -391,9 +392,21 @@ class DownloadEngine: (tracks, albums) tuple, or ``([], [])`` when every source in the chain is exhausted. + Quality-aware: a source is only accepted when it can deliver a + candidate that meets a real target in the user's quality profile. + A source that meets no target is skipped in favour of the next one + (source priority still wins among satisfying sources). RAW tracks are + returned — the per-source ``satisfied`` check is a coarse gate; + match-filtering and final quality ranking happen in the orchestrator. + When no source satisfies a target, the first source's results are + returned as a fallback (unless the profile disables fallback). + Replaces orchestrator's hand-rolled hybrid search loop. The chain is ordered (most-preferred first). """ + targets, fallback_enabled = load_profile_targets() + first_fallback = None # (tracks, albums) of the first source with hits + for i, source_name in enumerate(source_chain): plugin = self._plugins.get(source_name) if plugin is None: @@ -406,14 +419,43 @@ class DownloadEngine: try: logger.info(f"Trying {source_name} (priority {i+1}): {query}") tracks, albums = await plugin.search(query, timeout, progress_callback) - if tracks: - logger.info(f"{source_name} found {len(tracks)} tracks") + if not tracks: + continue + try: + _, satisfied = rank_with_targets( + tracks, targets, fallback_enabled=fallback_enabled, + ) + except Exception as rank_err: + # Fail open: a ranking bug must never drop a source's real + # download results. Accept them as-is. + logger.warning( + f"Quality ranking failed for {source_name} — accepting " + f"its results unranked: {rank_err}" + ) + satisfied = True + if satisfied: + logger.info( + f"{source_name} found {len(tracks)} tracks (meets quality target)" + ) return (tracks, albums) + logger.info( + f"{source_name} found {len(tracks)} tracks but none meet a " + f"quality target — trying next source" + ) + if first_fallback is None: + first_fallback = (tracks, albums) except Exception as e: logger.warning(f"{source_name} search failed: {e}") + if fallback_enabled and first_fallback is not None: + logger.info( + "Hybrid search: no source met a quality target — falling back to " + "first source's results" + ) + return first_fallback + logger.warning( - "Hybrid search: all sources (%s) found nothing for: %s", + "Hybrid search: all sources (%s) found nothing acceptable for: %s", ', '.join(source_chain), query, ) return ([], []) diff --git a/core/quality/selection.py b/core/quality/selection.py new file mode 100644 index 00000000..fc9e8ef4 --- /dev/null +++ b/core/quality/selection.py @@ -0,0 +1,80 @@ +"""Quality-aware candidate selection shared by the search engine and the +download orchestrator. + +``rank_with_targets`` is the pure core: it ranks candidates against a target +list and reports whether any candidate met a *real* target (strict, fallback +off). The engine uses that ``satisfied`` flag to decide whether the current +source is good enough or it should fall through to the next source in the +hybrid chain. + +``rank_for_profile`` is the thin DB-backed wrapper that loads the user's +quality profile (with v2->v3 migration) and delegates. +""" + +from __future__ import annotations + +from typing import List, Tuple + +from core.quality.model import ( + QualityTarget, + filter_and_rank, + v2_qualities_to_ranked_targets, +) + + +def rank_with_targets( + candidates: list, + targets: List[QualityTarget], + *, + fallback_enabled: bool = True, +) -> Tuple[list, bool]: + """Rank *candidates* against *targets*. + + Returns ``(ranked, satisfied)`` where ``satisfied`` is True when at least + one candidate meets a real target. When no targets are configured the + profile imposes no constraint, so any non-empty result counts as + satisfied (the first source wins, quality-sorted). + """ + if not candidates: + return [], False + + if not targets: + ranked = filter_and_rank(candidates, targets, fallback_enabled=True) + return ranked, bool(ranked) + + strict = filter_and_rank(candidates, targets, fallback_enabled=False) + if strict: + return strict, True + + if fallback_enabled: + return filter_and_rank(candidates, targets, fallback_enabled=True), False + return [], False + + +def load_profile_targets() -> Tuple[List[QualityTarget], bool]: + """Load the user's quality profile from the DB and return + ``(targets, fallback_enabled)`` with v2->v3 migration applied. + + Callers that rank across many sources should load once and reuse via + :func:`rank_with_targets` rather than calling :func:`rank_for_profile` + per source. + """ + from database.music_database import MusicDatabase + + profile = MusicDatabase().get_quality_profile() + raw_targets = profile.get('ranked_targets') + if not raw_targets and 'qualities' in profile: + raw_targets = v2_qualities_to_ranked_targets(profile['qualities']) + + targets = [QualityTarget.from_dict(t) for t in (raw_targets or [])] + fallback_enabled = profile.get('fallback_enabled', True) + return targets, fallback_enabled + + +def rank_for_profile(candidates: list) -> Tuple[list, bool]: + """Load the user's quality profile and rank *candidates* against it. + + Returns ``(ranked, satisfied)`` — see :func:`rank_with_targets`. + """ + targets, fallback_enabled = load_profile_targets() + return rank_with_targets(candidates, targets, fallback_enabled=fallback_enabled) diff --git a/tests/quality/test_engine_fallback.py b/tests/quality/test_engine_fallback.py new file mode 100644 index 00000000..e85af0c3 --- /dev/null +++ b/tests/quality/test_engine_fallback.py @@ -0,0 +1,108 @@ +"""Engine search_with_fallback is quality-aware: it falls through to the next +source when the current source can deliver no target-satisfying quality, but +returns RAW tracks (match-filtering happens later in the orchestrator). +""" + +import asyncio + +import pytest + +from core.download_engine import engine as engine_mod +from core.download_engine.engine import DownloadEngine +from core.quality.model import AudioQuality, QualityTarget + + +class _Cand: + def __init__(self, aq, name): + self.audio_quality = aq + self.name = name + + +class _FakePlugin: + def __init__(self, tracks): + self._tracks = tracks + self.searched = False + + def is_configured(self): + return True + + async def search(self, query, timeout=None, progress_callback=None): + self.searched = True + return (self._tracks, []) + + +FLAC = AudioQuality('flac', sample_rate=44100, bit_depth=16) +MP3 = AudioQuality('mp3', bitrate=320) +WANT_FLAC_ONLY = [QualityTarget(label='FLAC 16', format='flac', bit_depth=16)] + + +def _engine_with(plugins): + eng = object.__new__(DownloadEngine) + eng._plugins = plugins + return eng + + +def _patch_profile(monkeypatch, targets, fallback_enabled): + monkeypatch.setattr( + engine_mod, 'load_profile_targets', + lambda: (targets, fallback_enabled), + ) + + +def test_escalates_to_next_source_when_first_cannot_meet_target(monkeypatch): + _patch_profile(monkeypatch, WANT_FLAC_ONLY, True) + first = _FakePlugin([_Cand(MP3, 'a-mp3')]) + second = _FakePlugin([_Cand(FLAC, 'b-flac')]) + eng = _engine_with({'first': first, 'second': second}) + + tracks, _ = asyncio.run(eng.search_with_fallback('q', ['first', 'second'])) + + assert [t.name for t in tracks] == ['b-flac'] # escalated to FLAC source + assert second.searched is True + + +def test_stops_on_first_satisfying_source(monkeypatch): + _patch_profile(monkeypatch, WANT_FLAC_ONLY, True) + first = _FakePlugin([_Cand(FLAC, 'a-flac')]) + second = _FakePlugin([_Cand(FLAC, 'b-flac')]) + eng = _engine_with({'first': first, 'second': second}) + + tracks, _ = asyncio.run(eng.search_with_fallback('q', ['first', 'second'])) + + assert [t.name for t in tracks] == ['a-flac'] + assert second.searched is False # never queried — source priority king + + +def test_returns_raw_tracks_not_pruned(monkeypatch): + # A source satisfied by one candidate must still return ALL its tracks so + # the orchestrator's match filter can pick the correct one. + _patch_profile(monkeypatch, WANT_FLAC_ONLY, True) + first = _FakePlugin([_Cand(MP3, 'wrong-but-present'), _Cand(FLAC, 'flac')]) + eng = _engine_with({'first': first}) + + tracks, _ = asyncio.run(eng.search_with_fallback('q', ['first'])) + + names = {t.name for t in tracks} + assert names == {'wrong-but-present', 'flac'} # nothing pruned + + +def test_no_source_satisfies_fallback_on_returns_first_source(monkeypatch): + _patch_profile(monkeypatch, WANT_FLAC_ONLY, True) + first = _FakePlugin([_Cand(MP3, 'a-mp3')]) + second = _FakePlugin([_Cand(MP3, 'b-mp3')]) + eng = _engine_with({'first': first, 'second': second}) + + tracks, _ = asyncio.run(eng.search_with_fallback('q', ['first', 'second'])) + + assert [t.name for t in tracks] == ['a-mp3'] # source priority for fallback + + +def test_no_source_satisfies_fallback_off_returns_empty(monkeypatch): + _patch_profile(monkeypatch, WANT_FLAC_ONLY, False) + first = _FakePlugin([_Cand(MP3, 'a-mp3')]) + eng = _engine_with({'first': first}) + + tracks, albums = asyncio.run(eng.search_with_fallback('q', ['first'])) + + assert tracks == [] + assert albums == [] diff --git a/tests/quality/test_selection.py b/tests/quality/test_selection.py new file mode 100644 index 00000000..1bcb6bd9 --- /dev/null +++ b/tests/quality/test_selection.py @@ -0,0 +1,64 @@ +"""core.quality.selection — quality-aware ranking + the satisfied flag that +drives the engine's source fall-through. + +A source is "satisfied" when at least one of its candidates meets a real +target (strict, fallback off). The engine uses that to decide whether to +stop on the current source or escalate to the next. +""" + +import pytest + +from core.quality.model import AudioQuality, QualityTarget +from core.quality.selection import rank_with_targets + + +class _Cand: + """Minimal candidate: filter_and_rank only needs ``.audio_quality``.""" + def __init__(self, aq, name=""): + self.audio_quality = aq + self.name = name + + def __repr__(self): + return f"_Cand({self.name})" + + +FLAC_HIRES = AudioQuality('flac', sample_rate=96000, bit_depth=24) +FLAC_CD = AudioQuality('flac', sample_rate=44100, bit_depth=16) +MP3_320 = AudioQuality('mp3', bitrate=320) + +WANT_HIRES = [QualityTarget(label='FLAC 24', format='flac', bit_depth=24, min_sample_rate=96000)] +WANT_FLAC_ONLY = [QualityTarget(label='FLAC 16', format='flac', bit_depth=16)] + + +def test_satisfied_when_a_candidate_meets_a_target(): + cands = [_Cand(MP3_320, 'mp3'), _Cand(FLAC_HIRES, 'hires')] + ranked, satisfied = rank_with_targets(cands, WANT_HIRES, fallback_enabled=True) + assert satisfied is True + assert ranked[0].name == 'hires' # the matching candidate wins + + +def test_unsatisfied_but_fallback_returns_sorted_when_enabled(): + cands = [_Cand(MP3_320, 'mp3')] + ranked, satisfied = rank_with_targets(cands, WANT_FLAC_ONLY, fallback_enabled=True) + assert satisfied is False # no FLAC → no target met + assert [c.name for c in ranked] == ['mp3'] # but fallback keeps it + + +def test_unsatisfied_and_fallback_off_returns_empty(): + cands = [_Cand(MP3_320, 'mp3')] + ranked, satisfied = rank_with_targets(cands, WANT_FLAC_ONLY, fallback_enabled=False) + assert satisfied is False + assert ranked == [] + + +def test_empty_targets_accepts_everything_satisfied(): + cands = [_Cand(MP3_320, 'mp3'), _Cand(FLAC_CD, 'cd')] + ranked, satisfied = rank_with_targets(cands, [], fallback_enabled=True) + assert satisfied is True # no constraint → first source wins + assert ranked[0].name == 'cd' # still quality-sorted + + +def test_no_candidates_is_unsatisfied(): + ranked, satisfied = rank_with_targets([], WANT_FLAC_ONLY, fallback_enabled=True) + assert satisfied is False + assert ranked == [] diff --git a/tests/tools/test_amazon_download_client.py b/tests/tools/test_amazon_download_client.py index f831f268..75c6a54b 100644 --- a/tests/tools/test_amazon_download_client.py +++ b/tests/tools/test_amazon_download_client.py @@ -224,7 +224,10 @@ class TestSearch: assert t.artist == "Kendrick Lamar" assert t.title == "Track 0" assert t.album == "GNX" - assert t.quality == "Lossless" + # Quality is now stamped as a real format token (was the display + # label "Lossless", which broke audio_quality format derivation). + assert t.quality == "flac" + assert t.audio_quality.format == "flac" assert t.duration == 200_000 def test_track_source_metadata(self, tmp_path):