1use std::{
4 collections::{BTreeSet, HashMap, HashSet},
5 sync::Arc,
6};
7
8use anyhow::{Context, ensure};
9use chrono::{DateTime, Utc};
10use kcode_speaker_extract::ExtractionOutcome;
11use kcode_speaker_system::{Cohort, IdentifyEvidence, SpeechClassifier, TrainOutcome};
12pub use kcode_speaker_system::{FeatureRow, ObservationKey};
13use serde::{Deserialize, Serialize};
14use uuid::Uuid;
15
16pub const CLASSIFIER_PROVIDER: &str = "google";
18pub const CLASSIFIER_MODEL: &str = "gemini-3.1-pro-preview";
20pub const CLASSIFIER_PROMPT_VERSION: &str = "gemini-speaker-24-freeform/1";
22pub const CLASSIFIER_SCHEMA_VERSION: &str = "gemini-speaker-24-normalized/1";
24
25const READ_ONLY_IDENTIFY_THRESHOLD: f64 = 1e308;
26
27#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
29#[serde(deny_unknown_fields)]
30pub struct ParsedUtterance {
31 pub speaker: String,
33 pub language: String,
35 pub original_text: String,
37 pub english_translation: String,
39 pub corrected_natural_text: Option<String>,
41 pub coaching: Vec<String>,
43 pub annotations: Vec<String>,
45}
46
47#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
49#[serde(deny_unknown_fields)]
50pub struct ParsedSpeaker {
51 pub local_label: String,
53 pub primary_language: Option<String>,
55 pub feature_row: Option<FeatureRow>,
57}
58
59#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
61#[serde(deny_unknown_fields)]
62pub struct ParsedChunk {
63 pub utterances: Vec<ParsedUtterance>,
65 pub notes: Vec<String>,
67 pub clip_valid: bool,
69 pub clip_validity_reason: Option<String>,
71 pub speakers: Vec<ParsedSpeaker>,
73}
74
75#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
77pub struct CandidateMapping {
78 pub full_name: String,
80 pub cost: f64,
82 pub confidence: f64,
84 pub runner_up_full_name: Option<String>,
86 pub runner_up_cost: Option<f64>,
88 pub background_population_cost: f64,
90}
91
92#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
94pub struct CorrectionObservation {
95 pub local_label: String,
97 pub speaker_ordinal: u32,
99 pub observation_key: ObservationKey,
101 pub candidate: Option<CandidateMapping>,
103 pub identified_full_name: Option<String>,
105 pub confirmed_full_name: Option<String>,
107}
108
109#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
111pub struct CorrectionChunk {
112 pub chunk_index: usize,
114 pub chunk_count: usize,
116 pub audio_start_ms: u64,
118 pub audio_end_ms: u64,
120 pub raw_gemini_response: String,
122 pub parsed: ParsedChunk,
124 pub observations: Vec<CorrectionObservation>,
126 pub clean: bool,
128}
129
130#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
132#[serde(rename_all = "snake_case")]
133pub enum ConfirmationState {
134 Unconfirmed,
136 AutomaticallyTrained,
138 Confirmed,
140}
141
142#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
144pub struct CorrectionPacket {
145 pub recording_id: Uuid,
147 pub user_id: String,
149 pub sha256: String,
151 pub original_filename: String,
153 pub size_bytes: u64,
155 pub recorded_at: DateTime<Utc>,
157 pub clean: bool,
159 pub chunk_count: usize,
161 pub chunks: Vec<CorrectionChunk>,
163 pub confirmation_state: ConfirmationState,
165}
166
167#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
169pub struct ObservationConfirmation {
170 pub observation_key: ObservationKey,
172 pub confirmed_full_name: String,
174}
175
176#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
178pub struct RecordingConfirmation {
179 pub recording_id: Uuid,
181 pub observations: Vec<ObservationConfirmation>,
183}
184
185#[derive(Clone)]
186pub(crate) struct ClassificationContext {
187 pub(crate) recording_id: Uuid,
188 pub(crate) user_id: String,
189 pub(crate) sha256: String,
190 pub(crate) original_filename: String,
191 pub(crate) size_bytes: u64,
192 pub(crate) recorded_at: DateTime<Utc>,
193 pub(crate) classifier: Arc<SpeechClassifier>,
194}
195
196#[derive(Deserialize)]
197#[serde(deny_unknown_fields)]
198struct TranscriptChunk {
199 utterances: Vec<ParsedUtterance>,
200 notes: Vec<String>,
201 speakers: Vec<TranscriptSpeaker>,
202}
203
204#[derive(Deserialize)]
205#[serde(deny_unknown_fields)]
206struct TranscriptSpeaker {
207 local_label: String,
208 speaker_ordinal: u16,
209}
210
211pub(crate) fn parse_and_validate_chunk(
212 response: &str,
213 extraction: &ExtractionOutcome,
214 chunk_duration_seconds: f64,
215) -> anyhow::Result<ParsedChunk> {
216 ensure!(
217 !response.trim().is_empty(),
218 "GPT parser returned an empty response"
219 );
220 let transcript: TranscriptChunk = serde_json::from_str(response)
221 .context("GPT parser response is not one valid transcript JSON object")?;
222
223 let (profile_count, clip_valid, clip_validity_reason) = match extraction {
224 ExtractionOutcome::Scored(scored) if scored.additional_speakers.is_empty() => {
225 (scored.speakers.len(), true, None)
226 }
227 ExtractionOutcome::Scored(scored) => (
228 scored.speakers.len() + scored.additional_speakers.len(),
229 false,
230 Some("One or more speakers lacked a complete feature profile.".to_owned()),
231 ),
232 ExtractionOutcome::Unscorable {
233 reason,
234 additional_speakers,
235 } => (additional_speakers.len(), false, Some(reason.clone())),
236 };
237 ensure!(
238 transcript.speakers.len() == profile_count,
239 "transcript and normalized speaker counts differ"
240 );
241
242 let mut speakers = transcript.speakers;
243 speakers.sort_by_key(|speaker| speaker.speaker_ordinal);
244 ensure!(
245 speakers
246 .iter()
247 .enumerate()
248 .all(|(index, speaker)| usize::from(speaker.speaker_ordinal) == index),
249 "transcript speaker ordinals must be contiguous from zero"
250 );
251 let speakers = speakers
252 .into_iter()
253 .map(|speaker| {
254 let profile = match extraction {
255 ExtractionOutcome::Scored(scored) => scored
256 .speakers
257 .iter()
258 .find(|profile| profile.speaker_ordinal == speaker.speaker_ordinal),
259 ExtractionOutcome::Unscorable { .. } => None,
260 };
261 ParsedSpeaker {
262 local_label: speaker.local_label,
263 primary_language: profile.map(|value| value.primary_language.as_ref().to_owned()),
264 feature_row: profile.map(|value| value.features),
265 }
266 })
267 .collect();
268 let mut parsed = ParsedChunk {
269 utterances: transcript.utterances,
270 notes: transcript.notes,
271 clip_valid,
272 clip_validity_reason,
273 speakers,
274 };
275 validate_parsed_chunk(&mut parsed, chunk_duration_seconds)?;
276 Ok(parsed)
277}
278
279pub(crate) fn validate_parsed_chunk(
280 parsed: &mut ParsedChunk,
281 chunk_duration_seconds: f64,
282) -> anyhow::Result<()> {
283 ensure!(
284 chunk_duration_seconds.is_finite() && chunk_duration_seconds > 0.0,
285 "chunk duration must be finite and positive"
286 );
287 ensure!(
288 parsed.notes.iter().all(|note| !note.trim().is_empty()),
289 "chunk notes must not contain empty entries"
290 );
291 match (parsed.clip_valid, parsed.clip_validity_reason.as_deref()) {
292 (true, None) => {}
293 (false, Some(reason)) if !reason.trim().is_empty() => {}
294 (true, Some(_)) => anyhow::bail!("a valid clip must not carry an invalidity reason"),
295 (false, _) => anyhow::bail!("an invalid clip requires a brief reason"),
296 }
297
298 let mut labels = HashSet::new();
299 for speaker in &parsed.speakers {
300 ensure!(
301 !speaker.local_label.trim().is_empty(),
302 "speaker labels must not be empty"
303 );
304 ensure!(
305 labels.insert(speaker.local_label.clone()),
306 "duplicate speaker label {:?}",
307 speaker.local_label
308 );
309 ensure!(
310 speaker.primary_language.is_some() == speaker.feature_row.is_some(),
311 "speaker language and feature row must be present together"
312 );
313 if let Some(language) = speaker.primary_language.as_deref() {
314 validate_iso_639_3(language)
315 .with_context(|| format!("speaker {} primary language", speaker.local_label))?;
316 }
317 }
318
319 let mut referenced = HashSet::new();
320 for (index, utterance) in parsed.utterances.iter().enumerate() {
321 ensure!(
322 labels.contains(&utterance.speaker),
323 "utterance {index} references unknown speaker {:?}",
324 utterance.speaker
325 );
326 referenced.insert(utterance.speaker.clone());
327 validate_iso_639_3(&utterance.language)
328 .with_context(|| format!("utterance {index} language"))?;
329 ensure!(
330 !utterance.original_text.trim().is_empty(),
331 "utterance {index} original text must not be empty"
332 );
333 if utterance.language == "eng" {
334 ensure!(
335 utterance.english_translation.is_empty(),
336 "English utterance {index} must use an empty translation"
337 );
338 } else {
339 ensure!(
340 !utterance.english_translation.trim().is_empty(),
341 "non-English utterance {index} requires a complete English translation"
342 );
343 }
344 if let Some(corrected) = utterance.corrected_natural_text.as_deref() {
345 ensure!(
346 !corrected.trim().is_empty(),
347 "utterance {index} corrected text must not be empty"
348 );
349 }
350 ensure!(
351 utterance
352 .coaching
353 .iter()
354 .chain(&utterance.annotations)
355 .all(|entry| !entry.trim().is_empty()),
356 "utterance {index} notes must not contain empty entries"
357 );
358 }
359
360 ensure!(
361 parsed
362 .speakers
363 .iter()
364 .all(|speaker| referenced.contains(&speaker.local_label)),
365 "every speaker row must be referenced by at least one utterance"
366 );
367 if parsed.clip_valid {
368 ensure!(
369 !parsed.speakers.is_empty()
370 && parsed
371 .speakers
372 .iter()
373 .all(|speaker| speaker.feature_row.is_some()),
374 "a valid clip must contain complete speaker rows"
375 );
376 }
377 Ok(())
378}
379
380pub(crate) fn classify_speakers(
381 context: &ClassificationContext,
382 chunk_index: usize,
383 parsed: &ParsedChunk,
384) -> anyhow::Result<(Vec<CorrectionObservation>, bool)> {
385 let mut observations = Vec::with_capacity(parsed.speakers.len());
386 for (ordinal, speaker) in parsed.speakers.iter().enumerate() {
387 let (primary_language, feature_row) = match (
388 speaker.primary_language.as_deref(),
389 speaker.feature_row.as_ref(),
390 ) {
391 (Some(language), Some(row)) => (language, row),
392 (None, None) => continue,
393 _ => anyhow::bail!("speaker profile is only partially present"),
394 };
395 let speaker_ordinal = u32::try_from(ordinal)
396 .context("chunk has more speakers than the key schema supports")?;
397 let cohort = cohort(primary_language);
398 let probe_key = ObservationKey {
399 object_id: probe_object_id(context.recording_id, chunk_index),
400 piece_index: speaker_ordinal,
401 };
402 let outcome = context
403 .classifier
404 .identify(
405 probe_key.clone(),
406 cohort,
407 *feature_row,
408 READ_ONLY_IDENTIFY_THRESHOLD,
409 )
410 .with_context(|| {
411 format!(
412 "read-only identity scoring failed for chunk {chunk_index} speaker {}",
413 speaker.local_label
414 )
415 })?;
416 if outcome.speaker_id.is_some() {
417 context
418 .classifier
419 .delete(probe_key)
420 .context("removing an unexpectedly accepted read-only probe")?;
421 }
422 let candidate = outcome.evidence.as_ref().map(candidate_mapping);
423 let identified_full_name = candidate.as_ref().map(|value| value.full_name.clone());
424 observations.push(CorrectionObservation {
425 local_label: speaker.local_label.clone(),
426 speaker_ordinal,
427 observation_key: ObservationKey {
428 object_id: training_object_id(context.recording_id, chunk_index),
429 piece_index: speaker_ordinal,
430 },
431 candidate,
432 identified_full_name,
433 confirmed_full_name: None,
434 });
435 }
436 let clean = chunk_is_clean(parsed.clip_valid, &observations);
437 Ok((observations, clean))
438}
439
440pub(crate) fn unclassified_observations(
441 recording_id: Uuid,
442 chunk_index: usize,
443 parsed: &ParsedChunk,
444) -> anyhow::Result<Vec<CorrectionObservation>> {
445 parsed
446 .speakers
447 .iter()
448 .enumerate()
449 .filter(|(_, speaker)| speaker.feature_row.is_some())
450 .map(|(ordinal, speaker)| {
451 let speaker_ordinal = u32::try_from(ordinal)
452 .context("chunk has more speakers than the key schema supports")?;
453 Ok(CorrectionObservation {
454 local_label: speaker.local_label.clone(),
455 speaker_ordinal,
456 observation_key: ObservationKey {
457 object_id: training_object_id(recording_id, chunk_index),
458 piece_index: speaker_ordinal,
459 },
460 candidate: None,
461 identified_full_name: None,
462 confirmed_full_name: None,
463 })
464 })
465 .collect()
466}
467
468pub(crate) fn chunk_is_clean(clip_valid: bool, observations: &[CorrectionObservation]) -> bool {
469 if !clip_valid || observations.is_empty() {
470 return false;
471 }
472 let mut names = HashSet::new();
473 observations.iter().all(|observation| {
474 observation.candidate.as_ref().is_some_and(|candidate| {
475 candidate.confidence > 0.0
476 && candidate.cost < candidate.background_population_cost
477 && candidate
478 .runner_up_cost
479 .is_some_and(|cost| cost > candidate.background_population_cost)
480 && names.insert(candidate.full_name.as_str())
481 })
482 })
483}
484
485pub(crate) fn build_packet(
486 context: &ClassificationContext,
487 chunks: Vec<CorrectionChunk>,
488) -> anyhow::Result<CorrectionPacket> {
489 ensure!(!chunks.is_empty(), "correction packet has no chunks");
490 let chunk_count = chunks.len();
491 ensure!(
492 chunks.iter().enumerate().all(|(index, chunk)| {
493 chunk.chunk_index == index
494 && chunk.chunk_count == chunk_count
495 && chunk.audio_end_ms > chunk.audio_start_ms
496 && chunk.observations.len()
497 == chunk
498 .parsed
499 .speakers
500 .iter()
501 .filter(|speaker| speaker.feature_row.is_some())
502 .count()
503 }),
504 "correction packet chunks are not one complete chronological plan"
505 );
506 let clean = chunks.iter().all(|chunk| chunk.clean);
507 Ok(CorrectionPacket {
508 recording_id: context.recording_id,
509 user_id: context.user_id.clone(),
510 sha256: context.sha256.clone(),
511 original_filename: context.original_filename.clone(),
512 size_bytes: context.size_bytes,
513 recorded_at: context.recorded_at,
514 clean,
515 chunk_count,
516 chunks,
517 confirmation_state: ConfirmationState::Unconfirmed,
518 })
519}
520
521pub(crate) fn train_clean_packet(
522 classifier: &SpeechClassifier,
523 packet: &mut CorrectionPacket,
524) -> anyhow::Result<()> {
525 if !packet.clean {
526 ensure!(
527 packet.confirmation_state == ConfirmationState::Unconfirmed,
528 "unclean packet unexpectedly claims retained training"
529 );
530 return Ok(());
531 }
532
533 let mut added = Vec::new();
534 for chunk in &packet.chunks {
535 for observation in &chunk.observations {
536 let speaker = speaker_for_observation(chunk, observation)?;
537 let (primary_language, feature_row) = classifier_row(speaker)?;
538 let full_name = observation
539 .identified_full_name
540 .as_deref()
541 .context("clean observation omitted its identified full name")?;
542 match classifier.train(
543 observation.observation_key.clone(),
544 cohort(primary_language),
545 feature_row,
546 full_name.to_owned(),
547 ) {
548 Ok(TrainOutcome::Added) => added.push(observation.observation_key.clone()),
549 Ok(TrainOutcome::Unchanged | TrainOutcome::Corrected) => {}
550 Err(error) => {
551 let rollback_errors = rollback_added(classifier, &added);
552 if rollback_errors.is_empty() {
553 anyhow::bail!("automatic identity training failed: {error}");
554 }
555 anyhow::bail!(
556 "automatic identity training failed: {error}; rollback also failed: {}",
557 rollback_errors.join("; ")
558 );
559 }
560 }
561 }
562 }
563 packet.confirmation_state = ConfirmationState::AutomaticallyTrained;
564 Ok(())
565}
566
567pub(crate) fn validate_confirmation_coverage(
568 packet: &CorrectionPacket,
569 confirmation: &RecordingConfirmation,
570) -> Result<(), String> {
571 if confirmation.recording_id != packet.recording_id {
572 return Err("Confirmation recording ID does not match the packet.".into());
573 }
574
575 let known = packet
576 .chunks
577 .iter()
578 .flat_map(|chunk| &chunk.observations)
579 .map(|observation| key_tuple(&observation.observation_key))
580 .collect::<BTreeSet<_>>();
581 if known.is_empty() {
582 return Err("The correction packet contains no speaker observations.".into());
583 }
584
585 let mut supplied = BTreeSet::new();
586 for observation in &confirmation.observations {
587 if observation.confirmed_full_name.trim().is_empty()
588 || observation.confirmed_full_name.chars().count() > 512
589 {
590 return Err("Confirmed full names must contain between 1 and 512 characters.".into());
591 }
592 if !supplied.insert(key_tuple(&observation.observation_key)) {
593 return Err("Confirmation contains a duplicate observation key.".into());
594 }
595 }
596 if supplied != known {
597 return Err(
598 "Confirmation must cover every known observation exactly once, with no extras.".into(),
599 );
600 }
601 Ok(())
602}
603
604pub(crate) fn apply_confirmations(
605 classifier: &SpeechClassifier,
606 packet: &mut CorrectionPacket,
607 confirmation: &RecordingConfirmation,
608 legacy_observation_keys: &HashSet<(String, u32)>,
609) -> anyhow::Result<()> {
610 validate_confirmation_coverage(packet, confirmation).map_err(anyhow::Error::msg)?;
611 let assignments = confirmation
612 .observations
613 .iter()
614 .map(|entry| {
615 (
616 key_tuple(&entry.observation_key),
617 entry.confirmed_full_name.trim().to_owned(),
618 )
619 })
620 .collect::<HashMap<_, _>>();
621
622 #[derive(Clone)]
623 struct Target {
624 chunk_position: usize,
625 observation_position: usize,
626 key: ObservationKey,
627 cohort: Cohort,
628 row: FeatureRow,
629 new_name: String,
630 old_name: Option<String>,
631 }
632
633 let mut targets = Vec::new();
634 let mut legacy_assignments = Vec::new();
635 for (chunk_position, chunk) in packet.chunks.iter().enumerate() {
636 for (observation_position, observation) in chunk.observations.iter().enumerate() {
637 let new_name = assignments
638 .get(&key_tuple(&observation.observation_key))
639 .context("validated confirmation assignment disappeared")?
640 .clone();
641 if legacy_observation_keys.contains(&key_tuple(&observation.observation_key)) {
642 legacy_assignments.push((chunk_position, observation_position, new_name));
643 continue;
644 }
645 let speaker = speaker_for_observation(chunk, observation)?;
646 let (primary_language, feature_row) = classifier_row(speaker)?;
647 targets.push(Target {
648 chunk_position,
649 observation_position,
650 key: observation.observation_key.clone(),
651 cohort: cohort(primary_language),
652 row: feature_row,
653 new_name,
654 old_name: retained_name(packet.confirmation_state, observation),
655 });
656 }
657 }
658
659 let mut applied = Vec::<(Target, TrainOutcome)>::new();
660 for target in targets {
661 match classifier.train(
662 target.key.clone(),
663 target.cohort.clone(),
664 target.row,
665 target.new_name.clone(),
666 ) {
667 Ok(outcome) => applied.push((target, outcome)),
668 Err(error) => {
669 let mut rollback_errors = Vec::new();
670 for (previous, outcome) in applied.iter().rev() {
671 let rollback = if let Some(old_name) = &previous.old_name {
672 classifier
673 .train(
674 previous.key.clone(),
675 previous.cohort.clone(),
676 previous.row,
677 old_name.clone(),
678 )
679 .map(|_| ())
680 } else if *outcome == TrainOutcome::Added {
681 classifier.delete(previous.key.clone()).map(|_| ())
682 } else {
683 Ok(())
684 };
685 if let Err(rollback_error) = rollback {
686 rollback_errors.push(rollback_error.to_string());
687 }
688 }
689 if rollback_errors.is_empty() {
690 anyhow::bail!("applying identity confirmations failed: {error}");
691 }
692 anyhow::bail!(
693 "applying identity confirmations failed: {error}; rollback also failed: {}",
694 rollback_errors.join("; ")
695 );
696 }
697 }
698 }
699
700 for (target, _) in applied {
701 packet.chunks[target.chunk_position].observations[target.observation_position]
702 .confirmed_full_name = Some(target.new_name);
703 }
704 for (chunk_position, observation_position, name) in legacy_assignments {
705 packet.chunks[chunk_position].observations[observation_position].confirmed_full_name =
706 Some(name);
707 }
708 packet.confirmation_state = ConfirmationState::Confirmed;
709 Ok(())
710}
711
712pub(crate) fn restore_packet_training(
713 classifier: &SpeechClassifier,
714 packet: &CorrectionPacket,
715 legacy_observation_keys: &HashSet<(String, u32)>,
716) -> Vec<String> {
717 let mut errors = Vec::new();
718 for chunk in packet.chunks.iter().rev() {
719 for observation in chunk.observations.iter().rev() {
720 if legacy_observation_keys.contains(&key_tuple(&observation.observation_key)) {
721 continue;
722 }
723 let result = match retained_name(packet.confirmation_state, observation) {
724 Some(name) => speaker_for_observation(chunk, observation).and_then(|speaker| {
725 let (primary_language, feature_row) = classifier_row(speaker)?;
726 classifier
727 .train(
728 observation.observation_key.clone(),
729 cohort(primary_language),
730 feature_row,
731 name,
732 )
733 .map(|_| ())
734 .map_err(anyhow::Error::from)
735 }),
736 None => classifier
737 .delete(observation.observation_key.clone())
738 .map(|_| ())
739 .map_err(anyhow::Error::from),
740 };
741 if let Err(error) = result {
742 errors.push(error.to_string());
743 }
744 }
745 }
746 errors
747}
748
749pub(crate) fn training_object_id(recording_id: Uuid, chunk_index: usize) -> String {
750 format!("kcode-audio-ingress/recording/{recording_id}/chunk/{chunk_index}")
751}
752
753fn probe_object_id(recording_id: Uuid, chunk_index: usize) -> String {
754 format!("kcode-audio-ingress/probe/{recording_id}/chunk/{chunk_index}")
755}
756
757fn cohort(primary_language: &str) -> Cohort {
758 Cohort {
759 provider: CLASSIFIER_PROVIDER.into(),
760 model: CLASSIFIER_MODEL.into(),
761 prompt_version: CLASSIFIER_PROMPT_VERSION.into(),
762 schema_version: CLASSIFIER_SCHEMA_VERSION.into(),
763 primary_language: primary_language.into(),
764 }
765}
766
767fn candidate_mapping(evidence: &IdentifyEvidence) -> CandidateMapping {
768 CandidateMapping {
769 full_name: evidence.best.speaker_id.clone(),
770 cost: evidence.best.cost,
771 confidence: evidence.confidence_score,
772 runner_up_full_name: evidence
773 .runner_up
774 .as_ref()
775 .map(|candidate| candidate.speaker_id.clone()),
776 runner_up_cost: evidence.runner_up.as_ref().map(|candidate| candidate.cost),
777 background_population_cost: evidence.background_population_cost,
778 }
779}
780
781fn speaker_for_observation<'a>(
782 chunk: &'a CorrectionChunk,
783 observation: &CorrectionObservation,
784) -> anyhow::Result<&'a ParsedSpeaker> {
785 let speaker = chunk
786 .parsed
787 .speakers
788 .get(observation.speaker_ordinal as usize)
789 .context("observation ordinal is outside the parsed speaker rows")?;
790 ensure!(
791 speaker.local_label == observation.local_label,
792 "observation label does not match its parsed speaker row"
793 );
794 Ok(speaker)
795}
796
797fn classifier_row(speaker: &ParsedSpeaker) -> anyhow::Result<(&str, FeatureRow)> {
798 Ok((
799 speaker
800 .primary_language
801 .as_deref()
802 .context("speaker observation omitted its primary language")?,
803 speaker
804 .feature_row
805 .context("speaker observation omitted its feature row")?,
806 ))
807}
808
809fn retained_name(state: ConfirmationState, observation: &CorrectionObservation) -> Option<String> {
810 match state {
811 ConfirmationState::Unconfirmed => None,
812 ConfirmationState::AutomaticallyTrained => observation.identified_full_name.clone(),
813 ConfirmationState::Confirmed => observation.confirmed_full_name.clone(),
814 }
815}
816
817fn rollback_added(classifier: &SpeechClassifier, keys: &[ObservationKey]) -> Vec<String> {
818 let mut errors = Vec::new();
819 for key in keys.iter().rev() {
820 if let Err(error) = classifier.delete(key.clone()) {
821 errors.push(error.to_string());
822 }
823 }
824 errors
825}
826
827fn key_tuple(key: &ObservationKey) -> (String, u32) {
828 (key.object_id.clone(), key.piece_index)
829}
830
831fn validate_iso_639_3(value: &str) -> anyhow::Result<()> {
832 ensure!(
833 value.len() == 3 && value.bytes().all(|byte| byte.is_ascii_lowercase()),
834 "must be a lowercase ISO 639-3 code"
835 );
836 ensure!(
837 !matches!(value, "mis" | "mul" | "und" | "zxx"),
838 "must identify one primary spoken language"
839 );
840 Ok(())
841}
842
843#[cfg(test)]
844mod tests {
845 use super::*;
846 use kcode_speaker_extract::{AdditionalSpeaker, CompleteSpeaker, ScoredClip};
847 use kcode_speaker_system::DeleteOutcome;
848 use serde_json::{Value, json};
849 use std::{
850 fs,
851 path::{Path, PathBuf},
852 sync::atomic::{AtomicU64, Ordering},
853 };
854
855 static NEXT_PATH: AtomicU64 = AtomicU64::new(0);
856
857 fn database_path(label: &str) -> PathBuf {
858 std::env::temp_dir().join(format!(
859 "kcode-audio-ingress-identity-{}-{label}-{}.sqlite3",
860 std::process::id(),
861 NEXT_PATH.fetch_add(1, Ordering::Relaxed)
862 ))
863 }
864
865 fn remove_database(path: &Path) {
866 for suffix in ["", "-wal", "-shm"] {
867 let mut value = path.as_os_str().to_os_string();
868 value.push(suffix);
869 let _ = fs::remove_file(PathBuf::from(value));
870 }
871 }
872
873 fn row() -> FeatureRow {
874 FeatureRow::new(std::array::from_fn(|index| {
875 u8::try_from(index * 3 + 10).unwrap()
876 }))
877 .unwrap()
878 }
879
880 fn extraction() -> ExtractionOutcome {
881 ExtractionOutcome::Scored(ScoredClip {
882 speakers: vec![CompleteSpeaker {
883 speaker_ordinal: 0,
884 primary_language: kcode_speaker_extract::Key::parse("eng").unwrap(),
885 closest_dialect: "General American English".into(),
886 usable_speech_ms: 1_000,
887 features: row(),
888 }],
889 additional_speakers: Vec::new(),
890 })
891 }
892
893 fn parsed() -> ParsedChunk {
894 ParsedChunk {
895 utterances: vec![ParsedUtterance {
896 speaker: "Speaker A".into(),
897 language: "eng".into(),
898 original_text: "Hello.".into(),
899 english_translation: String::new(),
900 corrected_natural_text: None,
901 coaching: Vec::new(),
902 annotations: Vec::new(),
903 }],
904 notes: vec!["Clear recording.".into()],
905 clip_valid: true,
906 clip_validity_reason: None,
907 speakers: vec![ParsedSpeaker {
908 local_label: "Speaker A".into(),
909 primary_language: Some("eng".into()),
910 feature_row: Some(row()),
911 }],
912 }
913 }
914
915 fn observation(name: &str, confidence: f64, ordinal: u32) -> CorrectionObservation {
916 CorrectionObservation {
917 local_label: format!("Speaker {}", char::from(b'A' + ordinal as u8)),
918 speaker_ordinal: ordinal,
919 observation_key: ObservationKey {
920 object_id: "recording".into(),
921 piece_index: ordinal,
922 },
923 candidate: Some(CandidateMapping {
924 full_name: name.into(),
925 cost: 1.0,
926 confidence,
927 runner_up_full_name: Some("Runner Up".into()),
928 runner_up_cost: Some(4.0),
929 background_population_cost: 3.0,
930 }),
931 identified_full_name: Some(name.into()),
932 confirmed_full_name: None,
933 }
934 }
935
936 fn packet(clean: bool) -> CorrectionPacket {
937 CorrectionPacket {
938 recording_id: Uuid::nil(),
939 user_id: "user".into(),
940 sha256: "0".repeat(64),
941 original_filename: "audio.wav".into(),
942 size_bytes: 44,
943 recorded_at: DateTime::parse_from_rfc3339("2026-01-01T00:00:00Z")
944 .unwrap()
945 .with_timezone(&Utc),
946 clean,
947 chunk_count: 1,
948 chunks: vec![CorrectionChunk {
949 chunk_index: 0,
950 chunk_count: 1,
951 audio_start_ms: 0,
952 audio_end_ms: 1_000,
953 raw_gemini_response: "raw".into(),
954 parsed: parsed(),
955 observations: vec![observation("David Example", 2.0, 0)],
956 clean,
957 }],
958 confirmation_state: ConfirmationState::Unconfirmed,
959 }
960 }
961
962 #[test]
963 fn parser_merges_frozen_features_and_requires_matching_speakers() {
964 let valid = json!({
965 "utterances": [{
966 "speaker":"Speaker A",
967 "language":"eng",
968 "original_text":"Hello.",
969 "english_translation":"",
970 "corrected_natural_text":null,
971 "coaching":[],
972 "annotations":[]
973 }],
974 "notes":["Clear recording."],
975 "speakers":[{"local_label":"Speaker A", "speaker_ordinal":0}]
976 });
977 let restored = parse_and_validate_chunk(&valid.to_string(), &extraction(), 2.0).unwrap();
978 assert_eq!(restored, parsed());
979
980 let mut unknown = valid.clone();
981 unknown["utterances"][0]["speaker"] = Value::String("Speaker Z".into());
982 assert!(parse_and_validate_chunk(&unknown.to_string(), &extraction(), 2.0).is_err());
983
984 let mut wrong_ordinal = valid.clone();
985 wrong_ordinal["speakers"][0]["speaker_ordinal"] = json!(1);
986 assert!(parse_and_validate_chunk(&wrong_ordinal.to_string(), &extraction(), 2.0).is_err());
987
988 let incomplete = ExtractionOutcome::Scored(ScoredClip {
989 speakers: Vec::new(),
990 additional_speakers: vec![AdditionalSpeaker {
991 speaker_ordinal: 0,
992 description: "Too little speech.".into(),
993 }],
994 });
995 let restored = parse_and_validate_chunk(&valid.to_string(), &incomplete, 2.0).unwrap();
996 assert!(
997 !restored.clip_valid
998 && restored.speakers[0].primary_language.is_none()
999 && restored.speakers[0].feature_row.is_none()
1000 );
1001 }
1002
1003 #[test]
1004 fn deterministic_keys_are_unique_for_multi_speaker_chunks() {
1005 let recording = Uuid::new_v4();
1006 let first = training_object_id(recording, 0);
1007 let second = training_object_id(recording, 1);
1008 assert_ne!(first, second);
1009 let keys = [
1010 ObservationKey {
1011 object_id: first.clone(),
1012 piece_index: 0,
1013 },
1014 ObservationKey {
1015 object_id: first,
1016 piece_index: 1,
1017 },
1018 ObservationKey {
1019 object_id: second,
1020 piece_index: 0,
1021 },
1022 ];
1023 assert_eq!(
1024 keys.iter().map(key_tuple).collect::<BTreeSet<_>>().len(),
1025 keys.len()
1026 );
1027 }
1028
1029 #[test]
1030 fn clean_gate_requires_background_bracketing_and_unique_candidates() {
1031 let bracketed = observation("David Example", 1.0, 0);
1032 assert!(chunk_is_clean(true, std::slice::from_ref(&bracketed)));
1033 assert!(!chunk_is_clean(true, &[]));
1034
1035 let mut missing_candidate = bracketed.clone();
1036 missing_candidate.candidate = None;
1037 assert!(!chunk_is_clean(true, &[missing_candidate]));
1038
1039 let zero_confidence = observation("David Example", 0.0, 0);
1040 assert!(!chunk_is_clean(true, &[zero_confidence]));
1041 let negative_confidence = observation("David Example", -1.0, 0);
1042 assert!(!chunk_is_clean(true, &[negative_confidence]));
1043
1044 let mut best_equal = bracketed.clone();
1045 best_equal.candidate.as_mut().unwrap().cost = 3.0;
1046 assert!(!chunk_is_clean(true, &[best_equal]));
1047 let mut best_greater = bracketed.clone();
1048 best_greater.candidate.as_mut().unwrap().cost = 4.0;
1049 assert!(!chunk_is_clean(true, &[best_greater]));
1050
1051 let mut runner_up_absent = bracketed.clone();
1052 let candidate = runner_up_absent.candidate.as_mut().unwrap();
1053 candidate.runner_up_full_name = None;
1054 candidate.runner_up_cost = None;
1055 assert!(!chunk_is_clean(true, &[runner_up_absent]));
1056
1057 let mut runner_up_equal = bracketed.clone();
1058 runner_up_equal.candidate.as_mut().unwrap().runner_up_cost = Some(3.0);
1059 assert!(!chunk_is_clean(true, &[runner_up_equal]));
1060 let mut runner_up_below = bracketed.clone();
1061 runner_up_below.candidate.as_mut().unwrap().runner_up_cost = Some(2.0);
1062 assert!(!chunk_is_clean(true, &[runner_up_below]));
1063
1064 assert!(!chunk_is_clean(
1065 true,
1066 &[bracketed.clone(), observation("David Example", 2.0, 1),]
1067 ));
1068 assert!(!chunk_is_clean(false, &[bracketed]));
1069 }
1070
1071 #[test]
1072 fn recording_training_gate_trains_only_clean_packets() {
1073 let path = database_path("training-gate");
1074 let classifier = SpeechClassifier::open(&path).unwrap();
1075 let mut unclean = packet(false);
1076 train_clean_packet(&classifier, &mut unclean).unwrap();
1077 assert_eq!(
1078 classifier
1079 .delete(unclean.chunks[0].observations[0].observation_key.clone())
1080 .unwrap(),
1081 DeleteOutcome::NotFound
1082 );
1083
1084 let mut clean = packet(true);
1085 train_clean_packet(&classifier, &mut clean).unwrap();
1086 assert_eq!(
1087 clean.confirmation_state,
1088 ConfirmationState::AutomaticallyTrained
1089 );
1090 assert_eq!(
1091 classifier
1092 .delete(clean.chunks[0].observations[0].observation_key.clone())
1093 .unwrap(),
1094 DeleteOutcome::Deleted
1095 );
1096 drop(classifier);
1097 remove_database(&path);
1098 }
1099
1100 #[test]
1101 fn confirmation_requires_exact_coverage() {
1102 let packet = packet(false);
1103 let key = packet.chunks[0].observations[0].observation_key.clone();
1104 let exact = RecordingConfirmation {
1105 recording_id: packet.recording_id,
1106 observations: vec![ObservationConfirmation {
1107 observation_key: key.clone(),
1108 confirmed_full_name: "David Example".into(),
1109 }],
1110 };
1111 assert!(validate_confirmation_coverage(&packet, &exact).is_ok());
1112
1113 let duplicate = RecordingConfirmation {
1114 recording_id: packet.recording_id,
1115 observations: vec![
1116 ObservationConfirmation {
1117 observation_key: key.clone(),
1118 confirmed_full_name: "David Example".into(),
1119 },
1120 ObservationConfirmation {
1121 observation_key: key,
1122 confirmed_full_name: "David Example".into(),
1123 },
1124 ],
1125 };
1126 assert!(validate_confirmation_coverage(&packet, &duplicate).is_err());
1127
1128 let empty = RecordingConfirmation {
1129 recording_id: packet.recording_id,
1130 observations: Vec::new(),
1131 };
1132 assert!(validate_confirmation_coverage(&packet, &empty).is_err());
1133 }
1134}