use std::{
collections::{BTreeSet, HashMap, HashSet},
sync::Arc,
};
use anyhow::{Context, ensure};
use chrono::{DateTime, Utc};
pub use kcode_speech_classification::{Cefr, FeatureRow, ObservationKey};
use kcode_speech_classification::{Cohort, IdentifyEvidence, SpeechClassifier, TrainOutcome};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use uuid::Uuid;
pub const CLASSIFIER_PROVIDER: &str = "google";
pub const CLASSIFIER_MODEL: &str = "gemini-3.1-pro-preview";
pub const CLASSIFIER_PROMPT_VERSION: &str = "gemini-speaker-prompt-v0.1";
pub const CLASSIFIER_SCHEMA_VERSION: &str = "gemini-speaker-features-v0.1";
const READ_ONLY_IDENTIFY_THRESHOLD: f64 = 1e308;
const FEATURE_FIELDS: [&str; 24] = [
"accent_variety",
"perceived_age",
"vocal_gender_presentation",
"median_f0_hz",
"formant_dispersion_hz",
"vai",
"hypernasality",
"creaky_phonation_percent",
"rhotic_realization",
"word_initial_stressed_prevocalic_t_vot_ms",
"breathiness",
"roughness",
"f0_pitch_span_semitones",
"articulation_rate_syllables_per_second",
"npvi_v",
"cefr",
"foreign_accentedness",
"unstressed_vowel_reduction_percent",
"lateral_realization",
"filled_pauses_per_100_words",
"s_realization",
"lexical_stress_accuracy_percent",
"monophthongization_percent",
"consonant_cluster_reduction_percent",
];
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct ParsedUtterance {
pub speaker: String,
pub language: String,
pub original_text: String,
pub english_translation: String,
pub corrected_natural_text: Option<String>,
pub coaching: Vec<String>,
pub annotations: Vec<String>,
}
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct ParsedSpeaker {
pub local_label: String,
pub primary_language: String,
pub feature_row: FeatureRow,
}
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct ParsedChunk {
pub utterances: Vec<ParsedUtterance>,
pub notes: Vec<String>,
pub clip_valid: bool,
pub clip_validity_reason: Option<String>,
pub speakers: Vec<ParsedSpeaker>,
}
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
pub struct CandidateMapping {
pub full_name: String,
pub cost: f64,
pub confidence: f64,
pub runner_up_full_name: Option<String>,
pub runner_up_cost: Option<f64>,
pub background_population_cost: f64,
}
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
pub struct CorrectionObservation {
pub local_label: String,
pub speaker_ordinal: u32,
pub observation_key: ObservationKey,
pub candidate: Option<CandidateMapping>,
pub identified_full_name: Option<String>,
pub confirmed_full_name: Option<String>,
}
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
pub struct CorrectionChunk {
pub chunk_index: usize,
pub chunk_count: usize,
pub audio_start_ms: u64,
pub audio_end_ms: u64,
pub raw_gemini_response: String,
pub parsed: ParsedChunk,
pub observations: Vec<CorrectionObservation>,
pub clean: bool,
}
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum ConfirmationState {
Unconfirmed,
AutomaticallyTrained,
Confirmed,
}
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
pub struct CorrectionPacket {
pub recording_id: Uuid,
pub user_id: String,
pub sha256: String,
pub original_filename: String,
pub size_bytes: u64,
pub recorded_at: DateTime<Utc>,
pub clean: bool,
pub chunk_count: usize,
pub chunks: Vec<CorrectionChunk>,
pub confirmation_state: ConfirmationState,
}
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct ObservationConfirmation {
pub observation_key: ObservationKey,
pub confirmed_full_name: String,
}
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct RecordingConfirmation {
pub recording_id: Uuid,
pub observations: Vec<ObservationConfirmation>,
}
#[derive(Clone)]
pub(crate) struct ClassificationContext {
pub(crate) recording_id: Uuid,
pub(crate) user_id: String,
pub(crate) sha256: String,
pub(crate) original_filename: String,
pub(crate) size_bytes: u64,
pub(crate) recorded_at: DateTime<Utc>,
pub(crate) classifier: Arc<SpeechClassifier>,
}
pub(crate) fn parse_and_validate_chunk(
response: &str,
chunk_duration_seconds: f64,
) -> anyhow::Result<ParsedChunk> {
ensure!(
!response.trim().is_empty(),
"GPT parser returned an empty response"
);
let value: Value =
serde_json::from_str(response).context("GPT parser response is not one JSON value")?;
validate_feature_row_shapes(&value)?;
let mut parsed: ParsedChunk =
serde_json::from_value(value).context("GPT parser JSON has an invalid typed shape")?;
validate_parsed_chunk(&mut parsed, chunk_duration_seconds)?;
Ok(parsed)
}
pub(crate) fn validate_parsed_chunk(
parsed: &mut ParsedChunk,
chunk_duration_seconds: f64,
) -> anyhow::Result<()> {
ensure!(
chunk_duration_seconds.is_finite() && chunk_duration_seconds > 0.0,
"chunk duration must be finite and positive"
);
ensure!(
parsed.notes.iter().all(|note| !note.trim().is_empty()),
"chunk notes must not contain empty entries"
);
match (parsed.clip_valid, parsed.clip_validity_reason.as_deref()) {
(true, None) => {}
(false, Some(reason)) if !reason.trim().is_empty() => {}
(true, Some(_)) => anyhow::bail!("a valid clip must not carry an invalidity reason"),
(false, _) => anyhow::bail!("an invalid clip requires a brief reason"),
}
parsed
.speakers
.sort_by(|left, right| left.local_label.cmp(&right.local_label));
let mut labels = HashSet::new();
for speaker in &parsed.speakers {
ensure!(
!speaker.local_label.trim().is_empty(),
"speaker labels must not be empty"
);
ensure!(
labels.insert(speaker.local_label.clone()),
"duplicate speaker label {:?}",
speaker.local_label
);
validate_iso_639_3(&speaker.primary_language)
.with_context(|| format!("speaker {} primary language", speaker.local_label))?;
validate_feature_row(&speaker.feature_row)
.with_context(|| format!("speaker {} feature row", speaker.local_label))?;
}
let mut referenced = HashSet::new();
for (index, utterance) in parsed.utterances.iter().enumerate() {
ensure!(
labels.contains(&utterance.speaker),
"utterance {index} references unknown speaker {:?}",
utterance.speaker
);
referenced.insert(utterance.speaker.clone());
validate_iso_639_3(&utterance.language)
.with_context(|| format!("utterance {index} language"))?;
ensure!(
!utterance.original_text.trim().is_empty(),
"utterance {index} original text must not be empty"
);
if utterance.language == "eng" {
ensure!(
utterance.english_translation.is_empty(),
"English utterance {index} must use an empty translation"
);
} else {
ensure!(
!utterance.english_translation.trim().is_empty(),
"non-English utterance {index} requires a complete English translation"
);
}
if let Some(corrected) = utterance.corrected_natural_text.as_deref() {
ensure!(
!corrected.trim().is_empty(),
"utterance {index} corrected text must not be empty"
);
}
ensure!(
utterance
.coaching
.iter()
.chain(&utterance.annotations)
.all(|entry| !entry.trim().is_empty()),
"utterance {index} notes must not contain empty entries"
);
}
ensure!(
parsed
.speakers
.iter()
.all(|speaker| referenced.contains(&speaker.local_label)),
"every speaker row must be referenced by at least one utterance"
);
if parsed.clip_valid {
ensure!(
!parsed.speakers.is_empty(),
"a valid clip must contain at least one speaker row"
);
}
Ok(())
}
pub(crate) fn classify_speakers(
context: &ClassificationContext,
chunk_index: usize,
parsed: &ParsedChunk,
) -> anyhow::Result<(Vec<CorrectionObservation>, bool)> {
let mut observations = Vec::with_capacity(parsed.speakers.len());
for (ordinal, speaker) in parsed.speakers.iter().enumerate() {
let speaker_ordinal = u32::try_from(ordinal)
.context("chunk has more speakers than the key schema supports")?;
let cohort = cohort(&speaker.primary_language);
let probe_key = ObservationKey {
object_id: probe_object_id(context.recording_id, chunk_index),
piece_index: speaker_ordinal,
};
let outcome = context
.classifier
.identify(
probe_key.clone(),
cohort,
speaker.feature_row.clone(),
READ_ONLY_IDENTIFY_THRESHOLD,
)
.with_context(|| {
format!(
"read-only identity scoring failed for chunk {chunk_index} speaker {}",
speaker.local_label
)
})?;
if outcome.speaker_id.is_some() {
context
.classifier
.delete(probe_key)
.context("removing an unexpectedly accepted read-only probe")?;
}
let candidate = outcome.evidence.as_ref().map(candidate_mapping);
let identified_full_name = candidate.as_ref().map(|value| value.full_name.clone());
observations.push(CorrectionObservation {
local_label: speaker.local_label.clone(),
speaker_ordinal,
observation_key: ObservationKey {
object_id: training_object_id(context.recording_id, chunk_index),
piece_index: speaker_ordinal,
},
candidate,
identified_full_name,
confirmed_full_name: None,
});
}
let clean = chunk_is_clean(parsed.clip_valid, &observations);
Ok((observations, clean))
}
pub(crate) fn unclassified_observations(
recording_id: Uuid,
chunk_index: usize,
parsed: &ParsedChunk,
) -> anyhow::Result<Vec<CorrectionObservation>> {
parsed
.speakers
.iter()
.enumerate()
.map(|(ordinal, speaker)| {
let speaker_ordinal = u32::try_from(ordinal)
.context("chunk has more speakers than the key schema supports")?;
Ok(CorrectionObservation {
local_label: speaker.local_label.clone(),
speaker_ordinal,
observation_key: ObservationKey {
object_id: training_object_id(recording_id, chunk_index),
piece_index: speaker_ordinal,
},
candidate: None,
identified_full_name: None,
confirmed_full_name: None,
})
})
.collect()
}
pub(crate) fn chunk_is_clean(clip_valid: bool, observations: &[CorrectionObservation]) -> bool {
if !clip_valid || observations.is_empty() {
return false;
}
let mut names = HashSet::new();
observations.iter().all(|observation| {
observation.candidate.as_ref().is_some_and(|candidate| {
candidate.confidence > 0.0
&& candidate.cost < candidate.background_population_cost
&& candidate
.runner_up_cost
.is_some_and(|cost| cost > candidate.background_population_cost)
&& names.insert(candidate.full_name.as_str())
})
})
}
pub(crate) fn build_packet(
context: &ClassificationContext,
chunks: Vec<CorrectionChunk>,
) -> anyhow::Result<CorrectionPacket> {
ensure!(!chunks.is_empty(), "correction packet has no chunks");
let chunk_count = chunks.len();
ensure!(
chunks.iter().enumerate().all(|(index, chunk)| {
chunk.chunk_index == index
&& chunk.chunk_count == chunk_count
&& chunk.audio_end_ms > chunk.audio_start_ms
&& chunk.observations.len() == chunk.parsed.speakers.len()
}),
"correction packet chunks are not one complete chronological plan"
);
let clean = chunks.iter().all(|chunk| chunk.clean);
Ok(CorrectionPacket {
recording_id: context.recording_id,
user_id: context.user_id.clone(),
sha256: context.sha256.clone(),
original_filename: context.original_filename.clone(),
size_bytes: context.size_bytes,
recorded_at: context.recorded_at,
clean,
chunk_count,
chunks,
confirmation_state: ConfirmationState::Unconfirmed,
})
}
pub(crate) fn train_clean_packet(
classifier: &SpeechClassifier,
packet: &mut CorrectionPacket,
) -> anyhow::Result<()> {
if !packet.clean {
ensure!(
packet.confirmation_state == ConfirmationState::Unconfirmed,
"unclean packet unexpectedly claims retained training"
);
return Ok(());
}
let mut added = Vec::new();
for chunk in &packet.chunks {
for observation in &chunk.observations {
let speaker = speaker_for_observation(chunk, observation)?;
let full_name = observation
.identified_full_name
.as_deref()
.context("clean observation omitted its identified full name")?;
match classifier.train(
observation.observation_key.clone(),
cohort(&speaker.primary_language),
speaker.feature_row.clone(),
full_name.to_owned(),
) {
Ok(TrainOutcome::Added) => added.push(observation.observation_key.clone()),
Ok(TrainOutcome::Unchanged | TrainOutcome::Corrected) => {}
Err(error) => {
let rollback_errors = rollback_added(classifier, &added);
if rollback_errors.is_empty() {
anyhow::bail!("automatic identity training failed: {error}");
}
anyhow::bail!(
"automatic identity training failed: {error}; rollback also failed: {}",
rollback_errors.join("; ")
);
}
}
}
}
packet.confirmation_state = ConfirmationState::AutomaticallyTrained;
Ok(())
}
pub(crate) fn validate_confirmation_coverage(
packet: &CorrectionPacket,
confirmation: &RecordingConfirmation,
) -> Result<(), String> {
if confirmation.recording_id != packet.recording_id {
return Err("Confirmation recording ID does not match the packet.".into());
}
let known = packet
.chunks
.iter()
.flat_map(|chunk| &chunk.observations)
.map(|observation| key_tuple(&observation.observation_key))
.collect::<BTreeSet<_>>();
if known.is_empty() {
return Err("The correction packet contains no speaker observations.".into());
}
let mut supplied = BTreeSet::new();
for observation in &confirmation.observations {
if observation.confirmed_full_name.trim().is_empty()
|| observation.confirmed_full_name.chars().count() > 512
{
return Err("Confirmed full names must contain between 1 and 512 characters.".into());
}
if !supplied.insert(key_tuple(&observation.observation_key)) {
return Err("Confirmation contains a duplicate observation key.".into());
}
}
if supplied != known {
return Err(
"Confirmation must cover every known observation exactly once, with no extras.".into(),
);
}
Ok(())
}
pub(crate) fn apply_confirmations(
classifier: &SpeechClassifier,
packet: &mut CorrectionPacket,
confirmation: &RecordingConfirmation,
) -> anyhow::Result<()> {
validate_confirmation_coverage(packet, confirmation).map_err(anyhow::Error::msg)?;
let assignments = confirmation
.observations
.iter()
.map(|entry| {
(
key_tuple(&entry.observation_key),
entry.confirmed_full_name.trim().to_owned(),
)
})
.collect::<HashMap<_, _>>();
#[derive(Clone)]
struct Target {
chunk_position: usize,
observation_position: usize,
key: ObservationKey,
cohort: Cohort,
row: FeatureRow,
new_name: String,
old_name: Option<String>,
}
let mut targets = Vec::new();
for (chunk_position, chunk) in packet.chunks.iter().enumerate() {
for (observation_position, observation) in chunk.observations.iter().enumerate() {
let speaker = speaker_for_observation(chunk, observation)?;
targets.push(Target {
chunk_position,
observation_position,
key: observation.observation_key.clone(),
cohort: cohort(&speaker.primary_language),
row: speaker.feature_row.clone(),
new_name: assignments
.get(&key_tuple(&observation.observation_key))
.context("validated confirmation assignment disappeared")?
.clone(),
old_name: retained_name(packet.confirmation_state, observation),
});
}
}
let mut applied = Vec::<(Target, TrainOutcome)>::new();
for target in targets {
match classifier.train(
target.key.clone(),
target.cohort.clone(),
target.row.clone(),
target.new_name.clone(),
) {
Ok(outcome) => applied.push((target, outcome)),
Err(error) => {
let mut rollback_errors = Vec::new();
for (previous, outcome) in applied.iter().rev() {
let rollback = if let Some(old_name) = &previous.old_name {
classifier
.train(
previous.key.clone(),
previous.cohort.clone(),
previous.row.clone(),
old_name.clone(),
)
.map(|_| ())
} else if *outcome == TrainOutcome::Added {
classifier.delete(previous.key.clone()).map(|_| ())
} else {
Ok(())
};
if let Err(rollback_error) = rollback {
rollback_errors.push(rollback_error.to_string());
}
}
if rollback_errors.is_empty() {
anyhow::bail!("applying identity confirmations failed: {error}");
}
anyhow::bail!(
"applying identity confirmations failed: {error}; rollback also failed: {}",
rollback_errors.join("; ")
);
}
}
}
for (target, _) in applied {
packet.chunks[target.chunk_position].observations[target.observation_position]
.confirmed_full_name = Some(target.new_name);
}
packet.confirmation_state = ConfirmationState::Confirmed;
Ok(())
}
pub(crate) fn restore_packet_training(
classifier: &SpeechClassifier,
packet: &CorrectionPacket,
) -> Vec<String> {
let mut errors = Vec::new();
for chunk in packet.chunks.iter().rev() {
for observation in chunk.observations.iter().rev() {
let result = match retained_name(packet.confirmation_state, observation) {
Some(name) => speaker_for_observation(chunk, observation).and_then(|speaker| {
classifier
.train(
observation.observation_key.clone(),
cohort(&speaker.primary_language),
speaker.feature_row.clone(),
name,
)
.map(|_| ())
.map_err(anyhow::Error::from)
}),
None => classifier
.delete(observation.observation_key.clone())
.map(|_| ())
.map_err(anyhow::Error::from),
};
if let Err(error) = result {
errors.push(error.to_string());
}
}
}
errors
}
pub(crate) fn training_object_id(recording_id: Uuid, chunk_index: usize) -> String {
format!("kcode-audio-ingress/recording/{recording_id}/chunk/{chunk_index}")
}
fn probe_object_id(recording_id: Uuid, chunk_index: usize) -> String {
format!("kcode-audio-ingress/probe/{recording_id}/chunk/{chunk_index}")
}
fn cohort(primary_language: &str) -> Cohort {
Cohort {
provider: CLASSIFIER_PROVIDER.into(),
model: CLASSIFIER_MODEL.into(),
prompt_version: CLASSIFIER_PROMPT_VERSION.into(),
schema_version: CLASSIFIER_SCHEMA_VERSION.into(),
primary_language: primary_language.into(),
}
}
fn candidate_mapping(evidence: &IdentifyEvidence) -> CandidateMapping {
CandidateMapping {
full_name: evidence.best.speaker_id.clone(),
cost: evidence.best.cost,
confidence: evidence.confidence_score,
runner_up_full_name: evidence
.runner_up
.as_ref()
.map(|candidate| candidate.speaker_id.clone()),
runner_up_cost: evidence.runner_up.as_ref().map(|candidate| candidate.cost),
background_population_cost: evidence.background_population_cost,
}
}
fn speaker_for_observation<'a>(
chunk: &'a CorrectionChunk,
observation: &CorrectionObservation,
) -> anyhow::Result<&'a ParsedSpeaker> {
let speaker = chunk
.parsed
.speakers
.get(observation.speaker_ordinal as usize)
.context("observation ordinal is outside the parsed speaker rows")?;
ensure!(
speaker.local_label == observation.local_label,
"observation label does not match its parsed speaker row"
);
Ok(speaker)
}
fn retained_name(state: ConfirmationState, observation: &CorrectionObservation) -> Option<String> {
match state {
ConfirmationState::Unconfirmed => None,
ConfirmationState::AutomaticallyTrained => observation.identified_full_name.clone(),
ConfirmationState::Confirmed => observation.confirmed_full_name.clone(),
}
}
fn rollback_added(classifier: &SpeechClassifier, keys: &[ObservationKey]) -> Vec<String> {
let mut errors = Vec::new();
for key in keys.iter().rev() {
if let Err(error) = classifier.delete(key.clone()) {
errors.push(error.to_string());
}
}
errors
}
fn key_tuple(key: &ObservationKey) -> (String, u32) {
(key.object_id.clone(), key.piece_index)
}
fn validate_feature_row_shapes(value: &Value) -> anyhow::Result<()> {
let speakers = value
.get("speakers")
.and_then(Value::as_array)
.context("GPT parser JSON omitted the speakers array")?;
let expected = FEATURE_FIELDS.iter().copied().collect::<BTreeSet<_>>();
for (index, speaker) in speakers.iter().enumerate() {
let row = speaker
.get("feature_row")
.and_then(Value::as_object)
.with_context(|| format!("speaker row {index} omitted feature_row"))?;
let actual = row.keys().map(String::as_str).collect::<BTreeSet<_>>();
ensure!(
actual == expected,
"speaker row {index} must contain exactly the 24 feature fields"
);
}
Ok(())
}
fn validate_iso_639_3(value: &str) -> anyhow::Result<()> {
ensure!(
value.len() == 3 && value.bytes().all(|byte| byte.is_ascii_lowercase()),
"must be a lowercase ISO 639-3 code"
);
ensure!(
!matches!(value, "mis" | "mul" | "und" | "zxx"),
"must identify one primary spoken language"
);
Ok(())
}
fn validate_feature_row(row: &FeatureRow) -> anyhow::Result<()> {
validate_nonempty("accent_variety", &row.accent_variety)?;
validate_positive("perceived_age", row.perceived_age)?;
validate_range(
"vocal_gender_presentation",
row.vocal_gender_presentation,
0.0,
100.0,
)?;
validate_positive("median_f0_hz", row.median_f0_hz)?;
validate_positive("formant_dispersion_hz", row.formant_dispersion_hz)?;
validate_positive("vai", row.vai)?;
validate_range("hypernasality", row.hypernasality, 0.0, 4.0)?;
validate_range(
"creaky_phonation_percent",
row.creaky_phonation_percent,
0.0,
100.0,
)?;
validate_nonempty("rhotic_realization", &row.rhotic_realization)?;
validate_positive(
"word_initial_stressed_prevocalic_t_vot_ms",
row.word_initial_stressed_prevocalic_t_vot_ms,
)?;
validate_range("breathiness", row.breathiness, 0.0, 100.0)?;
validate_range("roughness", row.roughness, 0.0, 100.0)?;
validate_positive("f0_pitch_span_semitones", row.f0_pitch_span_semitones)?;
validate_positive(
"articulation_rate_syllables_per_second",
row.articulation_rate_syllables_per_second,
)?;
validate_nonnegative("npvi_v", row.npvi_v)?;
validate_range("foreign_accentedness", row.foreign_accentedness, 1.0, 9.0)?;
validate_range(
"unstressed_vowel_reduction_percent",
row.unstressed_vowel_reduction_percent,
0.0,
100.0,
)?;
validate_nonempty("lateral_realization", &row.lateral_realization)?;
validate_nonnegative(
"filled_pauses_per_100_words",
row.filled_pauses_per_100_words,
)?;
validate_nonempty("s_realization", &row.s_realization)?;
validate_range(
"lexical_stress_accuracy_percent",
row.lexical_stress_accuracy_percent,
0.0,
100.0,
)?;
validate_range(
"monophthongization_percent",
row.monophthongization_percent,
0.0,
100.0,
)?;
validate_range(
"consonant_cluster_reduction_percent",
row.consonant_cluster_reduction_percent,
0.0,
100.0,
)
}
fn validate_nonempty(field: &str, value: &str) -> anyhow::Result<()> {
ensure!(!value.trim().is_empty(), "{field} must not be empty");
Ok(())
}
fn validate_finite(field: &str, value: f64) -> anyhow::Result<()> {
ensure!(value.is_finite(), "{field} must be finite");
Ok(())
}
fn validate_positive(field: &str, value: f64) -> anyhow::Result<()> {
validate_finite(field, value)?;
ensure!(value > 0.0, "{field} must be positive");
Ok(())
}
fn validate_nonnegative(field: &str, value: f64) -> anyhow::Result<()> {
validate_finite(field, value)?;
ensure!(value >= 0.0, "{field} must be nonnegative");
Ok(())
}
fn validate_range(field: &str, value: f64, minimum: f64, maximum: f64) -> anyhow::Result<()> {
validate_finite(field, value)?;
ensure!(
(minimum..=maximum).contains(&value),
"{field} must be between {minimum} and {maximum} inclusive"
);
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use kcode_speech_classification::DeleteOutcome;
use std::{
fs,
path::{Path, PathBuf},
sync::atomic::{AtomicU64, Ordering},
};
static NEXT_PATH: AtomicU64 = AtomicU64::new(0);
fn database_path(label: &str) -> PathBuf {
std::env::temp_dir().join(format!(
"kcode-audio-ingress-identity-{}-{label}-{}.sqlite3",
std::process::id(),
NEXT_PATH.fetch_add(1, Ordering::Relaxed)
))
}
fn remove_database(path: &Path) {
for suffix in ["", "-wal", "-shm"] {
let mut value = path.as_os_str().to_os_string();
value.push(suffix);
let _ = fs::remove_file(PathBuf::from(value));
}
}
fn row() -> FeatureRow {
FeatureRow {
accent_variety: "stan1293 Standard American English".into(),
perceived_age: 36.0,
vocal_gender_presentation: 55.0,
median_f0_hz: 145.0,
formant_dispersion_hz: 1050.0,
vai: 1.1,
hypernasality: 0.0,
creaky_phonation_percent: 5.0,
rhotic_realization: "[ɹ] alveolar approximant".into(),
word_initial_stressed_prevocalic_t_vot_ms: 58.0,
breathiness: 8.0,
roughness: 4.0,
f0_pitch_span_semitones: 10.0,
articulation_rate_syllables_per_second: 4.1,
npvi_v: 48.0,
cefr: Cefr::C2,
foreign_accentedness: 1.0,
unstressed_vowel_reduction_percent: 75.0,
lateral_realization: "mixed".into(),
filled_pauses_per_100_words: 1.0,
s_realization: "laminal [sÌ»]".into(),
lexical_stress_accuracy_percent: 99.0,
monophthongization_percent: 2.0,
consonant_cluster_reduction_percent: 1.0,
}
}
fn parsed() -> ParsedChunk {
ParsedChunk {
utterances: vec![ParsedUtterance {
speaker: "Speaker A".into(),
language: "eng".into(),
original_text: "Hello.".into(),
english_translation: String::new(),
corrected_natural_text: None,
coaching: Vec::new(),
annotations: Vec::new(),
}],
notes: vec!["Clear recording.".into()],
clip_valid: true,
clip_validity_reason: None,
speakers: vec![ParsedSpeaker {
local_label: "Speaker A".into(),
primary_language: "eng".into(),
feature_row: row(),
}],
}
}
fn observation(name: &str, confidence: f64, ordinal: u32) -> CorrectionObservation {
CorrectionObservation {
local_label: format!("Speaker {}", char::from(b'A' + ordinal as u8)),
speaker_ordinal: ordinal,
observation_key: ObservationKey {
object_id: "recording".into(),
piece_index: ordinal,
},
candidate: Some(CandidateMapping {
full_name: name.into(),
cost: 1.0,
confidence,
runner_up_full_name: Some("Runner Up".into()),
runner_up_cost: Some(4.0),
background_population_cost: 3.0,
}),
identified_full_name: Some(name.into()),
confirmed_full_name: None,
}
}
fn packet(clean: bool) -> CorrectionPacket {
CorrectionPacket {
recording_id: Uuid::nil(),
user_id: "user".into(),
sha256: "0".repeat(64),
original_filename: "audio.wav".into(),
size_bytes: 44,
recorded_at: DateTime::parse_from_rfc3339("2026-01-01T00:00:00Z")
.unwrap()
.with_timezone(&Utc),
clean,
chunk_count: 1,
chunks: vec![CorrectionChunk {
chunk_index: 0,
chunk_count: 1,
audio_start_ms: 0,
audio_end_ms: 1_000,
raw_gemini_response: "raw".into(),
parsed: parsed(),
observations: vec![observation("David Example", 2.0, 0)],
clean,
}],
confirmation_state: ConfirmationState::Unconfirmed,
}
}
#[test]
fn parser_requires_exact_rows_and_known_utterance_speakers() {
let valid = serde_json::to_string(&parsed()).unwrap();
let restored = parse_and_validate_chunk(&valid, 2.0).unwrap();
assert_eq!(restored, parsed());
let mut missing: Value = serde_json::from_str(&valid).unwrap();
missing["speakers"][0]["feature_row"]
.as_object_mut()
.unwrap()
.remove("median_f0_hz");
assert!(parse_and_validate_chunk(&missing.to_string(), 2.0).is_err());
let mut unknown: Value = serde_json::from_str(&valid).unwrap();
unknown["utterances"][0]["speaker"] = Value::String("Speaker Z".into());
assert!(parse_and_validate_chunk(&unknown.to_string(), 2.0).is_err());
let mut invalid_language: Value = serde_json::from_str(&valid).unwrap();
invalid_language["speakers"][0]["primary_language"] = Value::String("EN".into());
assert!(parse_and_validate_chunk(&invalid_language.to_string(), 2.0).is_err());
let mut duplicate = parsed();
duplicate.speakers.push(duplicate.speakers[0].clone());
assert!(
parse_and_validate_chunk(&serde_json::to_string(&duplicate).unwrap(), 2.0).is_err()
);
}
#[test]
fn deterministic_keys_are_unique_for_multi_speaker_chunks() {
let recording = Uuid::new_v4();
let first = training_object_id(recording, 0);
let second = training_object_id(recording, 1);
assert_ne!(first, second);
let keys = [
ObservationKey {
object_id: first.clone(),
piece_index: 0,
},
ObservationKey {
object_id: first,
piece_index: 1,
},
ObservationKey {
object_id: second,
piece_index: 0,
},
];
assert_eq!(
keys.iter().map(key_tuple).collect::<BTreeSet<_>>().len(),
keys.len()
);
}
#[test]
fn clean_gate_requires_background_bracketing_and_unique_candidates() {
let bracketed = observation("David Example", 1.0, 0);
assert!(chunk_is_clean(true, std::slice::from_ref(&bracketed)));
assert!(!chunk_is_clean(true, &[]));
let mut missing_candidate = bracketed.clone();
missing_candidate.candidate = None;
assert!(!chunk_is_clean(true, &[missing_candidate]));
let zero_confidence = observation("David Example", 0.0, 0);
assert!(!chunk_is_clean(true, &[zero_confidence]));
let negative_confidence = observation("David Example", -1.0, 0);
assert!(!chunk_is_clean(true, &[negative_confidence]));
let mut best_equal = bracketed.clone();
best_equal.candidate.as_mut().unwrap().cost = 3.0;
assert!(!chunk_is_clean(true, &[best_equal]));
let mut best_greater = bracketed.clone();
best_greater.candidate.as_mut().unwrap().cost = 4.0;
assert!(!chunk_is_clean(true, &[best_greater]));
let mut runner_up_absent = bracketed.clone();
let candidate = runner_up_absent.candidate.as_mut().unwrap();
candidate.runner_up_full_name = None;
candidate.runner_up_cost = None;
assert!(!chunk_is_clean(true, &[runner_up_absent]));
let mut runner_up_equal = bracketed.clone();
runner_up_equal.candidate.as_mut().unwrap().runner_up_cost = Some(3.0);
assert!(!chunk_is_clean(true, &[runner_up_equal]));
let mut runner_up_below = bracketed.clone();
runner_up_below.candidate.as_mut().unwrap().runner_up_cost = Some(2.0);
assert!(!chunk_is_clean(true, &[runner_up_below]));
assert!(!chunk_is_clean(
true,
&[bracketed.clone(), observation("David Example", 2.0, 1),]
));
assert!(!chunk_is_clean(false, &[bracketed]));
}
#[test]
fn recording_training_gate_trains_only_clean_packets() {
let path = database_path("training-gate");
let classifier = SpeechClassifier::open(&path).unwrap();
let mut unclean = packet(false);
train_clean_packet(&classifier, &mut unclean).unwrap();
assert_eq!(
classifier
.delete(unclean.chunks[0].observations[0].observation_key.clone())
.unwrap(),
DeleteOutcome::NotFound
);
let mut clean = packet(true);
train_clean_packet(&classifier, &mut clean).unwrap();
assert_eq!(
clean.confirmation_state,
ConfirmationState::AutomaticallyTrained
);
assert_eq!(
classifier
.delete(clean.chunks[0].observations[0].observation_key.clone())
.unwrap(),
DeleteOutcome::Deleted
);
drop(classifier);
remove_database(&path);
}
#[test]
fn confirmation_requires_exact_coverage() {
let packet = packet(false);
let key = packet.chunks[0].observations[0].observation_key.clone();
let exact = RecordingConfirmation {
recording_id: packet.recording_id,
observations: vec![ObservationConfirmation {
observation_key: key.clone(),
confirmed_full_name: "David Example".into(),
}],
};
assert!(validate_confirmation_coverage(&packet, &exact).is_ok());
let duplicate = RecordingConfirmation {
recording_id: packet.recording_id,
observations: vec![
ObservationConfirmation {
observation_key: key.clone(),
confirmed_full_name: "David Example".into(),
},
ObservationConfirmation {
observation_key: key,
confirmed_full_name: "David Example".into(),
},
],
};
assert!(validate_confirmation_coverage(&packet, &duplicate).is_err());
let empty = RecordingConfirmation {
recording_id: packet.recording_id,
observations: Vec::new(),
};
assert!(validate_confirmation_coverage(&packet, &empty).is_err());
}
}