Skip to main content

kcode_k1_audio_classification_event_format/
lib.rs

1use kcode_k1_person_types::PersonId;
2use serde::{Deserialize, Serialize};
3
4pub use kcode_k1_audio_classification_format_error::FormatError;
5pub use kcode_k1_transaction_id::TxId;
6pub use kcode_speaker_v3_analysis::{ExecutedAnalysis, FeatureVector24, LocalSpeakerLabel};
7
8const EVENT_VERSION: u8 = 5;
9const QUEUE_TAG: u8 = 1;
10const TRANSCRIPTION_COMPLETE_TAG: u8 = 2;
11const FAILED_TAG: u8 = 3;
12const DISCARDED_TAG: u8 = 4;
13const LABEL_CONFIRMATION_TAG: u8 = 5;
14const PROGRESS_TAG: u8 = 6;
15
16#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
17pub enum FragmentStageV1 {
18    Queue,
19    Transcript,
20    SpeakerLabels,
21    SpeakerFeatures,
22    Structuring,
23    LabelConfirmation,
24}
25
26#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
27pub struct QueueV2 {
28    #[serde(with = "txid_serde")]
29    pub audio_object_id: TxId,
30}
31
32#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
33pub struct ProgressV1 {
34    #[serde(with = "txid_serde")]
35    pub fragment_id: TxId,
36    pub update: ProgressUpdateV1,
37}
38
39#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
40pub enum ProgressUpdateV1 {
41    LlmJobStarted {
42        sequence: u64,
43        stage: FragmentStageV1,
44        name: String,
45    },
46    LlmJobSucceeded {
47        sequence: u64,
48    },
49    LlmJobFailed {
50        sequence: u64,
51        error: String,
52    },
53    StageCompleted {
54        stage: FragmentStageV1,
55    },
56}
57
58#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
59pub struct TranscriptionCompleteV1 {
60    #[serde(with = "txid_serde")]
61    pub fragment_id: TxId,
62    pub analysis: ExecutedAnalysis,
63}
64
65#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
66pub struct FailedV2 {
67    #[serde(with = "txid_serde")]
68    pub fragment_id: TxId,
69    pub stage: FragmentStageV1,
70    pub llm_job_sequence: Option<u64>,
71    pub error: String,
72}
73
74#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
75pub struct DiscardedV2 {
76    #[serde(with = "txid_serde")]
77    pub fragment_id: TxId,
78}
79
80#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
81pub struct SpeakerLabelV1 {
82    pub speaker: LocalSpeakerLabel,
83    #[serde(with = "person_id_option_serde")]
84    pub person_id: Option<PersonId>,
85}
86
87#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
88pub struct LabelConfirmationV1 {
89    #[serde(with = "txid_serde")]
90    pub fragment_id: TxId,
91    #[serde(with = "txid_serde")]
92    pub interim_txid: TxId,
93    pub speakers: Vec<SpeakerLabelV1>,
94}
95
96#[allow(clippy::large_enum_variant)]
97#[derive(Debug, Clone, PartialEq)]
98pub enum AudioClassificationEventV3 {
99    Queue(QueueV2),
100    Progress(ProgressV1),
101    TranscriptionComplete(TranscriptionCompleteV1),
102    Failed(FailedV2),
103    Discarded(DiscardedV2),
104    LabelConfirmation(LabelConfirmationV1),
105}
106
107pub fn encode_event(event: &AudioClassificationEventV3) -> Result<Vec<u8>, FormatError> {
108    let (tag, body) = match event {
109        AudioClassificationEventV3::Queue(value) => (QUEUE_TAG, postcard::to_allocvec(value)),
110        AudioClassificationEventV3::TranscriptionComplete(value) => {
111            (TRANSCRIPTION_COMPLETE_TAG, postcard::to_allocvec(value))
112        }
113        AudioClassificationEventV3::Failed(value) => (FAILED_TAG, postcard::to_allocvec(value)),
114        AudioClassificationEventV3::Discarded(value) => {
115            (DISCARDED_TAG, postcard::to_allocvec(value))
116        }
117        AudioClassificationEventV3::LabelConfirmation(value) => {
118            (LABEL_CONFIRMATION_TAG, postcard::to_allocvec(value))
119        }
120        AudioClassificationEventV3::Progress(value) => (PROGRESS_TAG, postcard::to_allocvec(value)),
121    };
122    let body = body.map_err(|_| FormatError::InvalidEventBody)?;
123    let mut output = Vec::with_capacity(2 + body.len());
124    output.extend_from_slice(&[EVENT_VERSION, tag]);
125    output.extend_from_slice(&body);
126    Ok(output)
127}
128
129pub fn decode_event(bytes: &[u8]) -> Result<AudioClassificationEventV3, FormatError> {
130    if bytes.len() < 2 {
131        return Err(FormatError::Truncated);
132    }
133    if bytes[0] != EVENT_VERSION {
134        return Err(FormatError::UnsupportedVersion(bytes[0]));
135    }
136    match bytes[1] {
137        QUEUE_TAG => decode_body(&bytes[2..]).map(AudioClassificationEventV3::Queue),
138        TRANSCRIPTION_COMPLETE_TAG => {
139            decode_body(&bytes[2..]).map(AudioClassificationEventV3::TranscriptionComplete)
140        }
141        FAILED_TAG => decode_body(&bytes[2..]).map(AudioClassificationEventV3::Failed),
142        DISCARDED_TAG => decode_body(&bytes[2..]).map(AudioClassificationEventV3::Discarded),
143        LABEL_CONFIRMATION_TAG => {
144            decode_body(&bytes[2..]).map(AudioClassificationEventV3::LabelConfirmation)
145        }
146        PROGRESS_TAG => decode_body(&bytes[2..]).map(AudioClassificationEventV3::Progress),
147        tag => Err(FormatError::UnknownEventTag(tag)),
148    }
149}
150
151fn decode_body<T: for<'de> Deserialize<'de>>(bytes: &[u8]) -> Result<T, FormatError> {
152    let (value, remaining) =
153        postcard::take_from_bytes(bytes).map_err(|_| FormatError::InvalidEventBody)?;
154    if !remaining.is_empty() {
155        return Err(FormatError::TrailingBytes);
156    }
157    Ok(value)
158}
159
160mod txid_serde {
161    use super::TxId;
162    use serde::{Deserialize, Deserializer, Serialize, Serializer};
163
164    pub fn serialize<S: Serializer>(value: &TxId, serializer: S) -> Result<S::Ok, S::Error> {
165        value.into_bytes().serialize(serializer)
166    }
167
168    pub fn deserialize<'de, D: Deserializer<'de>>(deserializer: D) -> Result<TxId, D::Error> {
169        Ok(TxId::from_bytes(<[u8; 12]>::deserialize(deserializer)?))
170    }
171}
172
173mod person_id_option_serde {
174    use super::{PersonId, TxId};
175    use serde::{Deserialize, Deserializer, Serialize, Serializer};
176
177    pub fn serialize<S: Serializer>(
178        value: &Option<PersonId>,
179        serializer: S,
180    ) -> Result<S::Ok, S::Error> {
181        value
182            .map(|person_id| person_id.as_tx_id().into_bytes())
183            .serialize(serializer)
184    }
185
186    pub fn deserialize<'de, D: Deserializer<'de>>(
187        deserializer: D,
188    ) -> Result<Option<PersonId>, D::Error> {
189        Ok(Option::<[u8; 12]>::deserialize(deserializer)?
190            .map(|bytes| PersonId::from_tx_id(TxId::from_bytes(bytes))))
191    }
192}