1use 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
13pub const CLASSIFIER_PROVIDER: &str = "google";
15pub const CLASSIFIER_MODEL: &str = "gemini-3.1-pro-preview";
17pub const CLASSIFIER_PROMPT_VERSION: &str = "gemini-transcript-speaker-24-freeform/2";
19pub const CLASSIFIER_SCHEMA_VERSION: &str = "gemini-speaker-24-normalized/1";
21
22const READ_ONLY_IDENTIFY_THRESHOLD: f64 = 1e308;
23
24#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
26#[serde(deny_unknown_fields)]
27pub struct ParsedSpeaker {
28 pub local_label: String,
30 pub primary_language: Option<String>,
32 pub feature_row: Option<FeatureRow>,
34}
35
36#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
38pub struct ParsedChunk {
39 pub clip_valid: bool,
41 pub clip_validity_reason: Option<String>,
43 pub speakers: Vec<ParsedSpeaker>,
45}
46
47#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
49pub struct CandidateMapping {
50 pub full_name: String,
52 pub score: f64,
54 pub runner_up_score: Option<f64>,
56}
57
58#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
60#[serde(tag = "kind", rename_all = "snake_case")]
61pub enum SpeakerResolution {
62 Known {
64 full_name: String,
66 },
67 Unknown,
69}
70
71#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
73pub struct CorrectionObservation {
74 pub local_label: String,
76 pub speaker_ordinal: u32,
78 pub observation_key: ObservationKey,
80 pub candidate: Option<CandidateMapping>,
82 pub resolution: Option<SpeakerResolution>,
84}
85
86#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
88pub struct CorrectionChunk {
89 pub chunk_index: usize,
91 pub chunk_count: usize,
93 pub audio_start_ms: u64,
95 pub audio_end_ms: u64,
97 pub raw_gemini_response: String,
99 pub parsed: ParsedChunk,
101 pub observations: Vec<CorrectionObservation>,
103 pub signed_off: bool,
105}
106
107#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
109#[serde(rename_all = "snake_case")]
110pub enum ConfirmationState {
111 Unconfirmed,
113 AutomaticallyTrained,
115 Confirmed,
117}
118
119#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
121pub struct CorrectionPacket {
122 pub recording_id: Uuid,
124 pub user_id: String,
126 pub sha256: String,
128 pub original_filename: String,
130 pub size_bytes: u64,
132 pub recorded_at: DateTime<Utc>,
134 pub chunk_count: usize,
136 pub chunks: Vec<CorrectionChunk>,
138 pub confirmation_state: ConfirmationState,
140}
141
142#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
144pub struct ObservationConfirmation {
145 pub observation_key: ObservationKey,
147 pub resolution: SpeakerResolution,
149}
150
151#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
153pub struct ChunkConfirmation {
154 pub recording_id: Uuid,
156 pub chunk_index: usize,
158 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}