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 EVENT_VERSION_V6: u8 = 6;
10const FINAL_EVENT_VERSION_V7: u8 = 7;
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
13pub enum FragmentStageV1 {
14    Queue,
15    Transcript,
16    SpeakerLabels,
17    SpeakerFeatures,
18    Structuring,
19    LabelConfirmation,
20}
21
22#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
23pub struct QueueV2 {
24    #[serde(with = "txid_serde")]
25    pub audio_object_id: TxId,
26}
27
28#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
29pub struct ProgressV1 {
30    #[serde(with = "txid_serde")]
31    pub fragment_id: TxId,
32    pub update: ProgressUpdateV1,
33}
34
35#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
36pub enum ProgressUpdateV1 {
37    LlmJobStarted {
38        sequence: u64,
39        stage: FragmentStageV1,
40        name: String,
41    },
42    LlmJobSucceeded {
43        sequence: u64,
44    },
45    LlmJobFailed {
46        sequence: u64,
47        error: String,
48    },
49    StageCompleted {
50        stage: FragmentStageV1,
51    },
52}
53
54#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
55pub struct TranscriptionCompleteV1 {
56    #[serde(with = "txid_serde")]
57    pub fragment_id: TxId,
58    pub analysis: ExecutedAnalysis,
59}
60
61#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
62pub struct FailedV2 {
63    #[serde(with = "txid_serde")]
64    pub fragment_id: TxId,
65    pub stage: FragmentStageV1,
66    pub llm_job_sequence: Option<u64>,
67    pub error: String,
68}
69
70#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
71pub struct DiscardedV2 {
72    #[serde(with = "txid_serde")]
73    pub fragment_id: TxId,
74}
75
76#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
77pub struct SpeakerLabelV1 {
78    pub speaker: LocalSpeakerLabel,
79    #[serde(with = "person_id_option_serde")]
80    pub person_id: Option<PersonId>,
81}
82
83#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
84pub struct LabelConfirmationV1 {
85    #[serde(with = "txid_serde")]
86    pub fragment_id: TxId,
87    #[serde(with = "txid_serde")]
88    pub interim_txid: TxId,
89    pub speakers: Vec<SpeakerLabelV1>,
90}
91
92#[allow(clippy::large_enum_variant)]
93#[derive(Debug, Clone, PartialEq)]
94pub enum AudioClassificationEventV3 {
95    Queue(QueueV2),
96    Progress(ProgressV1),
97    TranscriptionComplete(TranscriptionCompleteV1),
98    Failed(FailedV2),
99    Discarded(DiscardedV2),
100    LabelConfirmation(LabelConfirmationV1),
101}
102
103#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
104pub struct AttemptStartedV1 {
105    #[serde(with = "txid_serde")]
106    pub fragment_id: TxId,
107}
108
109macro_rules! attempted_payload {
110    ($name:ident { $($field:ident: $field_type:ty),+ $(,)? }) => {
111        #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
112        pub struct $name {
113            #[serde(with = "txid_serde")]
114            pub fragment_id: TxId,
115            #[serde(with = "txid_serde")]
116            pub attempt_txid: TxId,
117            $(pub $field: $field_type),+
118        }
119    };
120}
121
122attempted_payload!(AttemptProgressV1 {
123    update: ProgressUpdateV1
124});
125attempted_payload!(GeminiTranscriptV1 { transcript: String });
126attempted_payload!(TerraSpeakerLabelsV1 { speakers: Vec<LocalSpeakerLabel> });
127attempted_payload!(GeminiFeatureBundleV1 {
128    speaker: LocalSpeakerLabel,
129    packets: [String; 3],
130});
131attempted_payload!(AttemptFailedV1 {
132    stage: FragmentStageV1,
133    llm_job_sequence: Option<u64>,
134    error: String,
135});
136
137#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
138pub enum AudioClassificationEventV6 {
139    AttemptStarted(AttemptStartedV1),
140    Progress(AttemptProgressV1),
141    GeminiTranscript(GeminiTranscriptV1),
142    TerraSpeakerLabels(TerraSpeakerLabelsV1),
143    GeminiFeatureBundle(GeminiFeatureBundleV1),
144    Failed(AttemptFailedV1),
145}
146
147#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
148pub struct AttemptFinalAnalysisV1 {
149    #[serde(with = "txid_serde")]
150    pub fragment_id: TxId,
151    #[serde(with = "txid_serde")]
152    pub attempt_txid: TxId,
153    pub completion_v5: Vec<u8>,
154}
155
156fn encode_frame<T: Serialize>(version: u8, tag: u8, value: &T) -> Result<Vec<u8>, FormatError> {
157    let body = postcard::to_allocvec(value).map_err(|_| FormatError::InvalidEventBody)?;
158    let mut output = Vec::with_capacity(2 + body.len());
159    output.extend_from_slice(&[version, tag]);
160    output.extend_from_slice(&body);
161    Ok(output)
162}
163
164fn frame_body(bytes: &[u8], version: u8) -> Result<(u8, &[u8]), FormatError> {
165    if bytes.len() < 2 {
166        return Err(FormatError::Truncated);
167    }
168    if bytes[0] != version {
169        return Err(FormatError::UnsupportedVersion(bytes[0]));
170    }
171    Ok((bytes[1], &bytes[2..]))
172}
173
174fn decode_body<T: for<'de> Deserialize<'de>>(bytes: &[u8]) -> Result<T, FormatError> {
175    let (value, remaining) =
176        postcard::take_from_bytes(bytes).map_err(|_| FormatError::InvalidEventBody)?;
177    if !remaining.is_empty() {
178        return Err(FormatError::TrailingBytes);
179    }
180    Ok(value)
181}
182
183macro_rules! event_codec {
184    ($encode:ident, $decode:ident, $event:ident, $version:expr, {
185        $($tag:literal => $variant:ident),+ $(,)?
186    }) => {
187        pub fn $encode(value: &$event) -> Result<Vec<u8>, FormatError> {
188            match value {
189                $($event::$variant(body) => encode_frame($version, $tag, body),)+
190            }
191        }
192
193        pub fn $decode(bytes: &[u8]) -> Result<$event, FormatError> {
194            let (tag, body) = frame_body(bytes, $version)?;
195            match tag {
196                $($tag => decode_body(body).map($event::$variant),)+
197                tag => Err(FormatError::UnknownEventTag(tag)),
198            }
199        }
200    };
201}
202
203event_codec!(encode_event, decode_event, AudioClassificationEventV3, EVENT_VERSION, {
204    1 => Queue,
205    2 => TranscriptionComplete,
206    3 => Failed,
207    4 => Discarded,
208    5 => LabelConfirmation,
209    6 => Progress,
210});
211
212event_codec!(encode_event_v6, decode_event_v6, AudioClassificationEventV6, EVENT_VERSION_V6, {
213    1 => AttemptStarted,
214    2 => Progress,
215    3 => GeminiTranscript,
216    4 => TerraSpeakerLabels,
217    5 => GeminiFeatureBundle,
218    6 => Failed,
219});
220
221pub fn encode_final_event_v7(value: &AttemptFinalAnalysisV1) -> Result<Vec<u8>, FormatError> {
222    encode_frame(FINAL_EVENT_VERSION_V7, 1, value)
223}
224
225pub fn decode_final_event_v7(bytes: &[u8]) -> Result<AttemptFinalAnalysisV1, FormatError> {
226    let (tag, body) = frame_body(bytes, FINAL_EVENT_VERSION_V7)?;
227    match tag {
228        1 => decode_body(body),
229        tag => Err(FormatError::UnknownEventTag(tag)),
230    }
231}
232
233mod txid_serde {
234    use super::TxId;
235    use serde::{Deserialize, Deserializer, Serialize, Serializer};
236
237    pub fn serialize<S: Serializer>(value: &TxId, serializer: S) -> Result<S::Ok, S::Error> {
238        value.into_bytes().serialize(serializer)
239    }
240
241    pub fn deserialize<'de, D: Deserializer<'de>>(deserializer: D) -> Result<TxId, D::Error> {
242        Ok(TxId::from_bytes(<[u8; 12]>::deserialize(deserializer)?))
243    }
244}
245
246mod person_id_option_serde {
247    use super::{PersonId, TxId};
248    use serde::{Deserialize, Deserializer, Serialize, Serializer};
249
250    pub fn serialize<S: Serializer>(
251        value: &Option<PersonId>,
252        serializer: S,
253    ) -> Result<S::Ok, S::Error> {
254        value
255            .map(|person_id| person_id.as_tx_id().into_bytes())
256            .serialize(serializer)
257    }
258
259    pub fn deserialize<'de, D: Deserializer<'de>>(
260        deserializer: D,
261    ) -> Result<Option<PersonId>, D::Error> {
262        Ok(Option::<[u8; 12]>::deserialize(deserializer)?
263            .map(|bytes| PersonId::from_tx_id(TxId::from_bytes(bytes))))
264    }
265}