spb/ai-atlas

1"""Unit tests for the pure API 1.1 services: Pareto frontier, cost arithmetic, search compiler v2, hardware fit breakdown. No database."""2from __future__ import annotations34from datetime import UTC, datetime56import pytest78from aiatlas.services import hardware_fit as hf9from aiatlas.services.cost import compute_cost, context_fill_cost10from aiatlas.services.pareto import is_dominated, pareto_frontier11from aiatlas.services.search import compile_query, params_bound_from_memory1213# ------------------------------------------------------------------------------------------------------------------ pareto141516def test_pareto_frontier_is_efficient_and_keeps_ties() -> None:17    pts = [18        {"id": "a", "x": 1.0, "y": 50},   # cheap, low quality → efficient19        {"id": "b", "x": 2.0, "y": 70},   # efficient20        {"id": "c", "x": 3.0, "y": 65},   # dominated by b (more expensive, lower)21        {"id": "d", "x": 5.0, "y": 90},   # efficient22        {"id": "e", "x": 5.0, "y": 90},   # exact tie with d → kept23        {"id": "f", "x": 6.0, "y": 90},   # dominated by d (same y, more expensive)24        {"id": "g", "x": None, "y": 99},  # ignored25    ]26    front = pareto_frontier(pts)27    assert front == ["a", "b", "d", "e"]28    valid = [p for p in pts if p["x"] is not None]29    for p in valid:30        assert is_dominated(p, valid) == (p["id"] not in front), p["id"]313233def test_pareto_minimise_y() -> None:34    pts = [{"id": "a", "x": 1, "y": 10}, {"id": "b", "x": 2, "y": 5}, {"id": "c", "x": 3, "y": 7}]35    assert pareto_frontier(pts, maximize_y=False) == ["a", "b"]363738# ------------------------------------------------------------------------------------------------------------------ cost394041def test_cost_arithmetic_standard() -> None:42    c = compute_cost({"input": 3.0, "output": 15.0}, input_tokens=1000, output_tokens=500, requests_per_day=1000)43    assert c["per_request"] == pytest.approx(1000 * 3 / 1e6 + 500 * 15 / 1e6)44    assert c["daily"] == pytest.approx(c["per_request"] * 1000)45    assert c["monthly"] == pytest.approx(c["daily"] * 30) and c["annual"] == pytest.approx(c["daily"] * 365)46    assert c["notes"] == []474849def test_cost_cached_and_batch_with_notes() -> None:50    c = compute_cost({"input": 3.0, "output": 15.0, "cached_input": 0.3}, input_tokens=1_000_000, output_tokens=0, cached_share=0.5)51    assert c["effective_input_per_mtok"] == pytest.approx(1.65) and c["per_request"] == pytest.approx(1.65)52    c2 = compute_cost({"input": 3.0, "output": 15.0}, input_tokens=1_000_000, output_tokens=0, cached_share=0.5)53    assert c2["per_request"] == pytest.approx(3.0) and any("cached input price not published" in n for n in c2["notes"])54    c3 = compute_cost({"input": 3.0, "output": 15.0, "batch_input": 1.5, "batch_output": 7.5}, input_tokens=1_000_000, output_tokens=1_000_000, batch=True)55    assert c3["per_request"] == pytest.approx(9.0)56    c4 = compute_cost({"input": 3.0, "output": 15.0}, input_tokens=1_000_000, output_tokens=1_000_000, batch=True)57    assert c4["per_request"] == pytest.approx(18.0) and any("batch prices not published" in n for n in c4["notes"])58    missing = compute_cost({"input": None, "output": 15.0}, input_tokens=10, output_tokens=10)59    assert missing["per_request"] is None and "input price unavailable" in missing["notes"]60    assert context_fill_cost(2.0, 1_000_000) == 2.0 and context_fill_cost(None, 10) is None616263# ------------------------------------------------------------------------------------------------------------------ search compiler v26465YEAR = datetime.now(UTC).year666768def _d(q: str) -> dict:69    d = compile_query(q).as_dict()70    d.pop("compiled"), d.pop("text"), d.pop("filters")71    return d727374def test_compile_open_models_over_100b_this_year() -> None:75    d = _d("open models over 100B released this year")76    assert d["entity_type"] == "model" and d["openness"] == "open" and d["params_min"] == 100_000_000_00077    assert d["year_from"] == YEAR and d["year_to"] == YEAR and d["residual"] == ""787980def test_compile_cheapest_1m_context() -> None:81    d = _d("cheapest models with 1M context")82    assert d["entity_type"] == "model" and d["context_min"] == 1_000_000 and d["sort"] == "cheapest" and d["residual"] == ""838485def test_compile_reasoning_under_price() -> None:86    d = _d("reasoning models under $1/M tokens")87    assert d["reasoning"] is True and d["max_output_price"] == 1.0 and d["entity_type"] == "model"888990def test_compile_fits_in_memory() -> None:91    d = _d("models that fit in 64GB")92    assert d["memory_gb"] == 64.0 and d["entity_type"] == "model"93    assert 0 < params_bound_from_memory(64.0) < 200e9949596def test_compile_org_since_year() -> None:97    d = _d("Anthropic models released since 2025")98    assert d["organization"] == "Anthropic" and d["year_from"] == 2025 and "year_to" not in d and d["entity_type"] == "model"99100101def test_compile_open_vision_apache() -> None:102    d = _d("open vision models with Apache license")103    assert d["openness"] == "open" and d["modalities"] == ["image"] and d["license_key"] == "Apache-2.0"104105106def test_compile_papers_free_text() -> None:107    d = _d("papers introducing MoE models")108    assert d["entity_type"] == "paper" and "organization" not in d and "MoE" in d["residual"]109110111def test_compile_misc_rules() -> None:112    assert _d("models between 7B and 70B")["params_min"] == 7_000_000_000 and _d("models between 7B and 70B")["params_max"] == 70_000_000_000113    assert _d("benchmark gpqa")["benchmark"] == "gpqa" and _d("provider groq")["provider"] == "groq"114    d = _d("open source models with commercial use released in the last 30 days")115    assert d["commercial_use"] is True and d["days_back"] == 30 and d["residual"] == ""116    assert _d("largest proprietary models")["sort"] == "largest" and _d("largest proprietary models")["openness"] == "proprietary"117    plain = compile_query("claude")118    assert plain.residual == "claude" and plain.unrecognised == [] and not plain.has_structure119    comp = compile_query("open models over 100B").compiled120    assert {c["filter"] for c in comp} >= {"openness", "params_min", "entity_type"} and all("label" in c and "source_span" in c for c in comp)121122123# ------------------------------------------------------------------------------------------------------------------ hardware fit124125126def test_fit_detailed_observed_vs_estimated_and_kv() -> None:127    est = hf.fit_detailed({"parameter_count": 70e9}, 64, quant="4bit", context=8192)128    assert est and est["estimated"] is True and est["breakdown"]["weights_source"] == "estimated" and est["breakdown"]["kv_cache_method"] == "heuristic"129    obs = hf.fit_detailed({"parameter_count": 70e9}, 64, quant="4bit", context=8192, observed_size_gb=40.0)130    assert obs and obs["breakdown"]["weights_gb"] == 40.0 and obs["breakdown"]["weights_source"] == "observed"131    arch = hf.fit_detailed({"parameter_count": 8e9, "num_hidden_layers": 32, "num_key_value_heads": 8, "head_dim": 128}, 24, quant="4bit", context=8192)132    assert arch and arch["breakdown"]["kv_cache_method"] == "architecture" and arch["breakdown"]["kv_cache_gb"] == pytest.approx(2 * 32 * 8 * 128 * 2 * 8192 / 1e9, rel=1e-3)133    multi = hf.fit_detailed({"parameter_count": 400e9}, 80, quant="8bit", gpu_count=8)134    assert multi and multi["device"]["total_memory_gb"] == 640 and "multi_gpu_note" in multi135    assert hf.fit_detailed({}, 64) is None  # nothing estimated from thin air136    assert hf.fit(70e9, 64, "4bit", 8192)["estimated"] is True  # v1 helper unchanged137