steeldb/

learn.rs

//! **Learn** — the third verb, and the one that works differently on purpose.
//!
//! `ingest` and `query` are offline, deterministic and free. `learn` is none of those: it calls a language
//! model, which means credentials, network, latency and a bill. Hiding that behind a method that looks like the
//! other two would be a trap, so this module makes all four facts visible in the shape of the API.
//!
//! ```no_run
//! # #[cfg(feature = "bedrock")]
//! # async fn demo() -> Result<(), Box<dyn std::error::Error>> {
//! use steeldb::{SteelDb, learn::Teacher};
//!
//! let mut db = SteelDb::ingest(["…documents…"])?;
//!
//! // credentials are checked when the teacher is built, not when it is used
//! let teacher = Teacher::bedrock("us.anthropic.claude-sonnet-4-5-20250929-v1:0")?;
//!
//! // a proposal is returned, NOT applied
//! let proposal = teacher.propose_categories(&db).await?;
//! println!("{proposal}");
//!
//! // you decide, and the same MECE test that gates local discovery gates this too
//! let adopted = db.adopt(&proposal);
//! println!("kept {} of {}", adopted.len(), proposal.candidates.len());
//! # Ok(()) }
//! ```
//!
//! ## Why a proposal instead of a mutation
//!
//! A model suggesting categories is a suggestion, not an authority. Returning a [`Proposal`] means you can
//! print it, diff it, log it, or reject it before your vocabulary changes — and it keeps the model advisory,
//! which is the same separation the query planner has. `adopt` then applies the *same* gate that local
//! discovery uses, so a model cannot sneak in a category that a deterministic test would have rejected.
//!
//! Without a network-capable feature enabled, this module still compiles: [`Proposal`] and [`SteelDb::adopt`]
//! work with candidates from any source, so the offline path is testable.

use crate::api::SteelDb;

/// A candidate category, from wherever.
#[derive(Debug, Clone, PartialEq)]
pub struct Candidate {
    /// the category name, which becomes its query stem
    pub name: String,
    /// words that should put a document in this category
    pub words: Vec<String>,
    /// why the proposer thinks it belongs, for a human reading the diff
    pub rationale: String,
}

/// What a teacher suggests. Inert until adopted.
#[derive(Debug, Clone, Default)]
pub struct Proposal {
    pub candidates: Vec<Candidate>,
    /// what produced this, so a log entry is traceable
    pub source: String,
}

impl std::fmt::Display for Proposal {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        writeln!(f, "proposal from {} — {} candidate(s)", self.source, self.candidates.len())?;
        for c in &self.candidates {
            writeln!(f, "  {} — {}", c.name, c.rationale)?;
            writeln!(f, "    words: {}", c.words.join(", "))?;
        }
        Ok(())
    }
}

/// The outcome of adopting one candidate. Reported per candidate so a rejection is explicable.
#[derive(Debug, Clone)]
pub struct Verdict {
    pub name: String,
    pub kept: bool,
    /// the gate's own words
    pub reason: String,
    /// share of documents the candidate covers
    pub coverage: f64,
    /// how much it duplicates a category already accepted
    pub overlap: f64,
}

impl SteelDb {
    /// Apply a proposal, keeping only what the MECE gate accepts.
    ///
    /// The gate is the same one [`SteelDb::ingest`] uses, so a model-proposed category has to earn its place on
    /// the same terms as a locally-discovered one: enough coverage, and not a near-duplicate of something
    /// already present. Candidates are judged in order, so the second of two similar suggestions is rejected
    /// against the first.
    ///
    /// Returns a verdict per candidate; the kept ones are queryable immediately.
    pub fn adopt(&mut self, proposal: &Proposal) -> Vec<Verdict> {
        let mut verdicts = Vec::new();
        for (round, cand) in proposal.candidates.iter().enumerate() {
            let c = crate::grow::Candidate {
                name: cand.name.clone(),
                parent: None,
                description: cand.rationale.clone(),
                examples: cand.words.clone(),
                worth_adding: true,
            };
            let docs = self.documents().to_vec();
            let spec = self.spec_snapshot();
            let scored = crate::grow::score_candidate_full(&spec, &docs, &c);
            let (score, dup) = match scored {
                Some((s, d)) => (Some(s), d),
                None => (None, None),
            };
            let ev = crate::grow::gate_full(&spec, &c, score.as_ref(), dup, self.min_gain(), round);
            if ev.kept {
                self.push_category(cand.name.clone(), cand.words.clone());
            }
            verdicts.push(Verdict {
                name: cand.name.clone(),
                kept: ev.kept,
                reason: ev.reason,
                coverage: ev.coverage,
                overlap: ev.maxcos,
            });
        }
        if verdicts.iter().any(|v| v.kept) {
            // the index must be rebuilt: a new category changes what every document projects to
            self.reproject();
        }
        verdicts
    }
}

/// A source of proposals.
///
/// Constructing one is where credentials are checked, so a missing configuration fails before any work is
/// queued rather than part-way through a corpus.
pub struct Teacher {
    kind: Kind,
}

enum Kind {
    /// proposals supplied by the caller — the offline path, and what the tests use
    Fixed(Proposal),
    /// any OpenAI-compatible endpoint: llama.cpp, vLLM, Ollama, LM Studio. No credentials, no bill.
    #[cfg(feature = "paddock")]
    Local { base_url: String, model: String },
    #[cfg(feature = "bedrock")]
    Bedrock { model_id: String },
}

/// Why learning could not proceed.
#[derive(Debug)]
pub enum LearnError {
    /// no credentials, or the region/model is not configured
    NotConfigured(String),
    /// the model answered, but not with something usable
    BadResponse(String),
    /// the call itself failed
    Transport(String),
}

impl std::fmt::Display for LearnError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            LearnError::NotConfigured(m) => write!(f, "not configured for learning: {m}"),
            LearnError::BadResponse(m) => write!(f, "unusable response: {m}"),
            LearnError::Transport(m) => write!(f, "call failed: {m}"),
        }
    }
}

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

impl Teacher {
    /// A teacher that returns a fixed proposal. Useful for tests, for replaying a recorded proposal, and for
    /// feeding in candidates from a source of your own.
    pub fn fixed(proposal: Proposal) -> Teacher {
        Teacher { kind: Kind::Fixed(proposal) }
    }

    /// A teacher backed by a **local** model behind an OpenAI-compatible endpoint.
    ///
    /// This is the preferred route. It needs no credentials, sends nothing off the machine, and costs nothing,
    /// which removes every objection to `learn` except quality. Works with llama.cpp, vLLM, LM Studio, Ollama —
    /// anything speaking `/v1/chat/completions`.
    ///
    /// A small model is a reasonable choice here *because* of the gate. Proposals are judged by the same
    /// deterministic MECE test as local discovery, so a weak model produces rejected candidates rather than a
    /// polluted vocabulary. The failure mode of choosing badly is wasted effort, not a wrong answer.
    ///
    /// ```no_run
    /// # #[cfg(feature = "paddock")]
    /// # fn demo() -> Result<(), Box<dyn std::error::Error>> {
    /// use steeldb::learn::Teacher;
    /// // a function-calling model small enough to run on a laptop
    /// let teacher = Teacher::local("http://localhost:11434/v1", "functiongemma-270m-it")?;
    /// # Ok(()) }
    /// ```
    #[cfg(feature = "paddock")]
    pub fn local(base_url: impl Into<String>, model: impl Into<String>) -> Result<Teacher, LearnError> {
        let base_url = base_url.into();
        if !base_url.starts_with("http") {
            return Err(LearnError::NotConfigured(format!(
                "base_url should be an http(s) endpoint, got {base_url:?}"
            )));
        }
        Ok(Teacher { kind: Kind::Local { base_url, model: model.into() } })
    }

    /// A teacher backed by a model served by Ollama on the default port.
    ///
    /// Shorthand for [`Teacher::local`] against `http://localhost:11434/v1`.
    #[cfg(feature = "paddock")]
    pub fn ollama(model: impl Into<String>) -> Result<Teacher, LearnError> {
        Teacher::local("http://localhost:11434/v1", model)
    }

    /// A teacher backed by Amazon Bedrock.
    ///
    /// Requires AWS credentials resolvable by the standard chain (environment, profile, or instance role) and
    /// `AWS_REGION`. Checked here so the failure is immediate and names what is missing.
    ///
    /// This costs money per call. The amount is small for vocabulary proposal — one request over a sample of
    /// documents — but it is a real charge and worth saying out loud.
    #[cfg(feature = "bedrock")]
    pub fn bedrock(model_id: impl Into<String>) -> Result<Teacher, LearnError> {
        if std::env::var("AWS_REGION").is_err() && std::env::var("AWS_DEFAULT_REGION").is_err() {
            return Err(LearnError::NotConfigured(
                "set AWS_REGION (or AWS_DEFAULT_REGION) to the region hosting the model".into(),
            ));
        }
        Ok(Teacher { kind: Kind::Bedrock { model_id: model_id.into() } })
    }

    /// Ask for categories the local discovery may have missed.
    ///
    /// Returns a [`Proposal`]; nothing changes until you [`SteelDb::adopt`] it.
    pub async fn propose_categories(&self, db: &SteelDb) -> Result<Proposal, LearnError> {
        // only the network paths read the corpus; without them this is deliberately unused
        let _ = db;
        match &self.kind {
            Kind::Fixed(p) => Ok(p.clone()),
            #[cfg(feature = "paddock")]
            Kind::Local { base_url, model } => {
                let cfg = crate::agent::config::ProviderConfig::Paddock {
                    base_url: base_url.clone(),
                    model: model.clone(),
                    api_key: None,
                };
                Self::propose_via(cfg, db).await
            }
            #[cfg(feature = "bedrock")]
            Kind::Bedrock { model_id } => {
                let cfg = crate::agent::config::ProviderConfig::Bedrock {
                    model_id: model_id.clone(),
                    region: std::env::var("AWS_REGION").ok(),
                };
                Self::propose_via(cfg, db).await
            }
        }
    }

    /// The shared path: whichever provider, the prompt, parsing and shaping are identical. Only the transport
    /// differs, which is what lets a local 270M model and a hosted frontier model be swapped freely.
    #[cfg(any(feature = "paddock", feature = "bedrock"))]
    async fn propose_via(
        cfg: crate::agent::config::ProviderConfig,
        db: &SteelDb,
    ) -> Result<Proposal, LearnError> {
        // label the proposal by its transport, so a logged proposal says which model produced it
        let label = match &cfg {
            #[cfg(feature = "paddock")]
            crate::agent::config::ProviderConfig::Paddock { model, .. } => format!("local:{model}"),
            #[cfg(feature = "bedrock")]
            crate::agent::config::ProviderConfig::Bedrock { model_id, .. } => format!("bedrock:{model_id}"),
            _ => "model".to_string(),
        };
        let provider = cfg.build().await.map_err(LearnError::Transport)?;
        let sample: Vec<String> = db.documents().iter().take(48).cloned().collect();
        let spec = crate::vocabulary::propose(provider.as_ref(), "documents", &sample)
            .await
            .map_err(LearnError::BadResponse)?;
        Ok(Proposal {
            source: label,
            candidates: spec
                .entity_facets
                .into_iter()
                .map(|f| Candidate { name: f.name, words: f.examples, rationale: f.description })
                .collect(),
        })
    }
}

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

    fn docs() -> Vec<String> {
        [
            "Morty Shade defeated Wallace Gale at Ecruteak City during the Indigo Invitational in 2025.",
            "Bea Strike defeated Iris Draco at Ecruteak City during the Indigo Invitational in 2025.",
            "A habitat survey recorded Aggron near Sootopolis City at an elevation of 1082 m.",
            "A habitat survey recorded Salamence near Sootopolis City at an elevation of 2369 m.",
            "Milotic is not permitted in Series 1 play for the 2025 season.",
        ]
        .iter()
        .map(|s| s.to_string())
        .collect()
    }

    #[test]
    fn a_proposal_changes_nothing_until_adopted() {
        let db = SteelDb::ingest(docs()).unwrap();
        let before = db.categories().len();
        let _p = Proposal {
            source: "test".into(),
            candidates: vec![Candidate {
                name: "trainer".into(),
                words: vec!["defeated".into(), "Shade".into()],
                rationale: "people who compete".into(),
            }],
        };
        // holding a proposal must not alter the database
        assert_eq!(db.categories().len(), before);
    }

    #[test]
    fn a_fixed_teacher_needs_no_credentials() {
        // deliberately driven without an async runtime: the offline path must not require tokio, which is an
        // optional dependency, so the default build can still test learning end to end
        let db = SteelDb::ingest(docs()).unwrap();
        let p = Proposal {
            source: "fixed".into(),
            candidates: vec![Candidate {
                name: "ruling".into(),
                words: vec!["permitted".into(), "Series".into(), "season".into()],
                rationale: "competition rules".into(),
            }],
        };
        let teacher = Teacher::fixed(p.clone());
        let got = block_on(teacher.propose_categories(&db)).unwrap();
        assert_eq!(got.candidates, p.candidates);
    }

    /// Drive a future to completion without a runtime. Sound here because the fixed teacher never yields.
    fn block_on<F: std::future::Future>(mut fut: F) -> F::Output {
        use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};
        fn noop(_: *const ()) {}
        fn clone(p: *const ()) -> RawWaker {
            RawWaker::new(p, &VTABLE)
        }
        static VTABLE: RawWakerVTable = RawWakerVTable::new(clone, noop, noop, noop);
        let waker = unsafe { Waker::from_raw(RawWaker::new(std::ptr::null(), &VTABLE)) };
        let mut cx = Context::from_waker(&waker);
        let mut fut = unsafe { std::pin::Pin::new_unchecked(&mut fut) };
        loop {
            match fut.as_mut().poll(&mut cx) {
                Poll::Ready(v) => return v,
                Poll::Pending => panic!("the fixed teacher must not yield"),
            }
        }
    }

    #[test]
    fn adoption_reports_a_verdict_per_candidate_and_can_reject() {
        let mut db = SteelDb::ingest(docs()).unwrap();
        let proposal = Proposal {
            source: "test".into(),
            candidates: vec![
                Candidate {
                    name: "ruling".into(),
                    words: vec!["permitted".into(), "Series".into()],
                    rationale: "rules".into(),
                },
                // deliberately a near-duplicate of the first, judged after it
                Candidate {
                    name: "ruling2".into(),
                    words: vec!["permitted".into(), "Series".into()],
                    rationale: "the same thing again".into(),
                },
            ],
        };
        let verdicts = db.adopt(&proposal);
        assert_eq!(verdicts.len(), 2, "one verdict per candidate");
        for v in &verdicts {
            assert!(!v.reason.is_empty(), "a rejection must be explicable: {v:?}");
        }
        // the gate is the point: a model cannot add what a deterministic test would refuse
        assert!(
            !verdicts[1].kept || verdicts[1].overlap < 0.99,
            "an exact duplicate should not be adopted unexamined: {:?}",
            verdicts[1]
        );
    }

    #[test]
    fn an_adopted_category_becomes_queryable() {
        let mut db = SteelDb::ingest(docs()).unwrap();
        let proposal = Proposal {
            source: "test".into(),
            candidates: vec![Candidate {
                name: "ruling".into(),
                words: vec!["permitted".into(), "season".into(), "Series".into()],
                rationale: "rules".into(),
            }],
        };
        let verdicts = db.adopt(&proposal);
        if verdicts[0].kept {
            let answer = db.query("ruling/*").expect("an adopted category must be queryable");
            assert!(!answer.is_empty(), "and must actually match documents");
        }
    }

    #[test]
    fn proposals_display_for_review_before_adoption() {
        let p = Proposal {
            source: "bedrock:test".into(),
            candidates: vec![Candidate {
                name: "trainer".into(),
                words: vec!["defeated".into()],
                rationale: "competitors".into(),
            }],
        };
        let shown = p.to_string();
        assert!(shown.contains("bedrock:test"), "{shown}");
        assert!(shown.contains("trainer"), "{shown}");
        assert!(shown.contains("competitors"), "the rationale must be reviewable: {shown}");
    }

    #[test]
    fn an_adopted_category_survives_into_an_artefact_and_is_followed_on_reload() {
        // The whole point of pairing `learn` with artefacts: the expensive, non-deterministic step runs once,
        // and every run afterwards is offline and identical. That only holds if what a teacher contributed is
        // actually written to the files that `ingest` follows — otherwise `learn` is a change you lose.
        let docs = docs();
        let mut db = SteelDb::ingest(docs.clone()).unwrap();
        let before: Vec<String> = db.categories().iter().map(|c| c.name.to_string()).collect();

        let proposal = Proposal {
            source: "test".into(),
            candidates: vec![Candidate {
                name: "ruling".into(),
                words: vec!["permitted".into(), "season".into(), "Series".into()],
                rationale: "competition rules".into(),
            }],
        };
        let verdicts = db.adopt(&proposal);
        if !verdicts[0].kept {
            // the gate is allowed to reject; there is then nothing to persist and nothing to assert
            return;
        }
        assert!(
            !before.contains(&"ruling".to_string()) && db.askable().contains(&"ruling/*".to_string()),
            "adoption should have added the category"
        );
        let expected = db.query("ruling/*").expect("adopted category must be queryable").len();

        let dir = std::env::temp_dir().join(format!("hsdb_learn_artifact_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        db.save(&dir).expect("save");

        // a fresh process, no teacher, no network: the reload must follow what learning produced
        let reloaded = SteelDb::ingest_using(docs, &dir).expect("reload");
        assert!(
            reloaded.askable().contains(&"ruling/*".to_string()),
            "the adopted category must come back: {:?}",
            reloaded.askable()
        );
        assert_eq!(
            reloaded.query("ruling/*").expect("still queryable").len(),
            expected,
            "and answer identically without the teacher"
        );
        assert_eq!(db.tags(), reloaded.tags(), "tag for tag");
        let _ = std::fs::remove_dir_all(&dir);
    }
}