Skip to main content

kcode_speech_classification/
lib.rs

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