use std::path::{Path, PathBuf};
use std::sync::Arc;
use kcode_k1_audio_classification::{
AudioClassification, FragmentId, FragmentStatus, OverallState,
};
use kcode_k1_full_audio_domain::{
ManifestEntry, decode_manifest, encode_manifest, stitch_transcripts, validate_fragment_geometry,
};
use kcode_k1_objects::{K1Objects, Object};
const MANIFEST_FILE_TYPE: &str = "k1-full-audio-manifest-v1";
pub type FullAudioId = kcode_k1_objects::TxId;
pub struct K1FullAudio {
ffmpeg_path: PathBuf,
objects: Arc<K1Objects>,
classification: Arc<AudioClassification>,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum FullAudioState {
Processing,
AwaitingLabels,
NeedsAttention,
Complete,
}
#[derive(Clone, Debug, PartialEq)]
pub struct FullAudioFragmentStatus {
pub fragment_id: FragmentId,
pub start_sample_48k: u64,
pub end_sample_48k: u64,
pub status: FragmentStatus,
}
#[derive(Clone, Debug, PartialEq)]
pub struct FullAudioStatus {
pub state: FullAudioState,
pub fragments: Vec<FullAudioFragmentStatus>,
pub final_transcript: Option<String>,
}
struct PendingFragment {
start_sample_48k: u64,
end_sample_48k: u64,
ogg_bytes: Vec<u8>,
}
impl K1FullAudio {
pub fn open(
ffmpeg_path: impl AsRef<Path>,
objects: Arc<K1Objects>,
classification: Arc<AudioClassification>,
) -> Result<Self, String> {
Ok(Self {
ffmpeg_path: validate_ffmpeg_path(ffmpeg_path.as_ref())?,
objects,
classification,
})
}
pub fn submit(&self, audio: &[u8]) -> Result<FullAudioId, String> {
submit_orchestration(
audio,
|input| kcode_audio_to_ogg_opus::convert_to_ogg_opus(&self.ffmpeg_path, input),
|ogg_bytes| {
kcode_ogg_opus_fragments::split_ogg_opus(ogg_bytes).map(|fragments| {
fragments
.into_iter()
.map(|fragment| PendingFragment {
start_sample_48k: fragment.start_sample_48k,
end_sample_48k: fragment.end_sample_48k,
ogg_bytes: fragment.ogg_bytes,
})
.collect()
})
},
|ogg_bytes| self.classification.submit(ogg_bytes),
|manifest| self.objects.save("", MANIFEST_FILE_TYPE, "", manifest),
)
}
pub fn status(&self, id: FullAudioId) -> Result<FullAudioStatus, String> {
status_orchestration(
id,
|manifest_id| self.objects.load(manifest_id),
|fragment_id| self.classification.status(fragment_id),
)
}
}
fn validate_ffmpeg_path(path: &Path) -> Result<PathBuf, String> {
if !path.is_absolute() {
return Err("ffmpeg path must be absolute".to_owned());
}
Ok(path.to_path_buf())
}
fn submit_orchestration<Convert, Split, Submit, Save>(
audio: &[u8],
convert: Convert,
split: Split,
mut submit: Submit,
save: Save,
) -> Result<FullAudioId, String>
where
Convert: FnOnce(&[u8]) -> Result<Vec<u8>, String>,
Split: FnOnce(&[u8]) -> Result<Vec<PendingFragment>, String>,
Submit: FnMut(&[u8]) -> Result<FragmentId, String>,
Save: FnOnce(&[u8]) -> Result<FullAudioId, String>,
{
let ogg_bytes = convert(audio)?;
let fragments = split(&ogg_bytes)?;
validate_fragment_geometry(
fragments
.iter()
.map(|fragment| (fragment.start_sample_48k, fragment.end_sample_48k)),
)?;
let mut entries = Vec::new();
entries
.try_reserve_exact(fragments.len())
.map_err(|error| format!("allocate manifest entries: {error}"))?;
for fragment in fragments {
let fragment_id = submit(&fragment.ogg_bytes)?;
entries.push(ManifestEntry {
fragment_id,
start_sample_48k: fragment.start_sample_48k,
end_sample_48k: fragment.end_sample_48k,
});
}
let manifest = encode_manifest(&entries)?;
save(&manifest)
}
fn status_orchestration<Load, Status>(
id: FullAudioId,
mut load: Load,
mut status: Status,
) -> Result<FullAudioStatus, String>
where
Load: FnMut(FullAudioId) -> Result<Option<Object>, String>,
Status: FnMut(FragmentId) -> Result<Option<FragmentStatus>, String>,
{
let object = load(id)?.ok_or_else(|| "unknown full-audio ID".to_owned())?;
if object.file_type != MANIFEST_FILE_TYPE
|| !object.filename.is_empty()
|| !object.description.is_empty()
{
return Err("object is not a full-audio manifest".to_owned());
}
let entries = decode_manifest(&object.data)?;
let mut fragments = Vec::new();
fragments
.try_reserve_exact(entries.len())
.map_err(|error| format!("allocate fragment statuses: {error}"))?;
for entry in entries {
let fragment_status = status(entry.fragment_id)?
.ok_or_else(|| format!("unknown audio fragment {}", entry.fragment_id))?;
fragments.push(FullAudioFragmentStatus {
fragment_id: entry.fragment_id,
start_sample_48k: entry.start_sample_48k,
end_sample_48k: entry.end_sample_48k,
status: fragment_status,
});
}
let state = derive_state(&fragments)?;
let final_transcript = if state == FullAudioState::Complete {
let transcripts = fragments
.iter()
.filter_map(|fragment| fragment.status.final_transcript.as_deref())
.collect::<Vec<_>>();
Some(stitch_transcripts(&transcripts))
} else {
None
};
Ok(FullAudioStatus {
state,
fragments,
final_transcript,
})
}
fn derive_state(fragments: &[FullAudioFragmentStatus]) -> Result<FullAudioState, String> {
if fragments.is_empty() {
return Err("manifest has no fragments".to_owned());
}
for fragment in fragments {
match (&fragment.status.state, &fragment.status.final_transcript) {
(OverallState::Confirmed, None) => {
return Err("confirmed fragment has no final transcript".to_owned());
}
(OverallState::Confirmed, Some(_)) => {}
(_, Some(_)) => {
return Err("non-confirmed fragment has a final transcript".to_owned());
}
(_, None) => {}
}
}
if fragments.iter().any(|fragment| {
matches!(
fragment.status.state,
OverallState::Failed | OverallState::Discarded
)
}) {
return Ok(FullAudioState::NeedsAttention);
}
if fragments.iter().any(|fragment| {
matches!(
fragment.status.state,
OverallState::Queued | OverallState::Running
)
}) {
return Ok(FullAudioState::Processing);
}
if fragments
.iter()
.all(|fragment| fragment.status.state == OverallState::Confirmed)
{
return Ok(FullAudioState::Complete);
}
if fragments.iter().all(|fragment| {
matches!(
fragment.status.state,
OverallState::Completed | OverallState::Confirmed
)
}) {
return Ok(FullAudioState::AwaitingLabels);
}
Err("classification returned an inconsistent state".to_owned())
}
#[cfg(test)]
mod tests {
use super::*;
use kcode_k1_audio_classification::{FragmentStageV1, StageState, StageStatus};
use std::cell::{Cell, RefCell};
use std::rc::Rc;
use std::sync::mpsc::{self, Receiver, SyncSender};
use std::thread;
use std::time::Duration;
fn id(value: u8) -> FragmentId {
FragmentId::from_bytes([value; 12])
}
fn entry(value: u8, start: u64, end: u64) -> ManifestEntry {
ManifestEntry {
fragment_id: id(value),
start_sample_48k: start,
end_sample_48k: end,
}
}
fn pending(value: u8, start: u64, end: u64) -> PendingFragment {
PendingFragment {
start_sample_48k: start,
end_sample_48k: end,
ogg_bytes: vec![value],
}
}
#[test]
fn submit_orders_work_and_returns_manifest_id_without_status_wait() {
let events = Rc::new(RefCell::new(Vec::new()));
let convert_events = events.clone();
let split_events = events.clone();
let submit_events = events.clone();
let save_events = events.clone();
let returned = submit_orchestration(
b"input",
move |audio| {
assert_eq!(audio, b"input");
convert_events.borrow_mut().push("convert".to_owned());
Ok(b"ogg".to_vec())
},
move |ogg| {
assert_eq!(ogg, b"ogg");
split_events.borrow_mut().push("split".to_owned());
Ok(vec![pending(10, 0, 100), pending(20, 80, 180)])
},
move |fragment| {
submit_events
.borrow_mut()
.push(format!("submit-{}", fragment[0]));
Ok(if fragment[0] == 10 { id(1) } else { id(2) })
},
move |manifest| {
save_events.borrow_mut().push("save".to_owned());
assert_eq!(
decode_manifest(manifest).expect("saved manifest"),
vec![entry(1, 0, 100), entry(2, 80, 180)]
);
Ok(id(9))
},
)
.expect("submit");
assert_eq!(returned, id(9));
assert_eq!(
events.borrow().as_slice(),
["convert", "split", "submit-10", "submit-20", "save"]
);
}
#[test]
fn partial_submit_leaves_no_manifest() {
let submissions = Rc::new(Cell::new(0));
let saves = Rc::new(Cell::new(0));
let submit_count = submissions.clone();
let save_count = saves.clone();
let result = submit_orchestration(
b"input",
|_| Ok(vec![1]),
|_| Ok(vec![pending(1, 0, 100), pending(2, 80, 180)]),
move |_| {
let count = submit_count.get() + 1;
submit_count.set(count);
if count == 2 {
Err("second submission failed".to_owned())
} else {
Ok(id(1))
}
},
move |_| {
save_count.set(save_count.get() + 1);
Ok(id(9))
},
);
assert!(result.is_err());
assert_eq!(submissions.get(), 2);
assert_eq!(saves.get(), 0);
}
fn stage(stage: FragmentStageV1) -> StageStatus {
StageStatus {
stage,
state: StageState::Pending,
}
}
fn classification_status(
state: OverallState,
final_transcript: Option<&str>,
) -> FragmentStatus {
FragmentStatus {
state,
queue: stage(FragmentStageV1::Queue),
transcript: stage(FragmentStageV1::Transcript),
speaker_labels: stage(FragmentStageV1::SpeakerLabels),
speaker_features: stage(FragmentStageV1::SpeakerFeatures),
structuring: stage(FragmentStageV1::Structuring),
label_confirmation: stage(FragmentStageV1::LabelConfirmation),
attempt_count: 0,
jobs: Vec::new(),
interim_txid: None,
analysis: None,
confirmed_labels: Vec::new(),
final_transcript: final_transcript.map(str::to_owned),
errors: Vec::new(),
errors_truncated: false,
}
}
fn object_with_data(data: Vec<u8>) -> Object {
Object {
filename: String::new(),
file_type: MANIFEST_FILE_TYPE.to_owned(),
description: String::new(),
data,
}
}
fn status_result(statuses: Vec<FragmentStatus>) -> Result<FullAudioStatus, String> {
let entries = statuses
.iter()
.enumerate()
.map(|(index, _)| {
let value = u8::try_from(index + 1).expect("test fragment ID");
let start = u64::try_from(index).expect("test index") * 80;
entry(value, start, start + 100)
})
.collect::<Vec<_>>();
let object = object_with_data(encode_manifest(&entries).expect("manifest"));
status_orchestration(
id(90),
|_| Ok(Some(object.clone())),
|fragment_id| {
Ok(entries
.iter()
.position(|entry| entry.fragment_id == fragment_id)
.map(|index| statuses[index].clone()))
},
)
}
#[test]
fn status_derives_all_states_with_precedence() {
let attention = status_result(vec![
classification_status(OverallState::Failed, None),
classification_status(OverallState::Queued, None),
])
.expect("attention");
assert_eq!(attention.state, FullAudioState::NeedsAttention);
assert_eq!(attention.final_transcript, None);
let processing = status_result(vec![
classification_status(OverallState::Completed, None),
classification_status(OverallState::Running, None),
])
.expect("processing");
assert_eq!(processing.state, FullAudioState::Processing);
let awaiting = status_result(vec![
classification_status(OverallState::Completed, None),
classification_status(OverallState::Confirmed, Some("ready")),
])
.expect("awaiting");
assert_eq!(awaiting.state, FullAudioState::AwaitingLabels);
assert_eq!(awaiting.final_transcript, None);
let complete = status_result(vec![
classification_status(
OverallState::Confirmed,
Some("[high] A: one two three four"),
),
classification_status(
OverallState::Confirmed,
Some("[medium] B: two three four five"),
),
])
.expect("complete");
assert_eq!(complete.state, FullAudioState::Complete);
assert_eq!(
complete.final_transcript.as_deref(),
Some("[high] A: one two three four\n[medium] B: five")
);
assert_eq!(complete.fragments[0].fragment_id, id(1));
assert_eq!(complete.fragments[1].start_sample_48k, 80);
}
#[test]
fn status_rejects_transcript_inconsistency() {
assert!(status_result(vec![classification_status(OverallState::Confirmed, None)]).is_err());
assert!(
status_result(vec![classification_status(
OverallState::Completed,
Some("not confirmed")
)])
.is_err()
);
}
#[test]
fn status_rejects_unknown_wrong_and_malformed_objects() {
assert!(status_orchestration(id(90), |_| Ok(None), |_| Ok(None)).is_err());
let entries = vec![entry(1, 0, 100)];
let valid = object_with_data(encode_manifest(&entries).expect("manifest"));
assert!(status_orchestration(id(90), |_| Ok(Some(valid.clone())), |_| Ok(None)).is_err());
let mut wrong_type = valid.clone();
wrong_type.file_type = "wrong".to_owned();
assert!(
status_orchestration(id(90), |_| Ok(Some(wrong_type.clone())), |_| Ok(None)).is_err()
);
let mut wrong_filename = valid;
wrong_filename.filename = "named".to_owned();
assert!(
status_orchestration(id(90), |_| Ok(Some(wrong_filename.clone())), |_| Ok(None))
.is_err()
);
let malformed = object_with_data(vec![1]);
assert!(
status_orchestration(id(90), |_| Ok(Some(malformed.clone())), |_| Ok(None)).is_err()
);
}
#[test]
fn repeated_complete_status_is_stable() {
let statuses = vec![
classification_status(OverallState::Confirmed, Some("Alpha, beta gamma delta")),
classification_status(OverallState::Confirmed, Some("BETA GAMMA DELTA epsilon")),
];
let first = status_result(statuses.clone()).expect("first status");
let second = status_result(statuses).expect("second status");
assert_eq!(first, second);
assert_eq!(
first.final_transcript.as_deref(),
Some("Alpha, beta gamma delta\nepsilon")
);
}
struct TestFacadeAdapter {
fragment_id: FragmentId,
manifest_id: FullAudioId,
entered: Option<SyncSender<()>>,
release: Option<Receiver<()>>,
}
impl TestFacadeAdapter {
fn submit(mut self) -> Result<FullAudioId, String> {
let mut entered = self.entered.take();
let mut release = self.release.take();
submit_orchestration(
b"input",
|_| Ok(vec![1]),
|_| Ok(vec![pending(1, 0, 100)]),
|_| {
if let Some(sender) = entered.take() {
sender
.send(())
.map_err(|error| format!("signal blocked submit: {error}"))?;
}
if let Some(receiver) = release.take() {
receiver
.recv()
.map_err(|error| format!("release blocked submit: {error}"))?;
}
Ok(self.fragment_id)
},
|_| Ok(self.manifest_id),
)
}
}
#[test]
fn blocked_operation_does_not_stall_an_independent_adapter() {
let (entered_sender, entered_receiver) = mpsc::sync_channel(1);
let (release_sender, release_receiver) = mpsc::channel();
let blocked = TestFacadeAdapter {
fragment_id: id(1),
manifest_id: id(10),
entered: Some(entered_sender),
release: Some(release_receiver),
};
let blocked_thread = thread::spawn(move || blocked.submit());
entered_receiver
.recv_timeout(Duration::from_secs(2))
.expect("blocked operation entered");
let independent = TestFacadeAdapter {
fragment_id: id(2),
manifest_id: id(20),
entered: None,
release: None,
};
assert_eq!(independent.submit().expect("independent submit"), id(20));
release_sender.send(()).expect("release blocked operation");
assert_eq!(
blocked_thread
.join()
.expect("join blocked operation")
.expect("blocked submit"),
id(10)
);
}
#[test]
fn absolute_path_validation_does_not_probe_the_filesystem() {
assert_eq!(
validate_ffmpeg_path(Path::new("/definitely/not/present")).expect("absolute path"),
PathBuf::from("/definitely/not/present")
);
assert!(validate_ffmpeg_path(Path::new("relative/ffmpeg")).is_err());
}
}