1use serde::{Deserialize, Serialize};
2
3pub use kcode_k1_audio_classification_format_error::FormatError;
4pub use kcode_k1_transaction_id::TxId;
5pub use kcode_speaker_v3_analysis::{ExecutedAnalysis, FeatureVector24, LocalSpeakerLabel};
6
7const EVENT_VERSION: u8 = 4;
8const QUEUE_TAG: u8 = 1;
9const TRANSCRIPTION_COMPLETE_TAG: u8 = 2;
10const FAILED_TAG: u8 = 3;
11const DISCARDED_TAG: u8 = 4;
12const LABEL_CONFIRMATION_TAG: u8 = 5;
13const PROGRESS_TAG: u8 = 6;
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
16pub enum FragmentStageV1 {
17 Queue,
18 Transcript,
19 SpeakerLabels,
20 SpeakerFeatures,
21 Structuring,
22 LabelConfirmation,
23}
24
25#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
26pub struct QueueV2 {
27 #[serde(with = "txid_serde")]
28 pub audio_object_id: TxId,
29}
30
31#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
32pub struct ProgressV1 {
33 #[serde(with = "txid_serde")]
34 pub fragment_id: TxId,
35 pub update: ProgressUpdateV1,
36}
37
38#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
39pub enum ProgressUpdateV1 {
40 LlmJobStarted {
41 sequence: u64,
42 stage: FragmentStageV1,
43 name: String,
44 },
45 LlmJobSucceeded {
46 sequence: u64,
47 },
48 LlmJobFailed {
49 sequence: u64,
50 error: String,
51 },
52 StageCompleted {
53 stage: FragmentStageV1,
54 },
55}
56
57#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
58pub struct TranscriptionCompleteV1 {
59 #[serde(with = "txid_serde")]
60 pub fragment_id: TxId,
61 pub analysis: ExecutedAnalysis,
62}
63
64#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
65pub struct FailedV2 {
66 #[serde(with = "txid_serde")]
67 pub fragment_id: TxId,
68 pub stage: FragmentStageV1,
69 pub llm_job_sequence: Option<u64>,
70 pub error: String,
71}
72
73#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
74pub struct DiscardedV2 {
75 #[serde(with = "txid_serde")]
76 pub fragment_id: TxId,
77}
78
79#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
80pub struct SpeakerLabelV1 {
81 pub speaker: LocalSpeakerLabel,
82 pub person_id: String,
83}
84
85#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
86pub struct LabelConfirmationV1 {
87 #[serde(with = "txid_serde")]
88 pub fragment_id: TxId,
89 #[serde(with = "txid_serde")]
90 pub interim_txid: TxId,
91 pub speakers: Vec<SpeakerLabelV1>,
92}
93
94#[allow(clippy::large_enum_variant)]
95#[derive(Debug, Clone, PartialEq)]
96pub enum AudioClassificationEventV3 {
97 Queue(QueueV2),
98 Progress(ProgressV1),
99 TranscriptionComplete(TranscriptionCompleteV1),
100 Failed(FailedV2),
101 Discarded(DiscardedV2),
102 LabelConfirmation(LabelConfirmationV1),
103}
104
105pub fn encode_event(event: &AudioClassificationEventV3) -> Result<Vec<u8>, FormatError> {
106 let (tag, body) = match event {
107 AudioClassificationEventV3::Queue(value) => (QUEUE_TAG, postcard::to_allocvec(value)),
108 AudioClassificationEventV3::TranscriptionComplete(value) => {
109 (TRANSCRIPTION_COMPLETE_TAG, postcard::to_allocvec(value))
110 }
111 AudioClassificationEventV3::Failed(value) => (FAILED_TAG, postcard::to_allocvec(value)),
112 AudioClassificationEventV3::Discarded(value) => {
113 (DISCARDED_TAG, postcard::to_allocvec(value))
114 }
115 AudioClassificationEventV3::LabelConfirmation(value) => {
116 (LABEL_CONFIRMATION_TAG, postcard::to_allocvec(value))
117 }
118 AudioClassificationEventV3::Progress(value) => (PROGRESS_TAG, postcard::to_allocvec(value)),
119 };
120 let body = body.map_err(|_| FormatError::InvalidEventBody)?;
121 let mut output = Vec::with_capacity(2 + body.len());
122 output.extend_from_slice(&[EVENT_VERSION, tag]);
123 output.extend_from_slice(&body);
124 Ok(output)
125}
126
127pub fn decode_event(bytes: &[u8]) -> Result<AudioClassificationEventV3, FormatError> {
128 if bytes.len() < 2 {
129 return Err(FormatError::Truncated);
130 }
131 if bytes[0] != EVENT_VERSION {
132 return Err(FormatError::UnsupportedVersion(bytes[0]));
133 }
134 match bytes[1] {
135 QUEUE_TAG => decode_body(&bytes[2..]).map(AudioClassificationEventV3::Queue),
136 TRANSCRIPTION_COMPLETE_TAG => {
137 decode_body(&bytes[2..]).map(AudioClassificationEventV3::TranscriptionComplete)
138 }
139 FAILED_TAG => decode_body(&bytes[2..]).map(AudioClassificationEventV3::Failed),
140 DISCARDED_TAG => decode_body(&bytes[2..]).map(AudioClassificationEventV3::Discarded),
141 LABEL_CONFIRMATION_TAG => {
142 decode_body(&bytes[2..]).map(AudioClassificationEventV3::LabelConfirmation)
143 }
144 PROGRESS_TAG => decode_body(&bytes[2..]).map(AudioClassificationEventV3::Progress),
145 tag => Err(FormatError::UnknownEventTag(tag)),
146 }
147}
148
149fn decode_body<T: for<'de> Deserialize<'de>>(bytes: &[u8]) -> Result<T, FormatError> {
150 let (value, remaining) =
151 postcard::take_from_bytes(bytes).map_err(|_| FormatError::InvalidEventBody)?;
152 if !remaining.is_empty() {
153 return Err(FormatError::TrailingBytes);
154 }
155 Ok(value)
156}
157
158mod txid_serde {
159 use super::TxId;
160 use serde::{Deserialize, Deserializer, Serialize, Serializer};
161
162 pub fn serialize<S: Serializer>(value: &TxId, serializer: S) -> Result<S::Ok, S::Error> {
163 value.into_bytes().serialize(serializer)
164 }
165
166 pub fn deserialize<'de, D: Deserializer<'de>>(deserializer: D) -> Result<TxId, D::Error> {
167 Ok(TxId::from_bytes(<[u8; 12]>::deserialize(deserializer)?))
168 }
169}
170
171#[cfg(test)]
172mod tests {
173 use super::*;
174 use kcode_speaker_v3_analysis::{
175 AnalysisEnvelope, GeminiCohort, OggAudioMetadata, StructuredAnalysis, StructuredSpeaker,
176 StructurerProvenance,
177 };
178 use serde::Serialize;
179
180 fn txid(seed: u8) -> TxId {
181 let mut bytes = [0; 12];
182 for (index, value) in bytes.iter_mut().enumerate() {
183 *value = seed.wrapping_add(index as u8);
184 }
185 TxId::from_bytes(bytes)
186 }
187
188 fn analysis() -> ExecutedAnalysis {
189 let mut ogg = vec![0; 29];
190 ogg[..4].copy_from_slice(b"OggS");
191 ogg[26] = 1;
192 ogg[27] = 1;
193 ExecutedAnalysis {
194 envelope: AnalysisEnvelope {
195 audio: OggAudioMetadata::from_bytes(&ogg, 1_250, Some("sample.ogg".into()))
196 .expect("Ogg metadata"),
197 analysis: StructuredAnalysis {
198 transcript: "Speaker 1".into(),
199 speakers: vec![StructuredSpeaker {
200 speaker: LocalSpeakerLabel::new(1).expect("speaker label"),
201 language: "English".into(),
202 features: FeatureVector24 {
203 median_f0_hz: Some(182.5),
204 ..FeatureVector24::default()
205 },
206 features_usable_for_training: true,
207 }],
208 },
209 gemini: GeminiCohort::new("gemini-3.1-pro"),
210 structurer: StructurerProvenance::new("terra-structurer"),
211 },
212 label_extractor: StructurerProvenance::new("terra-labeler"),
213 }
214 }
215
216 fn assert_golden<T: Serialize>(event: AudioClassificationEventV3, tag: u8, body: &T) {
217 let mut expected = vec![EVENT_VERSION, tag];
218 expected.extend(postcard::to_allocvec(body).expect("serialize golden body"));
219 assert_eq!(encode_event(&event), Ok(expected.clone()));
220 assert_eq!(decode_event(&expected), Ok(event));
221 }
222
223 #[test]
224 fn all_event_variants_match_v4_golden_envelopes() {
225 let queue = QueueV2 {
226 audio_object_id: txid(1),
227 };
228 assert_golden(
229 AudioClassificationEventV3::Queue(queue.clone()),
230 QUEUE_TAG,
231 &queue,
232 );
233
234 let complete = TranscriptionCompleteV1 {
235 fragment_id: txid(2),
236 analysis: analysis(),
237 };
238 assert_golden(
239 AudioClassificationEventV3::TranscriptionComplete(complete.clone()),
240 TRANSCRIPTION_COMPLETE_TAG,
241 &complete,
242 );
243
244 let failed = FailedV2 {
245 fragment_id: txid(3),
246 stage: FragmentStageV1::SpeakerFeatures,
247 llm_job_sequence: Some(7),
248 error: "failure".into(),
249 };
250 assert_golden(
251 AudioClassificationEventV3::Failed(failed.clone()),
252 FAILED_TAG,
253 &failed,
254 );
255
256 let discarded = DiscardedV2 {
257 fragment_id: txid(4),
258 };
259 assert_golden(
260 AudioClassificationEventV3::Discarded(discarded.clone()),
261 DISCARDED_TAG,
262 &discarded,
263 );
264
265 let confirmation = LabelConfirmationV1 {
266 fragment_id: txid(5),
267 interim_txid: txid(6),
268 speakers: vec![SpeakerLabelV1 {
269 speaker: LocalSpeakerLabel::new(1).expect("speaker label"),
270 person_id: "person-1".into(),
271 }],
272 };
273 assert_golden(
274 AudioClassificationEventV3::LabelConfirmation(confirmation.clone()),
275 LABEL_CONFIRMATION_TAG,
276 &confirmation,
277 );
278
279 let progress = ProgressV1 {
280 fragment_id: txid(7),
281 update: ProgressUpdateV1::LlmJobStarted {
282 sequence: 9,
283 stage: FragmentStageV1::Transcript,
284 name: "transcribe".into(),
285 },
286 };
287 assert_golden(
288 AudioClassificationEventV3::Progress(progress.clone()),
289 PROGRESS_TAG,
290 &progress,
291 );
292 }
293
294 #[test]
295 fn every_progress_update_roundtrips() {
296 let updates = [
297 ProgressUpdateV1::LlmJobStarted {
298 sequence: 1,
299 stage: FragmentStageV1::Transcript,
300 name: "transcript".into(),
301 },
302 ProgressUpdateV1::LlmJobSucceeded { sequence: 1 },
303 ProgressUpdateV1::LlmJobFailed {
304 sequence: 2,
305 error: "failed".into(),
306 },
307 ProgressUpdateV1::StageCompleted {
308 stage: FragmentStageV1::Structuring,
309 },
310 ];
311 for update in updates {
312 let event = AudioClassificationEventV3::Progress(ProgressV1 {
313 fragment_id: txid(8),
314 update,
315 });
316 let bytes = encode_event(&event).expect("encode progress");
317 assert_eq!(decode_event(&bytes), Ok(event));
318 }
319 }
320
321 #[test]
322 fn canonical_txid_encoding_is_twelve_raw_bytes() {
323 let id = txid(0);
324 let event = AudioClassificationEventV3::Queue(QueueV2 {
325 audio_object_id: id,
326 });
327 let mut expected = vec![EVENT_VERSION, QUEUE_TAG];
328 expected.extend_from_slice(id.as_bytes());
329 assert_eq!(encode_event(&event), Ok(expected));
330 }
331
332 #[test]
333 fn current_speaker_analysis_serialization_is_embedded_unchanged() {
334 let fragment_id = txid(12);
335 let analysis = analysis();
336 let expected_analysis = postcard::to_allocvec(&analysis).expect("serialize analysis");
337 let bytes = encode_event(&AudioClassificationEventV3::TranscriptionComplete(
338 TranscriptionCompleteV1 {
339 fragment_id,
340 analysis,
341 },
342 ))
343 .expect("encode complete event");
344 assert_eq!(&bytes[..2], &[EVENT_VERSION, TRANSCRIPTION_COMPLETE_TAG]);
345 assert_eq!(&bytes[2..14], fragment_id.as_bytes());
346 assert_eq!(&bytes[14..], expected_analysis);
347 }
348
349 #[test]
350 fn malformed_unknown_version_and_trailing_inputs_are_rejected() {
351 assert_eq!(decode_event(&[]), Err(FormatError::Truncated));
352 assert_eq!(decode_event(&[EVENT_VERSION]), Err(FormatError::Truncated));
353 assert_eq!(
354 decode_event(&[3, QUEUE_TAG]),
355 Err(FormatError::UnsupportedVersion(3))
356 );
357 assert_eq!(
358 decode_event(&[EVENT_VERSION, 7]),
359 Err(FormatError::UnknownEventTag(7))
360 );
361 assert_eq!(
362 decode_event(&[EVENT_VERSION, QUEUE_TAG]),
363 Err(FormatError::InvalidEventBody)
364 );
365
366 let event = AudioClassificationEventV3::Discarded(DiscardedV2 {
367 fragment_id: txid(9),
368 });
369 let mut bytes = encode_event(&event).expect("encode discarded");
370 bytes.push(0);
371 assert_eq!(decode_event(&bytes), Err(FormatError::TrailingBytes));
372 }
373}