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
21pub struct SpeechClassifier {
23 store: Store,
24}
25
26impl SpeechClassifier {
27 pub fn open(path: impl AsRef<Path>) -> Result<Self, Error> {
29 Ok(Self {
30 store: Store::open(path)?,
31 })
32 }
33
34 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 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 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}