kcode_k1_audio_classification_format/
lib.rs

1use serde::{Deserialize, Serialize};
2use std::fmt::{Display, Formatter};
3use std::path::{Component, Path, PathBuf};
4
5pub use kcode_k1_transaction_id::TxId;
6pub use kcode_speaker_v3_analysis::{ExecutedAnalysis, FeatureVector24, LocalSpeakerLabel};
7
8const EVENT_VERSION: u8 = 1;
9const QUEUE_TAG: u8 = 1;
10const PROCESSED_TAG: u8 = 2;
11const FAILED_TAG: u8 = 3;
12const DISCARDED_TAG: u8 = 4;
13const CONFIRMED_TAG: u8 = 5;
14const FRAGMENT_VERSION: u8 = 1;
15const STAGED_KIND: u8 = 1;
16const FINAL_KIND: u8 = 2;
17const HEADER_LEN: usize = 16;
18const ANALYSIS_SLOT_OFFSET: usize = 16;
19const CONFIRMATION_SLOT_OFFSET: usize = 32;
20const BODY_OFFSET: usize = 48;
21const BASE64_ALPHABET: &[u8; 64] =
22    b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_";
23
24#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
25pub enum AnalysisStageV1 {
26    GeminiTranscript,
27    TerraLabels,
28    GeminiFeatures,
29    TerraStructuring,
30}
31
32#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
33pub struct QueueV1 {
34    #[serde(with = "txid_serde")]
35    pub audio_object_id: TxId,
36    pub duration_ms: u64,
37    pub filename: Option<String>,
38    #[serde(with = "optional_txid_serde")]
39    pub retry_of: Option<TxId>,
40}
41
42#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
43pub struct ProcessedV1 {
44    #[serde(with = "txid_serde")]
45    pub queue_id: TxId,
46    #[serde(with = "txid_serde")]
47    pub audio_object_id: TxId,
48    pub analysis: ExecutedAnalysis,
49}
50
51#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
52pub struct FailedV1 {
53    #[serde(with = "txid_serde")]
54    pub queue_id: TxId,
55    #[serde(with = "txid_serde")]
56    pub audio_object_id: TxId,
57    pub final_stage: AnalysisStageV1,
58    pub final_error: String,
59}
60
61#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
62pub struct DiscardedV1 {
63    #[serde(with = "txid_serde")]
64    pub failed_queue_id: TxId,
65    #[serde(with = "txid_serde")]
66    pub audio_object_id: TxId,
67}
68
69#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
70pub struct ConfirmedSpeakerV1 {
71    pub speaker: LocalSpeakerLabel,
72    pub person_id: String,
73}
74
75#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
76pub struct ConfirmedV1 {
77    #[serde(with = "txid_serde")]
78    pub queue_id: TxId,
79    #[serde(with = "txid_serde")]
80    pub audio_object_id: TxId,
81    #[serde(with = "txid_serde")]
82    pub analysis_txid: TxId,
83    pub speakers: Vec<ConfirmedSpeakerV1>,
84}
85
86#[allow(clippy::large_enum_variant)]
87#[derive(Debug, Clone, PartialEq)]
88pub enum AudioClassificationEventV1 {
89    Queue(QueueV1),
90    Processed(ProcessedV1),
91    Failed(FailedV1),
92    Discarded(DiscardedV1),
93    Confirmed(ConfirmedV1),
94}
95
96#[derive(Debug, Clone, PartialEq)]
97pub struct StagedSpeakerV1 {
98    pub speaker: LocalSpeakerLabel,
99    pub language: String,
100    pub features: FeatureVector24,
101    pub usable_for_training: bool,
102}
103
104#[derive(Debug, Clone, PartialEq)]
105pub struct StagedFragmentV1 {
106    pub analysis_txid: TxId,
107    pub transcript: String,
108    pub speakers: Vec<StagedSpeakerV1>,
109}
110
111#[derive(Debug, Clone, PartialEq)]
112pub struct FinalSpeakerV1 {
113    pub speaker: LocalSpeakerLabel,
114    pub person_id: String,
115    pub language: String,
116    pub features: FeatureVector24,
117    pub usable_for_training: bool,
118}
119
120#[derive(Debug, Clone, PartialEq)]
121pub struct FinalFragmentV1 {
122    pub analysis_txid: TxId,
123    pub confirmation_txid: TxId,
124    pub transcript: String,
125    pub speakers: Vec<FinalSpeakerV1>,
126}
127
128#[derive(Debug, Clone, PartialEq, Eq)]
129pub enum FormatError {
130    Truncated,
131    LengthOverflow,
132    UnsupportedVersion(u8),
133    UnknownEventTag(u8),
134    InvalidEventBody,
135    InvalidFragmentKind(u8),
136    NonZeroReserved,
137    NonZeroPadding,
138    NonZeroStagedConfirmation,
139    InvalidTxIdSlotLength(usize),
140    InvalidUtf8,
141    InvalidSpeakerLabel,
142    InvalidFeatureBody,
143    InvalidBoolean(u8),
144    TrailingBytes,
145    InvalidPath,
146}
147
148impl Display for FormatError {
149    fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
150        match self {
151            Self::Truncated => formatter.write_str("truncated input"),
152            Self::LengthOverflow => formatter.write_str("encoded length overflow"),
153            Self::UnsupportedVersion(version) => {
154                write!(formatter, "unsupported format version {version}")
155            }
156            Self::UnknownEventTag(tag) => write!(formatter, "unknown event tag {tag}"),
157            Self::InvalidEventBody => formatter.write_str("invalid event body"),
158            Self::InvalidFragmentKind(kind) => {
159                write!(formatter, "invalid fragment kind {kind}")
160            }
161            Self::NonZeroReserved => formatter.write_str("nonzero reserved bytes"),
162            Self::NonZeroPadding => formatter.write_str("nonzero transaction ID slot padding"),
163            Self::NonZeroStagedConfirmation => {
164                formatter.write_str("nonzero staged confirmation slot")
165            }
166            Self::InvalidTxIdSlotLength(length) => {
167                write!(formatter, "invalid transaction ID slot length {length}")
168            }
169            Self::InvalidUtf8 => formatter.write_str("invalid UTF-8"),
170            Self::InvalidSpeakerLabel => formatter.write_str("invalid speaker label"),
171            Self::InvalidFeatureBody => formatter.write_str("invalid feature body"),
172            Self::InvalidBoolean(value) => write!(formatter, "invalid boolean byte {value}"),
173            Self::TrailingBytes => formatter.write_str("trailing bytes"),
174            Self::InvalidPath => formatter.write_str("invalid transaction ID path"),
175        }
176    }
177}
178
179impl std::error::Error for FormatError {}
180
181#[derive(Debug, Clone, Copy, PartialEq, Eq)]
182pub struct TxIdSlot {
183    txid: TxId,
184}
185
186impl TxIdSlot {
187    pub const LEN: usize = 16;
188    pub const PADDING_LEN: usize = 4;
189
190    pub const fn new(txid: TxId) -> Self {
191        Self { txid }
192    }
193
194    pub const fn txid(self) -> TxId {
195        self.txid
196    }
197
198    pub fn encode(self) -> [u8; Self::LEN] {
199        let mut encoded = [0_u8; Self::LEN];
200        encoded[..12].copy_from_slice(self.txid.as_bytes());
201        encoded
202    }
203
204    pub fn decode(bytes: &[u8]) -> Result<Self, FormatError> {
205        if bytes.len() != Self::LEN {
206            return Err(FormatError::InvalidTxIdSlotLength(bytes.len()));
207        }
208        if bytes[12..].iter().any(|byte| *byte != 0) {
209            return Err(FormatError::NonZeroPadding);
210        }
211        let mut txid_bytes = [0_u8; 12];
212        txid_bytes.copy_from_slice(&bytes[..12]);
213        Ok(Self::new(TxId::from_bytes(txid_bytes)))
214    }
215}
216
217pub fn encode_event(event: &AudioClassificationEventV1) -> Result<Vec<u8>, FormatError> {
218    let (tag, body) = match event {
219        AudioClassificationEventV1::Queue(value) => (QUEUE_TAG, postcard::to_allocvec(value)),
220        AudioClassificationEventV1::Processed(value) => {
221            (PROCESSED_TAG, postcard::to_allocvec(value))
222        }
223        AudioClassificationEventV1::Failed(value) => (FAILED_TAG, postcard::to_allocvec(value)),
224        AudioClassificationEventV1::Discarded(value) => {
225            (DISCARDED_TAG, postcard::to_allocvec(value))
226        }
227        AudioClassificationEventV1::Confirmed(value) => {
228            (CONFIRMED_TAG, postcard::to_allocvec(value))
229        }
230    };
231    let body = body.map_err(|_| FormatError::InvalidEventBody)?;
232    let mut encoded = Vec::with_capacity(2 + body.len());
233    encoded.push(EVENT_VERSION);
234    encoded.push(tag);
235    encoded.extend_from_slice(&body);
236    Ok(encoded)
237}
238
239pub fn decode_event(bytes: &[u8]) -> Result<AudioClassificationEventV1, FormatError> {
240    if bytes.len() < 2 {
241        return Err(FormatError::Truncated);
242    }
243    if bytes[0] != EVENT_VERSION {
244        return Err(FormatError::UnsupportedVersion(bytes[0]));
245    }
246    let body = &bytes[2..];
247    match bytes[1] {
248        QUEUE_TAG => decode_event_body(body).map(AudioClassificationEventV1::Queue),
249        PROCESSED_TAG => decode_event_body(body).map(AudioClassificationEventV1::Processed),
250        FAILED_TAG => decode_event_body(body).map(AudioClassificationEventV1::Failed),
251        DISCARDED_TAG => decode_event_body(body).map(AudioClassificationEventV1::Discarded),
252        CONFIRMED_TAG => decode_event_body(body).map(AudioClassificationEventV1::Confirmed),
253        tag => Err(FormatError::UnknownEventTag(tag)),
254    }
255}
256
257pub fn encode_staged_fragment(value: &StagedFragmentV1) -> Result<Vec<u8>, FormatError> {
258    let mut encoded = encode_fragment_prefix(STAGED_KIND, value.analysis_txid, None);
259    append_string(&mut encoded, &value.transcript)?;
260    append_length(&mut encoded, value.speakers.len())?;
261    for speaker in &value.speakers {
262        append_string(&mut encoded, &speaker.speaker.to_string())?;
263        append_string(&mut encoded, &speaker.language)?;
264        append_features(&mut encoded, &speaker.features)?;
265        encoded.push(u8::from(speaker.usable_for_training));
266    }
267    Ok(encoded)
268}
269
270pub fn decode_staged_fragment(bytes: &[u8]) -> Result<StagedFragmentV1, FormatError> {
271    let (analysis_txid, _) = decode_fragment_header(bytes, STAGED_KIND, false)?;
272    let mut decoder = BodyDecoder::new(&bytes[BODY_OFFSET..]);
273    let transcript = decoder.read_string()?;
274    let speaker_count = decoder.read_count()?;
275    let mut speakers = Vec::new();
276    for _ in 0..speaker_count {
277        speakers.push(StagedSpeakerV1 {
278            speaker: decoder.read_speaker_label()?,
279            language: decoder.read_string()?,
280            features: decoder.read_features()?,
281            usable_for_training: decoder.read_boolean()?,
282        });
283    }
284    decoder.finish()?;
285    Ok(StagedFragmentV1 {
286        analysis_txid,
287        transcript,
288        speakers,
289    })
290}
291
292pub fn encode_final_fragment(value: &FinalFragmentV1) -> Result<Vec<u8>, FormatError> {
293    let mut encoded = encode_fragment_prefix(
294        FINAL_KIND,
295        value.analysis_txid,
296        Some(value.confirmation_txid),
297    );
298    append_string(&mut encoded, &value.transcript)?;
299    append_length(&mut encoded, value.speakers.len())?;
300    for speaker in &value.speakers {
301        append_string(&mut encoded, &speaker.speaker.to_string())?;
302        append_string(&mut encoded, &speaker.person_id)?;
303        append_string(&mut encoded, &speaker.language)?;
304        append_features(&mut encoded, &speaker.features)?;
305        encoded.push(u8::from(speaker.usable_for_training));
306    }
307    Ok(encoded)
308}
309
310pub fn decode_final_fragment(bytes: &[u8]) -> Result<FinalFragmentV1, FormatError> {
311    let (analysis_txid, confirmation_txid) = decode_fragment_header(bytes, FINAL_KIND, true)?;
312    let mut decoder = BodyDecoder::new(&bytes[BODY_OFFSET..]);
313    let transcript = decoder.read_string()?;
314    let speaker_count = decoder.read_count()?;
315    let mut speakers = Vec::new();
316    for _ in 0..speaker_count {
317        speakers.push(FinalSpeakerV1 {
318            speaker: decoder.read_speaker_label()?,
319            person_id: decoder.read_string()?,
320            language: decoder.read_string()?,
321            features: decoder.read_features()?,
322            usable_for_training: decoder.read_boolean()?,
323        });
324    }
325    decoder.finish()?;
326    Ok(FinalFragmentV1 {
327        analysis_txid,
328        confirmation_txid: confirmation_txid.expect("final header has a confirmation slot"),
329        transcript,
330        speakers,
331    })
332}
333
334pub fn txid_path(txid: TxId) -> PathBuf {
335    let encoded = encode_txid_base64(txid);
336    let first = String::from_utf8(encoded[..1].to_vec()).expect("base64 is ASCII");
337    let remaining = String::from_utf8(encoded[1..].to_vec()).expect("base64 is ASCII");
338    PathBuf::from(first).join(format!("{remaining}.dat"))
339}
340
341pub fn txid_from_path(path: impl AsRef<Path>) -> Result<TxId, FormatError> {
342    let mut components = path.as_ref().components();
343    let first = normal_utf8_component(components.next())?;
344    let filename = normal_utf8_component(components.next())?;
345    if components.next().is_some() || first.len() != 1 {
346        return Err(FormatError::InvalidPath);
347    }
348    let remaining = filename
349        .strip_suffix(".dat")
350        .ok_or(FormatError::InvalidPath)?;
351    if remaining.len() != 15 {
352        return Err(FormatError::InvalidPath);
353    }
354    let mut encoded = [0_u8; 16];
355    encoded[0] = first.as_bytes()[0];
356    encoded[1..].copy_from_slice(remaining.as_bytes());
357    decode_txid_base64(encoded)
358}
359
360fn decode_event_body<T>(bytes: &[u8]) -> Result<T, FormatError>
361where
362    T: for<'de> Deserialize<'de>,
363{
364    let (value, remaining) =
365        postcard::take_from_bytes(bytes).map_err(|_| FormatError::InvalidEventBody)?;
366    if !remaining.is_empty() {
367        return Err(FormatError::TrailingBytes);
368    }
369    Ok(value)
370}
371
372fn encode_fragment_prefix(kind: u8, analysis: TxId, confirmation: Option<TxId>) -> Vec<u8> {
373    let mut encoded = vec![0_u8; BODY_OFFSET];
374    encoded[0] = FRAGMENT_VERSION;
375    encoded[1] = kind;
376    encoded[ANALYSIS_SLOT_OFFSET..ANALYSIS_SLOT_OFFSET + TxIdSlot::LEN]
377        .copy_from_slice(&TxIdSlot::new(analysis).encode());
378    if let Some(confirmation) = confirmation {
379        encoded[CONFIRMATION_SLOT_OFFSET..CONFIRMATION_SLOT_OFFSET + TxIdSlot::LEN]
380            .copy_from_slice(&TxIdSlot::new(confirmation).encode());
381    }
382    encoded
383}
384
385fn decode_fragment_header(
386    bytes: &[u8],
387    expected_kind: u8,
388    needs_confirmation: bool,
389) -> Result<(TxId, Option<TxId>), FormatError> {
390    if bytes.len() < BODY_OFFSET {
391        return Err(FormatError::Truncated);
392    }
393    if bytes[0] != FRAGMENT_VERSION {
394        return Err(FormatError::UnsupportedVersion(bytes[0]));
395    }
396    if bytes[1] != expected_kind {
397        return Err(FormatError::InvalidFragmentKind(bytes[1]));
398    }
399    if bytes[2..HEADER_LEN].iter().any(|byte| *byte != 0) {
400        return Err(FormatError::NonZeroReserved);
401    }
402    let analysis =
403        TxIdSlot::decode(&bytes[ANALYSIS_SLOT_OFFSET..ANALYSIS_SLOT_OFFSET + TxIdSlot::LEN])?
404            .txid();
405    let confirmation_bytes =
406        &bytes[CONFIRMATION_SLOT_OFFSET..CONFIRMATION_SLOT_OFFSET + TxIdSlot::LEN];
407    let confirmation = if needs_confirmation {
408        Some(TxIdSlot::decode(confirmation_bytes)?.txid())
409    } else {
410        if confirmation_bytes.iter().any(|byte| *byte != 0) {
411            return Err(FormatError::NonZeroStagedConfirmation);
412        }
413        None
414    };
415    Ok((analysis, confirmation))
416}
417
418fn append_length(encoded: &mut Vec<u8>, length: usize) -> Result<(), FormatError> {
419    let length = u64::try_from(length).map_err(|_| FormatError::LengthOverflow)?;
420    encoded.extend_from_slice(&length.to_le_bytes());
421    Ok(())
422}
423
424fn append_bytes(encoded: &mut Vec<u8>, bytes: &[u8]) -> Result<(), FormatError> {
425    append_length(encoded, bytes.len())?;
426    encoded.extend_from_slice(bytes);
427    Ok(())
428}
429
430fn append_string(encoded: &mut Vec<u8>, value: &str) -> Result<(), FormatError> {
431    append_bytes(encoded, value.as_bytes())
432}
433
434fn append_features(encoded: &mut Vec<u8>, features: &FeatureVector24) -> Result<(), FormatError> {
435    let feature_bytes =
436        postcard::to_allocvec(features).map_err(|_| FormatError::InvalidFeatureBody)?;
437    append_bytes(encoded, &feature_bytes)
438}
439
440struct BodyDecoder<'a> {
441    bytes: &'a [u8],
442    position: usize,
443}
444
445impl<'a> BodyDecoder<'a> {
446    fn new(bytes: &'a [u8]) -> Self {
447        Self { bytes, position: 0 }
448    }
449
450    fn take(&mut self, length: usize) -> Result<&'a [u8], FormatError> {
451        let end = self
452            .position
453            .checked_add(length)
454            .ok_or(FormatError::LengthOverflow)?;
455        let value = self
456            .bytes
457            .get(self.position..end)
458            .ok_or(FormatError::Truncated)?;
459        self.position = end;
460        Ok(value)
461    }
462
463    fn read_u64(&mut self) -> Result<u64, FormatError> {
464        let bytes = self.take(8)?;
465        let mut value = [0_u8; 8];
466        value.copy_from_slice(bytes);
467        Ok(u64::from_le_bytes(value))
468    }
469
470    fn read_length(&mut self) -> Result<usize, FormatError> {
471        usize::try_from(self.read_u64()?).map_err(|_| FormatError::LengthOverflow)
472    }
473
474    fn read_count(&mut self) -> Result<usize, FormatError> {
475        self.read_length()
476    }
477
478    fn read_bytes(&mut self) -> Result<&'a [u8], FormatError> {
479        let length = self.read_length()?;
480        self.take(length)
481    }
482
483    fn read_string(&mut self) -> Result<String, FormatError> {
484        let bytes = self.read_bytes()?;
485        let value = std::str::from_utf8(bytes).map_err(|_| FormatError::InvalidUtf8)?;
486        Ok(value.to_owned())
487    }
488
489    fn read_speaker_label(&mut self) -> Result<LocalSpeakerLabel, FormatError> {
490        self.read_string()?
491            .parse()
492            .map_err(|_| FormatError::InvalidSpeakerLabel)
493    }
494
495    fn read_features(&mut self) -> Result<FeatureVector24, FormatError> {
496        let bytes = self.read_bytes()?;
497        let (features, remaining) =
498            postcard::take_from_bytes(bytes).map_err(|_| FormatError::InvalidFeatureBody)?;
499        if !remaining.is_empty() {
500            return Err(FormatError::InvalidFeatureBody);
501        }
502        Ok(features)
503    }
504
505    fn read_boolean(&mut self) -> Result<bool, FormatError> {
506        match self.take(1)?[0] {
507            0 => Ok(false),
508            1 => Ok(true),
509            value => Err(FormatError::InvalidBoolean(value)),
510        }
511    }
512
513    fn finish(self) -> Result<(), FormatError> {
514        if self.position == self.bytes.len() {
515            Ok(())
516        } else {
517            Err(FormatError::TrailingBytes)
518        }
519    }
520}
521
522fn normal_utf8_component(component: Option<Component<'_>>) -> Result<&str, FormatError> {
523    match component {
524        Some(Component::Normal(value)) => value.to_str().ok_or(FormatError::InvalidPath),
525        _ => Err(FormatError::InvalidPath),
526    }
527}
528
529fn encode_txid_base64(txid: TxId) -> [u8; 16] {
530    let bytes = txid.into_bytes();
531    let mut encoded = [0_u8; 16];
532    for group in 0..4 {
533        let input = group * 3;
534        let output = group * 4;
535        encoded[output] = BASE64_ALPHABET[(bytes[input] >> 2) as usize];
536        encoded[output + 1] = BASE64_ALPHABET
537            [(((bytes[input] & 0b0000_0011) << 4) | (bytes[input + 1] >> 4)) as usize];
538        encoded[output + 2] = BASE64_ALPHABET
539            [(((bytes[input + 1] & 0b0000_1111) << 2) | (bytes[input + 2] >> 6)) as usize];
540        encoded[output + 3] = BASE64_ALPHABET[(bytes[input + 2] & 0b0011_1111) as usize];
541    }
542    encoded
543}
544
545fn decode_txid_base64(encoded: [u8; 16]) -> Result<TxId, FormatError> {
546    let mut bytes = [0_u8; 12];
547    for group in 0..4 {
548        let input = group * 4;
549        let output = group * 3;
550        let first = decode_base64_character(encoded[input])?;
551        let second = decode_base64_character(encoded[input + 1])?;
552        let third = decode_base64_character(encoded[input + 2])?;
553        let fourth = decode_base64_character(encoded[input + 3])?;
554        bytes[output] = (first << 2) | (second >> 4);
555        bytes[output + 1] = (second << 4) | (third >> 2);
556        bytes[output + 2] = (third << 6) | fourth;
557    }
558    Ok(TxId::from_bytes(bytes))
559}
560
561fn decode_base64_character(value: u8) -> Result<u8, FormatError> {
562    match value {
563        b'A'..=b'Z' => Ok(value - b'A'),
564        b'a'..=b'z' => Ok(value - b'a' + 26),
565        b'0'..=b'9' => Ok(value - b'0' + 52),
566        b'-' => Ok(62),
567        b'_' => Ok(63),
568        _ => Err(FormatError::InvalidPath),
569    }
570}
571
572mod txid_serde {
573    use super::TxId;
574    use serde::{Deserialize, Deserializer, Serialize, Serializer};
575
576    pub fn serialize<S>(value: &TxId, serializer: S) -> Result<S::Ok, S::Error>
577    where
578        S: Serializer,
579    {
580        value.into_bytes().serialize(serializer)
581    }
582
583    pub fn deserialize<'de, D>(deserializer: D) -> Result<TxId, D::Error>
584    where
585        D: Deserializer<'de>,
586    {
587        let bytes = <[u8; 12]>::deserialize(deserializer)?;
588        Ok(TxId::from_bytes(bytes))
589    }
590}
591
592mod optional_txid_serde {
593    use super::TxId;
594    use serde::{Deserialize, Deserializer, Serialize, Serializer};
595
596    pub fn serialize<S>(value: &Option<TxId>, serializer: S) -> Result<S::Ok, S::Error>
597    where
598        S: Serializer,
599    {
600        value.map(TxId::into_bytes).serialize(serializer)
601    }
602
603    pub fn deserialize<'de, D>(deserializer: D) -> Result<Option<TxId>, D::Error>
604    where
605        D: Deserializer<'de>,
606    {
607        let bytes = Option::<[u8; 12]>::deserialize(deserializer)?;
608        Ok(bytes.map(TxId::from_bytes))
609    }
610}
611
612#[cfg(test)]
613mod tests {
614    use super::*;
615    use kcode_speaker_v3_analysis::{
616        AnalysisEnvelope, GeminiCohort, OggAudioMetadata, StructuredAnalysis, StructuredSpeaker,
617        StructurerProvenance,
618    };
619
620    fn txid(seed: u8) -> TxId {
621        let mut bytes = [0_u8; 12];
622        for (index, byte) in bytes.iter_mut().enumerate() {
623            *byte = seed.wrapping_add(index as u8);
624        }
625        TxId::from_bytes(bytes)
626    }
627
628    fn features(seed: f64, name: Option<&str>) -> FeatureVector24 {
629        FeatureVector24 {
630            median_f0_hz: Some(seed),
631            high_front_vowel_f1_hz: None,
632            dominant_rhotic_realization: name.map(str::to_owned),
633            hypernasality_0_to_4: Some(seed / 100.0),
634            ..FeatureVector24::default()
635        }
636    }
637
638    fn executed_analysis() -> ExecutedAnalysis {
639        let mut ogg = vec![0_u8; 29];
640        ogg[..4].copy_from_slice(b"OggS");
641        ogg[4] = 0;
642        ogg[26] = 1;
643        ogg[27] = 1;
644        ogg[28] = 0;
645        let audio = OggAudioMetadata::from_bytes(&ogg, 1_250, Some("sample.ogg".to_owned()))
646            .expect("valid Ogg metadata");
647        let analysis = StructuredAnalysis {
648            transcript: "Speaker 1: Héllo from 東京".to_owned(),
649            speakers: vec![StructuredSpeaker {
650                speaker: LocalSpeakerLabel::new(1).expect("valid label"),
651                language: "English / 日本語".to_owned(),
652                features: features(182.5, Some("approximant")),
653                features_usable_for_training: true,
654            }],
655        };
656        ExecutedAnalysis {
657            envelope: AnalysisEnvelope {
658                audio,
659                analysis,
660                gemini: GeminiCohort::new("gemini-3.1-pro"),
661                structurer: StructurerProvenance::new("terra-structurer"),
662            },
663            label_extractor: StructurerProvenance::new("terra-labeler"),
664        }
665    }
666
667    fn staged_fragment() -> StagedFragmentV1 {
668        StagedFragmentV1 {
669            analysis_txid: txid(20),
670            transcript: "Élodie: bonjour 🌍\n話者 2: こんにちは".to_owned(),
671            speakers: vec![
672                StagedSpeakerV1 {
673                    speaker: LocalSpeakerLabel::new(1).expect("valid label"),
674                    language: "français".to_owned(),
675                    features: features(201.25, Some("uvulaire")),
676                    usable_for_training: true,
677                },
678                StagedSpeakerV1 {
679                    speaker: LocalSpeakerLabel::new(2).expect("valid label"),
680                    language: "日本語".to_owned(),
681                    features: FeatureVector24 {
682                        low_vowel_f2_hz: Some(1_234.5),
683                        dominant_lateral_realization: Some("明瞭".to_owned()),
684                        ..FeatureVector24::default()
685                    },
686                    usable_for_training: false,
687                },
688            ],
689        }
690    }
691
692    fn final_fragment() -> FinalFragmentV1 {
693        FinalFragmentV1 {
694            analysis_txid: txid(30),
695            confirmation_txid: txid(40),
696            transcript: "Élodie: bonjour 🌍\n話者 2: こんにちは".to_owned(),
697            speakers: vec![
698                FinalSpeakerV1 {
699                    speaker: LocalSpeakerLabel::new(1).expect("valid label"),
700                    person_id: "person-éloïse".to_owned(),
701                    language: "français".to_owned(),
702                    features: features(201.25, Some("uvulaire")),
703                    usable_for_training: false,
704                },
705                FinalSpeakerV1 {
706                    speaker: LocalSpeakerLabel::new(2).expect("valid label"),
707                    person_id: "人物-東京".to_owned(),
708                    language: "日本語".to_owned(),
709                    features: FeatureVector24::default(),
710                    usable_for_training: true,
711                },
712            ],
713        }
714    }
715
716    fn read_test_length(bytes: &[u8], position: &mut usize) -> usize {
717        let end = *position + 8;
718        let mut value = [0_u8; 8];
719        value.copy_from_slice(&bytes[*position..end]);
720        *position = end;
721        usize::try_from(u64::from_le_bytes(value)).expect("test length fits")
722    }
723
724    fn first_staged_feature_and_boolean_offsets(bytes: &[u8]) -> (usize, usize) {
725        let mut position = BODY_OFFSET;
726        let transcript_length = read_test_length(bytes, &mut position);
727        position += transcript_length;
728        let speaker_count = read_test_length(bytes, &mut position);
729        assert!(speaker_count > 0);
730        let label_length = read_test_length(bytes, &mut position);
731        position += label_length;
732        let language_length = read_test_length(bytes, &mut position);
733        position += language_length;
734        let feature_length = read_test_length(bytes, &mut position);
735        let feature_offset = position;
736        position += feature_length;
737        (feature_offset, position)
738    }
739
740    #[test]
741    fn txid_slot_has_exact_bytes_and_padding() {
742        let id = txid(3);
743        let encoded = TxIdSlot::new(id).encode();
744        assert_eq!(encoded.len(), TxIdSlot::LEN);
745        assert_eq!(TxIdSlot::PADDING_LEN, 4);
746        assert_eq!(&encoded[..12], id.as_bytes());
747        assert_eq!(&encoded[12..], &[0_u8; 4]);
748        assert_eq!(TxIdSlot::decode(&encoded).expect("slot").txid(), id);
749        assert_eq!(
750            TxIdSlot::decode(&encoded[..15]),
751            Err(FormatError::InvalidTxIdSlotLength(15))
752        );
753        let mut corrupt = encoded;
754        corrupt[12] = 1;
755        assert_eq!(TxIdSlot::decode(&corrupt), Err(FormatError::NonZeroPadding));
756    }
757
758    #[test]
759    fn synthetic_slots_start_on_sixteen_byte_boundaries() {
760        let mut bytes = [0_u8; 64];
761        for (index, offset) in (0..64).step_by(TxIdSlot::LEN).enumerate() {
762            assert_eq!(offset % TxIdSlot::LEN, 0);
763            let slot = TxIdSlot::new(txid(index as u8)).encode();
764            bytes[offset..offset + TxIdSlot::LEN].copy_from_slice(&slot);
765        }
766        for (index, offset) in (0..64).step_by(TxIdSlot::LEN).enumerate() {
767            let slot =
768                TxIdSlot::decode(&bytes[offset..offset + TxIdSlot::LEN]).expect("aligned slot");
769            assert_eq!(slot.txid(), txid(index as u8));
770        }
771    }
772
773    #[test]
774    fn fragment_slots_are_at_fixed_offsets() {
775        let staged = staged_fragment();
776        let staged_bytes = encode_staged_fragment(&staged).expect("encode staged");
777        assert_eq!(
778            &staged_bytes[16..32],
779            &TxIdSlot::new(staged.analysis_txid).encode()
780        );
781        assert_eq!(&staged_bytes[32..48], &[0_u8; 16]);
782
783        let final_value = final_fragment();
784        let final_bytes = encode_final_fragment(&final_value).expect("encode final");
785        assert_eq!(
786            &final_bytes[16..32],
787            &TxIdSlot::new(final_value.analysis_txid).encode()
788        );
789        assert_eq!(
790            &final_bytes[32..48],
791            &TxIdSlot::new(final_value.confirmation_txid).encode()
792        );
793    }
794
795    #[test]
796    fn every_shard_character_roundtrips() {
797        for shard in 0..64 {
798            let mut bytes = [0_u8; 12];
799            bytes[0] = (shard as u8) << 2;
800            bytes[11] = shard as u8;
801            let id = TxId::from_bytes(bytes);
802            let path = txid_path(id);
803            let first = path
804                .components()
805                .next()
806                .expect("first component")
807                .as_os_str()
808                .to_str()
809                .expect("UTF-8");
810            assert_eq!(first.as_bytes(), &BASE64_ALPHABET[shard..shard + 1]);
811            assert_eq!(txid_from_path(&path), Ok(id));
812        }
813    }
814
815    #[test]
816    fn malformed_paths_are_rejected() {
817        for path in [
818            "",
819            "A",
820            "A/B.dat",
821            "AA/AAAAAAAAAAAAAAA.dat",
822            "A/AAAAAAAAAAAAAAA.bin",
823            "A/AAAAAAAAAAAAAA!.dat",
824            "/A/AAAAAAAAAAAAAAA.dat",
825            "A/../B.dat",
826            "A/AAAAAAAAAAAAAAA.dat/extra",
827        ] {
828            assert_eq!(txid_from_path(path), Err(FormatError::InvalidPath));
829        }
830    }
831
832    #[test]
833    fn every_event_variant_roundtrips_with_stable_tags() {
834        let events = vec![
835            AudioClassificationEventV1::Queue(QueueV1 {
836                audio_object_id: txid(1),
837                duration_ms: 98_765,
838                filename: Some("réunion.ogg".to_owned()),
839                retry_of: Some(txid(2)),
840            }),
841            AudioClassificationEventV1::Processed(ProcessedV1 {
842                queue_id: txid(3),
843                audio_object_id: txid(4),
844                analysis: executed_analysis(),
845            }),
846            AudioClassificationEventV1::Failed(FailedV1 {
847                queue_id: txid(5),
848                audio_object_id: txid(6),
849                final_stage: AnalysisStageV1::GeminiFeatures,
850                final_error: "packet 2 unavailable".to_owned(),
851            }),
852            AudioClassificationEventV1::Discarded(DiscardedV1 {
853                failed_queue_id: txid(7),
854                audio_object_id: txid(8),
855            }),
856            AudioClassificationEventV1::Confirmed(ConfirmedV1 {
857                queue_id: txid(9),
858                audio_object_id: txid(10),
859                analysis_txid: txid(11),
860                speakers: vec![ConfirmedSpeakerV1 {
861                    speaker: LocalSpeakerLabel::new(1).expect("valid label"),
862                    person_id: "person-α".to_owned(),
863                }],
864            }),
865        ];
866        for (index, event) in events.into_iter().enumerate() {
867            let encoded = encode_event(&event).expect("encode event");
868            assert_eq!(encoded[0], EVENT_VERSION);
869            assert_eq!(encoded[1], index as u8 + 1);
870            assert_eq!(decode_event(&encoded), Ok(event));
871        }
872    }
873
874    #[test]
875    fn every_analysis_stage_roundtrips() {
876        for stage in [
877            AnalysisStageV1::GeminiTranscript,
878            AnalysisStageV1::TerraLabels,
879            AnalysisStageV1::GeminiFeatures,
880            AnalysisStageV1::TerraStructuring,
881        ] {
882            let event = AudioClassificationEventV1::Failed(FailedV1 {
883                queue_id: txid(1),
884                audio_object_id: txid(2),
885                final_stage: stage,
886                final_error: "failure".to_owned(),
887            });
888            let encoded = encode_event(&event).expect("encode event");
889            assert_eq!(decode_event(&encoded), Ok(event));
890        }
891    }
892
893    #[test]
894    fn event_corruption_is_rejected() {
895        let event = AudioClassificationEventV1::Queue(QueueV1 {
896            audio_object_id: txid(1),
897            duration_ms: 500,
898            filename: None,
899            retry_of: None,
900        });
901        let encoded = encode_event(&event).expect("encode event");
902        assert_eq!(decode_event(&[]), Err(FormatError::Truncated));
903        let mut wrong_version = encoded.clone();
904        wrong_version[0] = 2;
905        assert_eq!(
906            decode_event(&wrong_version),
907            Err(FormatError::UnsupportedVersion(2))
908        );
909        let mut wrong_tag = encoded.clone();
910        wrong_tag[1] = 99;
911        assert_eq!(
912            decode_event(&wrong_tag),
913            Err(FormatError::UnknownEventTag(99))
914        );
915        assert_eq!(
916            decode_event(&encoded[..2]),
917            Err(FormatError::InvalidEventBody)
918        );
919        let mut trailing = encoded;
920        trailing.push(0);
921        assert_eq!(decode_event(&trailing), Err(FormatError::TrailingBytes));
922    }
923
924    #[test]
925    fn staged_and_final_fragments_roundtrip_utf8_speakers() {
926        let staged = staged_fragment();
927        let staged_bytes = encode_staged_fragment(&staged).expect("encode staged");
928        assert_eq!(decode_staged_fragment(&staged_bytes), Ok(staged));
929
930        let final_value = final_fragment();
931        let final_bytes = encode_final_fragment(&final_value).expect("encode final");
932        assert_eq!(decode_final_fragment(&final_bytes), Ok(final_value));
933    }
934
935    #[test]
936    fn fragment_header_and_slot_corruption_is_rejected() {
937        let staged = staged_fragment();
938        let encoded = encode_staged_fragment(&staged).expect("encode staged");
939        for cut in [0, 15, 47, encoded.len() - 1] {
940            assert!(decode_staged_fragment(&encoded[..cut]).is_err());
941        }
942        let mut wrong_version = encoded.clone();
943        wrong_version[0] = 2;
944        assert_eq!(
945            decode_staged_fragment(&wrong_version),
946            Err(FormatError::UnsupportedVersion(2))
947        );
948        let mut wrong_kind = encoded.clone();
949        wrong_kind[1] = FINAL_KIND;
950        assert_eq!(
951            decode_staged_fragment(&wrong_kind),
952            Err(FormatError::InvalidFragmentKind(FINAL_KIND))
953        );
954        let mut reserved = encoded.clone();
955        reserved[2] = 1;
956        assert_eq!(
957            decode_staged_fragment(&reserved),
958            Err(FormatError::NonZeroReserved)
959        );
960        let mut analysis_padding = encoded.clone();
961        analysis_padding[28] = 1;
962        assert_eq!(
963            decode_staged_fragment(&analysis_padding),
964            Err(FormatError::NonZeroPadding)
965        );
966        let mut staged_confirmation = encoded;
967        staged_confirmation[32] = 1;
968        assert_eq!(
969            decode_staged_fragment(&staged_confirmation),
970            Err(FormatError::NonZeroStagedConfirmation)
971        );
972
973        let final_value = final_fragment();
974        let mut final_bytes = encode_final_fragment(&final_value).expect("encode final");
975        final_bytes[44] = 1;
976        assert_eq!(
977            decode_final_fragment(&final_bytes),
978            Err(FormatError::NonZeroPadding)
979        );
980    }
981
982    #[test]
983    fn fragment_body_corruption_is_rejected() {
984        let value = StagedFragmentV1 {
985            analysis_txid: txid(1),
986            transcript: "hello".to_owned(),
987            speakers: vec![StagedSpeakerV1 {
988                speaker: LocalSpeakerLabel::new(1).expect("valid label"),
989                language: "English".to_owned(),
990                features: FeatureVector24::default(),
991                usable_for_training: true,
992            }],
993        };
994        let encoded = encode_staged_fragment(&value).expect("encode staged");
995
996        let mut length_overflow = encoded.clone();
997        length_overflow[BODY_OFFSET..BODY_OFFSET + 8].copy_from_slice(&u64::MAX.to_le_bytes());
998        assert_eq!(
999            decode_staged_fragment(&length_overflow),
1000            Err(FormatError::LengthOverflow)
1001        );
1002
1003        let mut invalid_utf8 = encoded.clone();
1004        invalid_utf8[BODY_OFFSET + 8] = 0xff;
1005        assert_eq!(
1006            decode_staged_fragment(&invalid_utf8),
1007            Err(FormatError::InvalidUtf8)
1008        );
1009
1010        let (feature_offset, boolean_offset) = first_staged_feature_and_boolean_offsets(&encoded);
1011        let mut invalid_feature = encoded.clone();
1012        invalid_feature[feature_offset] = 2;
1013        assert_eq!(
1014            decode_staged_fragment(&invalid_feature),
1015            Err(FormatError::InvalidFeatureBody)
1016        );
1017
1018        let mut invalid_boolean = encoded.clone();
1019        invalid_boolean[boolean_offset] = 2;
1020        assert_eq!(
1021            decode_staged_fragment(&invalid_boolean),
1022            Err(FormatError::InvalidBoolean(2))
1023        );
1024
1025        let mut trailing = encoded;
1026        trailing.push(0);
1027        assert_eq!(
1028            decode_staged_fragment(&trailing),
1029            Err(FormatError::TrailingBytes)
1030        );
1031    }
1032
1033    #[test]
1034    fn malformed_label_and_internal_feature_trailing_bytes_are_rejected() {
1035        let value = StagedFragmentV1 {
1036            analysis_txid: txid(1),
1037            transcript: "hello".to_owned(),
1038            speakers: vec![StagedSpeakerV1 {
1039                speaker: LocalSpeakerLabel::new(1).expect("valid label"),
1040                language: "English".to_owned(),
1041                features: FeatureVector24::default(),
1042                usable_for_training: true,
1043            }],
1044        };
1045        let encoded = encode_staged_fragment(&value).expect("encode staged");
1046        let mut position = BODY_OFFSET;
1047        let transcript_length = read_test_length(&encoded, &mut position);
1048        position += transcript_length;
1049        let _ = read_test_length(&encoded, &mut position);
1050        let label_length = read_test_length(&encoded, &mut position);
1051        assert_eq!(label_length, "Speaker 1".len());
1052        let mut invalid_label = encoded.clone();
1053        invalid_label[position] = b'X';
1054        assert_eq!(
1055            decode_staged_fragment(&invalid_label),
1056            Err(FormatError::InvalidSpeakerLabel)
1057        );
1058
1059        let (feature_offset, boolean_offset) = first_staged_feature_and_boolean_offsets(&encoded);
1060        let mut feature_trailing = encoded;
1061        feature_trailing.insert(boolean_offset, 0);
1062        let feature_length_offset = feature_offset - 8;
1063        let feature_length = boolean_offset - feature_offset + 1;
1064        feature_trailing[feature_length_offset..feature_length_offset + 8]
1065            .copy_from_slice(&(feature_length as u64).to_le_bytes());
1066        assert_eq!(
1067            decode_staged_fragment(&feature_trailing),
1068            Err(FormatError::InvalidFeatureBody)
1069        );
1070    }
1071}
kcode_k1_audio_classification_format/lib.rs

kcode_k1_audio_classification_format/
lib.rs
