"""
catalog — load and merge the bundled model catalog with the user's cache.
The catalog lives in two layers:
1. The **bundled seed** (`models.yaml` in the package root). Hand-maintained,
always present, never deleted by auto-refresh.
2. The **user cache** (`~/.best-engine-ai-helper/catalog_cache.yaml`). Written
by `catalog update`; entries keyed by `id` overwrite matching seed entries.
Absent on first run; load_catalog silently skips it in that case.
Keeping the seed immutable and the cache additive means offline machines always
have a usable catalog and updates never lose hand-curated data.
Author
------
Warith Harchaoui <warithmetics@deraison.ai>
"""
from __future__ import annotations
import os
import tempfile
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
import os_helper as osh
import yaml
# Root of the installed package; models.yaml sits next to pyproject.toml
# Data lives INSIDE the package. These files used to sit beside it, next to
# pyproject.toml, where `package-data` patterns — which resolve relative to a
# package directory — could never reach them: the wheel shipped none of them,
# and `load_catalog()` returned an EMPTY catalogue on any machine without a
# populated user cache. A repo checkout and a populated cache both hid it.
_PACKAGE_ROOT = Path(__file__).resolve().parent / "data"
# Default path for the bundled seed catalog
_SEED_PATH = _PACKAGE_ROOT / "models.yaml"
# User-writable runtime directory; created on first run
_USER_DIR = Path.home() / ".best-engine-ai-helper"
_CACHE_PATH = _USER_DIR / "catalog_cache.yaml"
# Public alias so callers (the CLI, tests) can name the cache without reaching
# into a private attribute.
CACHE_PATH = _CACHE_PATH
# Quant-specific overhead factors for estimating peak inference RAM.
# Values are empirically derived: KV cache at 4K context adds ~10-20%
# beyond on-disk size depending on quantization precision.
_QUANT_OVERHEAD: dict[str, float] = {
"Q4_K_M": 1.12,
"Q8_0": 1.10,
"FP16": 1.05,
"Q2_K": 1.20,
}
_DEFAULT_OVERHEAD = 1.15 # fallback for unknown quant identifiers
#: Sources whose ``id`` is a directory slug rather than a registry pull tag.
#: The catalog's own contract says ``id`` IS the Ollama pull tag, and the
#: hand-maintained seed honours it; the ApXML adapter cannot, because the
#: directory it scrapes names models after their paper (``qwen36-35b-a3b``).
#: Such an entry stays useful — it carries a ``vllm_id`` — but offering it to
#: Ollama yields a descriptor naming a model that cannot be pulled, and the
#: caller finds out at its first inference call rather than at resolution.
_SLUG_SOURCES: frozenset[str] = frozenset({"apxml"})
#: Backends this gate knows how to rule on. vLLM is deliberately absent: an
#: entry without a ``vllm_id`` already falls back to the raw tag with a
#: warning, which is a documented behaviour and a reasonable one. The gate
#: exists for the cases that are *not* recoverable — a name Ollama cannot pull,
#: and (for ``lightweight``) an entry naming no HuggingFace repository at all, since
#: layer streaming reads safetensors from the Hub and has no registry to fall
#: back to.
_GATED_BACKENDS: frozenset[str] = frozenset({"ollama", "lightweight"})
[docs]
def servable(entry: dict[str, Any], backend: str) -> bool:
"""
Whether a backend can actually serve this catalog entry.
Parameters
----------
entry : dict[str, Any]
A catalog entry.
backend : str
``ollama``, ``vllm``, ``lightweight``, or anything else (unknown backends are
trusted: this is a gate against a known mismatch, not a whitelist).
Returns
-------
bool
True when the entry names something the backend can load.
An entry may say so itself with ``backends: [...]``; otherwise the rule is
the source. A bare name is **not** treated as suspect — ``gemma3`` is a
perfectly legal pull tag that Ollama resolves to ``gemma3:latest`` — so the
discriminator is where the entry came from, not how it is spelled.
Examples
--------
>>> servable({"id": "qwen3-vl:8b"}, "ollama")
True
>>> servable({"id": "gemma3"}, "ollama") # bare tags are legal
True
>>> servable({"id": "qwen36-35b-a3b", "source": "apxml"}, "ollama")
False
>>> servable({"id": "x", "source": "apxml"}, "vllm") # vLLM is not gated
True
>>> servable({"id": "x", "backends": ["vllm"]}, "ollama")
False
>>> servable({"id": "whatever"}, "some-future-backend")
True
>>> servable({"id": "llama3.3:70b", "vllm_id": "meta-llama/x"}, "lightweight")
True
>>> servable({"id": "gemma3:12b"}, "lightweight") # no Hub repo to stream
False
"""
declared = entry.get("backends")
if isinstance(declared, list) and declared:
return backend in declared
if backend not in _GATED_BACKENDS:
return True
if not str(entry.get("id") or "").strip():
return False
if backend == "lightweight":
# A pull tag is useless to a streaming backend: lightweight downloads
# safetensors straight out of the Hub repository, so an
# entry that names no repository names nothing it can load. The slug
# rule below is the opposite case and does not apply — a slug source
# like ApXML carries exactly the `vllm_id` this needs.
from . import offload as _offload
return _offload.hf_id(entry) is not None
return str(entry.get("source") or "").strip() not in _SLUG_SOURCES
[docs]
def servable_only(rows: list[dict[str, Any]], backend: str) -> list[dict[str, Any]]:
"""
Keep the entries a backend can serve, and say what was dropped.
Parameters
----------
rows : list[dict[str, Any]]
Catalog entries.
backend : str
The backend about to be resolved for.
Returns
-------
list[dict[str, Any]]
The servable subset, in the original order.
Examples
--------
>>> rows = [{"id": "gemma3:12b"}, {"id": "llama-3-70b", "source": "apxml"}]
>>> [r["id"] for r in servable_only(rows, "ollama")]
['gemma3:12b']
>>> len(servable_only(rows, "vllm")) # vLLM keeps both
2
"""
kept = [row for row in rows if servable(row, backend)]
dropped = len(rows) - len(kept)
if dropped:
osh.debug(
f"{dropped} catalog entr{'y' if dropped == 1 else 'ies'} cannot be served "
f"by {backend} and were left out of the recommendation."
)
return kept
[docs]
def estimate_ram(disk_gb: float, quant: str) -> float:
"""
Estimate peak inference RAM from on-disk model size.
The estimate covers the model weights plus KV cache at the default context
length (4K tokens). For models with very large context windows (256K+),
actual RAM may exceed this estimate significantly; treat it as a lower bound.
Parameters
----------
disk_gb : float
On-disk footprint of the model in GB.
quant : str
Quantization identifier, e.g. 'Q4_K_M', 'Q8_0', 'FP16', 'Q2_K'.
Returns
-------
float
Estimated peak RAM in GB.
Examples
--------
>>> estimate_ram(6.1, 'Q4_K_M')
6.832
>>> estimate_ram(10.0, 'FP16')
10.5
"""
overhead = _QUANT_OVERHEAD.get(quant, _DEFAULT_OVERHEAD)
return round(disk_gb * overhead, 3)
def _load_yaml_file(path: Path) -> list[dict[str, Any]]:
"""
Load a YAML file and return its contents as a list of dicts.
Returns an empty list on FileNotFoundError so callers can treat the
absence of the cache file as a no-op rather than an error.
Parameters
----------
path : Path
Absolute path to the YAML file.
Returns
-------
list[dict[str, Any]]
Parsed entries, or [] if the file is absent or empty.
"""
if not osh.file_exists(str(path)):
osh.debug(f"YAML file absent, treating as empty:\n\t{path}")
return []
try:
raw = yaml.safe_load(path.read_text(encoding="utf-8"))
except yaml.YAMLError as exc:
osh.warning(f"Malformed YAML, ignoring:\n\t{path}\n\t{exc}")
return []
# yaml.safe_load returns None for empty files
if raw is None:
osh.debug(f"YAML file empty, treating as empty:\n\t{path}")
return []
if not isinstance(raw, list):
# A syntactically valid YAML document whose top level isn't a list (a
# bare scalar, or a mapping) is not a usable catalog: `list(a_dict)`
# would silently return only the dict's KEYS as "entries", and
# `list(a_string)` would silently return one-character "entries" --
# either way, garbage that breaks every downstream `entry.get(...)`
# call instead of a clear, loggable error at the source.
osh.warning(f"YAML file does not contain a list at the top level, ignoring:\n\t{path}")
return []
return list(raw)
def _overlay(base: dict[str, Any] | None, incoming: dict[str, Any]) -> dict[str, Any]:
"""
Lay a refreshed entry over an existing one without erasing what it knows.
A refresh from a specs directory is a claim about *specs*, not about
*scores*. Sources like the Ollama library publish sizes and capabilities
and no benchmarks at all, so their entries carry null scores — and a plain
dict replacement would hand those nulls to a seed entry that had real,
hand-curated numbers. Widening the catalog would then quietly degrade it,
which is the opposite of the point.
The bug stayed hidden while the only refresh source was ApXML, whose ids
are directory slugs that never collide with the seed's pull tags. A source
that emits real tags collides constantly, and surfaced it.
So: the incoming entry wins field by field, except that a null never
replaces a value that is already there.
Parameters
----------
base : dict or None
The entry already held, usually from the seed. None when the incoming
entry is new, in which case it is returned unchanged.
incoming : dict
The refreshed entry.
Returns
-------
dict
The merged entry. Neither input is mutated.
Examples
--------
>>> seed = {"id": "m", "benchmarks": {"general": 77}, "notes": "curated"}
>>> fresh = {"id": "m", "benchmarks": {"general": None}, "disk_gb": 8.0}
>>> merged = _overlay(seed, fresh)
>>> merged["benchmarks"]["general"] # the real score survives
77
>>> merged["disk_gb"] # the fresh spec lands
8.0
>>> merged["notes"] # untouched fields stay
'curated'
"""
if not base:
return incoming
result = dict(base)
for key, value in incoming.items():
if key == "benchmarks":
continue
if value is None and result.get(key) is not None:
continue
result[key] = value
# Benchmarks merge per axis rather than wholesale: a source may know one
# axis and nothing about the others, and replacing the block would discard
# every axis it was silent on.
scores = dict(base.get("benchmarks") or {})
for axis, value in (incoming.get("benchmarks") or {}).items():
if value is not None:
scores[axis] = value
if scores or "benchmarks" in incoming:
result["benchmarks"] = scores
return result
[docs]
def load_catalog(catalog_path: Path | None = None) -> list[dict[str, Any]]:
"""
Load the bundled seed catalog merged with the user's local cache.
Cache entries whose `id` matches a seed entry overwrite the seed entry.
New cache entries (no matching seed id) are appended. The seed is never
modified on disk.
Parameters
----------
catalog_path : Path or None
Path to the seed `models.yaml`. Defaults to the bundled file next to
`pyproject.toml`. Pass an explicit path in tests to use a fixture.
Returns
-------
list[dict[str, Any]]
Merged model entries. Each entry is guaranteed to have at minimum:
``id``, ``kind``, ``ram_gb``, ``benchmarks``.
Raises
------
FileNotFoundError
If ``catalog_path`` is given explicitly and does not exist.
Examples
--------
>>> entries = load_catalog()
>>> len(entries) > 0
True
>>> all('id' in e for e in entries)
True
"""
seed_path = catalog_path if catalog_path is not None else _SEED_PATH
# Start with the bundled seed; this must exist
seed = _load_yaml_file(seed_path)
if not seed and catalog_path is not None:
osh.error(f"Catalog not found:\n\t{seed_path}")
raise FileNotFoundError(f"Catalog not found: {seed_path}")
osh.info(f"Loaded {len(seed)} model(s) from seed catalog:\n\t{seed_path}")
# Build a lookup by id so cache entries can overwrite in O(1)
merged: dict[str, dict[str, Any]] = {e["id"]: e for e in seed}
# Overlay user cache entries; absent cache is silently ignored
cache_entries = _load_yaml_file(_CACHE_PATH)
if cache_entries:
osh.info(f"Overlaying {len(cache_entries)} cached entry(ies):\n\t{_CACHE_PATH}")
for entry in cache_entries:
merged[entry["id"]] = _overlay(merged.get(entry["id"]), entry)
# Preserve seed ordering; append cache-only entries at the end
result: list[dict[str, Any]] = []
seen: set[str] = set()
for e in seed:
result.append(merged[e["id"]])
seen.add(e["id"])
for eid, entry in merged.items():
if eid not in seen:
result.append(entry)
return result
# ---------------------------------------------------------------------------
# Cache refresh (`catalog update`)
# ---------------------------------------------------------------------------
# The default quant we pull, and the quant whose VRAM ApXML's headline figure
# mirrors, so disk_gb is estimated back out of ram_gb through its overhead.
_REFRESH_QUANT = "Q4_K_M"
# Empty benchmark block: ApXML's static pages carry specs and memory-fit
# figures but no numeric leaderboard scores, so refreshed entries rank at the
# bottom until a scored source fills these axes.
_EMPTY_BENCHMARKS: dict[str, float | None] = {
"general": None,
"vision": None,
"ocr": None,
"code": None,
"math": None,
}
def _atomic_write_text(path: Path, content: str) -> None:
"""Write ``content`` to ``path`` atomically: no reader ever sees a partial write.
Writes to a sibling temp file in the same directory (so the final
``os.replace`` is a same-filesystem rename, not a cross-filesystem copy)
then renames it over the destination. A crash or a concurrent
:func:`load_catalog` mid-write sees either the old cache or the fully
-written new one, never a truncated YAML document that would otherwise
be silently discarded as malformed (see :func:`_load_yaml_file`).
"""
fd, tmp_name = tempfile.mkstemp(dir=str(path.parent), prefix=f".{path.name}.")
try:
with os.fdopen(fd, "w", encoding="utf-8") as f:
f.write(content)
os.replace(tmp_name, path)
except BaseException:
Path(tmp_name).unlink(missing_ok=True)
raise
def _today() -> str:
"""Return today's date as an ISO 8601 string (UTC), for the ``fetched_at`` stamp."""
return datetime.now(timezone.utc).date().isoformat()
[docs]
def normalize_apxml_spec(spec: dict[str, Any], fetched_at: str) -> dict[str, Any] | None:
"""
Map one ApXML spec dict onto a catalog entry.
The ApXML adapter (:mod:`best_engine_ai_helper.sources.apxml`) yields spec
and memory-fit metadata but no numeric benchmarks, so the ``benchmarks``
block is left null. ``disk_gb`` is estimated from the Q4 VRAM figure by
dividing out the quant overhead — the inverse of :func:`estimate_ram`.
Parameters
----------
spec : dict[str, Any]
A normalized spec as returned by ``apxml.parse_model_page``.
fetched_at : str
ISO 8601 date recorded on the entry as its refresh timestamp.
Returns
-------
dict[str, Any] or None
A catalog entry, or None when the spec lacks the ``slug`` needed to key
it (such an entry could never be merged or pulled).
"""
slug = spec.get("slug")
if not slug:
osh.debug(f"ApXML spec without a slug, skipping:\n\t{spec.get('name')!r}")
return None
ram_gb = spec.get("ram_gb")
if ram_gb is None and spec.get("size_b") is None:
# ApXML does not carry a Q4 VRAM figure for every model, and a record
# with neither that nor a parameter count cannot be sized against this
# machine's memory -- which is the catalog's entire job. Skipped rather
# than stored with a null, exactly as `ollama_library.to_specs` skips a
# family that lists no sizes.
osh.debug(f"ApXML spec with no size information, skipping:\n\t{slug!r}")
return None
# Invert estimate_ram: disk ≈ ram / overhead. Estimate only, like the seed.
disk_gb = (
round(float(ram_gb) / _QUANT_OVERHEAD[_REFRESH_QUANT], 2) if ram_gb is not None else None
)
return {
"id": slug,
"kind": spec.get("kind", "llm"),
"size_b": spec.get("size_b"),
"quant": _REFRESH_QUANT,
"disk_gb": disk_gb,
"ram_gb": ram_gb,
"benchmarks": dict(_EMPTY_BENCHMARKS),
# Say it rather than leave it to be inferred. ApXML keys its directory
# by the paper's slug (`qwen36-35b-a3b`), which is a HuggingFace repo
# name and not an `ollama pull` tag, so these entries are servable by
# whatever loads from the Hub and by nothing else. `servable()` can
# deduce as much from `source`, and still does for entries cached by
# older versions, but a filter that infers a fact the writer already
# knew is a net under a gap rather than a fix for it.
#
# Two backends load from the Hub, not one: vLLM and layer streaming.
# This directory is in fact the best source of streaming candidates
# there is — it catalogues the 200B-and-up open-weight models that no
# machine here can load and that streaming is the only way to run — so
# omitting `lightweight` here hid exactly the entries the backend exists for.
"backends": ["vllm", "lightweight"],
# Mixture-of-experts, when the directory says so. Load-bearing for the
# streaming footprint: without it, a MoE model's decoder depth is
# interpolated from its parameter count, and that relation does not
# hold for MoE (DeepSeek-V3 is 671B across 61 layers where the dense
# trend predicts twice as many). Flagged here, `offload.layer_count`
# refuses to guess instead of guessing wrong by a factor of two.
**(
{"moe": True}
if (spec.get("num_of_experts") or 0) and float(spec["num_of_experts"]) > 1
else {}
),
"vllm_id": spec.get("vllm_id"),
"context_length": spec.get("context_length"),
"license": spec.get("license"),
"url": spec.get("url"),
"source": "apxml",
"fetched_at": fetched_at,
}
[docs]
def normalize_ollama_spec(spec: dict[str, Any], fetched_at: str) -> dict[str, Any] | None:
"""
Map one Ollama-library spec onto a catalog entry.
Unlike the ApXML path, ``id`` here is already a real ``ollama pull`` tag,
so the entry carries no ``backends`` restriction and :func:`servable` lets
it through for Ollama. That is the whole point of this source: it widens
the pool of models that can actually be served, where ApXML widens a pool
that only vLLM can load.
Benchmarks are left null. The library publishes descriptions and sizes, not
leaderboard results, and an invented score would be worse than an absent
one — a null ranks the entry low, which is honest, whereas a guess would
rank it wrongly and look authoritative.
Parameters
----------
spec : dict[str, Any]
A spec from ``sources.ollama_library.to_specs``.
fetched_at : str
ISO 8601 date recorded as the entry's refresh timestamp.
Returns
-------
dict[str, Any] or None
A catalog entry, or None when the spec has no id or no usable size —
an entry whose footprint is unknown cannot be sized against memory,
and one that silently became zero would "fit" every machine.
"""
ident = str(spec.get("id") or "").strip()
if not ident:
osh.debug(f"Ollama-library spec without an id, skipping:\n\t{spec!r}")
return None
disk_gb = spec.get("disk_gb")
if not isinstance(disk_gb, (int, float)) or disk_gb <= 0:
osh.debug(f"Ollama-library spec without a usable size, skipping:\n\t{ident}")
return None
return {
"id": ident,
"kind": spec.get("kind", "llm"),
"size_b": spec.get("size_b"),
"quant": _REFRESH_QUANT,
"disk_gb": round(float(disk_gb), 2),
"ram_gb": estimate_ram(float(disk_gb), _REFRESH_QUANT),
"benchmarks": dict(_EMPTY_BENCHMARKS),
"url": spec.get("url"),
"source": "ollama-library",
"fetched_at": fetched_at,
}
[docs]
def normalize_ollama_specs(
specs: list[dict[str, Any]], fetched_at: str | None = None
) -> list[dict[str, Any]]:
"""
Normalize a batch of Ollama-library specs, dropping unusable ones.
Parameters
----------
specs : list[dict[str, Any]]
Specs as returned by ``ollama_library.fetch_pullable_models``.
fetched_at : str or None
Refresh timestamp for every entry; defaults to today (UTC).
Returns
-------
list[dict[str, Any]]
Catalog entries, in input order, minus specs that could not be sized.
"""
stamp = fetched_at or _today()
entries = [normalize_ollama_spec(spec, stamp) for spec in specs]
return [e for e in entries if e is not None]
[docs]
def normalize_apxml_specs(
specs: list[dict[str, Any]], fetched_at: str | None = None
) -> list[dict[str, Any]]:
"""
Normalize a batch of ApXML specs into catalog entries, dropping unusable ones.
Parameters
----------
specs : list[dict[str, Any]]
Specs as returned by ``apxml.fetch_open_weight_models``.
fetched_at : str or None
Refresh timestamp for every entry; defaults to today (UTC).
Returns
-------
list[dict[str, Any]]
Catalog entries, in input order, minus specs with no slug.
"""
stamp = fetched_at or _today()
entries = [normalize_apxml_spec(spec, stamp) for spec in specs]
return [e for e in entries if e is not None]
[docs]
def fetch_arena_safely(timeout: float = 30.0) -> tuple[dict[str, Any] | None, str]:
"""
Fetch Arena ratings without ever raising, and say what happened.
The "a dead mirror must not sink a good refresh" contract lives here, in
one place, because two callers need it: :func:`merge_arena_axes`, which
scores entries, and the cloud refresh, which also *adopts* models from the
same snapshot. Threading one fetch through both beats fetching twice, and
beats a second copy of the contract more.
Parameters
----------
timeout : float
Per-request network timeout, in seconds.
Returns
-------
(result, message)
``result`` is a :func:`sources.arena.fetch_arena_ratings` payload, or
None when no mirror answered. ``message`` is always worth showing.
"""
from .sources import arena
try:
result = arena.fetch_arena_ratings(timeout=timeout)
except Exception as exc: # a dead mirror must not sink a good refresh
return None, f"Arena ratings unavailable: {exc}"
return result, str(result.get("note") or "no reason given")
[docs]
def merge_arena_axes(
entries: list[dict[str, Any]],
timeout: float = 30.0,
ratings: dict[str, Any] | None = None,
) -> tuple[list[dict[str, Any]], str, bool]:
"""
Add LM Arena conversational axes to ``entries``, and report what happened.
Lives here rather than in the CLI because the command line and the HTTP
ends mirror each other 1:1; a copy in each would be two places to drift.
The callers do the printing, this does the deciding.
Never raises on a dead mirror. The refresh that precedes this call has
already succeeded, and throwing away a good catalog because a third-party
scrape rotted would punish the user for someone else's outage. What it
refuses to do is return quietly: ``merged`` is False and ``message`` says
why whenever no rating was applied, because an axis the user believes is
populated is worse than one they know is empty.
Coverage skews to hosted models. Arena rates what its voters can chat
with, which is frontier endpoints rather than small quantised local
builds, so this axis serves the cloud catalog far more than the local one.
Parameters
----------
entries : list[dict[str, Any]]
Catalog entries to enrich. Not mutated; ``arena.merge_into`` copies.
timeout : float
Per-request network timeout, in seconds. Ignored when *ratings* is
supplied.
ratings : dict or None
A snapshot already fetched by :func:`fetch_arena_safely`, so a caller
that needs the same ratings twice fetches once. None fetches here.
Returns
-------
(entries, message, merged)
``entries`` enriched when ``merged`` is True, returned unchanged
otherwise. ``message`` is always worth showing the user.
"""
from .sources import arena
result = ratings
if result is None:
result, note = fetch_arena_safely(timeout)
if result is None:
return list(entries), note, False
merged, report = arena.merge_into(entries, result)
note = str(report.get("note") or "no reason given")
matched = report.get("matched") or []
if not matched:
return list(entries), f"No Arena rating merged: {note}", False
unmatched = report.get("unmatched") or []
total = len(matched) + len(unmatched)
return merged, f"Arena: matched {len(matched)}/{total} model(s) — {note}", True
[docs]
def write_cache(entries: list[dict[str, Any]], cache_path: Path | None = None) -> Path:
"""
Merge ``entries`` into the user catalog cache by ``id`` and write it to disk.
Existing cache entries are preserved; an incoming entry whose ``id`` matches
one already cached overwrites it, so a refresh is idempotent and never loses
previously cached models. The bundled seed is untouched.
Parameters
----------
entries : list[dict[str, Any]]
Catalog entries to add or update, e.g. from :func:`normalize_apxml_specs`.
cache_path : Path or None
Destination cache file. Defaults to :data:`CACHE_PATH`; override in tests.
Returns
-------
Path
The path written.
"""
path = cache_path if cache_path is not None else _CACHE_PATH
merged: dict[str, dict[str, Any]] = {e["id"]: e for e in _load_yaml_file(path)}
for entry in entries:
merged[entry["id"]] = entry
osh.make_directory(str(path.parent))
_atomic_write_text(
path, yaml.safe_dump(list(merged.values()), sort_keys=False, allow_unicode=True)
)
osh.info(f"Wrote {len(merged)} cached model(s):\n\t{path}")
return path