Skip to main content

kcode_speech_classification/
lib.rs

1#![deny(unsafe_code)]
2
3mod model;
4mod protocol;
5mod scoring;
6mod store;
7
8pub use model::{
9    CandidateEvidence, Cefr, Cohort, DeleteOutcome, Error, FeatureRow, IdentifyEvidence,
10    IdentifyOutcome, ObservationKey, TrainOutcome,
11};
12pub use protocol::{ProtocolError, ProtocolRequest, ProtocolResponse, ProtocolResult};
13
14use std::path::Path;
15use store::Store;
16
17/// A thread-safe classifier backed by one SQLite database.
18pub struct SpeechClassifier {
19    store: Store,
20}
21
22impl SpeechClassifier {
23    /// Opens or creates a classifier database and rebuilds current state from its event log.
24    pub fn open(path: impl AsRef<Path>) -> Result<Self, Error> {
25        Ok(Self {
26            store: Store::open(path)?,
27        })
28    }
29
30    /// Scores an observation and atomically retains it when the threshold is met.
31    pub fn identify(
32        &self,
33        key: ObservationKey,
34        cohort: Cohort,
35        row: FeatureRow,
36        threshold: f64,
37    ) -> Result<IdentifyOutcome, Error> {
38        key.validate()?;
39        cohort.validate()?;
40        row.validate()?;
41        if !threshold.is_finite() {
42            return Err(Error::validation("threshold", "must be finite"));
43        }
44        self.store.identify(&key, &cohort, &row, threshold)
45    }
46
47    /// Adds known training data or atomically corrects the current assignment.
48    pub fn train(
49        &self,
50        key: ObservationKey,
51        cohort: Cohort,
52        row: FeatureRow,
53        speaker_id: String,
54    ) -> Result<TrainOutcome, Error> {
55        key.validate()?;
56        cohort.validate()?;
57        row.validate()?;
58        if speaker_id.trim().is_empty() {
59            return Err(Error::validation("speaker_id", "must not be empty"));
60        }
61        self.store.train(&key, &cohort, &row, speaker_id.as_str())
62    }
63
64    /// Idempotently deletes the current observation for a key.
65    pub fn delete(&self, key: ObservationKey) -> Result<DeleteOutcome, Error> {
66        key.validate()?;
67        self.store.delete(&key)
68    }
69}
70
71#[cfg(test)]
72mod tests {
73    use super::*;
74    use rusqlite::Connection;
75    use std::fs;
76    use std::path::{Path, PathBuf};
77    use std::sync::atomic::{AtomicU64, Ordering};
78
79    static NEXT_PATH: AtomicU64 = AtomicU64::new(0);
80
81    fn database_path(label: &str) -> PathBuf {
82        std::env::temp_dir().join(format!(
83            "kcode-speech-classification-lib-{}-{label}-{}.sqlite3",
84            std::process::id(),
85            NEXT_PATH.fetch_add(1, Ordering::Relaxed)
86        ))
87    }
88
89    fn key(object_id: &str, piece_index: u32) -> ObservationKey {
90        ObservationKey {
91            object_id: object_id.to_owned(),
92            piece_index,
93        }
94    }
95
96    fn cohort() -> Cohort {
97        Cohort {
98            provider: "google".to_owned(),
99            model: "gemini-example".to_owned(),
100            prompt_version: "prompt-1".to_owned(),
101            schema_version: "features-1".to_owned(),
102            primary_language: "eng".to_owned(),
103        }
104    }
105
106    fn row(age: f64, accent: &str) -> FeatureRow {
107        FeatureRow {
108            accent_variety: accent.to_owned(),
109            perceived_age: age,
110            vocal_gender_presentation: 45.0,
111            median_f0_hz: 155.0 + age,
112            formant_dispersion_hz: 1100.0,
113            vai: 1.1,
114            hypernasality: 0.5,
115            creaky_phonation_percent: 8.0,
116            rhotic_realization: format!("{accent}-rhotic"),
117            word_initial_stressed_prevocalic_t_vot_ms: 62.0,
118            breathiness: 0.2,
119            roughness: 0.1,
120            f0_pitch_span_semitones: 9.0,
121            articulation_rate_syllables_per_second: 4.2,
122            npvi_v: 48.0,
123            cefr: Cefr::C1,
124            foreign_accentedness: 2.0,
125            unstressed_vowel_reduction_percent: 72.0,
126            lateral_realization: "alveolar".to_owned(),
127            filled_pauses_per_100_words: 2.5,
128            s_realization: "alveolar".to_owned(),
129            lexical_stress_accuracy_percent: 92.0,
130            monophthongization_percent: 4.0,
131            consonant_cluster_reduction_percent: 3.0,
132        }
133    }
134
135    fn event_count(path: &Path) -> i64 {
136        let connection = Connection::open(path).unwrap();
137        connection
138            .query_row("SELECT COUNT(*) FROM events", [], |record| record.get(0))
139            .unwrap()
140    }
141
142    fn sidecar(path: &Path, suffix: &str) -> PathBuf {
143        let mut value = path.as_os_str().to_os_string();
144        value.push(suffix);
145        PathBuf::from(value)
146    }
147
148    fn remove_database(path: &Path) {
149        for candidate in [
150            path.to_path_buf(),
151            sidecar(path, "-wal"),
152            sidecar(path, "-shm"),
153        ] {
154            let _ = fs::remove_file(candidate);
155        }
156    }
157
158    #[test]
159    fn public_identify_is_cohort_isolated_thresholded_and_idempotent() {
160        let path = database_path("identify");
161        let classifier = SpeechClassifier::open(&path).unwrap();
162        let cohort = cohort();
163
164        classifier
165            .train(
166                key("a-1", 0),
167                cohort.clone(),
168                row(30.0, "alpha"),
169                "a".into(),
170            )
171            .unwrap();
172        classifier
173            .train(
174                key("a-2", 0),
175                cohort.clone(),
176                row(32.0, "alpha"),
177                "a".into(),
178            )
179            .unwrap();
180        classifier
181            .train(key("b-1", 0), cohort.clone(), row(70.0, "beta"), "b".into())
182            .unwrap();
183        classifier
184            .train(key("b-2", 0), cohort.clone(), row(72.0, "beta"), "b".into())
185            .unwrap();
186        assert_eq!(event_count(&path), 4);
187
188        let query = row(31.0, "alpha");
189        let rejected = classifier
190            .identify(key("rejected", 0), cohort.clone(), query.clone(), f64::MAX)
191            .unwrap();
192        assert_eq!(rejected.speaker_id, None);
193        assert!(rejected.evidence.is_some());
194        assert_eq!(event_count(&path), 4);
195
196        let mut separate_cohort = cohort.clone();
197        separate_cohort.prompt_version = "prompt-2".to_owned();
198        let isolated = classifier
199            .identify(
200                key("isolated", 0),
201                separate_cohort,
202                query.clone(),
203                -1_000_000.0,
204            )
205            .unwrap();
206        assert_eq!(isolated.speaker_id, None);
207        assert_eq!(isolated.evidence, None);
208        assert_eq!(event_count(&path), 4);
209
210        let accepted = classifier
211            .identify(
212                key("accepted", 0),
213                cohort.clone(),
214                query.clone(),
215                -1_000_000.0,
216            )
217            .unwrap();
218        assert_eq!(accepted.speaker_id.as_deref(), Some("a"));
219        assert_eq!(event_count(&path), 5);
220
221        let repeated = classifier
222            .identify(
223                key("accepted", 0),
224                cohort.clone(),
225                query.clone(),
226                -1_000_000.0,
227            )
228            .unwrap();
229        assert_eq!(repeated.speaker_id.as_deref(), Some("a"));
230        assert_eq!(event_count(&path), 5);
231
232        let conflict = classifier.identify(
233            key("accepted", 0),
234            cohort,
235            row(60.0, "different"),
236            -1_000_000.0,
237        );
238        assert!(matches!(conflict, Err(Error::Conflict { .. })));
239        assert_eq!(event_count(&path), 5);
240
241        drop(classifier);
242        remove_database(&path);
243    }
244
245    #[test]
246    fn public_boundary_rejects_malformed_keys_speakers_and_thresholds() {
247        let path = database_path("validation");
248        let classifier = SpeechClassifier::open(&path).unwrap();
249
250        assert!(matches!(
251            classifier.delete(key(" ", 0)),
252            Err(Error::Validation { .. })
253        ));
254        assert!(matches!(
255            classifier.train(key("known", 0), cohort(), row(30.0, "alpha"), " ".into()),
256            Err(Error::Validation { .. })
257        ));
258        assert!(matches!(
259            classifier.identify(key("query", 0), cohort(), row(30.0, "alpha"), f64::NAN),
260            Err(Error::Validation { .. })
261        ));
262        assert_eq!(event_count(&path), 0);
263
264        drop(classifier);
265        remove_database(&path);
266    }
267}