Skip to main content

kcode_audio_ingress/
identity.rs

1//! Provider parsing and read-only speaker classification for audio ingress.
2
3use std::sync::Arc;
4
5use anyhow::{Context, ensure};
6use chrono::{DateTime, Utc};
7use kcode_audio_speaker_review::{
8    CLASSIFIER_MODEL, CLASSIFIER_PROMPT_VERSION, CLASSIFIER_PROVIDER, CLASSIFIER_SCHEMA_VERSION,
9    CandidateMapping, ConfirmationState, ObservationKey, ParsedSpeaker, training_object_id,
10};
11pub(crate) use kcode_audio_speaker_review::{
12    CorrectionChunk, CorrectionObservation, CorrectionPacket, ParsedChunk,
13};
14use kcode_speaker_extract::ExtractionOutcome;
15use kcode_speaker_system::{Cohort, IdentifyEvidence, SpeechClassifier};
16use uuid::Uuid;
17
18const READ_ONLY_IDENTIFY_THRESHOLD: f64 = 1e308;
19
20#[derive(Clone)]
21pub(crate) struct ClassificationContext {
22    pub(crate) recording_id: Uuid,
23    pub(crate) user_id: String,
24    pub(crate) sha256: String,
25    pub(crate) original_filename: String,
26    pub(crate) size_bytes: u64,
27    pub(crate) recorded_at: DateTime<Utc>,
28    pub(crate) classifier: Arc<SpeechClassifier>,
29}
30
31pub(crate) fn parsed_chunk_from_extraction(
32    extraction: &ExtractionOutcome,
33) -> anyhow::Result<ParsedChunk> {
34    let (count, clip_valid, clip_validity_reason) = match extraction {
35        ExtractionOutcome::Scored(scored) if scored.additional_speakers.is_empty() => {
36            (scored.speakers.len(), true, None)
37        }
38        ExtractionOutcome::Scored(scored) => (
39            scored.speakers.len() + scored.additional_speakers.len(),
40            false,
41            Some("One or more speakers lacked a complete feature profile.".to_owned()),
42        ),
43        ExtractionOutcome::Unscorable {
44            reason,
45            additional_speakers,
46        } => (additional_speakers.len(), false, Some(reason.clone())),
47    };
48    let speakers = (0..count)
49        .map(|ordinal| {
50            let profile = match extraction {
51                ExtractionOutcome::Scored(scored) => scored
52                    .speakers
53                    .iter()
54                    .find(|profile| usize::from(profile.speaker_ordinal) == ordinal),
55                ExtractionOutcome::Unscorable { .. } => None,
56            };
57            ParsedSpeaker {
58                local_label: format!("Speaker {}", ordinal + 1),
59                primary_language: profile.map(|value| value.primary_language.as_ref().to_owned()),
60                feature_row: profile.map(|value| value.features),
61            }
62        })
63        .collect();
64    Ok(ParsedChunk {
65        clip_valid,
66        clip_validity_reason,
67        speakers,
68    })
69}
70
71pub(crate) fn classify_speakers(
72    context: &ClassificationContext,
73    chunk_index: usize,
74    parsed: &ParsedChunk,
75) -> anyhow::Result<Vec<CorrectionObservation>> {
76    let mut observations = Vec::with_capacity(parsed.speakers.len());
77    for (ordinal, speaker) in parsed.speakers.iter().enumerate() {
78        let speaker_ordinal = u32::try_from(ordinal)
79            .context("chunk has more speakers than the key schema supports")?;
80        let candidate = match (
81            speaker.primary_language.as_deref(),
82            speaker.feature_row.as_ref(),
83        ) {
84            (Some(primary_language), Some(feature_row)) => {
85                let probe_key = ObservationKey {
86                    object_id: probe_object_id(context.recording_id, chunk_index),
87                    piece_index: speaker_ordinal,
88                };
89                let outcome = context
90                    .classifier
91                    .identify(
92                        probe_key.clone(),
93                        cohort(primary_language),
94                        *feature_row,
95                        READ_ONLY_IDENTIFY_THRESHOLD,
96                    )
97                    .with_context(|| {
98                        format!(
99                            "read-only identity scoring failed for chunk {chunk_index} speaker {}",
100                            speaker.local_label
101                        )
102                    })?;
103                if outcome.speaker_id.is_some() {
104                    context
105                        .classifier
106                        .delete(probe_key)
107                        .context("removing an unexpectedly accepted read-only probe")?;
108                }
109                outcome.evidence.as_ref().map(candidate_mapping)
110            }
111            (None, None) => None,
112            _ => anyhow::bail!("speaker profile is only partially present"),
113        };
114        observations.push(CorrectionObservation {
115            local_label: speaker.local_label.clone(),
116            speaker_ordinal,
117            observation_key: ObservationKey {
118                object_id: training_object_id(context.recording_id, chunk_index),
119                piece_index: speaker_ordinal,
120            },
121            candidate,
122            resolution: None,
123        });
124    }
125    Ok(observations)
126}
127
128pub(crate) fn unclassified_observations(
129    recording_id: Uuid,
130    chunk_index: usize,
131    parsed: &ParsedChunk,
132) -> anyhow::Result<Vec<CorrectionObservation>> {
133    parsed
134        .speakers
135        .iter()
136        .enumerate()
137        .map(|(ordinal, speaker)| {
138            let speaker_ordinal = u32::try_from(ordinal)
139                .context("chunk has more speakers than the key schema supports")?;
140            Ok(CorrectionObservation {
141                local_label: speaker.local_label.clone(),
142                speaker_ordinal,
143                observation_key: ObservationKey {
144                    object_id: training_object_id(recording_id, chunk_index),
145                    piece_index: speaker_ordinal,
146                },
147                candidate: None,
148                resolution: None,
149            })
150        })
151        .collect()
152}
153
154pub(crate) fn build_packet(
155    context: &ClassificationContext,
156    chunks: Vec<CorrectionChunk>,
157) -> anyhow::Result<CorrectionPacket> {
158    ensure!(!chunks.is_empty(), "correction packet has no chunks");
159    let chunk_count = chunks.len();
160    ensure!(
161        chunks.iter().enumerate().all(|(index, chunk)| {
162            chunk.chunk_index == index
163                && chunk.chunk_count == chunk_count
164                && chunk.audio_end_ms > chunk.audio_start_ms
165                && chunk.observations.len() == chunk.parsed.speakers.len()
166                && !chunk.signed_off
167        }),
168        "correction packet chunks are not one complete chronological plan"
169    );
170    Ok(CorrectionPacket {
171        recording_id: context.recording_id,
172        user_id: context.user_id.clone(),
173        sha256: context.sha256.clone(),
174        original_filename: context.original_filename.clone(),
175        size_bytes: context.size_bytes,
176        recorded_at: context.recorded_at,
177        chunk_count,
178        chunks,
179        confirmation_state: ConfirmationState::Unconfirmed,
180    })
181}
182
183fn probe_object_id(recording_id: Uuid, chunk_index: usize) -> String {
184    format!("kcode-audio-ingress/probe/{recording_id}/chunk/{chunk_index}")
185}
186
187fn cohort(primary_language: &str) -> Cohort {
188    Cohort {
189        provider: CLASSIFIER_PROVIDER.into(),
190        model: CLASSIFIER_MODEL.into(),
191        prompt_version: CLASSIFIER_PROMPT_VERSION.into(),
192        schema_version: CLASSIFIER_SCHEMA_VERSION.into(),
193        primary_language: primary_language.into(),
194    }
195}
196
197fn candidate_mapping(evidence: &IdentifyEvidence) -> CandidateMapping {
198    CandidateMapping {
199        full_name: evidence.best.speaker_id.clone(),
200        score: -evidence.best.cost,
201        runner_up_score: evidence.runner_up.as_ref().map(|candidate| -candidate.cost),
202    }
203}