Skip to main content

kcode_audio_ingress/
identity.rs

1//! Typed speaker-analysis validation and classifier orchestration.
2
3use 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
16/// Exact classifier provider cohort component.
17pub const CLASSIFIER_PROVIDER: &str = "google";
18/// Exact classifier model cohort component.
19pub const CLASSIFIER_MODEL: &str = "gemini-3.1-pro-preview";
20/// Exact classifier prompt-version cohort component.
21pub const CLASSIFIER_PROMPT_VERSION: &str = "gemini-speaker-24-freeform/1";
22/// Exact classifier feature-schema cohort component.
23pub const CLASSIFIER_SCHEMA_VERSION: &str = "gemini-speaker-24-normalized/1";
24
25const READ_ONLY_IDENTIFY_THRESHOLD: f64 = 1e308;
26
27/// One validated chunk-local utterance parsed from the raw Gemini response.
28#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
29#[serde(deny_unknown_fields)]
30pub struct ParsedUtterance {
31    /// Exact chunk-local speaker label.
32    pub speaker: String,
33    /// Lowercase ISO 639-3 language code for this utterance.
34    pub language: String,
35    /// Complete utterance in the language spoken.
36    pub original_text: String,
37    /// Complete English translation, or an empty string for English.
38    pub english_translation: String,
39    /// Corrected natural version for audibly non-native speech, when applicable.
40    pub corrected_natural_text: Option<String>,
41    /// Concise grammar, vocabulary, pronunciation, stress, and rhythm coaching.
42    pub coaching: Vec<String>,
43    /// Concise audible annotations that belong to this utterance.
44    pub annotations: Vec<String>,
45}
46
47/// One typed speaker row parsed from a raw Gemini response.
48#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
49#[serde(deny_unknown_fields)]
50pub struct ParsedSpeaker {
51    /// Exact chunk-local label used by the utterances.
52    pub local_label: String,
53    /// Lowercase ISO 639-3 code for the speaker's primary spoken language.
54    pub primary_language: Option<String>,
55    /// A validated 24-value row, absent when the extractor withheld a profile.
56    pub feature_row: Option<FeatureRow>,
57}
58
59/// Complete validated structure parsed from one raw Gemini response.
60#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
61#[serde(deny_unknown_fields)]
62pub struct ParsedChunk {
63    /// Complete ordered utterances for the chunk.
64    pub utterances: Vec<ParsedUtterance>,
65    /// Useful whole-chunk notes.
66    pub notes: Vec<String>,
67    /// Whether Gemini's whole-clip assessment was valid.
68    pub clip_valid: bool,
69    /// Brief invalidity reason, present exactly when `clip_valid` is false.
70    pub clip_validity_reason: Option<String>,
71    /// Exactly one typed row for each chunk-local speaker.
72    pub speakers: Vec<ParsedSpeaker>,
73}
74
75/// Classifier evidence retained for one chunk-local speaker.
76#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
77pub struct CandidateMapping {
78    /// Best candidate's caller-owned full name.
79    pub full_name: String,
80    /// Best candidate's raw classifier cost.
81    pub cost: f64,
82    /// Raw classifier confidence evidence, not a probability.
83    pub confidence: f64,
84    /// Optional runner-up full name.
85    pub runner_up_full_name: Option<String>,
86    /// Optional runner-up raw cost.
87    pub runner_up_cost: Option<f64>,
88    /// Raw background-population cost.
89    pub background_population_cost: f64,
90}
91
92/// One deterministic classifier observation in a correction packet.
93#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
94pub struct CorrectionObservation {
95    /// Chunk-local speaker label.
96    pub local_label: String,
97    /// Stable zero-based ordinal after labels are sorted.
98    pub speaker_ordinal: u32,
99    /// Deterministic persisted training and correction key.
100    pub observation_key: ObservationKey,
101    /// Best available read-only classifier evidence.
102    pub candidate: Option<CandidateMapping>,
103    /// Best candidate's full name, even when identity quality is insufficient.
104    pub identified_full_name: Option<String>,
105    /// Human-confirmed full name, when the confirmation API has been applied.
106    pub confirmed_full_name: Option<String>,
107}
108
109/// One complete chunk in a recording-level correction packet.
110#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
111pub struct CorrectionChunk {
112    /// Zero-based chronological chunk index.
113    pub chunk_index: usize,
114    /// Total recording chunk count.
115    pub chunk_count: usize,
116    /// Source-audio start in milliseconds.
117    pub audio_start_ms: u64,
118    /// Source-audio end in milliseconds.
119    pub audio_end_ms: u64,
120    /// Complete raw Gemini response without normalization.
121    pub raw_gemini_response: String,
122    /// Validated GPT-parsed structure.
123    pub parsed: ParsedChunk,
124    /// Read-only classifier mappings for every parsed speaker.
125    pub observations: Vec<CorrectionObservation>,
126    /// Whether validity, confidence, background-bracketing, and one-to-one checks passed.
127    pub clean: bool,
128}
129
130/// Durable identity-confirmation lifecycle for a correction packet.
131#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
132#[serde(rename_all = "snake_case")]
133pub enum ConfirmationState {
134    /// No classifier observations from this packet were intentionally retained.
135    Unconfirmed,
136    /// The recording was clean and all observations were automatically trained.
137    AutomaticallyTrained,
138    /// Exact observation-level human confirmations were applied.
139    Confirmed,
140}
141
142/// Complete transport-neutral correction packet for one recording.
143#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
144pub struct CorrectionPacket {
145    /// Stable recording UUID.
146    pub recording_id: Uuid,
147    /// Stable application user identifier associated with provider usage.
148    pub user_id: String,
149    /// Lowercase SHA-256 identity of the retained original bytes.
150    pub sha256: String,
151    /// Sanitized original filename.
152    pub original_filename: String,
153    /// Original retained file size in bytes.
154    pub size_bytes: u64,
155    /// Instant at which the recording began.
156    pub recorded_at: DateTime<Utc>,
157    /// Whether every chunk passed the recording-level clean gate.
158    pub clean: bool,
159    /// Total chronological chunk count.
160    pub chunk_count: usize,
161    /// Every raw response, parsed structure, row, mapping, key, and interval.
162    pub chunks: Vec<CorrectionChunk>,
163    /// Current durable confirmation and training state.
164    pub confirmation_state: ConfirmationState,
165}
166
167/// One observation-level full-name confirmation.
168#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
169pub struct ObservationConfirmation {
170    /// Exact deterministic observation key from the correction packet.
171    pub observation_key: ObservationKey,
172    /// Caller-confirmed full name.
173    pub confirmed_full_name: String,
174}
175
176/// Exact confirmations for one completed recording.
177#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
178pub struct RecordingConfirmation {
179    /// Completed recording receiving the confirmations.
180    pub recording_id: Uuid,
181    /// One confirmation for every known observation, with no extras.
182    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}