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}