Skip to main content

kcode_audio_ingress/
identity.rs

1//! Typed speaker-analysis validation and classifier orchestration.
2
3use std::{collections::HashSet, sync::Arc};
4
5use anyhow::{Context, ensure};
6use chrono::{DateTime, Utc};
7use kcode_speaker_extract::ExtractionOutcome;
8use kcode_speaker_system::{Cohort, IdentifyEvidence, SpeechClassifier};
9pub use kcode_speaker_system::{FeatureRow, ObservationKey};
10use serde::{Deserialize, Serialize};
11use uuid::Uuid;
12
13/// Exact classifier provider cohort component.
14pub const CLASSIFIER_PROVIDER: &str = "google";
15/// Exact classifier model cohort component.
16pub const CLASSIFIER_MODEL: &str = "gemini-3.1-pro-preview";
17/// Exact classifier prompt-version cohort component.
18pub const CLASSIFIER_PROMPT_VERSION: &str = "gemini-transcript-speaker-24-freeform/2";
19/// Exact classifier feature-schema cohort component.
20pub const CLASSIFIER_SCHEMA_VERSION: &str = "gemini-speaker-24-normalized/1";
21
22const READ_ONLY_IDENTIFY_THRESHOLD: f64 = 1e308;
23
24/// One typed speaker row normalized from a raw Gemini response.
25#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
26#[serde(deny_unknown_fields)]
27pub struct ParsedSpeaker {
28    /// Exact chunk-local label, such as `Speaker 1`.
29    pub local_label: String,
30    /// Primary language identifier, absent when Gemini withheld a complete profile.
31    pub primary_language: Option<String>,
32    /// Validated 24-value row, absent when Gemini withheld a complete profile.
33    pub feature_row: Option<FeatureRow>,
34}
35
36/// Complete normalized speaker structure for one raw Gemini result.
37#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
38pub struct ParsedChunk {
39    /// Whether every substantive speaker supplied a complete feature profile.
40    pub clip_valid: bool,
41    /// Brief incompleteness reason, present exactly when `clip_valid` is false.
42    pub clip_validity_reason: Option<String>,
43    /// One row for every substantive chunk-local speaker.
44    pub speakers: Vec<ParsedSpeaker>,
45}
46
47/// Classifier evidence retained for one chunk-local speaker.
48#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
49pub struct CandidateMapping {
50    /// Best candidate's caller-owned full name.
51    pub full_name: String,
52    /// Best-candidate log-likelihood score relative to background at zero.
53    pub score: f64,
54    /// Runner-up log-likelihood score relative to background at zero.
55    pub runner_up_score: Option<f64>,
56}
57
58/// Human resolution for one chunk-local speaker.
59#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
60#[serde(tag = "kind", rename_all = "snake_case")]
61pub enum SpeakerResolution {
62    /// A known or newly entered speaker name that may be trained.
63    Known {
64        /// Exact human-approved full name.
65        full_name: String,
66    },
67    /// A deliberately unidentified speaker that must not be trained.
68    Unknown,
69}
70
71/// One deterministic classifier observation in a correction packet.
72#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
73pub struct CorrectionObservation {
74    /// Chunk-local speaker label.
75    pub local_label: String,
76    /// Stable zero-based ordinal.
77    pub speaker_ordinal: u32,
78    /// Deterministic persisted training and correction key.
79    pub observation_key: ObservationKey,
80    /// Best available read-only classifier evidence.
81    pub candidate: Option<CandidateMapping>,
82    /// Human-approved resolution, absent until this chunk is signed off.
83    pub resolution: Option<SpeakerResolution>,
84}
85
86/// One complete chunk in a recording-level correction packet.
87#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
88pub struct CorrectionChunk {
89    /// Zero-based chronological chunk index.
90    pub chunk_index: usize,
91    /// Total recording chunk count.
92    pub chunk_count: usize,
93    /// Source-audio start in milliseconds.
94    pub audio_start_ms: u64,
95    /// Source-audio end in milliseconds.
96    pub audio_end_ms: u64,
97    /// Complete raw Gemini response without normalization.
98    pub raw_gemini_response: String,
99    /// Normalized feature structure from the single recording-wide GPT pass.
100    pub parsed: ParsedChunk,
101    /// Read-only classifier mappings for every substantive speaker.
102    pub observations: Vec<CorrectionObservation>,
103    /// Whether a human approved every resolution in this chunk.
104    pub signed_off: bool,
105}
106
107/// Durable identity-confirmation lifecycle for a correction packet.
108#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
109#[serde(rename_all = "snake_case")]
110pub enum ConfirmationState {
111    /// One or more chunks still await human signoff.
112    Unconfirmed,
113    /// Legacy state retained only so old packets remain decodable.
114    AutomaticallyTrained,
115    /// Every chunk has explicit human signoff.
116    Confirmed,
117}
118
119/// Complete transport-neutral correction packet for one recording.
120#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
121pub struct CorrectionPacket {
122    /// Stable recording UUID.
123    pub recording_id: Uuid,
124    /// Stable application user identifier associated with provider usage.
125    pub user_id: String,
126    /// Lowercase SHA-256 identity of the retained original bytes.
127    pub sha256: String,
128    /// Sanitized original filename.
129    pub original_filename: String,
130    /// Original retained file size in bytes.
131    pub size_bytes: u64,
132    /// Instant at which the recording began.
133    pub recorded_at: DateTime<Utc>,
134    /// Total chronological chunk count.
135    pub chunk_count: usize,
136    /// Every raw response, feature row, mapping, key, resolution, and interval.
137    pub chunks: Vec<CorrectionChunk>,
138    /// Current durable confirmation state.
139    pub confirmation_state: ConfirmationState,
140}
141
142/// One observation-level human resolution.
143#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
144pub struct ObservationConfirmation {
145    /// Exact deterministic observation key from the correction packet.
146    pub observation_key: ObservationKey,
147    /// Caller-confirmed known or unknown resolution.
148    pub resolution: SpeakerResolution,
149}
150
151/// Exact speaker resolutions and signoff for one review chunk.
152#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
153pub struct ChunkConfirmation {
154    /// Recording receiving the chunk signoff.
155    pub recording_id: Uuid,
156    /// Exact zero-based chunk index receiving signoff.
157    pub chunk_index: usize,
158    /// One resolution for every chunk-local observation, with no extras.
159    pub observations: Vec<ObservationConfirmation>,
160}
161
162#[derive(Clone)]
163pub(crate) struct ClassificationContext {
164    pub(crate) recording_id: Uuid,
165    pub(crate) user_id: String,
166    pub(crate) sha256: String,
167    pub(crate) original_filename: String,
168    pub(crate) size_bytes: u64,
169    pub(crate) recorded_at: DateTime<Utc>,
170    pub(crate) classifier: Arc<SpeechClassifier>,
171}
172
173pub(crate) fn parsed_chunk_from_extraction(
174    extraction: &ExtractionOutcome,
175) -> anyhow::Result<ParsedChunk> {
176    let (count, clip_valid, clip_validity_reason) = match extraction {
177        ExtractionOutcome::Scored(scored) if scored.additional_speakers.is_empty() => {
178            (scored.speakers.len(), true, None)
179        }
180        ExtractionOutcome::Scored(scored) => (
181            scored.speakers.len() + scored.additional_speakers.len(),
182            false,
183            Some("One or more speakers lacked a complete feature profile.".to_owned()),
184        ),
185        ExtractionOutcome::Unscorable {
186            reason,
187            additional_speakers,
188        } => (additional_speakers.len(), false, Some(reason.clone())),
189    };
190    let speakers = (0..count)
191        .map(|ordinal| {
192            let profile = match extraction {
193                ExtractionOutcome::Scored(scored) => scored
194                    .speakers
195                    .iter()
196                    .find(|profile| usize::from(profile.speaker_ordinal) == ordinal),
197                ExtractionOutcome::Unscorable { .. } => None,
198            };
199            ParsedSpeaker {
200                local_label: format!("Speaker {}", ordinal + 1),
201                primary_language: profile.map(|value| value.primary_language.as_ref().to_owned()),
202                feature_row: profile.map(|value| value.features),
203            }
204        })
205        .collect();
206    Ok(ParsedChunk {
207        clip_valid,
208        clip_validity_reason,
209        speakers,
210    })
211}
212
213pub(crate) fn classify_speakers(
214    context: &ClassificationContext,
215    chunk_index: usize,
216    parsed: &ParsedChunk,
217) -> anyhow::Result<Vec<CorrectionObservation>> {
218    let mut observations = Vec::with_capacity(parsed.speakers.len());
219    for (ordinal, speaker) in parsed.speakers.iter().enumerate() {
220        let speaker_ordinal = u32::try_from(ordinal)
221            .context("chunk has more speakers than the key schema supports")?;
222        let candidate = match (
223            speaker.primary_language.as_deref(),
224            speaker.feature_row.as_ref(),
225        ) {
226            (Some(primary_language), Some(feature_row)) => {
227                let probe_key = ObservationKey {
228                    object_id: probe_object_id(context.recording_id, chunk_index),
229                    piece_index: speaker_ordinal,
230                };
231                let outcome = context
232                    .classifier
233                    .identify(
234                        probe_key.clone(),
235                        cohort(primary_language),
236                        *feature_row,
237                        READ_ONLY_IDENTIFY_THRESHOLD,
238                    )
239                    .with_context(|| {
240                        format!(
241                            "read-only identity scoring failed for chunk {chunk_index} speaker {}",
242                            speaker.local_label
243                        )
244                    })?;
245                if outcome.speaker_id.is_some() {
246                    context
247                        .classifier
248                        .delete(probe_key)
249                        .context("removing an unexpectedly accepted read-only probe")?;
250                }
251                outcome.evidence.as_ref().map(candidate_mapping)
252            }
253            (None, None) => None,
254            _ => anyhow::bail!("speaker profile is only partially present"),
255        };
256        observations.push(CorrectionObservation {
257            local_label: speaker.local_label.clone(),
258            speaker_ordinal,
259            observation_key: ObservationKey {
260                object_id: training_object_id(context.recording_id, chunk_index),
261                piece_index: speaker_ordinal,
262            },
263            candidate,
264            resolution: None,
265        });
266    }
267    Ok(observations)
268}
269
270pub(crate) fn unclassified_observations(
271    recording_id: Uuid,
272    chunk_index: usize,
273    parsed: &ParsedChunk,
274) -> anyhow::Result<Vec<CorrectionObservation>> {
275    parsed
276        .speakers
277        .iter()
278        .enumerate()
279        .map(|(ordinal, speaker)| {
280            let speaker_ordinal = u32::try_from(ordinal)
281                .context("chunk has more speakers than the key schema supports")?;
282            Ok(CorrectionObservation {
283                local_label: speaker.local_label.clone(),
284                speaker_ordinal,
285                observation_key: ObservationKey {
286                    object_id: training_object_id(recording_id, chunk_index),
287                    piece_index: speaker_ordinal,
288                },
289                candidate: None,
290                resolution: None,
291            })
292        })
293        .collect()
294}
295
296pub(crate) fn build_packet(
297    context: &ClassificationContext,
298    chunks: Vec<CorrectionChunk>,
299) -> anyhow::Result<CorrectionPacket> {
300    ensure!(!chunks.is_empty(), "correction packet has no chunks");
301    let chunk_count = chunks.len();
302    ensure!(
303        chunks.iter().enumerate().all(|(index, chunk)| {
304            chunk.chunk_index == index
305                && chunk.chunk_count == chunk_count
306                && chunk.audio_end_ms > chunk.audio_start_ms
307                && chunk.observations.len() == chunk.parsed.speakers.len()
308                && !chunk.signed_off
309        }),
310        "correction packet chunks are not one complete chronological plan"
311    );
312    Ok(CorrectionPacket {
313        recording_id: context.recording_id,
314        user_id: context.user_id.clone(),
315        sha256: context.sha256.clone(),
316        original_filename: context.original_filename.clone(),
317        size_bytes: context.size_bytes,
318        recorded_at: context.recorded_at,
319        chunk_count,
320        chunks,
321        confirmation_state: ConfirmationState::Unconfirmed,
322    })
323}
324
325pub(crate) fn validate_confirmation_coverage(
326    packet: &CorrectionPacket,
327    confirmation: &ChunkConfirmation,
328) -> Result<(), String> {
329    if confirmation.recording_id != packet.recording_id {
330        return Err("Confirmation recording ID does not match the packet.".into());
331    }
332    let chunk = packet
333        .chunks
334        .get(confirmation.chunk_index)
335        .filter(|chunk| chunk.chunk_index == confirmation.chunk_index)
336        .ok_or_else(|| "Confirmation chunk does not exist.".to_owned())?;
337    let known = chunk
338        .observations
339        .iter()
340        .map(|observation| key_tuple(&observation.observation_key))
341        .collect::<HashSet<_>>();
342    let mut supplied = HashSet::new();
343    for observation in &confirmation.observations {
344        if let SpeakerResolution::Known { full_name } = &observation.resolution
345            && (full_name.trim().is_empty() || full_name.chars().count() > 512)
346        {
347            return Err("Known speaker names must contain between 1 and 512 characters.".into());
348        }
349        if !supplied.insert(key_tuple(&observation.observation_key)) {
350            return Err("Confirmation contains a duplicate observation key.".into());
351        }
352    }
353    if supplied != known {
354        return Err(
355            "Chunk signoff must resolve every speaker exactly once, with no extras.".into(),
356        );
357    }
358    Ok(())
359}
360
361pub(crate) fn apply_confirmations(
362    classifier: &SpeechClassifier,
363    packet: &mut CorrectionPacket,
364    confirmation: &ChunkConfirmation,
365    legacy_observation_keys: &HashSet<(String, u32)>,
366) -> anyhow::Result<()> {
367    validate_confirmation_coverage(packet, confirmation).map_err(anyhow::Error::msg)?;
368    let chunk_index = confirmation.chunk_index;
369    if packet.chunks[chunk_index].signed_off {
370        let exact_retry = confirmation.observations.iter().all(|entry| {
371            packet.chunks[chunk_index]
372                .observations
373                .iter()
374                .find(|observation| observation.observation_key == entry.observation_key)
375                .and_then(|observation| observation.resolution.as_ref())
376                == Some(&entry.resolution)
377        });
378        ensure!(exact_retry, "signed-off chunk cannot be changed");
379        return Ok(());
380    }
381
382    let assignments = confirmation
383        .observations
384        .iter()
385        .map(|entry| {
386            (
387                key_tuple(&entry.observation_key),
388                normalized(&entry.resolution),
389            )
390        })
391        .collect::<std::collections::HashMap<_, _>>();
392    let mut applied: Vec<ObservationKey> = Vec::new();
393    for position in 0..packet.chunks[chunk_index].observations.len() {
394        let observation = &packet.chunks[chunk_index].observations[position];
395        let resolution = assignments
396            .get(&key_tuple(&observation.observation_key))
397            .context("validated confirmation assignment disappeared")?
398            .clone();
399        if !legacy_observation_keys.contains(&key_tuple(&observation.observation_key)) {
400            let speaker = &packet.chunks[chunk_index].parsed.speakers[position];
401            let result = match (
402                &resolution,
403                speaker.primary_language.as_deref(),
404                speaker.feature_row,
405            ) {
406                (SpeakerResolution::Known { full_name }, Some(language), Some(row)) => classifier
407                    .train(
408                        observation.observation_key.clone(),
409                        cohort(language),
410                        row,
411                        full_name.clone(),
412                    )
413                    .map(|_| ()),
414                _ => classifier
415                    .delete(observation.observation_key.clone())
416                    .map(|_| ()),
417            };
418            if let Err(error) = result {
419                let rollback = applied
420                    .iter()
421                    .rev()
422                    .filter_map(|key| classifier.delete(key.clone()).err().map(|e| e.to_string()))
423                    .collect::<Vec<_>>();
424                if rollback.is_empty() {
425                    anyhow::bail!("applying identity confirmations failed: {error}");
426                }
427                anyhow::bail!(
428                    "applying identity confirmations failed: {error}; rollback also failed: {}",
429                    rollback.join("; ")
430                );
431            }
432            applied.push(observation.observation_key.clone());
433        }
434        packet.chunks[chunk_index].observations[position].resolution = Some(resolution);
435    }
436    packet.chunks[chunk_index].signed_off = true;
437    if packet.chunks.iter().all(|chunk| chunk.signed_off) {
438        packet.confirmation_state = ConfirmationState::Confirmed;
439    }
440    Ok(())
441}
442
443pub(crate) fn restore_packet_training(
444    classifier: &SpeechClassifier,
445    packet: &CorrectionPacket,
446    legacy_observation_keys: &HashSet<(String, u32)>,
447) -> Vec<String> {
448    let mut errors = Vec::new();
449    for chunk in packet.chunks.iter().rev() {
450        for observation in chunk.observations.iter().rev() {
451            if legacy_observation_keys.contains(&key_tuple(&observation.observation_key)) {
452                continue;
453            }
454            let speaker = chunk
455                .parsed
456                .speakers
457                .get(observation.speaker_ordinal as usize);
458            let result = match (observation.resolution.as_ref(), speaker) {
459                (
460                    Some(SpeakerResolution::Known { full_name }),
461                    Some(ParsedSpeaker {
462                        primary_language: Some(language),
463                        feature_row: Some(row),
464                        ..
465                    }),
466                ) => classifier
467                    .train(
468                        observation.observation_key.clone(),
469                        cohort(language),
470                        *row,
471                        full_name.clone(),
472                    )
473                    .map(|_| ()),
474                _ => classifier
475                    .delete(observation.observation_key.clone())
476                    .map(|_| ()),
477            };
478            if let Err(error) = result {
479                errors.push(error.to_string());
480            }
481        }
482    }
483    errors
484}
485
486pub(crate) fn training_object_id(recording_id: Uuid, chunk_index: usize) -> String {
487    format!("kcode-audio-ingress/recording/{recording_id}/chunk/{chunk_index}")
488}
489
490fn probe_object_id(recording_id: Uuid, chunk_index: usize) -> String {
491    format!("kcode-audio-ingress/probe/{recording_id}/chunk/{chunk_index}")
492}
493
494fn cohort(primary_language: &str) -> Cohort {
495    Cohort {
496        provider: CLASSIFIER_PROVIDER.into(),
497        model: CLASSIFIER_MODEL.into(),
498        prompt_version: CLASSIFIER_PROMPT_VERSION.into(),
499        schema_version: CLASSIFIER_SCHEMA_VERSION.into(),
500        primary_language: primary_language.into(),
501    }
502}
503
504fn candidate_mapping(evidence: &IdentifyEvidence) -> CandidateMapping {
505    CandidateMapping {
506        full_name: evidence.best.speaker_id.clone(),
507        score: -evidence.best.cost,
508        runner_up_score: evidence.runner_up.as_ref().map(|candidate| -candidate.cost),
509    }
510}
511
512fn normalized(resolution: &SpeakerResolution) -> SpeakerResolution {
513    match resolution {
514        SpeakerResolution::Known { full_name } => SpeakerResolution::Known {
515            full_name: full_name.trim().to_owned(),
516        },
517        SpeakerResolution::Unknown => SpeakerResolution::Unknown,
518    }
519}
520
521fn key_tuple(key: &ObservationKey) -> (String, u32) {
522    (key.object_id.clone(), key.piece_index)
523}
524
525#[cfg(test)]
526mod tests {
527    use super::*;
528
529    fn chunk(recording_id: Uuid, index: usize) -> CorrectionChunk {
530        let key = ObservationKey {
531            object_id: training_object_id(recording_id, index),
532            piece_index: 0,
533        };
534        CorrectionChunk {
535            chunk_index: index,
536            chunk_count: 2,
537            audio_start_ms: index as u64 * 1_000,
538            audio_end_ms: (index as u64 + 1) * 1_000,
539            raw_gemini_response: "Speaker 1: hello".into(),
540            parsed: ParsedChunk {
541                clip_valid: true,
542                clip_validity_reason: None,
543                speakers: vec![ParsedSpeaker {
544                    local_label: "Speaker 1".into(),
545                    primary_language: Some("eng".into()),
546                    feature_row: Some(FeatureRow::new([50; 24]).unwrap()),
547                }],
548            },
549            observations: vec![CorrectionObservation {
550                local_label: "Speaker 1".into(),
551                speaker_ordinal: 0,
552                observation_key: key,
553                candidate: None,
554                resolution: None,
555            }],
556            signed_off: false,
557        }
558    }
559
560    #[test]
561    fn per_chunk_signoff_trains_known_and_skips_unknown_speakers() {
562        let recording_id = Uuid::new_v4();
563        let path = std::env::temp_dir().join(format!("audio-ingress-{recording_id}.sqlite3"));
564        let classifier = SpeechClassifier::open(&path).unwrap();
565        let mut packet = CorrectionPacket {
566            recording_id,
567            user_id: "user".into(),
568            sha256: "a".repeat(64),
569            original_filename: "voice.wav".into(),
570            size_bytes: 1,
571            recorded_at: Utc::now(),
572            chunk_count: 2,
573            chunks: vec![chunk(recording_id, 0), chunk(recording_id, 1)],
574            confirmation_state: ConfirmationState::Unconfirmed,
575        };
576        let keys = packet
577            .chunks
578            .iter()
579            .map(|chunk| chunk.observations[0].observation_key.clone())
580            .collect::<Vec<_>>();
581        apply_confirmations(
582            &classifier,
583            &mut packet,
584            &ChunkConfirmation {
585                recording_id,
586                chunk_index: 0,
587                observations: vec![ObservationConfirmation {
588                    observation_key: keys[0].clone(),
589                    resolution: SpeakerResolution::Known {
590                        full_name: "David Example".into(),
591                    },
592                }],
593            },
594            &HashSet::new(),
595        )
596        .unwrap();
597        assert_eq!(packet.confirmation_state, ConfirmationState::Unconfirmed);
598        apply_confirmations(
599            &classifier,
600            &mut packet,
601            &ChunkConfirmation {
602                recording_id,
603                chunk_index: 1,
604                observations: vec![ObservationConfirmation {
605                    observation_key: keys[1].clone(),
606                    resolution: SpeakerResolution::Unknown,
607                }],
608            },
609            &HashSet::new(),
610        )
611        .unwrap();
612        assert_eq!(packet.confirmation_state, ConfirmationState::Confirmed);
613        assert_eq!(classifier.known_speakers().unwrap(), vec!["David Example"]);
614        drop(classifier);
615        let _ = std::fs::remove_file(path);
616    }
617}