kcode_k1_full_audio/
lib.rs

1use std::path::{Path, PathBuf};
2use std::sync::Arc;
3
4use kcode_k1_audio_classification::{
5    AudioClassification, FragmentId, FragmentStatus, OverallState,
6};
7use kcode_k1_full_audio_domain::{
8    ManifestEntry, decode_manifest, encode_manifest, stitch_transcripts, validate_fragment_geometry,
9};
10use kcode_k1_objects::{K1Objects, Object};
11
12const MANIFEST_FILE_TYPE: &str = "k1-full-audio-manifest-v1";
13const FRAGMENT_FILE_TYPE: &str = "audio/ogg";
14
15pub type FullAudioId = kcode_k1_objects::TxId;
16
17pub struct K1FullAudio {
18    ffmpeg_path: PathBuf,
19    objects: Arc<K1Objects>,
20    classification: Arc<AudioClassification>,
21}
22
23#[derive(Clone, Debug, Eq, PartialEq)]
24pub enum FullAudioState {
25    Processing,
26    AwaitingLabels,
27    NeedsAttention,
28    Complete,
29}
30
31#[derive(Clone, Debug, PartialEq)]
32pub struct FullAudioFragmentStatus {
33    pub fragment_id: FragmentId,
34    pub start_sample_48k: u64,
35    pub end_sample_48k: u64,
36    pub status: FragmentStatus,
37}
38
39#[derive(Clone, Debug, PartialEq)]
40pub struct FullAudioStatus {
41    pub state: FullAudioState,
42    pub fragments: Vec<FullAudioFragmentStatus>,
43    pub final_transcript: Option<String>,
44}
45
46struct PendingFragment {
47    start_sample_48k: u64,
48    end_sample_48k: u64,
49    ogg_bytes: Vec<u8>,
50}
51
52impl K1FullAudio {
53    pub fn open(
54        ffmpeg_path: impl AsRef<Path>,
55        objects: Arc<K1Objects>,
56        classification: Arc<AudioClassification>,
57    ) -> Result<Self, String> {
58        Ok(Self {
59            ffmpeg_path: validate_ffmpeg_path(ffmpeg_path.as_ref())?,
60            objects,
61            classification,
62        })
63    }
64
65    pub fn submit(&self, audio: &[u8]) -> Result<FullAudioId, String> {
66        submit_orchestration(
67            audio,
68            |input| kcode_audio_to_ogg_opus::convert_to_ogg_opus(&self.ffmpeg_path, input),
69            |ogg_bytes| {
70                kcode_ogg_opus_fragments::split_ogg_opus(ogg_bytes).map(|fragments| {
71                    fragments
72                        .into_iter()
73                        .map(|fragment| PendingFragment {
74                            start_sample_48k: fragment.start_sample_48k,
75                            end_sample_48k: fragment.end_sample_48k,
76                            ogg_bytes: fragment.ogg_bytes,
77                        })
78                        .collect()
79                })
80            },
81            |ogg_bytes| self.classification.submit(ogg_bytes),
82            |manifest| self.objects.save("", MANIFEST_FILE_TYPE, "", manifest),
83        )
84    }
85
86    pub fn status(&self, id: FullAudioId) -> Result<FullAudioStatus, String> {
87        status_orchestration(
88            id,
89            |object_id| self.objects.load(object_id),
90            |fragment_id| self.classification.status(fragment_id),
91        )
92    }
93
94    pub fn fragment_audio(
95        &self,
96        id: FullAudioId,
97        fragment_id: FragmentId,
98    ) -> Result<Vec<u8>, String> {
99        fragment_audio_orchestration(id, fragment_id, |object_id| self.objects.load(object_id))
100    }
101}
102
103fn validate_ffmpeg_path(path: &Path) -> Result<PathBuf, String> {
104    if !path.is_absolute() {
105        return Err("ffmpeg path must be absolute".to_owned());
106    }
107    Ok(path.to_path_buf())
108}
109
110fn submit_orchestration<Convert, Split, Submit, Save>(
111    audio: &[u8],
112    convert: Convert,
113    split: Split,
114    mut submit: Submit,
115    save: Save,
116) -> Result<FullAudioId, String>
117where
118    Convert: FnOnce(&[u8]) -> Result<Vec<u8>, String>,
119    Split: FnOnce(&[u8]) -> Result<Vec<PendingFragment>, String>,
120    Submit: FnMut(&[u8]) -> Result<FragmentId, String>,
121    Save: FnOnce(&[u8]) -> Result<FullAudioId, String>,
122{
123    let ogg_bytes = convert(audio)?;
124    let fragments = split(&ogg_bytes)?;
125    validate_fragment_geometry(
126        fragments
127            .iter()
128            .map(|fragment| (fragment.start_sample_48k, fragment.end_sample_48k)),
129    )?;
130    let mut entries = Vec::new();
131    entries
132        .try_reserve_exact(fragments.len())
133        .map_err(|error| format!("allocate manifest entries: {error}"))?;
134    for fragment in fragments {
135        let fragment_id = submit(&fragment.ogg_bytes)?;
136        entries.push(ManifestEntry {
137            fragment_id,
138            start_sample_48k: fragment.start_sample_48k,
139            end_sample_48k: fragment.end_sample_48k,
140        });
141    }
142    let manifest = encode_manifest(&entries)?;
143    save(&manifest)
144}
145
146fn load_manifest_entries<Load>(
147    id: FullAudioId,
148    load: &mut Load,
149) -> Result<Vec<ManifestEntry>, String>
150where
151    Load: FnMut(FullAudioId) -> Result<Option<Object>, String>,
152{
153    let object = load(id)?.ok_or_else(|| "unknown full-audio ID".to_owned())?;
154    if object.file_type != MANIFEST_FILE_TYPE
155        || !object.filename.is_empty()
156        || !object.description.is_empty()
157    {
158        return Err("object is not a full-audio manifest".to_owned());
159    }
160    decode_manifest(&object.data)
161}
162
163fn status_orchestration<Load, Status>(
164    id: FullAudioId,
165    mut load: Load,
166    mut status: Status,
167) -> Result<FullAudioStatus, String>
168where
169    Load: FnMut(FullAudioId) -> Result<Option<Object>, String>,
170    Status: FnMut(FragmentId) -> Result<Option<FragmentStatus>, String>,
171{
172    let entries = load_manifest_entries(id, &mut load)?;
173    let mut fragments = Vec::new();
174    fragments
175        .try_reserve_exact(entries.len())
176        .map_err(|error| format!("allocate fragment statuses: {error}"))?;
177    for entry in entries {
178        let fragment_status = status(entry.fragment_id)?
179            .ok_or_else(|| format!("unknown audio fragment {}", entry.fragment_id))?;
180        fragments.push(FullAudioFragmentStatus {
181            fragment_id: entry.fragment_id,
182            start_sample_48k: entry.start_sample_48k,
183            end_sample_48k: entry.end_sample_48k,
184            status: fragment_status,
185        });
186    }
187    let state = derive_state(&fragments)?;
188    let final_transcript = if state == FullAudioState::Complete {
189        let transcripts = fragments
190            .iter()
191            .filter_map(|fragment| fragment.status.final_transcript.as_deref())
192            .collect::<Vec<_>>();
193        Some(stitch_transcripts(&transcripts))
194    } else {
195        None
196    };
197    Ok(FullAudioStatus {
198        state,
199        fragments,
200        final_transcript,
201    })
202}
203
204fn fragment_audio_orchestration<Load>(
205    id: FullAudioId,
206    fragment_id: FragmentId,
207    mut load: Load,
208) -> Result<Vec<u8>, String>
209where
210    Load: FnMut(FullAudioId) -> Result<Option<Object>, String>,
211{
212    let entries = load_manifest_entries(id, &mut load)?;
213    if !entries.iter().any(|entry| entry.fragment_id == fragment_id) {
214        return Err("fragment does not belong to full audio".to_owned());
215    }
216    let object =
217        load(fragment_id)?.ok_or_else(|| format!("unknown audio fragment {fragment_id}"))?;
218    if object.file_type != FRAGMENT_FILE_TYPE
219        || !object.filename.is_empty()
220        || !object.description.is_empty()
221    {
222        return Err("object is not an audio fragment".to_owned());
223    }
224    Ok(object.data)
225}
226
227fn derive_state(fragments: &[FullAudioFragmentStatus]) -> Result<FullAudioState, String> {
228    if fragments.is_empty() {
229        return Err("manifest has no fragments".to_owned());
230    }
231    for fragment in fragments {
232        let confirmed = fragment.status.state == OverallState::Confirmed;
233        let has_transcript = fragment.status.final_transcript.is_some();
234        if confirmed && !has_transcript {
235            return Err("confirmed fragment has no final transcript".to_owned());
236        }
237        if !confirmed && has_transcript {
238            return Err("non-confirmed fragment has a final transcript".to_owned());
239        }
240    }
241    if has_state(fragments, &[OverallState::Failed, OverallState::Discarded]) {
242        return Ok(FullAudioState::NeedsAttention);
243    }
244    if has_state(fragments, &[OverallState::Queued, OverallState::Running]) {
245        return Ok(FullAudioState::Processing);
246    }
247    if fragments
248        .iter()
249        .all(|fragment| fragment.status.state == OverallState::Confirmed)
250    {
251        return Ok(FullAudioState::Complete);
252    }
253    if fragments.iter().all(|fragment| {
254        [OverallState::Completed, OverallState::Confirmed].contains(&fragment.status.state)
255    }) {
256        return Ok(FullAudioState::AwaitingLabels);
257    }
258    Err("classification returned an inconsistent state".to_owned())
259}
260
261fn has_state(fragments: &[FullAudioFragmentStatus], states: &[OverallState]) -> bool {
262    fragments
263        .iter()
264        .any(|fragment| states.contains(&fragment.status.state))
265}
266
267#[cfg(test)]
268mod tests;
kcode_k1_full_audio/lib.rs

kcode_k1_full_audio/
lib.rs
