steeldb/

artifact.rs

//! **Artefacts** — what `learn` leaves behind, and what `ingest` and `query` follow.
//!
//! `learn` is the only step that needs a model. Running it on every startup would make the system slow,
//! non-deterministic, and dependent on a service being reachable. So it runs **once** and writes its result to
//! a directory of small JSON files; every run after that reads those files and is offline and reproducible.
//!
//! ```text
//! .hypersteeldb/
//!   manifest.json      what produced this set, when, and a digest of each part
//!   vocabulary.json    the categories and the words each claims
//!   gazetteer.json     multi-word mentions kept whole
//!   relations.json     relation verbs, with polarity
//!   motifs.json        latent motifs, so `motif/*` tokens survive a reload
//! ```
//!
//! ```no_run
//! use steeldb::SteelDb;
//!
//! // once — discovers (or learns) and records the result
//! let db = SteelDb::ingest(["…documents…"])?;
//! db.save(".hypersteeldb")?;
//!
//! // every time after — same vocabulary, no model, no network, byte-identical results
//! let db = SteelDb::ingest_using(["…documents…"], ".hypersteeldb")?;
//! # Ok::<(), steeldb::Error>(())
//! ```
//!
//! # What is deliberately *not* in here
//!
//! **No document text.** An artefact set is derived vocabulary: category names, the words that signal them,
//! mention surfaces, relation verbs. It is not a copy of the corpus, and it is not an index. Two reasons that
//! matters: an artefact set is safe to commit to a repository next to code, and shipping one does not ship the
//! documents it was learned from. A test asserts the files contain no sentence from the corpus.
//!
//! You still need the documents to `ingest`. The artefacts say *how* to read them, not *what* they said.

use std::collections::BTreeMap;
use std::path::{Path, PathBuf};

/// Current on-disk layout version. Bumped when a field's meaning changes, so an old set is refused rather
/// than silently misread.
///
/// v3 added `motifs.json`. Without it a reload produced no `motif/*` tokens, so the same documents indexed
/// differently depending on whether they went through discovery or an artefact — the exact drift these files
/// exist to prevent.
///
/// v2 changed `gazetteer` from a list of surfaces to a surface-to-token mapping. v1 lost entity normalisation
/// on reload: "the Sootopolis City outage" and "Sootopolis City" both register as `entity/sootopolis-city`, but
/// with only the surfaces recorded a later ingest re-derived the token by slugging the raw text and the two
/// stopped merging.
pub const SCHEMA: u32 = 3;

/// File names inside an artefact directory.
pub const MANIFEST: &str = "manifest.json";
pub const VOCABULARY: &str = "vocabulary.json";
pub const GAZETTEER: &str = "gazetteer.json";
pub const RELATIONS: &str = "relations.json";
/// Latent motifs: the themes a situation can join without sharing a keyword (the paper's sixth dimension).
pub const MOTIFS: &str = "motifs.json";
/// Subdirectory for the finetuning set. Separate because, unlike everything else here, it necessarily
/// contains document text.
pub const TRAINING_DIR: &str = "training";
pub const TRAINING_FILE: &str = "spans.jsonl";
pub const IGNORE_FILE: &str = ".gitignore";

/// A registered entity mention: how it appears, and what it resolves to.
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct Registration {
    /// the surface form as found in the text
    pub surface: String,
    /// the canonical token, e.g. `entity/sootopolis-city`
    pub token: String,
}

/// One category and the words that put a document in it.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct CategoryRecord {
    pub name: String,
    pub words: Vec<String>,
}

/// The complete set. Small, human-readable, and diffable on purpose — a vocabulary change should show up in
/// review like any other change.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct Artifacts {
    pub schema: u32,
    /// what produced this set: `discovery` for the offline path, or `learn:<model>` when a model proposed it
    pub producer: String,
    /// RFC3339 timestamp, or empty where no clock is available
    pub created: String,
    pub categories: Vec<CategoryRecord>,
    /// Registered mentions: the surface as it appears in text, and the canonical token it resolves to.
    ///
    /// A mapping rather than a list because registration NORMALISES — leading determiners and trailing status
    /// words are stripped, so several surfaces can share one token and their postings merge. Recording only the
    /// surfaces would drop that, and two variants of one entity would become two entities.
    pub gazetteer: Vec<Registration>,
    /// relation verbs, canonical form
    pub relations: Vec<String>,
    /// latent motifs, in the same shape as a category: a name and the terms that join it
    #[serde(default)]
    pub motifs: Vec<CategoryRecord>,
    /// byte lengths of each written part, so a truncated set is detected on load
    pub digests: BTreeMap<String, usize>,
    /// Number of labelled examples written under `training/`, if any.
    ///
    /// The count lives in the manifest but the examples do not, because they contain corpus text and the
    /// manifest is meant to be readable and committable.
    #[serde(default)]
    pub training_examples: usize,
    /// Parts of this set that contain verbatim document text. Always empty for the vocabulary files; lists
    /// `training/` when a finetuning set was written.
    ///
    /// Recorded explicitly so a tool, a reviewer, or a CI check can tell which parts are safe to publish
    /// without having to know the layout by heart.
    #[serde(default)]
    pub contains_document_text: Vec<String>,
    /// The finetuning set, held in memory. Never serialised into the manifest: it goes to its own file.
    #[serde(skip)]
    training: Option<String>,
}

/// Why an artefact set could not be read or written.
#[derive(Debug)]
pub enum ArtifactError {
    Io(std::io::Error),
    /// the directory exists but does not hold an artefact set
    NotAnArtifactDir(PathBuf),
    /// written by a newer or older layout than this build understands
    SchemaMismatch { found: u32, expected: u32 },
    /// a part is present but its size does not match the manifest
    Corrupt(String),
    Parse(String),
}

impl std::fmt::Display for ArtifactError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ArtifactError::Io(e) => write!(f, "{e}"),
            ArtifactError::NotAnArtifactDir(p) => write!(
                f,
                "{} is not an artefact directory (no {MANIFEST}); run save() first",
                p.display()
            ),
            ArtifactError::SchemaMismatch { found, expected } => write!(
                f,
                "artefact schema {found} cannot be read by this build (expects {expected}); re-run learn"
            ),
            ArtifactError::Corrupt(what) => write!(f, "artefact set is incomplete: {what}"),
            ArtifactError::Parse(e) => write!(f, "could not parse artefact: {e}"),
        }
    }
}

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

impl From<std::io::Error> for ArtifactError {
    fn from(e: std::io::Error) -> Self {
        ArtifactError::Io(e)
    }
}

impl Artifacts {
    /// Build a set from the parts an ingest produced.
    pub fn new(
        producer: impl Into<String>,
        categories: Vec<CategoryRecord>,
        gazetteer: Vec<Registration>,
        relations: Vec<String>,
    ) -> Artifacts {
        Artifacts {
            schema: SCHEMA,
            producer: producer.into(),
            created: now_rfc3339(),
            categories,
            gazetteer,
            relations,
            motifs: Vec::new(),
            digests: BTreeMap::new(),
            training_examples: 0,
            contains_document_text: Vec::new(),
            training: None,
        }
    }

    /// Attach a finetuning set: JSONL lines, one labelled passage each.
    ///
    /// This is the one part of an artefact set that carries document text — a span label is meaningless
    /// without the words it points at. It is therefore written to a subdirectory, recorded in
    /// `contains_document_text`, and excluded by a `.gitignore` that [`Artifacts::save`] writes, so committing
    /// the artefact directory does not commit the corpus by accident.
    pub fn with_training(mut self, jsonl: impl Into<String>) -> Artifacts {
        let jsonl = jsonl.into();
        self.training_examples = jsonl.lines().filter(|l| !l.trim().is_empty()).count();
        self.contains_document_text = vec![format!("{TRAINING_DIR}/")];
        self.training = Some(jsonl);
        self
    }

    /// The finetuning set, if this build of the artefacts carries one in memory.
    pub fn training(&self) -> Option<&str> {
        self.training.as_deref()
    }

    /// Write the set to `dir`, creating it if needed.
    ///
    /// Each part is written separately so a reviewer sees a small, meaningful diff rather than one large blob.
    /// Record the discovered motifs, so a reload projects the same `motif/*` tokens discovery did.
    pub fn with_motifs(mut self, motifs: Vec<CategoryRecord>) -> Artifacts {
        self.motifs = motifs;
        self
    }

    pub fn save(&self, dir: impl AsRef<Path>) -> Result<(), ArtifactError> {
        let dir = dir.as_ref();
        std::fs::create_dir_all(dir)?;

        let vocab = serde_json::to_vec_pretty(&self.categories).map_err(|e| ArtifactError::Parse(e.to_string()))?;
        let gaz = serde_json::to_vec_pretty(&self.gazetteer).map_err(|e| ArtifactError::Parse(e.to_string()))?;
        let rel = serde_json::to_vec_pretty(&self.relations).map_err(|e| ArtifactError::Parse(e.to_string()))?;
        let mot = serde_json::to_vec_pretty(&self.motifs).map_err(|e| ArtifactError::Parse(e.to_string()))?;

        let mut manifest = self.clone();
        manifest.digests.clear();
        manifest.digests.insert(VOCABULARY.into(), vocab.len());
        manifest.digests.insert(GAZETTEER.into(), gaz.len());
        manifest.digests.insert(RELATIONS.into(), rel.len());
        manifest.digests.insert(MOTIFS.into(), mot.len());
        // the manifest carries only counts and provenance; the parts carry the content
        let mut slim = manifest.clone();
        slim.training = None;
        slim.categories = Vec::new();
        slim.gazetteer = Vec::new();
        slim.relations = Vec::new();
        slim.motifs = Vec::new();

        // Write the ignore file FIRST, so the training set can never exist in the directory without the rule
        // that keeps it out of a commit.
        std::fs::write(
            dir.join(IGNORE_FILE),
            format!(
                "# Written by hypersteeldb. The finetuning set under {TRAINING_DIR}/ contains verbatim document\n\
                 # text, because a span label is meaningless without the words it points at. Everything else in\n\
                 # this directory is derived vocabulary and is safe to commit.\n\
                 {TRAINING_DIR}/\n"
            ),
        )?;
        if let Some(jsonl) = &self.training {
            let tdir = dir.join(TRAINING_DIR);
            std::fs::create_dir_all(&tdir)?;
            std::fs::write(tdir.join(TRAINING_FILE), jsonl.as_bytes())?;
        }

        std::fs::write(dir.join(VOCABULARY), &vocab)?;
        std::fs::write(dir.join(GAZETTEER), &gaz)?;
        std::fs::write(dir.join(RELATIONS), &rel)?;
        std::fs::write(dir.join(MOTIFS), &mot)?;
        std::fs::write(
            dir.join(MANIFEST),
            serde_json::to_vec_pretty(&slim).map_err(|e| ArtifactError::Parse(e.to_string()))?,
        )?;
        Ok(())
    }

    /// Read a set from `dir`, refusing anything this build cannot interpret correctly.
    pub fn load(dir: impl AsRef<Path>) -> Result<Artifacts, ArtifactError> {
        let dir = dir.as_ref();
        let mpath = dir.join(MANIFEST);
        if !mpath.exists() {
            return Err(ArtifactError::NotAnArtifactDir(dir.to_path_buf()));
        }
        let mut set: Artifacts = serde_json::from_slice(&std::fs::read(&mpath)?)
            .map_err(|e| ArtifactError::Parse(e.to_string()))?;
        if set.schema != SCHEMA {
            return Err(ArtifactError::SchemaMismatch { found: set.schema, expected: SCHEMA });
        }

        let vocab = std::fs::read(dir.join(VOCABULARY))?;
        let gaz = std::fs::read(dir.join(GAZETTEER))?;
        let rel = std::fs::read(dir.join(RELATIONS))?;
        let mot = std::fs::read(dir.join(MOTIFS))?;

        // A truncated part would otherwise load as a smaller vocabulary and answer questions with it — a
        // quietly wrong result rather than a failure.
        for (name, actual) in
            [(VOCABULARY, vocab.len()), (GAZETTEER, gaz.len()), (RELATIONS, rel.len()), (MOTIFS, mot.len())]
        {
            if let Some(expected) = set.digests.get(name) {
                if *expected != actual {
                    return Err(ArtifactError::Corrupt(format!(
                        "{name} is {actual} bytes, manifest says {expected}"
                    )));
                }
            }
        }

        set.categories = serde_json::from_slice(&vocab).map_err(|e| ArtifactError::Parse(e.to_string()))?;
        set.gazetteer = serde_json::from_slice(&gaz).map_err(|e| ArtifactError::Parse(e.to_string()))?;
        set.relations = serde_json::from_slice(&rel).map_err(|e| ArtifactError::Parse(e.to_string()))?;
        set.motifs = serde_json::from_slice(&mot).map_err(|e| ArtifactError::Parse(e.to_string()))?;
        // the training set is optional and may legitimately be absent — it is gitignored, so a cloned
        // repository will have the vocabulary without it, and ingest/query do not need it
        let tpath = dir.join(TRAINING_DIR).join(TRAINING_FILE);
        set.training = std::fs::read_to_string(&tpath).ok();
        Ok(set)
    }

    /// True when `dir` holds a readable artefact set.
    pub fn exists(dir: impl AsRef<Path>) -> bool {
        dir.as_ref().join(MANIFEST).exists()
    }
}

/// An RFC3339 timestamp, or empty on targets without a clock (wasm has none).
fn now_rfc3339() -> String {
    #[cfg(not(target_arch = "wasm32"))]
    {
        use std::time::{SystemTime, UNIX_EPOCH};
        let Ok(d) = SystemTime::now().duration_since(UNIX_EPOCH) else { return String::new() };
        let secs = d.as_secs() as i64;
        // civil-from-days, so no date dependency is needed for a timestamp
        let days = secs.div_euclid(86_400);
        let tod = secs.rem_euclid(86_400);
        let z = days + 719_468;
        let era = z.div_euclid(146_097);
        let doe = z.rem_euclid(146_097);
        let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365;
        let y = yoe + era * 400;
        let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
        let mp = (5 * doy + 2) / 153;
        let d_ = doy - (153 * mp + 2) / 5 + 1;
        let m = if mp < 10 { mp + 3 } else { mp - 9 };
        let y = if m <= 2 { y + 1 } else { y };
        format!(
            "{y:04}-{m:02}-{d_:02}T{:02}:{:02}:{:02}Z",
            tod / 3600,
            (tod % 3600) / 60,
            tod % 60
        )
    }
    #[cfg(target_arch = "wasm32")]
    {
        String::new()
    }
}

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

    fn tmp(name: &str) -> PathBuf {
        let p = std::env::temp_dir().join(format!("hsdb_artifact_{}_{name}", std::process::id()));
        let _ = std::fs::remove_dir_all(&p);
        p
    }

    fn sample() -> Artifacts {
        Artifacts::new(
            "discovery",
            vec![
                CategoryRecord { name: "battle".into(), words: vec!["defeated".into(), "faced".into()] },
                CategoryRecord { name: "survey".into(), words: vec!["elevation".into()] },
            ],
            vec![
                Registration { surface: "Sootopolis City".into(), token: "entity/sootopolis-city".into() },
                Registration { surface: "Indigo Invitational".into(), token: "entity/indigo-invitational".into() },
            ],
            vec!["defeated".into(), "documented".into()],
        )
        .with_motifs(vec![CategoryRecord {
            name: "series".into(),
            words: vec!["series".into(), "season".into()],
        }])
    }

    #[test]
    fn a_set_round_trips_exactly() {
        let dir = tmp("roundtrip");
        let a = sample();
        a.save(&dir).unwrap();
        let b = Artifacts::load(&dir).unwrap();
        assert_eq!(b.schema, SCHEMA);
        assert_eq!(b.producer, "discovery");
        assert_eq!(b.categories.len(), 2);
        assert_eq!(b.categories[0].name, "battle");
        assert_eq!(b.gazetteer, a.gazetteer);
        assert_eq!(b.relations, a.relations);
        // motifs travel too, or a reload projects no motif/* tokens and indexes differently from discovery
        assert_eq!(b.motifs.len(), 1);
        assert_eq!(b.motifs[0].name, "series");
        assert_eq!(b.motifs[0].words, a.motifs[0].words);
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn the_files_contain_no_document_text() {
        // The leakage guard. An artefact set is derived vocabulary, so it is safe to commit and safe to ship;
        // if a sentence from the corpus could appear here, neither would be true.
        let dir = tmp("leak");
        let sentence = "Morty Shade defeated Wallace Gale at Ecruteak City during the Indigo Invitational.";
        sample().save(&dir).unwrap();
        for f in [MANIFEST, VOCABULARY, GAZETTEER, RELATIONS, MOTIFS] {
            let text = std::fs::read_to_string(dir.join(f)).unwrap();
            assert!(!text.contains(sentence), "{f} contains a corpus sentence");
            assert!(!text.contains("Morty Shade defeated"), "{f} contains a document fragment");
        }
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn a_missing_set_says_what_to_do() {
        let dir = tmp("missing");
        std::fs::create_dir_all(&dir).unwrap();
        let err = Artifacts::load(&dir).unwrap_err();
        let msg = err.to_string();
        assert!(msg.contains("not an artefact directory"), "{msg}");
        assert!(msg.contains("save()"), "must say how to create one: {msg}");
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn a_truncated_part_is_refused_not_silently_loaded() {
        // A shorter vocabulary file would otherwise load as a smaller vocabulary and answer questions with
        // it — wrong, and invisible.
        let dir = tmp("truncated");
        sample().save(&dir).unwrap();
        std::fs::write(dir.join(VOCABULARY), b"[]").unwrap();
        let err = Artifacts::load(&dir).unwrap_err();
        assert!(matches!(err, ArtifactError::Corrupt(_)), "{err}");
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn a_future_schema_is_refused() {
        let dir = tmp("schema");
        let mut a = sample();
        a.save(&dir).unwrap();
        a.schema = SCHEMA + 1;
        let mut slim = a.clone();
        slim.categories = Vec::new();
        slim.gazetteer = Vec::new();
        slim.relations = Vec::new();
        std::fs::write(dir.join(MANIFEST), serde_json::to_vec_pretty(&slim).unwrap()).unwrap();
        assert!(matches!(
            Artifacts::load(&dir).unwrap_err(),
            ArtifactError::SchemaMismatch { .. }
        ));
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn the_manifest_records_provenance() {
        let dir = tmp("provenance");
        Artifacts::new("learn:local:qwen2.5-0.5b", vec![], vec![], vec![]).save(&dir).unwrap();
        let loaded = Artifacts::load(&dir).unwrap();
        assert_eq!(loaded.producer, "learn:local:qwen2.5-0.5b");
        assert!(loaded.created.contains('T') || loaded.created.is_empty());
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn the_finetuning_set_is_separated_and_gitignored() {
        // A span label is meaningless without the words it points at, so this ONE part must carry document
        // text. It therefore has to be impossible to commit it by accident while committing the vocabulary.
        let dir = tmp("training");
        let jsonl = "{\"text\":\"Morty Shade defeated Wallace Gale.\",\"spans\":[]}\n\
                     {\"text\":\"A survey recorded Aggron at 1082 m.\",\"spans\":[]}\n";
        sample().with_training(jsonl).save(&dir).unwrap();

        // the ignore rule exists and covers the training directory
        let ignore = std::fs::read_to_string(dir.join(IGNORE_FILE)).unwrap();
        assert!(ignore.contains(&format!("{TRAINING_DIR}/")), "{ignore}");

        // the text is in the training file
        let train = std::fs::read_to_string(dir.join(TRAINING_DIR).join(TRAINING_FILE)).unwrap();
        assert!(train.contains("Morty Shade defeated"));

        // and NOT in any of the committable parts
        for f in [MANIFEST, VOCABULARY, GAZETTEER, RELATIONS, MOTIFS] {
            let text = std::fs::read_to_string(dir.join(f)).unwrap();
            assert!(!text.contains("Morty Shade defeated"), "{f} leaked document text");
        }

        // the manifest declares which parts are unsafe to publish, so a CI check need not know the layout
        let loaded = Artifacts::load(&dir).unwrap();
        assert_eq!(loaded.training_examples, 2);
        assert_eq!(loaded.contains_document_text, vec![format!("{TRAINING_DIR}/")]);
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn a_vocabulary_only_set_declares_no_text_and_still_loads() {
        // The common case after cloning a repository: the vocabulary is committed, the gitignored training set
        // is absent. ingest and query must work from that alone.
        let dir = tmp("novocab");
        sample().save(&dir).unwrap();
        let loaded = Artifacts::load(&dir).unwrap();
        assert!(loaded.contains_document_text.is_empty(), "nothing here carries document text");
        assert_eq!(loaded.training_examples, 0);
        assert!(loaded.training().is_none());
        assert_eq!(loaded.categories.len(), 2, "the vocabulary is complete without the training set");
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn the_ignore_rule_is_written_even_without_a_training_set() {
        // Otherwise a later run that DOES produce one would drop it into an unguarded directory.
        let dir = tmp("ignorefirst");
        sample().save(&dir).unwrap();
        assert!(dir.join(IGNORE_FILE).exists(), "the rule must exist before the data can");
        let _ = std::fs::remove_dir_all(&dir);
    }
}