steeldb/

models.rs

//! **Where the weights come from.**
//!
//! A published crate cannot carry its models. crates.io caps a package at 10 MB and the full model set is
//! ~388 MB, so the crate ships what genuinely fits and resolves the rest at runtime.
//!
//! The split is not arbitrary, and one half of it is a licensing constraint rather than a size one:
//!
//! | what | size | how it arrives |
//! |---|---|---|
//! | trained heads and the reasoning core | ~1.5 MB | **compiled into the crate** |
//! | relation head | ~2 MB | **compiled into the crate** |
//! | the `bert-tiny` embedding table | 15.2 MB | downloaded — too large to package |
//! | our tagger / SPLADE / LoRA weights | 100s of MB | downloaded from our own model repo |
//! | third-party models (needle3, model2vec, PP-OCR) | 100s of MB | fetched from **their** repos, never re-hosted |
//!
//! That last row is the important one. We trained the tagger, the facet heads and the LoRA, so we may
//! redistribute them. needle3 belongs to Cactus, `bert-tiny` to Google, and the others to their authors —
//! re-hosting someone else's weights under our name would be wrong regardless of whether it is convenient, so
//! those are fetched from the repository that published them and the licence is recorded here.
//!
//! Nothing downloads without being asked: [`resolve`] searches locally first and returns
//! [`ModelError::NotFound`] with the exact fetch instruction if it comes up empty.

use std::path::PathBuf;

/// The trained parameters small enough to travel with the crate.
///
/// These are the parts trained from scratch — the two-timescale core and the classification heads. The
/// embedding table they sit on top of is 10x larger than everything else combined and is fetched separately.
pub mod bundled {
    /// HRM reasoning core + the typed-span and epistemic heads, without the embedding table.
    ///
    /// Empty until `tools/split_checkpoint` writes it; the constant exists so the resolution path is the same
    /// whether or not a build has produced it.
    pub const HEADS: &[u8] = include_bytes!("../assets/heads.safetensors");

    /// True when this build actually carries head weights rather than a placeholder.
    pub fn have_heads() -> bool {
        HEADS.len() > 1024
    }
}

/// Where a set of weights comes from, and what we are permitted to do with it.
///
/// An earlier version of this enum forbade mirroring third-party weights on principle. That was the wrong
/// rule: every model this engine uses turns out to be permissively licensed (Apache-2.0 or MIT), so mirroring
/// is allowed. The obligations that *do* bind are attribution, carrying the licence text, and stating
/// modifications — none of which is satisfied by simply refusing to host a copy.
///
/// The more important protection is the pinned revision. A mirror without a pin still drifts when we refresh
/// it; a pin without a mirror still breaks when upstream deletes the repo. Both, where permitted.
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Source {
    /// Compiled into this crate. Small, ours, offline by construction.
    Bundled,
    /// Trained by this project and published to our own model repository.
    Ours { repo: &'static str, rev: &'static str },
    /// Someone else's weights, fetched from the repository that published them.
    Upstream { repo: &'static str, license: &'static str, rev: &'static str },
    /// Someone else's weights that we also host, as the licence permits, so a build does not depend on an
    /// upstream repository staying where it is. `upstream` is retained because attribution does not transfer.
    Mirrored {
        ours: &'static str,
        upstream: &'static str,
        license: &'static str,
        rev: &'static str,
    },
}

impl Source {
    /// The licence a redistribution must carry, if this is not our own work.
    pub fn license(&self) -> Option<&'static str> {
        match self {
            Source::Bundled | Source::Ours { .. } => None,
            Source::Upstream { license, .. } | Source::Mirrored { license, .. } => Some(license),
        }
    }
    /// Who to credit.
    pub fn attribution(&self) -> Option<&'static str> {
        match self {
            Source::Bundled | Source::Ours { .. } => None,
            Source::Upstream { repo, .. } => Some(repo),
            Source::Mirrored { upstream, .. } => Some(upstream),
        }
    }
    /// The revision this build expects. Pinning is what makes a fetch reproducible.
    pub fn revision(&self) -> &'static str {
        match self {
            Source::Bundled => "bundled",
            Source::Ours { rev, .. } | Source::Upstream { rev, .. } | Source::Mirrored { rev, .. } => rev,
        }
    }
    /// Where to fetch from, preferring our mirror when one exists.
    pub fn fetch_repo(&self) -> Option<&'static str> {
        match self {
            Source::Bundled => None,
            Source::Ours { repo, .. } => Some(repo),
            Source::Upstream { repo, .. } => Some(repo),
            Source::Mirrored { ours, .. } => Some(ours),
        }
    }
}

/// One resolvable set of weights.
#[derive(Debug, Clone, Copy)]
pub struct Artifact {
    /// stable name used by [`resolve`]
    pub name: &'static str,
    /// environment variable that overrides the search, if set
    pub env: &'static str,
    /// directory name under the model root
    pub dir: &'static str,
    /// a file that must exist for the directory to count as present
    pub marker: &'static str,
    pub source: Source,
    pub approx_mb: u32,
    pub purpose: &'static str,
}

/// Every artifact the engine knows how to find.
pub const ARTIFACTS: &[Artifact] = &[
    Artifact {
        name: "embeddings",
        env: "STEELDB_EMBEDDINGS",
        dir: "bert-tiny",
        marker: "model.safetensors",
        source: Source::Mirrored {
            ours: "cp500/steeldb-models",
            upstream: "google/bert_uncased_L-2_H-128_A-2",
            license: "Apache-2.0",
            rev: "main",
        },
        approx_mb: 16,
        purpose: "the embedding table the bundled heads sit on; 15.2 MB of it is the vocabulary alone",
    },
    Artifact {
        name: "spo-tagger",
        env: "STEELDB_ML_BUNDLE",
        dir: "step0_bundle_ml",
        marker: "spo.onnx",
        source: Source::Ours { repo: "cp500/steeldb-models", rev: "main" },
        approx_mb: 168,
        purpose: "typed span tagger: text becomes ENT/REL/GEO/TIME/QTY spans",
    },
    Artifact {
        name: "splade",
        env: "STEELDB_SPLADE_DIR",
        dir: "splade",
        marker: "splade.onnx",
        source: Source::Ours { repo: "cp500/steeldb-models", rev: "main" },
        approx_mb: 107,
        purpose: "learned facet heads: the projection that builds the bitmap",
    },
    Artifact {
        name: "model2vec",
        env: "STEELDB_MODEL2VEC",
        dir: "model2vec",
        marker: "potion.f32",
        source: Source::Mirrored {
            ours: "cp500/steeldb-models",
            upstream: "minishlab/potion-base-4M",
            license: "MIT",
            rev: "main",
        },
        approx_mb: 15,
        purpose: "static embeddings for optimal-transport ontology discovery",
    },
    Artifact {
        name: "needle3",
        env: "STEELDB_NEEDLE_DIR",
        dir: "needle3",
        marker: "needle3.cact",
        // Apache-2.0, confirmed from the LICENSE file in the model repository — so mirroring is permitted
        // provided the licence and attribution travel with it.
        source: Source::Mirrored {
            ours: "cp500/steeldb-models",
            upstream: "Cactus-Compute/needle3",
            license: "Apache-2.0",
            rev: "b274efcb211a9eef48c9a88da4b43bd569696a39",
        },
        approx_mb: 242,
        purpose: "121M tool-calling model used as the query planner",
    },
];

/// Why a set of weights could not be produced.
#[derive(Debug, Clone)]
pub enum ModelError {
    /// no such artifact name
    Unknown(String),
    /// searched everywhere and found nothing; carries the instruction to fix it
    NotFound { name: String, searched: Vec<PathBuf>, hint: String },
}

impl std::fmt::Display for ModelError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ModelError::Unknown(n) => write!(f, "unknown model '{n}'"),
            ModelError::NotFound { name, searched, hint } => {
                writeln!(f, "model '{name}' not found. Looked in:")?;
                for p in searched {
                    writeln!(f, "  {}", p.display())?;
                }
                write!(f, "{hint}")
            }
        }
    }
}

impl std::error::Error for ModelError {}

/// Look up an artifact by name.
pub fn artifact(name: &str) -> Option<&'static Artifact> {
    ARTIFACTS.iter().find(|a| a.name == name)
}

/// Directories searched for models, in order.
///
/// An explicit environment variable wins, then a user-level cache, then a `models/` directory beside the
/// working tree — which is what makes a checkout work without configuration.
pub fn search_roots() -> Vec<PathBuf> {
    let mut roots = Vec::new();
    if let Ok(home) = std::env::var("HOME") {
        roots.push(PathBuf::from(home).join(".steeldb").join("models"));
    }
    if let Ok(dir) = std::env::var("STEELDB_MODELS") {
        roots.insert(0, PathBuf::from(dir));
    }
    roots.push(PathBuf::from("models"));
    roots
}

/// Find the directory holding an artifact, without downloading anything.
///
/// Fetching is left to the caller on purpose: a library that reaches for the network on its own is a library
/// that surprises someone in production. The error says exactly what to run.
pub fn resolve(name: &str) -> Result<PathBuf, ModelError> {
    let art = artifact(name).ok_or_else(|| ModelError::Unknown(name.to_string()))?;

    if let Ok(dir) = std::env::var(art.env) {
        let p = PathBuf::from(dir);
        if p.join(art.marker).exists() {
            return Ok(p);
        }
    }

    let mut searched = Vec::new();
    for root in search_roots() {
        let cand = root.join(art.dir);
        if cand.join(art.marker).exists() {
            return Ok(cand);
        }
        searched.push(cand);
    }

    Err(ModelError::NotFound { name: name.to_string(), searched, hint: fetch_hint(art) })
}

/// The instruction that would make a missing artifact present.
pub fn fetch_hint(art: &Artifact) -> String {
    // The destination is the models ROOT, with `--include` selecting the artifact's subdirectory.
    //
    // An earlier version pointed `--local-dir` at `<root>/<dir>`, which downloads the whole repository INTO a
    // directory already named after one artifact: the file lands at `<root>/step0_bundle_ml/step0_bundle_ml/
    // spo.onnx` and resolution still fails. That could not be caught until the repository actually existed with
    // subdirectories in it — the command looked plausible against an empty repo.
    let root = search_roots().first().cloned().unwrap_or_else(|| PathBuf::from("models"));
    match art.source.fetch_repo() {
        None => "this artifact ships with the crate; the build is incomplete".to_string(),
        Some(repo) => {
            let credit = match art.source.attribution() {
                Some(up) => format!("\n  {up} — {}", art.source.license().unwrap_or("see model card")),
                None => String::new(),
            };
            format!(
                "Fetch it ({} MB) with:\n  \
                 huggingface-cli download {repo} --revision {} --include '{}/*' --local-dir {}\n\
                 or set {}=/path/to/{}{credit}",
                art.approx_mb,
                art.source.revision(),
                art.dir,
                root.display(),
                art.env,
                art.dir
            )
        }
    }
}

/// A short report of what is present and what is missing — useful in a CLI or a bug report.
pub fn status() -> Vec<(&'static str, Option<PathBuf>)> {
    ARTIFACTS.iter().map(|a| (a.name, resolve(a.name).ok())).collect()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn third_party_weights_keep_their_licence_and_credit() {
        // The obligations that actually bind when mirroring is permitted: the licence travels with the copy,
        // and attribution points at whoever trained it — not at us, mirror or no mirror.
        for a in ARTIFACTS {
            match a.source {
                Source::Ours { repo, .. } => {
                    assert!(repo.starts_with("cp500/"), "{} claims ours but points at {repo}", a.name);
                    assert!(a.source.license().is_none(), "our own work needs no third-party licence");
                }
                Source::Upstream { license, .. } | Source::Mirrored { license, .. } => {
                    assert!(!license.is_empty(), "{} must record a licence", a.name);
                    let credit = a.source.attribution().expect("third-party work needs attribution");
                    assert!(
                        !credit.starts_with("cp500/"),
                        "{} credits us for someone else's work",
                        a.name
                    );
                }
                Source::Bundled => {}
            }
        }
    }

    #[test]
    fn every_fetched_artifact_pins_a_revision() {
        // A mirror without a pin still drifts; a pin without a mirror still breaks on deletion. The pin is the
        // part that makes a build reproducible, so it is required even where we host the copy.
        for a in ARTIFACTS {
            if a.source.fetch_repo().is_some() {
                assert!(!a.source.revision().is_empty(), "{} must pin a revision", a.name);
            }
        }
        // the one we depend on most specifically is pinned to a commit rather than a moving branch
        let n = artifact("needle3").unwrap();
        assert_eq!(n.source.revision().len(), 40, "needle3 should pin an exact commit");
    }

    #[test]
    fn the_fetch_hint_downloads_into_the_models_ROOT_not_the_artifact_directory() {
        // The hint used to point `--local-dir` at `<root>/<artifact-dir>`, which downloads the repository INTO a
        // directory already named after one artifact — the file lands at
        // `<root>/step0_bundle_ml/step0_bundle_ml/spo.onnx` and resolution still fails. Verified against the real
        // huggingface-cli: with `--include '<dir>/*' --local-dir <root>` the subdirectory lands directly under
        // the root, which is where `resolve` looks.
        std::env::set_var("STEELDB_MODELS", "/models-root");
        let hint = fetch_hint(artifact("spo-tagger").unwrap());
        assert!(hint.contains("--include 'step0_bundle_ml/*'"), "must select the subdirectory: {hint}");
        assert!(hint.contains("--local-dir /models-root"), "must target the ROOT: {hint}");
        assert!(
            !hint.contains("--local-dir /models-root/step0_bundle_ml"),
            "must not nest the artifact directory inside itself: {hint}"
        );
        std::env::remove_var("STEELDB_MODELS");
    }

    #[test]
    fn a_missing_model_explains_how_to_get_it() {
        // resolution must never panic, and the error has to be actionable rather than just negative
        std::env::set_var("STEELDB_MODELS", "/nonexistent-steeldb-test-root");
        let err = resolve("needle3").unwrap_err();
        let msg = err.to_string();
        assert!(msg.contains("not found"), "{msg}");
        assert!(msg.contains("Cactus-Compute/needle3"), "must credit the author: {msg}");
        assert!(msg.contains("Apache-2.0"), "must state the licence: {msg}");
        assert!(msg.contains("--revision"), "must pin a revision: {msg}");
        assert!(msg.contains("242 MB"), "must state the size: {msg}");
        std::env::remove_var("STEELDB_MODELS");
    }

    #[test]
    fn an_unknown_name_is_an_error_not_a_panic() {
        assert!(matches!(resolve("no-such-model"), Err(ModelError::Unknown(_))));
    }

    #[test]
    fn every_artifact_is_uniquely_named_and_documented() {
        let mut seen = std::collections::HashSet::new();
        for a in ARTIFACTS {
            assert!(seen.insert(a.name), "duplicate artifact name {}", a.name);
            assert!(!a.purpose.is_empty(), "{} needs a purpose", a.name);
            assert!(a.env.starts_with("STEELDB_"), "{} env var should be namespaced", a.name);
        }
    }
}