distinct_agent/models.py

"""Known llama.cpp model manifests and deterministic local discovery."""

from __future__ import annotations

import hashlib
import re
from collections.abc import Iterable, Iterator, Mapping, Sequence
from dataclasses import dataclass
from pathlib import Path

from distinct_protocol import ModelManifest


class ModelVerificationError(ValueError):
    """A local model file could not be verified against its pinned digest."""


class UnpinnedModelError(ModelVerificationError):
    """No weights digest has been recorded for this model.

    This is deliberately a distinct, louder failure from a digest *mismatch*.
    An absent digest is not a passing check; it is an unanswered question, and
    the caller has to decide explicitly whether to proceed without identity.
    """


# A model is advertised only after its file is found locally. Presence in this
# catalogue never means usable.
#
# PUBLISHER POLICY. Weights are taken from the organisation that trained them,
# and from nobody else. The earlier policy defaulted to unsloth's GGUF
# republications on the argument that a smaller file for the same nominal
# quantisation is less to download and less to hold in memory. That argument is
# still true and it is still not the deciding one: a digest attests to whoever
# built the file, so a third-party requantisation pins the repackager rather
# than the trainer, and this catalogue's whole claim is that you can check who
# made what you are running. Ai2 publish their own GGUFs, which is what makes
# the entries below first party end to end -- the same organisation trained the
# model, measured its training, quantised it and published the digest.
#
# DIGESTS. A digest is meaningless without the repository and revision it came
# from: the same nominal quantisation published twice is two different files.
# Entries carrying a digest name all three. Entries without one are still
# pending, and an empty digest is an unanswered question rather than a passing
# check; `discover_models(verify_hashes=True)` refuses them loudly.
#
# EVERY ENTRY HERE IS RUNNABLE.
#
# There was one that was not: Llama 2 7B Chat, kept to be named rather than
# run, because Meta publish no GGUF and the source repository is gated, so
# there was no file to pin and the worker withheld it on every start. It has
# been removed. A catalogue whose entries can all be fetched and verified is
# easier to reason about than one carrying a permanent exception, and nothing
# was lost from the assessment surface: describing what a publisher did or did
# not disclose never required shipping their weights.
#
# WHAT AN UNPINNED ENTRY WOULD MEAN. Weights are fetched on demand rather than
# placed on disk by the operator, so a manifest is only usable when its
# repository, revision *and* digest are all recorded: without a revision there
# is no exact file to fetch, and without a digest there is no way to say what
# was run. :func:`distinct_agent.weights.fetchable_manifests` filters this list
# to the entries that satisfy all three, and a worker offers nothing else.
#
# The unpinned entry stays here rather than being deleted. It is the record of
# what is missing and for which model, it keeps the assessment surface able to
# name a release it has assessed but cannot run, and completing it is then a
# data change (three strings) rather than an archaeology exercise. The worker
# prints every withheld model and its missing part on every start, so an
# absence is never silent.
#
# WHY THESE THREE. Every entry here has a published, model-specific assessment
# of its training, recorded in ``model_assessments.json`` with its primary
# source. Models with no such assessment were removed rather than shown with an
# empty record: this project's only claim is about disclosed and measured
# environmental cost, and a model nobody has published anything about is one it
# has no business shipping. What was removed, and why:
#
#   Qwen3 0.6B and 4B   Qwen has published nothing for any release, in any
#                       version. Previously kept as the worked example of zero
#                       disclosure; that example now lives in the assessment
#                       surface's own copy, which does not require shipping the
#                       weights to make the point.
#   gpt-oss-20b         No disclosure of any kind on the model card or
#                       elsewhere.
#   OLMo 2 0425 1B      Released a month after Ai2's assessment closed. Ai2 have
#                       published nothing for it since, and a figure for a
#                       different OLMo checkpoint is not a figure for this one.
#   OLMoE 0125          Superseded here by OLMoE 0924, which is the checkpoint
#                       Ai2 actually measured. Same architecture, same parameter
#                       count, four months apart -- and by this catalogue's own
#                       rule a different version is a different model, so the
#                       date stamp is the entire difference between an entry
#                       with evidence and an entry without.
KNOWN_MODEL_MANIFESTS: tuple[ModelManifest, ...] = (
    # DEFAULT. Chosen on published, measured environmental impact, not on
    # capability, and not on file size.
    #
    # Ai2 measured this run: GPU power sampled at sub-second intervals rather
    # than inferred from rated wattage, giving 54 MWh, 18 tCO2e and 70 kL of
    # water for the pretraining run. That is the lowest published figure of any
    # assessed release at a size this network can run, and it is roughly a third
    # of the dense OLMo 2 7B measured by the same team on the same cluster.
    # Mixture of experts: 6.9B parameters resident, about 1.3B active per token.
    #
    # See model_assessments.json, record ``olmoe-1b-7b-0924``, and
    # https://arxiv.org/abs/2503.05804 for the measurement.
    #
    # NOTE ON COVERAGE. The published figures cover the 0924 *pretraining* run.
    # This entry is the Instruct checkpoint, whose tuning stage the source does
    # not separately cost. The assessment record says so; do not quietly promote
    # it to complete coverage.
    ModelManifest(
        id="olmoe-1b-7b-0924-instruct",
        label="OLMoE 1B-7B 0924 Instruct · Ai2 · Q4_K_M",
        filename="olmoe-1b-7b-0924-instruct-q4_k_m.gguf",
        min_ram_gb=8.0,
        sha256="8c310f1435a1222338fd2d3d974975be9cd908180b644bab0c2a94da1ac32f3f",
        source_repo="allenai/OLMoE-1B-7B-0924-Instruct-GGUF",
        source_revision="02ab6ea6894a8418eb14f6d8ee1bfb08bd298080",
        context_length=4096,
    ),
    # The dense comparison, and the fallback if the MoE architecture misbehaves
    # under a given llama.cpp build. Same team, same cluster, same measurement
    # method, same year, so the gap between this and OLMoE above is one of the
    # few genuinely like-for-like comparisons in this literature: 157 MWh,
    # 52 tCO2e, 202 kL. Roughly three times OLMoE on every disclosed axis.
    #
    # See model_assessments.json, record ``olmo-2-1124-7b``.
    ModelManifest(
        id="olmo-2-1124-7b-instruct",
        label="OLMo 2 1124 7B Instruct · Ai2 · Q4_K_M",
        filename="olmo-2-1124-7B-instruct-Q4_K_M.gguf",
        min_ram_gb=8.0,
        sha256="e08112e5f84aab7c05fa6e713c58e5214cd5d8e32ed773ff3354b006eed41b95",
        source_repo="allenai/OLMo-2-1124-7B-Instruct-GGUF",
        source_revision="410e0069f64869e4b1d17d8de04810b881fd824b",
        context_length=4096,
    ),
)


_FILENAME_ALIASES: Mapping[str, tuple[str, ...]] = {
    "olmoe-1b-7b-0924-instruct": (
        "olmoe-1b-7b-0924-instruct-q4_k_m.gguf",
        "OLMoE-1B-7B-0924-Instruct-Q4_K_M.gguf",
    ),
    "olmo-2-1124-7b-instruct": (
        "olmo-2-1124-7B-instruct-Q4_K_M.gguf",
        "OLMo-2-1124-7B-Instruct-Q4_K_M.gguf",
    ),
}


VERIFIED = "verified"
UNPINNED = "identity not verified: no digest recorded"


@dataclass(frozen=True)
class DiscoveredModel:
    manifest: ModelManifest
    path: Path
    size_bytes: int
    # ``verified`` is True only when a recorded digest was checked and matched.
    # It is never True merely because no digest existed to check.
    verified: bool = False

    @property
    def verification_status(self) -> str:
        return VERIFIED if self.verified else UNPINNED

    def to_dict(self) -> dict:
        value = self.manifest.to_dict()
        value.update(
            {
                "path": str(self.path),
                "size_bytes": self.size_bytes,
                "verified": self.verified,
                "verification_status": self.verification_status,
            }
        )
        return value


def model_manifest(model_id: str) -> ModelManifest:
    for manifest in KNOWN_MODEL_MANIFESTS:
        if manifest.id == model_id:
            return manifest
    raise KeyError(model_id)


def discover_models(
    search_paths: Iterable[Path | str],
    *,
    manifests: Sequence[ModelManifest] = KNOWN_MODEL_MANIFESTS,
    recursive: bool = True,
    max_files: int = 10_000,
    verify_hashes: bool = True,
    require_pinned_digest: bool = False,
) -> tuple[DiscoveredModel, ...]:
    """Find recognized, non-empty GGUF files beneath explicit roots.

    Discovery never downloads files and never imports model-provided Python.

    Verification is fail-closed, and the two negative cases are kept distinct
    because they mean different things:

    * A recorded digest that does **not** match the file on disk raises
      ``ModelVerificationError``.  A file presenting itself as a pinned model
      while being a different file is not something to skip quietly.
    * A manifest with **no** recorded digest yields a model whose ``verified``
      flag is ``False`` and whose ``verification_status`` says so.  Callers
      must surface that; it is not a passing check.  Pass
      ``require_pinned_digest=True`` to turn it into an ``UnpinnedModelError``.

    ``verify_hashes=False`` skips digest comparison entirely and is only for
    tests that deliberately construct mismatching fixtures; it still reports
    ``verified=False``, never ``True``.
    """

    if isinstance(max_files, bool) or max_files < 1:
        raise ValueError("max_files must be a positive integer")
    catalog = {manifest.id: manifest for manifest in manifests}
    aliases: dict[str, str] = {}
    for manifest in manifests:
        names = _FILENAME_ALIASES.get(manifest.id, (manifest.filename,))
        for name in names + (manifest.filename,):
            aliases[_normal_name(name)] = manifest.id

    candidates: list[Path] = []
    for raw_root in search_paths:
        root = Path(raw_root).expanduser()
        if root.is_file():
            candidates.append(root)
            continue
        if not root.is_dir():
            continue
        iterator: Iterator[Path]
        iterator = root.rglob("*.gguf") if recursive else root.glob("*.gguf")
        for path in iterator:
            candidates.append(path)
            if len(candidates) > max_files:
                raise ValueError(f"model discovery exceeded {max_files} files")

    discovered: dict[str, DiscoveredModel] = {}
    for path in sorted(candidates, key=lambda item: str(item).casefold()):
        model_id = aliases.get(_normal_name(path.name))
        if model_id is None or model_id in discovered:
            continue
        try:
            resolved = path.resolve(strict=True)
            stat = resolved.stat()
        except OSError:
            continue
        if not resolved.is_file() or stat.st_size <= 0:
            continue
        manifest = catalog[model_id]
        verified = False
        if manifest.digest_recorded:
            if require_pinned_digest or verify_hashes:
                actual = sha256_file(resolved)
                if actual != manifest.sha256:
                    raise ModelVerificationError(
                        f"{resolved} does not match the pinned SHA-256 for {model_id!r}: "
                        f"expected {manifest.sha256}, computed {actual}"
                    )
                verified = True
        elif require_pinned_digest:
            raise UnpinnedModelError(
                f"no SHA-256 digest is recorded for {model_id!r}, so {resolved} cannot be "
                "verified; record the upstream repository, revision and digest first"
            )
        discovered[model_id] = DiscoveredModel(manifest, resolved, stat.st_size, verified)

    return tuple(
        discovered[manifest.id]
        for manifest in manifests
        if manifest.id in discovered
    )


def index_discovered(models: Iterable[DiscoveredModel]) -> dict[str, DiscoveredModel]:
    result: dict[str, DiscoveredModel] = {}
    for model in models:
        if model.manifest.id in result:
            raise ValueError(f"duplicate discovered model {model.manifest.id!r}")
        result[model.manifest.id] = model
    return result


def sha256_file(path: Path | str, *, chunk_size: int = 1024 * 1024) -> str:
    if chunk_size < 1:
        raise ValueError("chunk_size must be positive")
    digest = hashlib.sha256()
    with Path(path).open("rb") as handle:
        while True:
            chunk = handle.read(chunk_size)
            if not chunk:
                break
            digest.update(chunk)
    return digest.hexdigest()


def _normal_name(value: str) -> str:
    return re.sub(r"[^a-z0-9]+", "", value.casefold())