tests/test_recommend.py

import re
from dataclasses import replace
from datetime import datetime, timezone
from pathlib import Path

from app.catalogue import build_record
from app.model_detail import ModelDetail
from app.recommend import CommunitySignal, HeuristicEngine, Query, sort_results

from .test_parsing import QWEN3_8B_CONFIG

NOW = datetime(2026, 9, 1, tzinfo=timezone.utc)
ENGINE = HeuristicEngine(now=NOW)


def rec(name, total, bits, downloads=1000, modified="2026-08-01T00:00:00Z", tags=None):
    return build_record({
        "id": f"mlx-community/{name}", "downloads": downloads, "likes": 1, "lastModified": modified,
        "pipeline_tag": "text-generation", "safetensors": {"total": total},
        "tags": (tags or []) + [f"{bits}-bit"],
    })


def family():
    return [
        rec("Fam-8B-2bit", 8e9, 2), rec("Fam-8B-3bit", 8e9, 3), rec("Fam-8B-4bit", 8e9, 4),
        rec("Fam-8B-6bit", 8e9, 6), rec("Fam-8B-8bit", 8e9, 8),
    ]


def ranked(models, q, details=None, community=None):
    return [s.model.name for s in sort_results(ENGINE.score(models, q, details or {}, community or {}), "recommended")]


def test_priority_changes_order():
    ms = family()
    quality = ranked(ms, Query(ram_gb=64, priority="quality"))
    memory = ranked(ms, Query(ram_gb=64, priority="memory"))
    assert quality[0] == "Fam-8B-8bit"
    assert memory[0] == "Fam-8B-2bit"
    assert ranked(ms, Query(ram_gb=64, priority="balanced"))[0] == "Fam-8B-4bit"


def test_ram_changes_order():
    ms = family()
    big = ranked(ms, Query(ram_gb=64, priority="quality"))
    small = ranked(ms, Query(ram_gb=8, priority="quality"))
    assert big[0] == "Fam-8B-8bit"
    assert small[0] != "Fam-8B-8bit"  # 8-bit 8B (~8.5 GB) can't load on an 8 GB Mac
    scored = {s.model.name: s for s in ENGINE.score(ms, Query(ram_gb=8, priority="quality"))}
    assert scored["Fam-8B-8bit"].memory.fit == "Unlikely"


def test_context_changes_score():
    m = rec("Fam-8B-4bit", 8e9, 4)
    d = {m.id: ModelDetail(m.id, QWEN3_8B_CONFIG, [], None, 0)}
    short = ENGINE.score([m], Query(ram_gb=16, context=4096), d)[0]
    long = ENGINE.score([m], Query(ram_gb=16, context=131072), d)[0]
    assert long.memory.total_gb > short.memory.total_gb
    assert long.score < short.score
    assert any("shorter maximum context" in r for r in long.reasons)


def test_quant_changes_score_for_same_model():
    a, b = rec("Fam-8B-4bit", 8e9, 4), rec("Fam-8B-3bit", 8e9, 3)
    q = Query(ram_gb=64, priority="balanced")
    s = {x.model.name: x.score for x in ENGINE.score([a, b], q)}
    assert s["Fam-8B-4bit"] > s["Fam-8B-3bit"]


def test_reasons_present_and_no_perf_claims():
    for s in ENGINE.score(family(), Query(ram_gb=16, priority="speed")):
        assert s.reasons
        text = " ".join(s.reasons).lower()
        assert "tok/s" not in text  # no invented throughput without community data
        assert "depends" in text


def test_unknown_ram_and_unknown_memory_degrade_gracefully():
    m = build_record({"id": "mlx-community/mystery", "tags": []})
    s = ENGINE.score([m], Query())[0]
    assert s.memory.total_gb is None and s.memory.fit is None
    assert any("isn't available" in r for r in s.reasons)


def test_neutral_to_names_and_brands():
    base = rec("Qwen3-8B-4bit", 8e9, 4, downloads=5000)
    renamed = replace(base, id="mlx-community/Qwen3-8B-OptiQ-4bit", name="Qwen3-8B-OptiQ-4bit")
    other_org = replace(base, id="someone/Qwen3-8B-4bit", name="Qwen3-8B-4bit")
    q = Query(ram_gb=24, priority="balanced", context=32768)
    scores = [ENGINE.score([m], q, {}, {})[0].score for m in (base, renamed, other_org)]
    assert scores[0] == scores[1] == scores[2]


def test_engine_source_mentions_no_product_names():
    src = (Path(__file__).resolve().parent.parent / "app" / "recommend.py").read_text().lower()
    code = re.sub(r'"""[\s\S]*?"""', "", src)  # the module docstring may explain the rule
    for word in ("optiq", "unsloth", "lmstudio", "bartowski", "dwq"):
        assert word not in code


def test_popularity_and_community_break_ties():
    a = rec("A-8B-4bit", 8e9, 4, downloads=100)
    b = rec("B-8B-4bit", 8e9, 4, downloads=100_000)
    assert ranked([a, b], Query(ram_gb=32))[0] == "B-8B-4bit"
    comm = {a.id: CommunitySignal(feedback_count=30, positive_share=1.0, benchmark_count=20,
                                  median_generation_tps=40)}
    s = {x.model.name: x for x in ENGINE.score([a, b], Query(ram_gb=32), {}, comm)}
    assert s["A-8B-4bit"].components["community"] > 0
    assert any("MLX benchmark" in r for r in s["A-8B-4bit"].reasons)


def test_sorts():
    old = rec("Old-8B-4bit", 8e9, 4, downloads=10**6, modified="2024-01-01T00:00:00Z")
    new = rec("New-8B-4bit", 8e9, 4, downloads=10, modified="2026-08-30T00:00:00Z")
    scored = ENGINE.score([old, new], Query(ram_gb=32))
    assert sort_results(scored, "popular")[0].model.name == "Old-8B-4bit"
    assert sort_results(scored, "recent")[0].model.name == "New-8B-4bit"
    assert [s.model.name for s in sort_results(scored, "all")] == ["New-8B-4bit", "Old-8B-4bit"]


def test_config_quant_overrides_misleading_name():
    m = rec("Fam-9B-4bit", 9e9, 4)
    cfg = {"model_type": "x", "num_hidden_layers": 4, "num_attention_heads": 8, "head_dim": 64}
    d = {m.id: ModelDetail(m.id, cfg, [], 18_000_000_000, 0)}
    s = ENGINE.score([m], Query(ram_gb=36), d)[0]
    assert s.quant["bits"] == 16 and s.to_dict()["quant"]["label"] == "16-bit"
    assert any("config.json says 16-bit" in r for r in s.reasons)


def test_likes_count_toward_popularity():
    a = rec("A-8B-4bit", 8e9, 4, downloads=20_000)
    b = replace(rec("B-8B-4bit", 8e9, 4, downloads=20_000), likes=150)
    assert ranked([a, b], Query(ram_gb=32))[0] == "B-8B-4bit"


def test_capacity_prefers_larger_models_that_fit():
    small = rec("Fam-1B-4bit", 1e9, 4, downloads=5000)
    mid = rec("Fam-8B-4bit", 8e9, 4, downloads=5000)
    huge = rec("Fam-70B-4bit", 70e9, 4, downloads=5000)
    order = ranked([small, mid, huge], Query(ram_gb=36, priority="balanced"))
    assert order[0] == "Fam-8B-4bit" and order[-1] == "Fam-70B-4bit"  # 70B doesn't fit 36 GB
    assert ranked([small, mid, huge], Query(ram_gb=128, priority="quality"))[0] == "Fam-70B-4bit"


def test_norms_make_scores_page_independent():
    ms = [rec(f"M{i}-8B-4bit", 8e9, 4, downloads=10 ** i) for i in range(1, 6)]
    q = Query(ram_gb=32)
    norms = ENGINE.norms(ms, q)
    full = {s.model.id: s.score for s in ENGINE.score(ms, q, {}, {}, norms)}
    page = {s.model.id: s.score for s in ENGINE.score(ms[:2], q, {}, {}, norms)}
    assert all(abs(full[k] - v) < 1e-9 for k, v in page.items())



def _refs(comm, models):
    from app.recommend import perplexity_refs
    return perplexity_refs(comm, {m.id: m.base_model for m in models}, {m.id: m.quant_bucket for m in models})


def test_measured_perplexity_compares_same_width_peers_only():
    a = rec("Base-4B-4bit", 4e9, 4, downloads=20_000, tags=["base_model:Org/Base-4B"])
    b = rec("Base-4B-Mixed-4bit", 4e9, 4, downloads=20_000, tags=["base_model:Org/Base-4B"])
    c = rec("Base-4B-8bit", 4e9, 8, downloads=20_000, tags=["base_model:Org/Base-4B"])
    u = rec("Base-4B-Unmeasured-4bit", 4e9, 4, downloads=20_000, tags=["base_model:Org/Base-4B"])
    comm = {a.id: CommunitySignal(median_perplexity=20.5, perplexity_count=1),
            b.id: CommunitySignal(median_perplexity=18.5, perplexity_count=1),
            c.id: CommunitySignal(median_perplexity=17.2, perplexity_count=1)}
    q = Query(ram_gb=36, priority="balanced")
    norms = ENGINE.norms([a, b, c, u], q)
    norms["ppl_refs"] = _refs(comm, [a, b, c, u])
    s = {x.model.name: x for x in ENGINE.score([a, b, c, u], q, {}, comm, norms)}
    assert s["Base-4B-Mixed-4bit"].components["measured_quality"] > 0 > s["Base-4B-4bit"].components["measured_quality"]
    # measuring is neutral in expectation: the better 4-bit rises above the unmeasured one, the worse falls below
    assert s["Base-4B-Mixed-4bit"].score > s["Base-4B-Unmeasured-4bit"].score > s["Base-4B-4bit"].score
    # the 8-bit has no measured 8-bit peer, so its score isn't moved by the 4-bit measurements
    assert "measured_quality" not in s["Base-4B-8bit"].components
    assert any("lower than the 1 other measured 4-bit" in r for r in s["Base-4B-Mixed-4bit"].reasons)


def test_single_measurement_does_not_change_score():
    a = rec("Solo-4B-4bit", 4e9, 4, tags=["base_model:Org/Solo"])
    comm = {a.id: CommunitySignal(median_perplexity=12.0, perplexity_count=1)}
    q = Query(ram_gb=36)
    norms = ENGINE.norms([a], q)
    base = ENGINE.score([a], q, {}, {}, dict(norms, ppl_refs={}))[0]
    norms["ppl_refs"] = _refs(comm, [a])
    s = ENGINE.score([a], q, {}, comm, norms)[0]
    assert "measured_quality" not in s.components
    assert any("no other 4-bit quantization" in r for r in s.reasons)
    assert abs(s.components["quant"] - base.components["quant"]) < 1e-9