kcode-speech-classification 0.1.2

#![deny(unsafe_code)]

mod ktool;
mod model;
mod protocol;
mod scoring;
mod store;

pub use ktool::{
    DELETE_TOOL, IDENTIFY_TOOL, KTOOLS, KtoolCall, KtoolError, TRAIN_TOOL, decode_ktool,
};
pub use model::{
    CandidateEvidence, Cefr, Cohort, DeleteOutcome, Error, FeatureRow, IdentifyEvidence,
    IdentifyOutcome, ObservationKey, TrainOutcome,
};
pub use protocol::{ProtocolError, ProtocolRequest, ProtocolResponse, ProtocolResult};

use std::path::Path;
use store::Store;

/// A thread-safe classifier backed by one SQLite database.
pub struct SpeechClassifier {
    store: Store,
}

impl SpeechClassifier {
    /// Opens or creates a classifier database and rebuilds current state from its event log.
    pub fn open(path: impl AsRef<Path>) -> Result<Self, Error> {
        Ok(Self {
            store: Store::open(path)?,
        })
    }

    /// Scores an observation and atomically retains it when the threshold is met.
    pub fn identify(
        &self,
        key: ObservationKey,
        cohort: Cohort,
        row: FeatureRow,
        threshold: f64,
    ) -> Result<IdentifyOutcome, Error> {
        key.validate()?;
        cohort.validate()?;
        row.validate()?;
        if !threshold.is_finite() {
            return Err(Error::validation("threshold", "must be finite"));
        }
        self.store.identify(&key, &cohort, &row, threshold)
    }

    /// Adds known training data or atomically corrects the current assignment.
    pub fn train(
        &self,
        key: ObservationKey,
        cohort: Cohort,
        row: FeatureRow,
        speaker_id: String,
    ) -> Result<TrainOutcome, Error> {
        key.validate()?;
        cohort.validate()?;
        row.validate()?;
        if speaker_id.trim().is_empty() {
            return Err(Error::validation("speaker_id", "must not be empty"));
        }
        self.store.train(&key, &cohort, &row, speaker_id.as_str())
    }

    /// Idempotently deletes the current observation for a key.
    pub fn delete(&self, key: ObservationKey) -> Result<DeleteOutcome, Error> {
        key.validate()?;
        self.store.delete(&key)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use rusqlite::Connection;
    use std::fs;
    use std::path::{Path, PathBuf};
    use std::sync::atomic::{AtomicU64, Ordering};

    static NEXT_PATH: AtomicU64 = AtomicU64::new(0);

    fn database_path(label: &str) -> PathBuf {
        std::env::temp_dir().join(format!(
            "kcode-speech-classification-lib-{}-{label}-{}.sqlite3",
            std::process::id(),
            NEXT_PATH.fetch_add(1, Ordering::Relaxed)
        ))
    }

    fn key(object_id: &str, piece_index: u32) -> ObservationKey {
        ObservationKey {
            object_id: object_id.to_owned(),
            piece_index,
        }
    }

    fn cohort() -> Cohort {
        Cohort {
            provider: "google".to_owned(),
            model: "gemini-example".to_owned(),
            prompt_version: "prompt-1".to_owned(),
            schema_version: "features-1".to_owned(),
            primary_language: "eng".to_owned(),
        }
    }

    fn row(age: f64, accent: &str) -> FeatureRow {
        FeatureRow {
            accent_variety: accent.to_owned(),
            perceived_age: age,
            vocal_gender_presentation: 45.0,
            median_f0_hz: 155.0 + age,
            formant_dispersion_hz: 1100.0,
            vai: 1.1,
            hypernasality: 0.5,
            creaky_phonation_percent: 8.0,
            rhotic_realization: format!("{accent}-rhotic"),
            word_initial_stressed_prevocalic_t_vot_ms: 62.0,
            breathiness: 0.2,
            roughness: 0.1,
            f0_pitch_span_semitones: 9.0,
            articulation_rate_syllables_per_second: 4.2,
            npvi_v: 48.0,
            cefr: Cefr::C1,
            foreign_accentedness: 2.0,
            unstressed_vowel_reduction_percent: 72.0,
            lateral_realization: "alveolar".to_owned(),
            filled_pauses_per_100_words: 2.5,
            s_realization: "alveolar".to_owned(),
            lexical_stress_accuracy_percent: 92.0,
            monophthongization_percent: 4.0,
            consonant_cluster_reduction_percent: 3.0,
        }
    }

    fn event_count(path: &Path) -> i64 {
        let connection = Connection::open(path).unwrap();
        connection
            .query_row("SELECT COUNT(*) FROM events", [], |record| record.get(0))
            .unwrap()
    }

    fn sidecar(path: &Path, suffix: &str) -> PathBuf {
        let mut value = path.as_os_str().to_os_string();
        value.push(suffix);
        PathBuf::from(value)
    }

    fn remove_database(path: &Path) {
        for candidate in [
            path.to_path_buf(),
            sidecar(path, "-wal"),
            sidecar(path, "-shm"),
        ] {
            let _ = fs::remove_file(candidate);
        }
    }

    #[test]
    fn public_identify_is_cohort_isolated_thresholded_and_idempotent() {
        let path = database_path("identify");
        let classifier = SpeechClassifier::open(&path).unwrap();
        let cohort = cohort();

        classifier
            .train(
                key("a-1", 0),
                cohort.clone(),
                row(30.0, "alpha"),
                "a".into(),
            )
            .unwrap();
        classifier
            .train(
                key("a-2", 0),
                cohort.clone(),
                row(32.0, "alpha"),
                "a".into(),
            )
            .unwrap();
        classifier
            .train(key("b-1", 0), cohort.clone(), row(70.0, "beta"), "b".into())
            .unwrap();
        classifier
            .train(key("b-2", 0), cohort.clone(), row(72.0, "beta"), "b".into())
            .unwrap();
        assert_eq!(event_count(&path), 4);

        let query = row(31.0, "alpha");
        let rejected = classifier
            .identify(key("rejected", 0), cohort.clone(), query.clone(), f64::MAX)
            .unwrap();
        assert_eq!(rejected.speaker_id, None);
        assert!(rejected.evidence.is_some());
        assert_eq!(event_count(&path), 4);

        let mut separate_cohort = cohort.clone();
        separate_cohort.prompt_version = "prompt-2".to_owned();
        let isolated = classifier
            .identify(
                key("isolated", 0),
                separate_cohort,
                query.clone(),
                -1_000_000.0,
            )
            .unwrap();
        assert_eq!(isolated.speaker_id, None);
        assert_eq!(isolated.evidence, None);
        assert_eq!(event_count(&path), 4);

        let accepted = classifier
            .identify(
                key("accepted", 0),
                cohort.clone(),
                query.clone(),
                -1_000_000.0,
            )
            .unwrap();
        assert_eq!(accepted.speaker_id.as_deref(), Some("a"));
        assert_eq!(event_count(&path), 5);

        let repeated = classifier
            .identify(
                key("accepted", 0),
                cohort.clone(),
                query.clone(),
                -1_000_000.0,
            )
            .unwrap();
        assert_eq!(repeated.speaker_id.as_deref(), Some("a"));
        assert_eq!(event_count(&path), 5);

        let conflict = classifier.identify(
            key("accepted", 0),
            cohort,
            row(60.0, "different"),
            -1_000_000.0,
        );
        assert!(matches!(conflict, Err(Error::Conflict { .. })));
        assert_eq!(event_count(&path), 5);

        drop(classifier);
        remove_database(&path);
    }

    #[test]
    fn public_boundary_rejects_malformed_keys_speakers_and_thresholds() {
        let path = database_path("validation");
        let classifier = SpeechClassifier::open(&path).unwrap();

        assert!(matches!(
            classifier.delete(key(" ", 0)),
            Err(Error::Validation { .. })
        ));
        assert!(matches!(
            classifier.train(key("known", 0), cohort(), row(30.0, "alpha"), " ".into()),
            Err(Error::Validation { .. })
        ));
        assert!(matches!(
            classifier.identify(key("query", 0), cohort(), row(30.0, "alpha"), f64::NAN),
            Err(Error::Validation { .. })
        ));
        assert_eq!(event_count(&path), 0);

        drop(classifier);
        remove_database(&path);
    }
}