kcode-audio-ingress 0.7.3

//! Provider parsing and read-only speaker classification for audio ingress.

use std::sync::Arc;

use anyhow::{Context, ensure};
use chrono::{DateTime, Utc};
use kcode_audio_speaker_review::{
    CLASSIFIER_MODEL, CLASSIFIER_PROMPT_VERSION, CLASSIFIER_PROVIDER, CLASSIFIER_SCHEMA_VERSION,
    CandidateMapping, ConfirmationState, ObservationKey, ParsedSpeaker, training_object_id,
};
pub(crate) use kcode_audio_speaker_review::{
    CorrectionChunk, CorrectionObservation, CorrectionPacket, ParsedChunk,
};
use kcode_speaker_extract::ExtractionOutcome;
use kcode_speaker_system::{Cohort, IdentifyEvidence, SpeechClassifier};
use uuid::Uuid;

const READ_ONLY_IDENTIFY_THRESHOLD: f64 = 1e308;

#[derive(Clone)]
pub(crate) struct ClassificationContext {
    pub(crate) recording_id: Uuid,
    pub(crate) user_id: String,
    pub(crate) sha256: String,
    pub(crate) original_filename: String,
    pub(crate) size_bytes: u64,
    pub(crate) recorded_at: DateTime<Utc>,
    pub(crate) classifier: Arc<SpeechClassifier>,
}

pub(crate) fn parsed_chunk_from_extraction(
    extraction: &ExtractionOutcome,
) -> anyhow::Result<ParsedChunk> {
    let (count, clip_valid, clip_validity_reason) = match extraction {
        ExtractionOutcome::Scored(scored) if scored.additional_speakers.is_empty() => {
            (scored.speakers.len(), true, None)
        }
        ExtractionOutcome::Scored(scored) => (
            scored.speakers.len() + scored.additional_speakers.len(),
            false,
            Some("One or more speakers lacked a complete feature profile.".to_owned()),
        ),
        ExtractionOutcome::Unscorable {
            reason,
            additional_speakers,
        } => (additional_speakers.len(), false, Some(reason.clone())),
    };
    let speakers = (0..count)
        .map(|ordinal| {
            let profile = match extraction {
                ExtractionOutcome::Scored(scored) => scored
                    .speakers
                    .iter()
                    .find(|profile| usize::from(profile.speaker_ordinal) == ordinal),
                ExtractionOutcome::Unscorable { .. } => None,
            };
            ParsedSpeaker {
                local_label: format!("Speaker {}", ordinal + 1),
                primary_language: profile.map(|value| value.primary_language.as_ref().to_owned()),
                feature_row: profile.map(|value| value.features),
            }
        })
        .collect();
    Ok(ParsedChunk {
        clip_valid,
        clip_validity_reason,
        speakers,
    })
}

pub(crate) fn classify_speakers(
    context: &ClassificationContext,
    chunk_index: usize,
    parsed: &ParsedChunk,
) -> anyhow::Result<Vec<CorrectionObservation>> {
    let mut observations = Vec::with_capacity(parsed.speakers.len());
    for (ordinal, speaker) in parsed.speakers.iter().enumerate() {
        let speaker_ordinal = u32::try_from(ordinal)
            .context("chunk has more speakers than the key schema supports")?;
        let candidate = match (
            speaker.primary_language.as_deref(),
            speaker.feature_row.as_ref(),
        ) {
            (Some(primary_language), Some(feature_row)) => {
                let probe_key = ObservationKey {
                    object_id: probe_object_id(context.recording_id, chunk_index),
                    piece_index: speaker_ordinal,
                };
                let outcome = context
                    .classifier
                    .identify(
                        probe_key.clone(),
                        cohort(primary_language),
                        *feature_row,
                        READ_ONLY_IDENTIFY_THRESHOLD,
                    )
                    .with_context(|| {
                        format!(
                            "read-only identity scoring failed for chunk {chunk_index} speaker {}",
                            speaker.local_label
                        )
                    })?;
                if outcome.speaker_id.is_some() {
                    context
                        .classifier
                        .delete(probe_key)
                        .context("removing an unexpectedly accepted read-only probe")?;
                }
                outcome.evidence.as_ref().map(candidate_mapping)
            }
            (None, None) => None,
            _ => anyhow::bail!("speaker profile is only partially present"),
        };
        observations.push(CorrectionObservation {
            local_label: speaker.local_label.clone(),
            speaker_ordinal,
            observation_key: ObservationKey {
                object_id: training_object_id(context.recording_id, chunk_index),
                piece_index: speaker_ordinal,
            },
            candidate,
            resolution: None,
        });
    }
    Ok(observations)
}

pub(crate) fn unclassified_observations(
    recording_id: Uuid,
    chunk_index: usize,
    parsed: &ParsedChunk,
) -> anyhow::Result<Vec<CorrectionObservation>> {
    parsed
        .speakers
        .iter()
        .enumerate()
        .map(|(ordinal, speaker)| {
            let speaker_ordinal = u32::try_from(ordinal)
                .context("chunk has more speakers than the key schema supports")?;
            Ok(CorrectionObservation {
                local_label: speaker.local_label.clone(),
                speaker_ordinal,
                observation_key: ObservationKey {
                    object_id: training_object_id(recording_id, chunk_index),
                    piece_index: speaker_ordinal,
                },
                candidate: None,
                resolution: None,
            })
        })
        .collect()
}

pub(crate) fn build_packet(
    context: &ClassificationContext,
    chunks: Vec<CorrectionChunk>,
) -> anyhow::Result<CorrectionPacket> {
    ensure!(!chunks.is_empty(), "correction packet has no chunks");
    let chunk_count = chunks.len();
    ensure!(
        chunks.iter().enumerate().all(|(index, chunk)| {
            chunk.chunk_index == index
                && chunk.chunk_count == chunk_count
                && chunk.audio_end_ms > chunk.audio_start_ms
                && chunk.observations.len() == chunk.parsed.speakers.len()
                && !chunk.signed_off
        }),
        "correction packet chunks are not one complete chronological plan"
    );
    Ok(CorrectionPacket {
        recording_id: context.recording_id,
        user_id: context.user_id.clone(),
        sha256: context.sha256.clone(),
        original_filename: context.original_filename.clone(),
        size_bytes: context.size_bytes,
        recorded_at: context.recorded_at,
        chunk_count,
        chunks,
        confirmation_state: ConfirmationState::Unconfirmed,
    })
}

fn probe_object_id(recording_id: Uuid, chunk_index: usize) -> String {
    format!("kcode-audio-ingress/probe/{recording_id}/chunk/{chunk_index}")
}

fn cohort(primary_language: &str) -> Cohort {
    Cohort {
        provider: CLASSIFIER_PROVIDER.into(),
        model: CLASSIFIER_MODEL.into(),
        prompt_version: CLASSIFIER_PROMPT_VERSION.into(),
        schema_version: CLASSIFIER_SCHEMA_VERSION.into(),
        primary_language: primary_language.into(),
    }
}

fn candidate_mapping(evidence: &IdentifyEvidence) -> CandidateMapping {
    CandidateMapping {
        full_name: evidence.best.speaker_id.clone(),
        score: -evidence.best.cost,
        runner_up_score: evidence.runner_up.as_ref().map(|candidate| -candidate.cost),
    }
}